diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 2819a99d..9c065b73 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -359,7 +359,9 @@ def _head_object( The result is cached under the path, or under the version-qualified path for an explicit version. An explicitly requested ``"null"`` - version is not cached. A missing object evicts its entry. + version is not cached. A missing object evicts its entry and, unless + a version was requested, the cached listing of its parent that still + lists it. Args: path: The object path, optionally with a versionId query. @@ -394,6 +396,13 @@ def _head_object( ) except FileNotFoundError: self._evict_cache(path) + if not version_id: + # Evict the cached listing of the parent only if it still + # lists the path. + parent_key = (self._parent(path), "/") + files = self.dircache.get(parent_key) + if files and any(f.name == path for f in files): + self._evict_cache(parent_key) return None if self.version_aware and not version_id: # Pin the version of the object so that subsequent reads see @@ -628,12 +637,18 @@ def info(self, path: str, **kwargs) -> S3Object: """Return information about an S3 path. Uses the directory cache first: a cached entry for the path is - returned, a cached listing of the path itself makes it a directory, - and a cached listing of its parent without it means it does not exist. + returned, the entry of the path in a cached listing of its parent is + returned, preferring an object to a key prefix of the same name as + HeadObject does, and a cached listing of its parent without it means + it does not exist. + The cached bucket listing holds only the buckets that the caller owns, + so a bucket missing from it is looked up with HeadBucket. Otherwise, a key path is looked up with HeadObject and, if no object exists, with a ListObjectsV2 request (``Delimiter="/"``, ``MaxKeys=1``) that checks whether it is a key prefix; a bucket path - is looked up with HeadBucket. With ``version_aware``, a cached file + is looked up with HeadBucket. If these requests find a listed object + missing, or find a key prefix, the cached listing of the parent is + removed. With ``version_aware``, a cached file entry without a version ID is looked up again. With an explicit version, the cached entries of the path are skipped, and the HeadObject result is cached under the version-qualified path apart @@ -647,16 +662,16 @@ def info(self, path: str, **kwargs) -> S3Object: version_id: The version ID to look up when the path has none. Returns: - S3Object describing the bucket, directory, or file. + S3Object describing the bucket, directory, or file. The root path + (``""``, ``"/"`` or ``"s3://"``) is a directory. Raises: FileNotFoundError: If the path does not exist. """ refresh = kwargs.pop("refresh", False) path = self._strip_protocol(path) - bucket, key, path_version_id = self.parse_path(path) - version_id = path_version_id if path_version_id else kwargs.pop("version_id", None) if path in ["/", ""]: + # parse_path rejects the root path. return S3Object( init={ "ContentLength": 0, @@ -666,17 +681,24 @@ def info(self, path: str, **kwargs) -> S3Object: "LastModified": None, }, type=S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, - bucket=bucket, + bucket="", key=None, version_id=None, ) + bucket, key, path_version_id = self.parse_path(path) + version_id = path_version_id if path_version_id else kwargs.pop("version_id", None) # Cached entries describe the current version of a path, so an # explicit version uses only the HeadObject cache of that version. if not refresh and not version_id: caches: list[S3Object] | S3Object | None = self._ls_from_cache(path) if caches is not None: if isinstance(caches, list): - cache = next((c for c in caches if c.name == path), None) + matches = [c for c in caches if c.name == path] + # A key can be both an object and a key prefix. + cache = next( + (c for c in matches if c.type == S3ObjectType.S3_OBJECT_TYPE_FILE), + next(iter(matches), None), + ) elif caches.name == path: cache = caches else: @@ -720,6 +742,9 @@ def info(self, path: str, **kwargs) -> S3Object: or response.get("Contents", []) or response.get("CommonPrefixes", []) ): + # Nothing caches the key prefix, and the cached listing of the + # parent may predate it. + self._evict_cache((self._parent(path), "/")) return self._directory_object(bucket, key.rstrip("/") if key else None, version_id) raise FileNotFoundError(path) @@ -898,7 +923,8 @@ def exists(self, path: str, **kwargs) -> bool: refresh: If True, bypass the cache and query S3. Returns: - True if the path exists, False otherwise. + True if the path exists, False otherwise. A bucket that HeadBucket + denies access to (403) exists. Example: >>> fs = S3FileSystem() @@ -919,15 +945,14 @@ def exists(self, path: str, **kwargs) -> bool: return bool(info) except FileNotFoundError: return False - if not refresh: - if self.dircache.get(bucket, False): - return True - try: - if self._ls_from_cache(bucket): - return True - except FileNotFoundError: - pass - file = self._head_bucket(bucket, refresh=refresh) + if not refresh and self._ls_from_cache(bucket): + return True + try: + file = self._head_bucket(bucket, refresh=refresh) + except PermissionError: + # HeadBucket answers 403 for a bucket that exists but that the + # caller may not access. + return True return bool(file) def rm_file(self, path: str, **kwargs) -> None: @@ -1231,8 +1256,8 @@ def mkdir(self, path: str, create_parents: bool = True, **kwargs) -> None: ) except botocore.exceptions.ParamValidationError as e: raise ValueError(f"Bucket create failed {bucket!r}: {e}") from e - # invalidate_cache walks parent paths and never pops the root - # entry itself, so evict the cached bucket listing directly. + # invalidate_cache of the bucket keeps the cached bucket + # listing, so evict it directly. self._evict_cache("") self.invalidate_cache(bucket) else: @@ -1304,8 +1329,8 @@ def rmdir(self, path: str) -> None: Bucket=bucket, ) self.invalidate_cache(bucket) - # invalidate_cache walks parent paths and never pops the root - # entry itself, so evict the cached bucket listing directly. + # invalidate_cache of the bucket keeps the cached bucket listing, + # so evict it directly. self._evict_cache("") def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]: @@ -2222,6 +2247,8 @@ def modified(self, path: str) -> datetime: def invalidate_cache(self, path: str | None = None) -> None: """Remove the cached entries of the path and its parent paths. + The cached bucket listing is removed only by the root path (``""``, + ``"/"`` or ``"s3://"``), not by the paths of buckets or keys. A version-qualified path invalidates the version under every query spelling that ``parse_path`` accepts, and also the object path without the version, because deleting or copying a version can change the @@ -2234,6 +2261,8 @@ def invalidate_cache(self, path: str | None = None) -> None: self.dircache.clear() else: path = self._strip_protocol(path) + if not path: + self._evict_cache("") while path: # parse_path does not accept "?" in keys, so it starts the # versionId query. @@ -2272,17 +2301,38 @@ def _evict_cache(self, key: str | tuple[str, str]) -> None: def _ls_from_cache(self, path: str) -> list[S3Object] | S3Object | None: """Check the dircache for a cached entry of the path. - fsspec's implementation assumes every dircache value is a listing, - but S3FileSystem also caches a single S3Object under the object's own - path (HeadObject/HeadBucket results). Guard the parent lookup so that - looking up a child path of a cached object does not fail, and fall - through to the S3 API instead. + fsspec's implementation looks up listings under the path itself, but + S3FileSystem caches a single S3Object under the path of an object or + a bucket (HeadObject/HeadBucket results), the bucket listing under + ``""``, and the other listings under ``(path, delimiter)`` (see + ``_ls_dirs``). + + Args: + path: The path without the protocol. + + Returns: + The cached entry of the path, the entries of a cached parent + listing named as the path, or None if no cached entry describes + the path. A listing of the path itself is not used, because it + cannot tell whether an object of the same name exists. A + version-qualified path uses only its own entry, because listings + describe the current versions. + + Raises: + FileNotFoundError: If a cached listing of the parent directory of + a key path does not contain the path. """ cache = self.dircache.get(path.rstrip("/")) if cache is not None: return cast("list[S3Object] | S3Object", cache) - parent_cache = self.dircache.get(self._parent(path)) - if isinstance(parent_cache, list): + _, key, version_id = self.parse_path(path) + if version_id: + return None + if key: + parent_cache = self.dircache.get((self._parent(path), "/")) + else: + parent_cache = self.dircache.get("") + if parent_cache is not None: files = [ f for f in parent_cache @@ -2294,7 +2344,10 @@ def _ls_from_cache(self, path: str) -> list[S3Object] | S3Object | None: ] if files: return files - raise FileNotFoundError(path) + if key: + raise FileNotFoundError(path) + # The bucket listing holds only the buckets that the caller owns, + # so a bucket missing from it is looked up with HeadBucket. return None def _open( @@ -2613,9 +2666,15 @@ def __init__( self._details = info if size is None: size = info.get("size") - elif "a" in mode and fs.exists(path): - append_info = fs.info(path) - if append_info.get("size", 0) < fs.MULTIPART_UPLOAD_MIN_PART_SIZE: + elif "a" in mode: + # The rewritten object keeps the metadata of the existing one, + # which a cached listing entry lacks, so look up the object. + with contextlib.suppress(FileNotFoundError): + append_info = fs.info(path, refresh=True) + if ( + append_info is not None + and append_info.get("size", 0) < fs.MULTIPART_UPLOAD_MIN_PART_SIZE + ): # Too small to be a part of a multipart upload: rewritten # from the buffer. append_data = fs.cat(path) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 76253929..0e6f262a 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -227,6 +227,175 @@ def test_ls_from_cache_with_cached_object(self): # every cache value is a listing and raises TypeError here). assert fs._ls_from_cache("bucket/key/child") is None + def test_info_uses_cached_listings(self): + # GH-965: listings cached under (path, delimiter) answer info() and + # exists() without HeadObject or ListObjectsV2 requests. + fs = self._make_fs() + fs._call.return_value = { + "CommonPrefixes": [{"Prefix": "d/sub/"}], + "Contents": [{"Key": "d/direct", "Size": 4, "ETag": '"etag"'}], + } + fs.ls("s3://bucket/d") + fs._call.reset_mock() + + file = fs.info("s3://bucket/d/direct") + assert (file.type, file.size, file.etag) == (S3ObjectType.S3_OBJECT_TYPE_FILE, 4, '"etag"') + assert fs.isdir("s3://bucket/d/sub") + with pytest.raises(FileNotFoundError): + fs.info("s3://bucket/d/missing") + assert not fs.exists("s3://bucket/d/missing") + fs._call.assert_not_called() + + def test_info_prefers_listed_object_to_prefix_of_same_name(self): + # A key that is both an object and a key prefix is an object, as the + # uncached lookup with HeadObject finds. + fs = self._make_fs() + fs.dircache[("bucket", "/")] = [ + fs._directory_object("bucket", "d"), + self._file_object("d"), + ] + + assert fs.isfile("s3://bucket/d") + fs._call.assert_not_called() + + def test_info_does_not_use_listing_of_path(self): + # The listing of the path cannot tell whether an object of the same + # name exists. + fs = self._make_fs() + fs.dircache[("bucket/d", "/")] = [self._file_object("d/direct")] + fs._call.return_value = {"ContentLength": 4} + + assert fs.isfile("s3://bucket/d") + fs._call.assert_called_once_with(fs._client.head_object, Bucket="bucket", Key="d") + + @pytest.mark.parametrize( + ("version_aware", "lookup"), + [ + pytest.param( + False, lambda fs: fs.exists("s3://bucket/d/key", refresh=True), id="exists" + ), + pytest.param(False, lambda fs: fs.ls("s3://bucket/d/key", refresh=True), id="ls"), + # The listed entry has no version, so it is looked up again. + pytest.param(True, lambda fs: fs.isfile("s3://bucket/d/key"), id="version_aware"), + ], + ) + def test_missing_object_drops_cached_parent_listing(self, version_aware, lookup): + # A lookup that finds a listed object deleted is not contradicted by + # the listing afterwards. + fs = self._make_fs() + fs.version_aware = version_aware + fs.dircache[("bucket/d", "/")] = [self._file_object("d/key")] + + def call(method, **kwargs): + if method == fs._client.head_object: + raise FileNotFoundError + return {} + + fs._call.side_effect = call + + lookup(fs) + assert ("bucket/d", "/") not in fs.dircache + assert not fs.exists("s3://bucket/d/key") + + def test_refreshed_prefix_drops_cached_parent_listing(self): + # A key prefix created after the parent was listed is not reported + # missing by the listing after a refreshed lookup finds it. + fs = self._make_fs() + fs.dircache[("bucket/d", "/")] = [self._file_object("d/key")] + + def call(method, **kwargs): + if method == fs._client.head_object: + raise FileNotFoundError + return {"KeyCount": 1} + + fs._call.side_effect = call + + assert fs.isdir("s3://bucket/d/new") is False + assert ( + fs.info("s3://bucket/d/new", refresh=True).type == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY + ) + assert fs.isdir("s3://bucket/d/new") + + def test_missing_object_keeps_cached_parent_listing_without_it(self): + fs = self._make_fs() + fs.dircache[("bucket/d", "/")] = [self._file_object("d/key")] + fs._call.side_effect = FileNotFoundError + + assert fs._head_object("bucket/d/other", refresh=True) is None + assert ("bucket/d", "/") in fs.dircache + + def test_info_version_aware_heads_listed_file(self): + fs = self._make_fs() + fs.version_aware = True + fs.dircache[("bucket/d", "/")] = [self._file_object("d/direct")] + fs._call.return_value = {"ContentLength": 4, "VersionId": "v1"} + + # The listed entry has no version to pin. + assert fs.info("s3://bucket/d/direct").version_id == "v1" + fs._call.assert_called_once_with(fs._client.head_object, Bucket="bucket", Key="d/direct") + + def test_exists_version_ignores_cached_parent_listing(self): + # The listing describes the current versions, so a version missing + # from it is looked up with HeadObject. + fs = self._make_fs() + fs.dircache[("bucket/d", "/")] = [self._file_object("d/other")] + fs._call.return_value = {"ContentLength": 4} + + assert fs.exists("s3://bucket/d/direct?versionId=v1") + fs._call.assert_called_once_with( + fs._client.head_object, Bucket="bucket", Key="d/direct", VersionId="v1" + ) + + def test_info_bucket_missing_from_bucket_listing(self): + # GH-980: the bucket listing holds only the buckets of the caller. + fs = self._make_fs() + fs.dircache[""] = [fs._directory_object("mine", None)] + fs._call.return_value = {} + + info = fs.info("s3://other-account-bucket") + assert info.storage_class == S3StorageClass.S3_STORAGE_CLASS_BUCKET + fs._call.assert_called_once_with(fs._client.head_bucket, Bucket="other-account-bucket") + fs._call.reset_mock() + assert fs.isdir("s3://other-account-bucket") + assert fs.isdir("s3://mine") + fs._call.assert_not_called() + + @pytest.mark.parametrize("path", ["", "/", "s3://"]) + def test_info_root(self, path): + fs = self._make_fs() + + info = fs.info(path) + assert (info.name, info.type, info.size) == ("", S3ObjectType.S3_OBJECT_TYPE_DIRECTORY, 0) + assert fs.isdir(path) + assert not fs.isfile(path) + assert fs.size(path) == 0 + fs._call.assert_not_called() + + @pytest.mark.parametrize("path", ["", "/", "s3://"]) + def test_invalidate_cache_root_drops_bucket_listing(self, path): + fs = self._make_fs() + fs.dircache[""] = [fs._directory_object("bucket", None)] + fs.dircache["bucket"] = fs._directory_object("bucket", None) + + fs.invalidate_cache(path) + assert list(fs.dircache) == ["bucket"] + + def test_exists_bucket_access_denied(self): + # GH-980: HeadBucket answers 403 for a bucket that exists but that + # the caller may not access. + fs = self._make_fs() + fs._call.side_effect = PermissionError + + assert fs.exists("s3://not-my-bucket") + fs.makedirs("s3://not-my-bucket/prefix", exist_ok=True) + assert ( + fs._call.call_args_list + == [ + mock.call(fs._client.head_bucket, Bucket="not-my-bucket"), + ] + * 2 + ) + def test_invalidate_cache_drops_listings_of_path_and_parents(self): fs = self._make_fs() invalidated = [ @@ -889,18 +1058,42 @@ def test_open_invalid_for_writing(self, monkeypatch, mode, path, block_size, mat assert unraisable == [] + def test_open_append_keeps_metadata_of_listed_object(self): + # A cached listing entry lacks the metadata that the rewritten object + # keeps, so the append looks up the object. + fs = self._make_fs() + fs.default_cache_type = "bytes" + fs.dircache[("bucket", "/")] = [self._file_object("key")] + fs._call.return_value = { + "ContentLength": 2, + "ContentType": "text/plain", + "Metadata": {"k": "v"}, + } + fs.cat = mock.MagicMock(return_value=b"aa") + fs._put_object = mock.MagicMock() + + with fs.open("s3://bucket/key", "ab") as f: + f.write(b"bb") + fs._call.assert_called_once_with(fs._client.head_object, Bucket="bucket", Key="key") + request = fs._put_object.call_args.kwargs + assert (request["body"], request["ContentType"], request["Metadata"]) == ( + b"aabb", + "text/plain", + {"k": "v"}, + ) + def test_open_append_lookup_failure(self, monkeypatch): # GH-976: an append whose lookup of the existing object fails leaves # no half-initialized file, whose garbage collection would close it. fs = self._make_fs() fs.default_cache_type = "bytes" - def exists(path): + def info(path, **kwargs): # A new exception each time: one kept by a mock would keep its # traceback, and the file, alive. raise PermissionError("denied") - fs.exists = exists + fs.info = info unraisable = [] monkeypatch.setattr(sys, "unraisablehook", unraisable.append) @@ -1555,14 +1748,11 @@ def test_ls_versions_object_path_falls_back_to_the_key(self): def test_dir_filesystem(self): # DirFileSystem copies every entry with copy() before renaming it. fs = self._make_fs() - fs._call.side_effect = [ - { - "CommonPrefixes": [{"Prefix": "path/dir/"}], - "Contents": [{"Key": "path/key", "Size": 4}], - "IsTruncated": False, - }, - {"ContentLength": 4, "ETag": '"etag"'}, - ] + fs._call.return_value = { + "CommonPrefixes": [{"Prefix": "path/dir/"}], + "Contents": [{"Key": "path/key", "Size": 4}], + "IsTruncated": False, + } dir_fs = DirFileSystem(path="bucket/path", fs=fs) actual = dir_fs.ls("", detail=True) @@ -1577,7 +1767,8 @@ def test_dir_filesystem(self): "bucket/path/key", ] assert fs.info("bucket/path/key").name == "bucket/path/key" - assert fs._call.call_count == 2 + # info() answers from the cached listing. + fs._call.assert_called_once() def test_metadata_with_version_id(self): fs = self._make_fs() @@ -2802,7 +2993,6 @@ def _make_append_fs(existing: bytes): fs.MULTIPART_UPLOAD_MIN_PART_SIZE = 4 fs.MULTIPART_UPLOAD_MAX_PART_SIZE = 64 fs.MULTIPART_UPLOAD_MAX_PARTS = S3FileSystem.MULTIPART_UPLOAD_MAX_PARTS - fs.exists.return_value = True fs.info.return_value = S3Object( init={"ContentLength": len(existing)}, type=S3ObjectType.S3_OBJECT_TYPE_FILE,