diff --git a/gcsfs/core.py b/gcsfs/core.py index 918c74d76..3a14f3eba 100644 --- a/gcsfs/core.py +++ b/gcsfs/core.py @@ -3007,7 +3007,7 @@ def _convert_fixed_key_metadata(metadata, *, from_google=False): async def upload_chunk(fs, location, data, offset, size, content_type): """ Uploads a chunk of data. This function has a conditional path to support - experimental features for Zonal buckets to append data using gRPC. + experimental features for Zonal and RCU buckets to append data using gRPC. """ from google.cloud.storage.asyncio.async_appendable_object_writer import ( AsyncAppendableObjectWriter, @@ -3117,7 +3117,7 @@ async def initiate_upload( ): """ Initiates a resumable upload. This function has a conditional path to support - experimental features for Zonal buckets to append data using gRPC, returning an + experimental features for Zonal and RCU buckets to append data using gRPC, returning an "AsyncAppendableObjectWriter" instance as location. """ from .extended_gcsfs import ExtendedGcsFileSystem @@ -3125,7 +3125,7 @@ async def initiate_upload( # Explicit type checking is used to ensure only the ExtendedGcsFileSystem # enters this path, ruling out false positives from mocks or coincidentally matching attributes. - if isinstance(fs, ExtendedGcsFileSystem) and await fs._is_zonal_bucket(bucket): + if isinstance(fs, ExtendedGcsFileSystem) and await fs._uses_grpc_data_path(bucket): return await ext_initiate_upload( fs, @@ -3178,14 +3178,14 @@ async def simple_upload( ): """ Performs a simple, single-request upload. This function has a conditional path to support - experimental features for Zonal buckets to upload data using gRPC. + experimental features for Zonal and RCU buckets to upload data using gRPC. """ from .extended_gcsfs import ExtendedGcsFileSystem from .extended_gcsfs import simple_upload as ext_simple_upload # Explicit type checking is used to ensure only the ExtendedGcsFileSystem # enters this path, ruling out false positives from mocks or coincidentally matching attributes. - if isinstance(fs, ExtendedGcsFileSystem) and await fs._is_zonal_bucket(bucket): + if isinstance(fs, ExtendedGcsFileSystem) and await fs._uses_grpc_data_path(bucket): return await ext_simple_upload( fs, diff --git a/gcsfs/extended_gcsfs.py b/gcsfs/extended_gcsfs.py index 4a2ce04b8..5b4863efb 100644 --- a/gcsfs/extended_gcsfs.py +++ b/gcsfs/extended_gcsfs.py @@ -159,7 +159,7 @@ class ExtendedGcsFileSystem(HnsDirCacheUpdater, GCSFileSystem): def __init__( self, *args, - finalize_on_close=False, + finalize_on_close=None, mrd_pool_cache_size=16, max_mrd_pool_cache_queue_size=16, **kwargs, @@ -167,8 +167,11 @@ def __init__( """ Parameters ---------- - finalize_on_close : bool, default False - By default, files in zonal buckets are left unfinalized to allow appends. + finalize_on_close : bool, optional + If set, overrides the default finalization behavior when closing + appendable gRPC writers. When ``None`` (the default), files in + Zonal buckets default to ``False`` (left unfinalized to allow + appends) and files in RCU buckets default to ``True``. mrd_pool_cache_size : int, default 16 Maximum number of idle pools to retain in the cache. max_mrd_pool_cache_queue_size : int, default 16 @@ -189,7 +192,7 @@ def __init__( if k.startswith("retry_") and k[6:] in valid_keys and v is not None } super().__init__(*args, **kwargs) - # By default, files in zonal buckets are left unfinalized to allow appends. + # When None, defaults to False for Zonal buckets and True for RCU buckets. self.finalize_on_close = finalize_on_close self._grpc_client = None self._storage_control_client = None @@ -221,7 +224,7 @@ def __init__( ) async def _get_threshold_for_disk_reads(self, bucket): - if await self._is_zonal_bucket(bucket): + if await self._uses_grpc_data_path(bucket): return ( 5 * 1024 * 1024 ) # Thanks to our in house, zero copy DirectMemmoveBuffer @@ -384,9 +387,20 @@ async def _get_bucket_layout(self, bucket): else: namespace_layout = NamespaceLayout.NON_HIERARCHICAL + rapid_cache_info = getattr(response, "rapid_cache_info", None) + cache_type = ( + getattr(rapid_cache_info, "cache_type", "") if rapid_cache_info else "" + ) + if cache_type == "rapid-cache-ultra": + caching_tier = CachingTier.RAPID_CACHE_ULTRA + elif cache_type == "rapid-cache": + caching_tier = CachingTier.RAPID_CACHE + else: + caching_tier = CachingTier.NONE + return BucketLayout( storage_class=storage_class, - caching_tier=CachingTier.NONE, + caching_tier=caching_tier, namespace_layout=namespace_layout, ) except api_exceptions.NotFound: @@ -421,7 +435,11 @@ def _open( """ bucket, _, _ = self.split_path(path) bucket_layout = self._sync_lookup_bucket_layout(bucket) - file_cls = ZonalFile if bucket_layout.is_zonal else GCSFile + use_zonal_file = "+" not in mode and bucket_layout.uses_grpc_data_path + file_cls = ZonalFile if use_zonal_file else GCSFile + finalize_on_close = self._resolve_finalize_on_close( + bucket_layout, kwargs.pop("finalize_on_close", None) + ) return file_cls( self, @@ -435,7 +453,7 @@ def _open( autocommit=autocommit, fixed_key_metadata=fixed_key_metadata, generation=generation, - finalize_on_close=kwargs.pop("finalize_on_close", self.finalize_on_close), + finalize_on_close=finalize_on_close, **kwargs, ) @@ -504,6 +522,43 @@ async def _uses_grpc_data_path(self, bucket): bucket_layout = await self._lookup_bucket_layout(bucket) return bucket_layout.uses_grpc_data_path + def _resolve_finalize_on_close(self, bucket_layout, finalize_on_close=None): + """Return the per-call value, else the filesystem default, else True for RCU.""" + if finalize_on_close is not None: + return finalize_on_close + fs_default = getattr(self, "finalize_on_close", None) + if fs_default is not None: + return fs_default + return bucket_layout.is_rcu + + async def _init_aaow( + self, + bucket, + key, + generation=None, + flush_interval_bytes=None, + if_generation_match=None, + ): + """Open an AsyncAppendableObjectWriter for ``bucket/key``. + + RCU buckets default new objects to the ``STANDARD`` storage class, + which appendable writes reject, so writes to them request + ``storage_class="RAPID"``: the object lands in the zonal cache and the + server later copies it to regional storage. Zonal buckets already + default to ``RAPID``, so their writes leave ``storage_class`` unset. + """ + await self._get_grpc_client() + storage_class = "RAPID" if await self._is_rcu_bucket(bucket) else None + return await zb_hns_utils.init_aaow( + self.grpc_client, + bucket, + key, + generation, + flush_interval_bytes, + if_generation_match=if_generation_match, + storage_class=storage_class, + ) + @staticmethod def _resolve_cache_config(kwargs): """Resolves cache_type and cache_source from kwargs if not already provided.""" @@ -677,16 +732,17 @@ async def _cat_file( mrd=None, **kwargs, ): - """Fetch a file's contents as bytes, with an optimized path for Zonal buckets. + """Fetch a file's contents as bytes, with an optimized path for Zonal and RCU buckets. - This method overrides the parent `_cat_file` to read objects in Zonal buckets using gRPC. + This method overrides the parent `_cat_file` to read objects in Zonal and RCU buckets using gRPC. Args: path (str): The full GCS path to the file (e.g., "bucket/object"). start (int, optional): The starting byte position to read from. end (int, optional): The ending byte position to read to. mrd (AsyncMultiRangeDownloader, MRDPool, optional): An existing multi-range - downloader instance or a pool of MRD. If not provided, a new one will be created for Zonal buckets. + downloader instance or a pool of MRD. If not provided, a new one will be + created for Zonal and RCU buckets. Returns: bytes: The content of the file or file range. @@ -701,8 +757,8 @@ async def _cat_file( if mrd is None: bucket, object_name, generation = self.split_path(path) - if not await self._is_zonal_bucket(bucket): - # Fall back to default implementation if not a zonal bucket + if not await self._uses_grpc_data_path(bucket): + # Fall back to default implementation if not a Zonal or RCU bucket return await super()._cat_file( path, start=start, end=end, concurrency=concurrency, **kwargs ) @@ -968,19 +1024,19 @@ async def _cat_ranges( ): """Get the contents of byte ranges from one or more files. - For Zonal buckets, ranges are grouped per object and served by a single + For Zonal and RCU buckets, ranges are grouped per object and served by a single MRD pool per object. Each object's ranges are split into byte-balanced groups, one per pooled MRD, and every group is sent as one multi-range ``download_ranges`` call. ``batch_size`` bounds the total number of concurrent MRD streams across all objects, subject to the per-object cap (``-1`` means at most one per range; ``None`` or ``0`` select the - fsspec default). Non-zonal paths, and any call with ``max_gap`` set, + fsspec default). Paths in other buckets, and any call with ``max_gap`` set, use the default fsspec implementation. - Zonal-only keyword arguments: + Zonal and RCU keyword arguments: concurrency (int, optional): Maximum MRD streams per object, instead of ``MAX_ZONAL_STREAMS_PER_OBJECT``. Also forwarded to - the fsspec path for non-zonal files. + the fsspec path for files in other buckets. With ``on_error="return"``, a failure while downloading a group of ranges is returned for every range in that group, and a failure to @@ -1015,7 +1071,7 @@ async def _cat_ranges( zonal, other = {}, [] for p, idxs in by_path.items(): bucket, _, _ = self.split_path(p) - if await self._is_zonal_bucket(bucket): + if await self._uses_grpc_data_path(bucket): zonal[p] = idxs else: other.extend(idxs) @@ -2039,10 +2095,14 @@ async def _put_file( ): """Upload a local file. - This method is optimized for Zonal buckets, using gRPC for uploads. - In zonal buckets, file is left *unfinalized* by default unless - `finalize_on_close` is set to True. - For non-Zonal buckets, it delegates to the parent class's implementation. + This method is optimized for Zonal and RCU buckets, using gRPC for uploads. + By default, files in Zonal buckets are left *unfinalized* (unless + `finalize_on_close` is set to True), whereas files in RCU buckets are + finalized on close by default. + On Zonal and RCU buckets an overwrite replaces the existing object as + soon as the writer opens, so a failed write leaves the partially written + object in its place, unfinalized. + For other buckets, it delegates to the parent class's implementation. Parameters ---------- @@ -2051,24 +2111,24 @@ async def _put_file( rpath: str Path on GCS to upload the file to. metadata: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. consistency: str, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. content_type: str, optional - Unsupported for Zonal buckets and will be ignored except for the default. + Unsupported for Zonal and RCU buckets and will be ignored except for the default. chunksize: int, optional The size of chunks to upload data in. callback: fsspec.callbacks.Callback, optional Callback to monitor the upload progress. fixed_key_metadata: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. mode: str, optional The write mode, either 'overwrite' or 'create'. if_generation_match: int, optional Only write if the object's current generation matches this value. """ bucket, key, generation = self.split_path(rpath) - if not await self._is_zonal_bucket(bucket): + if not await self._uses_grpc_data_path(bucket): return await super()._put_file( lpath, rpath, @@ -2096,26 +2156,35 @@ async def _put_file( or (content_type and content_type != "application/octet-stream") ): logger.warning( - "Zonal buckets do not support content_type, metadata, " + "Zonal and RCU buckets do not support content_type, metadata, " "fixed_key_metadata or consistency during upload. " "These parameters will be ignored." ) - await self._get_grpc_client() - writer = await zb_hns_utils.init_aaow( - self.grpc_client, - bucket, - key, - if_generation_match=_resolve_generation_match(mode, if_generation_match), + finalize_on_close = self._resolve_finalize_on_close( + await self._lookup_bucket_layout(bucket), kwargs.get("finalize_on_close") ) - - try: - with open(lpath, "rb") as f: - await writer.append_from_file(f, block_size=chunksize) - finally: - finalize_on_close = kwargs.get("finalize_on_close", self.finalize_on_close) - await zb_hns_utils.close_aaow(writer, finalize_on_close=finalize_on_close) - - await self._write_file_cache_update(rpath) + with open(lpath, "rb") as f: + writer = await self._init_aaow( + bucket, + key, + if_generation_match=_resolve_generation_match( + mode, if_generation_match + ), + ) + # The open writer has already replaced the object, so update the + # cache even if the write fails. + try: + try: + await writer.append_from_file(f, block_size=chunksize) + except BaseException: + # Leave a partially written object unfinalized. + await zb_hns_utils.close_aaow(writer, finalize_on_close=False) + raise + await zb_hns_utils.close_aaow( + writer, finalize_on_close=finalize_on_close + ) + finally: + await self._write_file_cache_update(rpath) async def _pipe_file( self, @@ -2132,10 +2201,14 @@ async def _pipe_file( ): """Upload bytes to a file. - This method is optimized for Zonal buckets, using gRPC for uploads. - In zonal buckets, file is left *unfinalized* by default unless - `finalize_on_close` is set to True. - For non-Zonal buckets, it delegates to the parent class's implementation. + This method is optimized for Zonal and RCU buckets, using gRPC for uploads. + By default, files in Zonal buckets are left *unfinalized* (unless + `finalize_on_close` is set to True), whereas files in RCU buckets are + finalized on close by default. + On Zonal and RCU buckets an overwrite replaces the existing object as + soon as the writer opens, so a failed write leaves the partially written + object in its place, unfinalized. + For other buckets, it delegates to the parent class's implementation. Parameters ---------- @@ -2144,13 +2217,13 @@ async def _pipe_file( data: bytes The content to write to the file. metadata: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. consistency: str, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. content_type: str, optional - Unsupported for Zonal buckets and will be ignored, except for the default. + Unsupported for Zonal and RCU buckets and will be ignored, except for the default. fixed_key_metadata: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. chunksize: int, optional The size of chunks to upload data in. mode: str, optional @@ -2159,7 +2232,7 @@ async def _pipe_file( Only write if the object's current generation matches this value. """ bucket, key, generation = self.split_path(path) - if not await self._is_zonal_bucket(bucket): + if not await self._uses_grpc_data_path(bucket): return await super()._pipe_file( path, data, @@ -2178,25 +2251,31 @@ async def _pipe_file( or (content_type and content_type != "application/octet-stream") ): logger.warning( - "Zonal buckets do not support content_type, metadata or " + "Zonal and RCU buckets do not support content_type, metadata or " "fixed_key_metadata during upload. These parameters will be ignored." ) - await self._get_grpc_client() - writer = await zb_hns_utils.init_aaow( - self.grpc_client, + finalize_on_close = self._resolve_finalize_on_close( + await self._lookup_bucket_layout(bucket), kwargs.get("finalize_on_close") + ) + writer = await self._init_aaow( bucket, key, if_generation_match=_resolve_generation_match(mode, if_generation_match), ) + # The open writer has already replaced the object, so update the cache + # even if the write fails. try: - with memoryview(data) as data_view: - for i in range(0, len(data_view), chunksize): - await writer.append(data_view[i : i + chunksize]) - finally: - finalize_on_close = kwargs.get("finalize_on_close", self.finalize_on_close) + try: + with memoryview(data) as data_view: + for i in range(0, len(data_view), chunksize): + await writer.append(data_view[i : i + chunksize]) + except BaseException: + # Leave a partially written object unfinalized. + await zb_hns_utils.close_aaow(writer, finalize_on_close=False) + raise await zb_hns_utils.close_aaow(writer, finalize_on_close=finalize_on_close) - - await self._write_file_cache_update(path) + finally: + await self._write_file_cache_update(path) async def _get_file_request( self, rpath, lpath, *args, headers=None, callback=None, **kwargs @@ -2204,8 +2283,8 @@ async def _get_file_request( """ Downloads a file from GCS to a local path. - For Zonal buckets, it uses gRPC client for optimized downloads. - For Standard buckets, it delegates to the parent class implementation. + For Zonal and RCU buckets, it uses gRPC client for optimized downloads. + For other buckets, it delegates to the parent class implementation. Parameters ---------- @@ -2216,12 +2295,12 @@ async def _get_file_request( callback: fsspec.callbacks.Callback, optional Callback to monitor the download progress. **kwargs: - For Zonal buckets, `chunksize` bytes (int) can be provided to control + For Zonal and RCU buckets, `chunksize` bytes (int) can be provided to control the download chunk size (default is 128KB). """ bucket, key, path_generation = self.split_path(rpath) - if not await self._is_zonal_bucket(bucket): + if not await self._uses_grpc_data_path(bucket): return await super()._get_file_request( rpath, lpath, *args, headers=headers, callback=callback, **kwargs ) @@ -2304,7 +2383,7 @@ async def _get_file_concurrent( bucket, key, path_generation = self.split_path(rpath) # Delegate standard buckets to the original implementation - if not await self._is_zonal_bucket(bucket): + if not await self._uses_grpc_data_path(bucket): return await super()._get_file_concurrent( rpath, lpath, @@ -2459,9 +2538,9 @@ async def _merge(self, path, paths, acl=None): async def upload_chunk(fs, location, data, offset, size, content_type): """ - Uploads a chunk of data using AsyncAppendableObjectWriter for zonal buckets. + Uploads a chunk of data using AsyncAppendableObjectWriter for Zonal and RCU buckets. Finalizes the upload when the total uploaded data size reaches the specified size. - Delegates to core upload_chunk implementation for Non-Zonal buckets. + Delegates to core upload_chunk implementation for other buckets. """ # If `location` is an HTTP resumable-upload URL (string), delegate to core upload_chunk # for Standard buckets. @@ -2472,12 +2551,13 @@ async def upload_chunk(fs, location, data, offset, size, content_type): if not isinstance(location, AsyncAppendableObjectWriter): raise TypeError( - "upload_chunk for Zonal buckets expects an AsyncAppendableObjectWriter instance." + "upload_chunk for Zonal and RCU buckets expects an AsyncAppendableObjectWriter instance." ) if offset or size or content_type: logger.warning( - "Zonal buckets do not support offset, or content_type during upload. These parameters will be ignored." + "Zonal and RCU buckets do not support offset, or content_type during upload. " + "These parameters will be ignored." ) if not location._is_stream_open: @@ -2510,8 +2590,8 @@ async def initiate_upload( if_generation_match=None, ): """ - Initiates an upload for Zonal buckets by creating an AsyncAppendableObjectWriter. - Delegates to core initiate_upload implementation for Non-Zonal buckets. + Initiates an upload for Zonal and RCU buckets by creating an AsyncAppendableObjectWriter. + Delegates to core initiate_upload implementation for other buckets. Parameters ---------- @@ -2522,19 +2602,19 @@ async def initiate_upload( key: str The target object key. content_type: str, optional - Unsupported for Zonal buckets and will be ignored, except for the default. + Unsupported for Zonal and RCU buckets and will be ignored, except for the default. metadata: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. fixed_key_metadata: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. mode: str, optional The write mode, either 'overwrite' or 'create'. kms_key_name: str, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. if_generation_match: int, optional Only write if the object's current generation matches this value. """ - if not await fs._is_zonal_bucket(bucket): + if not await fs._uses_grpc_data_path(bucket): from gcsfs.core import initiate_upload as core_initiate_upload return await core_initiate_upload( @@ -2556,13 +2636,11 @@ async def initiate_upload( or (content_type and content_type != "application/octet-stream") ): logger.warning( - "Zonal buckets do not support content_type, metadata, fixed_key_metadata, " + "Zonal and RCU buckets do not support content_type, metadata, fixed_key_metadata, " "or kms_key_name during upload. These parameters will be ignored." ) - await fs._get_grpc_client() - return await zb_hns_utils.init_aaow( - fs.grpc_client, + return await fs._init_aaow( bucket, key, if_generation_match=_resolve_generation_match(mode, if_generation_match), @@ -2584,10 +2662,14 @@ async def simple_upload( **kwargs, ): """ - Performs a simple, single-request upload to Zonal bucket using gRPC. - In zonal buckets, file is left *unfinalized* by default unless - `finalize_on_close` is set to True. - Delegates to core simple_upload implementation for Non-Zonal buckets. + Performs a simple, single-request upload to Zonal and RCU buckets using gRPC. + By default, files in Zonal buckets are left *unfinalized* (unless + `finalize_on_close` is set to True), whereas files in RCU buckets are + finalized on close by default. A failed upload is never finalized. + On Zonal and RCU buckets an overwrite replaces the existing object as soon + as the writer opens, so a failed write leaves the partially written object + in its place, unfinalized. + Delegates to core simple_upload implementation for other buckets. Parameters ---------- @@ -2600,21 +2682,21 @@ async def simple_upload( datain: bytes The data to be uploaded. metadatain: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. consistency: str, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. content_type: str, optional - Unsupported for Zonal buckets and will be ignored, except for the default. + Unsupported for Zonal and RCU buckets and will be ignored, except for the default. fixed_key_metadata: dict, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. mode: str, optional The write mode, either 'overwrite' or 'create'. kms_key_name: str, optional - Unsupported for Zonal buckets and will be ignored. + Unsupported for Zonal and RCU buckets and will be ignored. if_generation_match: int, optional Only write if the object's current generation matches this value. """ - if not await fs._is_zonal_bucket(bucket): + if not await fs._uses_grpc_data_path(bucket): from gcsfs.core import simple_upload as core_simple_upload return await core_simple_upload( @@ -2639,19 +2721,21 @@ async def simple_upload( or (content_type and content_type != "application/octet-stream") ): logger.warning( - "Zonal buckets do not support content_type, metadatain, fixed_key_metadata, " + "Zonal and RCU buckets do not support content_type, metadatain, fixed_key_metadata, " "consistency or kms_key_name during upload. These parameters will be ignored." ) - await fs._get_grpc_client() - writer = await zb_hns_utils.init_aaow( - fs.grpc_client, + finalize_on_close = fs._resolve_finalize_on_close( + await fs._lookup_bucket_layout(bucket), kwargs.get("finalize_on_close") + ) + writer = await fs._init_aaow( bucket, key, if_generation_match=_resolve_generation_match(mode, if_generation_match), ) try: await writer.append(datain) - finally: - default_finalize = getattr(fs, "finalize_on_close", False) - finalize_on_close = kwargs.get("finalize_on_close", default_finalize) - await zb_hns_utils.close_aaow(writer, finalize_on_close=finalize_on_close) + except BaseException: + # Leave a partially written object unfinalized. + await zb_hns_utils.close_aaow(writer, finalize_on_close=False) + raise + await zb_hns_utils.close_aaow(writer, finalize_on_close=finalize_on_close) diff --git a/gcsfs/tests/conftest.py b/gcsfs/tests/conftest.py index b7803e7bc..4aff89529 100644 --- a/gcsfs/tests/conftest.py +++ b/gcsfs/tests/conftest.py @@ -18,6 +18,7 @@ ) from gcsfs import GCSFileSystem +from gcsfs.extended_gcsfs import CachingTier from gcsfs.tests import layouts from gcsfs.tests.settings import ( TEST_BUCKET, @@ -546,11 +547,27 @@ def flat_bucket(flat_bucket_name, gcs_factory): _close_gcs(gcs) +# Bucket layouts whose appendable writes use the gRPC writer; pass them to +# ``zonal_write_mocks`` as indirect params. +RAPID_WRITE_LAYOUTS = [ + layouts.ZONAL_HIERARCHICAL, + layouts.STANDARD_RCU, + layouts.HIERARCHICAL_RCU, +] + + @pytest.fixture -def zonal_write_mocks(): - """A fixture for mocking Zonal bucket write functionality.""" +def zonal_write_mocks(request): + """A fixture for mocking Zonal and RCU bucket write functionality. + + Accepts an indirect ``BucketLayout`` param (default ``ZONAL_HIERARCHICAL``). + On real GCS it yields ``None`` for the zonal layout and skips the others. + """ + layout = getattr(request, "param", layouts.ZONAL_HIERARCHICAL) if is_real_gcs(): + if layout != layouts.ZONAL_HIERARCHICAL: + pytest.skip("Live write tests only run against the zonal test bucket") yield None return @@ -563,7 +580,16 @@ def zonal_write_mocks(): mock_aaow = mock.AsyncMock(spec=AsyncAppendableObjectWriter) mock_aaow.offset = 0 mock_aaow._is_stream_open = True - mock_init_aaow = mock.AsyncMock(return_value=mock_aaow) + + async def init_aaow_side_effect(grpc_client, bucket_name, object_name, *a, **k): + # Every open gets an open writer for its own object. + mock_aaow._is_stream_open = True + mock_aaow.offset = 0 + mock_aaow.bucket_name = bucket_name + mock_aaow.object_name = object_name + return mock_aaow + + mock_init_aaow = mock.AsyncMock(side_effect=init_aaow_side_effect) mock_gcsfs_info = mock.AsyncMock( return_value={"generation": "12345", "type": "file", "name": "mock_file"} ) @@ -585,10 +611,7 @@ async def finalize_side_effect(): mock_aaow.finalize.side_effect = finalize_side_effect with ( - mock.patch( - patch_target_get_bucket_layout, - return_value=layouts.ZONAL_HIERARCHICAL, - ), + mock.patch(patch_target_get_bucket_layout, return_value=layout), mock.patch(patch_target_gcsfs_info, mock_gcsfs_info), mock.patch(patch_target_init_aaow, mock_init_aaow), ): @@ -596,6 +619,11 @@ async def finalize_side_effect(): "aaow": mock_aaow, "init_aaow": mock_init_aaow, "_gcsfs_info": mock_gcsfs_info, + "storage_class": ( + "RAPID" + if layout.caching_tier == CachingTier.RAPID_CACHE_ULTRA + else None + ), } yield mocks @@ -631,6 +659,18 @@ def skip_if_zonal(gcs): ) +@pytest.fixture +def skip_if_grpc_read(gcs): + """Skips tests of HTTP read mechanics; Zonal and RCU buckets read over gRPC.""" + if ( + hasattr(gcs, "_sync_lookup_bucket_layout") + and gcs._sync_lookup_bucket_layout(TEST_BUCKET).uses_grpc_data_path + ): + pytest.skip( + "Test validates HTTP read mechanics; Zonal and RCU buckets read over gRPC" + ) + + def pytest_addoption(parser): parser.addoption( "--run-benchmarks", diff --git a/gcsfs/tests/test_core.py b/gcsfs/tests/test_core.py index f0a484e3a..e8d744700 100644 --- a/gcsfs/tests/test_core.py +++ b/gcsfs/tests/test_core.py @@ -3427,7 +3427,7 @@ def test_mv_file_raises_error_for_specific_generation(gcs): gcs.version_aware = original_version_aware -def test_cat_file_routing_and_thresholds(gcs, skip_if_zonal, monkeypatch): +def test_cat_file_routing_and_thresholds(gcs, skip_if_grpc_read, monkeypatch): monkeypatch.setattr(gcs, "MIN_CHUNK_SIZE_FOR_CONCURRENCY", 5) fn = f"{TEST_BUCKET}/core_routing.txt" data = b"0123456789abcdefghijk" @@ -3575,7 +3575,7 @@ def test_cat_file_default_concurrency(gcs, skip_if_zonal): assert mock_conc.call_count == 0 -def test_cat_file_explicit_concurrency(gcs, skip_if_zonal): +def test_cat_file_explicit_concurrency(gcs, skip_if_grpc_read): # Arrange fn = f"{TEST_BUCKET}/test_cat_explicit_concurrency.txt" gcs.pipe(fn, b"cat test data") @@ -4451,14 +4451,9 @@ def test_get_headers_includes_cache_type(gcs): assert headers_preset["User-Agent"] == "custom-ua" -def test_user_agent_includes_cache_type_and_source_in_read(gcs): +def test_user_agent_includes_cache_type_and_source_in_read(gcs, skip_if_grpc_read): # TODO(lankita): Remove this skip when User-Agent is updated to include cache_type for gRPC/rapid buckets. - # Zonal buckets currently use the gRPC path which does not use the HTTP client. - if hasattr(gcs, "_is_zonal_bucket") and sync( - gcs.loop, gcs._is_zonal_bucket, TEST_BUCKET - ): - pytest.skip("Zonal buckets use gRPC, skipping HTTP User-Agent test") - + # Zonal and RCU buckets currently read over gRPC, which does not use the HTTP client. fn = TEST_BUCKET + "/2014-01-01.csv" # 1. Explicit cache_type="mmap" -> User-Agent contains cache_type/mmap:e diff --git a/gcsfs/tests/test_hns_unit.py b/gcsfs/tests/test_hns_unit.py index 32e0a47b4..86cb86f26 100644 --- a/gcsfs/tests/test_hns_unit.py +++ b/gcsfs/tests/test_hns_unit.py @@ -1970,11 +1970,21 @@ async def test_is_bucket_hns_enabled_true_zonal_hierarchical(self): ): assert await fs._is_bucket_hns_enabled("my-bucket") is True + @pytest.mark.asyncio + async def test_is_bucket_hns_enabled_true_hns_rcu(self): + """Verifies HNS is enabled when bucket layout is HIERARCHICAL_RCU.""" + fs = ExtendedGcsFileSystem(token="anon") + with mock.patch.object( + fs, "_lookup_bucket_layout", return_value=layouts.HIERARCHICAL_RCU + ): + assert await fs._is_bucket_hns_enabled("my-bucket") is True + @pytest.mark.asyncio @pytest.mark.parametrize( "bucket_layout", [ layouts.STANDARD, + layouts.STANDARD_RCU, BucketLayout.UNKNOWN, ], ) @@ -3320,6 +3330,47 @@ def extended_gcsfs(self): layouts.STANDARD, ), ({"location_type": "region"}, layouts.STANDARD), + ( + {"location_type": "region", "rapid_cache_info": {"cache_type": ""}}, + layouts.STANDARD, + ), + ( + { + "location_type": "region", + "rapid_cache_info": {"cache_type": "unknown-cache-type"}, + }, + layouts.STANDARD, + ), + ( + { + "location_type": "region", + "rapid_cache_info": {"cache_type": "rapid-cache-ultra"}, + }, + layouts.STANDARD_RCU, + ), + ( + { + "location_type": "region", + "hierarchical_namespace": {"enabled": True}, + "rapid_cache_info": {"cache_type": "rapid-cache-ultra"}, + }, + layouts.HIERARCHICAL_RCU, + ), + ( + { + "location_type": "region", + "rapid_cache_info": {"cache_type": "rapid-cache"}, + }, + layouts.STANDARD_RAPID_CACHE, + ), + ( + { + "location_type": "region", + "hierarchical_namespace": {"enabled": True}, + "rapid_cache_info": {"cache_type": "rapid-cache"}, + }, + layouts.HIERARCHICAL_RAPID_CACHE, + ), ], ) async def test_get_bucket_layout_detection( diff --git a/gcsfs/tests/test_zb_hns_utils.py b/gcsfs/tests/test_zb_hns_utils.py index 669968fd4..68c6d9d63 100644 --- a/gcsfs/tests/test_zb_hns_utils.py +++ b/gcsfs/tests/test_zb_hns_utils.py @@ -126,6 +126,32 @@ async def test_init_aaow_if_generation_match(if_generation_match, expected_kwarg ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "storage_class, expected_kwargs", + [("RAPID", {"storage_class": "RAPID"}), (None, {})], +) +async def test_init_aaow_storage_class(storage_class, expected_kwargs): + """storage_class is forwarded to the writer only when it is set.""" + mock_writer_instance = mock.AsyncMock() + with mock.patch( + "gcsfs.zb_hns_utils.AsyncAppendableObjectWriter", + new_callable=mock.Mock, + return_value=mock_writer_instance, + ) as mock_writer_class: + await zb_hns_utils.init_aaow( + mock_grpc_client, "bucket", "object", storage_class=storage_class + ) + mock_writer_class.assert_called_once_with( + client=mock_grpc_client, + bucket_name="bucket", + object_name="object", + generation=None, + writer_options={}, + **expected_kwargs, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("if_generation_match", [0, 7]) async def test_init_aaow_rejects_generation_with_if_generation_match( @@ -294,7 +320,7 @@ async def test_init_mrd_not_found(): @pytest.mark.asyncio async def test_close_aaow(caplog): - """Tests all graceful closing scenarios for AsyncAppendableObjectWriter.""" + """Tests the closing scenarios for AsyncAppendableObjectWriter.""" # 1. Handles None gracefully await zb_hns_utils.close_aaow(None) @@ -318,6 +344,21 @@ async def test_close_aaow(caplog): in caplog.text ) + # 4. Logs and re-raises a failed finalize, so the write is not reported as successful + mock_aaow.reset_mock() + mock_aaow.close.side_effect = NotFound("Object deleted before finalize") + caplog.clear() + + with caplog.at_level(logging.WARNING, logger="gcsfs"): + with pytest.raises(NotFound, match="Object deleted before finalize"): + await zb_hns_utils.close_aaow(mock_aaow, finalize_on_close=True) + + mock_aaow.close.assert_awaited_once_with(finalize_on_close=True) + assert ( + "Error closing AsyncAppendableObjectWriter for test-bucket/test-object" + in caplog.text + ) + @pytest.mark.asyncio async def test_close_mrd(caplog): diff --git a/gcsfs/tests/test_zonal.py b/gcsfs/tests/test_zonal.py index a2a90dd3c..b6b741d00 100644 --- a/gcsfs/tests/test_zonal.py +++ b/gcsfs/tests/test_zonal.py @@ -34,7 +34,13 @@ upload_chunk, ) from gcsfs.tests import layouts -from gcsfs.tests.conftest import csv_files, files, requires_rapid, text_files +from gcsfs.tests.conftest import ( + RAPID_WRITE_LAYOUTS, + csv_files, + files, + requires_rapid, + text_files, +) from gcsfs.tests.settings import TEST_BUCKET, TEST_ZONAL_BUCKET from gcsfs.tests.utils import is_real_gcs, tempdir, tmpfile from gcsfs.zb_hns_utils import MRDPool @@ -767,6 +773,7 @@ def test_multiprocess_error_handling_zb(extended_gcsfs): # =========================== Zonal Upload Tests =========================== @pytest.mark.asyncio +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) async def test_simple_upload_zonal(async_gcs, zonal_write_mocks, file_path): """Test simple_upload for Zonal buckets calls the correct writer methods.""" data = b"test data for simple_upload" @@ -781,7 +788,13 @@ async def test_simple_upload_zonal(async_gcs, zonal_write_mocks, file_path): ) if zonal_write_mocks: zonal_write_mocks["init_aaow"].assert_awaited_once_with( - async_gcs.grpc_client, bucket, object_name, if_generation_match=None + async_gcs.grpc_client, + bucket, + object_name, + None, + None, + if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) zonal_write_mocks["aaow"].append.assert_awaited_once_with(data) zonal_write_mocks["aaow"].close.assert_awaited_once_with(finalize_on_close=True) @@ -790,13 +803,20 @@ async def test_simple_upload_zonal(async_gcs, zonal_write_mocks, file_path): @pytest.mark.asyncio +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) async def test_initiate_upload_zonal(async_gcs, zonal_write_mocks, file_path): """Test initiate_upload for Zonal buckets returns a writer instance.""" bucket, object_name, _ = async_gcs.split_path(file_path) writer = await initiate_upload(fs=async_gcs, bucket=bucket, key=object_name) if zonal_write_mocks: zonal_write_mocks["init_aaow"].assert_awaited_once_with( - async_gcs.grpc_client, bucket, object_name, if_generation_match=None + async_gcs.grpc_client, + bucket, + object_name, + None, + None, + if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) assert writer is zonal_write_mocks["aaow"] else: @@ -888,6 +908,7 @@ async def test_upload_chunk_zonal_final_chunk(async_gcs, zonal_write_mocks, file @pytest.mark.asyncio +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) async def test_put_file_zonal(async_gcs, zonal_write_mocks, file_path): """Test _put_file for Zonal buckets.""" data = b"test data for put_file" @@ -902,7 +923,13 @@ async def test_put_file_zonal(async_gcs, zonal_write_mocks, file_path): if zonal_write_mocks: grpc_client = await async_gcs._get_grpc_client() zonal_write_mocks["init_aaow"].assert_awaited_once_with( - grpc_client, bucket, object_name, if_generation_match=None + grpc_client, + bucket, + object_name, + None, + None, + if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) zonal_write_mocks["aaow"].append_from_file.assert_awaited_once() args, kwargs = zonal_write_mocks["aaow"].append_from_file.call_args @@ -921,6 +948,7 @@ async def test_put_file_zonal(async_gcs, zonal_write_mocks, file_path): @pytest.mark.asyncio +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) async def test_pipe_file_zonal(async_gcs, zonal_write_mocks, file_path): """Test _pipe_file for Zonal buckets.""" data = b"test data for pipe_file" @@ -931,7 +959,13 @@ async def test_pipe_file_zonal(async_gcs, zonal_write_mocks, file_path): if zonal_write_mocks: grpc_client = await async_gcs._get_grpc_client() zonal_write_mocks["init_aaow"].assert_awaited_once_with( - grpc_client, bucket, object_name, if_generation_match=None + grpc_client, + bucket, + object_name, + None, + None, + if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) zonal_write_mocks["aaow"].append.assert_awaited_once_with(data) zonal_write_mocks["aaow"].close.assert_awaited_once_with(finalize_on_close=True) @@ -990,7 +1024,7 @@ def test_get_list_from_zonal_bucket(extended_gcsfs): mock_create_mrd = mock.AsyncMock(side_effect=create_mrd_side_effect) with ( mock.patch( - "gcsfs.extended_gcsfs.ExtendedGcsFileSystem._is_zonal_bucket", + "gcsfs.extended_gcsfs.ExtendedGcsFileSystem._uses_grpc_data_path", return_value=True, ), mock.patch( @@ -1035,7 +1069,7 @@ def test_get_directory_from_zonal_bucket(extended_gcsfs): with ( mock.patch( - "gcsfs.extended_gcsfs.ExtendedGcsFileSystem._is_zonal_bucket", + "gcsfs.extended_gcsfs.ExtendedGcsFileSystem._uses_grpc_data_path", return_value=True, ), mock.patch( @@ -1339,9 +1373,6 @@ async def test_fetch_range_split_concurrent_success(extended_gcsfs): extended_gcsfs._mrd_pool_cache.get = mock.AsyncMock(return_value=mock_pool) with contextlib.ExitStack() as stack: - stack.enter_context( - mock.patch.object(extended_gcsfs, "_is_zonal_bucket", return_value=True) - ) stack.enter_context( mock.patch.object(extended_gcsfs, "_info", return_value={"size": 100}) ) @@ -1382,9 +1413,6 @@ async def test_fetch_range_split_concurrent_exception(extended_gcsfs): extended_gcsfs._mrd_pool_cache.get = mock.AsyncMock(return_value=mock_pool) with contextlib.ExitStack() as stack: - stack.enter_context( - mock.patch.object(extended_gcsfs, "_is_zonal_bucket", return_value=True) - ) stack.enter_context( mock.patch.object(extended_gcsfs, "_info", return_value={"size": 100}) ) @@ -1413,7 +1441,7 @@ async def test_cat_file_zero_length_read(extended_gcsfs): with contextlib.ExitStack() as stack: stack.enter_context( - mock.patch.object(extended_gcsfs, "_is_zonal_bucket", return_value=True) + mock.patch.object(extended_gcsfs, "_uses_grpc_data_path", return_value=True) ) stack.enter_context( mock.patch("gcsfs.extended_gcsfs._get_mrd_size", return_value=100) @@ -1454,7 +1482,7 @@ async def test_cat_file_delegates_resolved_range_to_mrd_fetch( with contextlib.ExitStack() as stack: stack.enter_context( - mock.patch.object(extended_gcsfs, "_is_zonal_bucket", return_value=True) + mock.patch.object(extended_gcsfs, "_uses_grpc_data_path", return_value=True) ) stack.enter_context( mock.patch("gcsfs.extended_gcsfs._get_mrd_size", return_value=5000) @@ -1556,7 +1584,9 @@ async def test_cat_file_non_zonal_fallback(extended_gcsfs): """Tests that _cat_file delegates to the parent class for non-zonal buckets.""" with contextlib.ExitStack() as stack: stack.enter_context( - mock.patch.object(extended_gcsfs, "_is_zonal_bucket", return_value=False) + mock.patch.object( + extended_gcsfs, "_uses_grpc_data_path", return_value=False + ) ) mock_super_cat = stack.enter_context( mock.patch( diff --git a/gcsfs/tests/test_zonal_file.py b/gcsfs/tests/test_zonal_file.py index 965e80b40..49a408bf7 100644 --- a/gcsfs/tests/test_zonal_file.py +++ b/gcsfs/tests/test_zonal_file.py @@ -16,14 +16,16 @@ import fsspec.asyn import pytest from fsspec.asyn import FSTimeoutError +from google.api_core.exceptions import NotFound from google.cloud.storage.asyncio.async_appendable_object_writer import ( _DEFAULT_FLUSH_INTERVAL_BYTES, ) from gcsfs import zb_hns_utils -from gcsfs.extended_gcsfs import ExtendedGcsFileSystem +from gcsfs.extended_gcsfs import CachingTier, ExtendedGcsFileSystem from gcsfs.retry import GenerationMismatchError, HttpError -from gcsfs.tests.conftest import requires_rapid +from gcsfs.tests import layouts +from gcsfs.tests.conftest import RAPID_WRITE_LAYOUTS, requires_rapid from gcsfs.tests.settings import TEST_ZONAL_BUCKET from gcsfs.tests.utils import is_real_gcs, tempdir, tmpfile from gcsfs.zb_hns_utils import DEFAULT_TEARDOWN_TIMEOUT_SECONDS @@ -36,7 +38,10 @@ class TestZonalGcsFileSystemWriteFileCacheUpdate: @pytest.mark.asyncio - async def test_pipe_file_zonal_routes_through_write_cache_update(self): + @pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) + async def test_pipe_file_zonal_routes_through_write_cache_update( + self, bucket_layout + ): """The zonal pipe_file path routes through the shared cache updater.""" fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) fs._grpc_client = mock.MagicMock() @@ -46,17 +51,22 @@ async def test_pipe_file_zonal_routes_through_write_cache_update(self): writer.append = mock.AsyncMock() with ( - mock.patch.object(fs, "_is_zonal_bucket", return_value=True), + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), mock.patch( "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", new_callable=mock.AsyncMock, return_value=writer, - ), + ) as mock_init_aaow, mock.patch( "gcsfs.extended_gcsfs.zb_hns_utils.close_aaow", new_callable=mock.AsyncMock, - ), + ) as mock_close_aaow, mock.patch.object( fs, "_write_file_cache_update", new_callable=mock.AsyncMock ) as mock_update, @@ -64,9 +74,17 @@ async def test_pipe_file_zonal_routes_through_write_cache_update(self): await fs._pipe_file(path, b"some-data") mock_update.assert_awaited_once_with(path) + expected_rcu = bucket_layout.caching_tier == CachingTier.RAPID_CACHE_ULTRA + assert mock_init_aaow.await_args.kwargs["storage_class"] == ( + "RAPID" if expected_rcu else None + ) + mock_close_aaow.assert_awaited_once_with(writer, finalize_on_close=expected_rcu) @pytest.mark.asyncio - async def test_put_file_zonal_routes_through_write_cache_update(self): + @pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) + async def test_put_file_zonal_routes_through_write_cache_update( + self, bucket_layout + ): """The zonal put_file path routes through the shared cache updater.""" fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) fs._grpc_client = mock.MagicMock() @@ -80,17 +98,22 @@ async def test_put_file_zonal_routes_through_write_cache_update(self): f.write(b"some-data") with ( - mock.patch.object(fs, "_is_zonal_bucket", return_value=True), + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), mock.patch( "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", new_callable=mock.AsyncMock, return_value=writer, - ), + ) as mock_init_aaow, mock.patch( "gcsfs.extended_gcsfs.zb_hns_utils.close_aaow", new_callable=mock.AsyncMock, - ), + ) as mock_close_aaow, mock.patch.object( fs, "_write_file_cache_update", new_callable=mock.AsyncMock ) as mock_update, @@ -98,6 +121,132 @@ async def test_put_file_zonal_routes_through_write_cache_update(self): await fs._put_file(lpath, rpath) mock_update.assert_awaited_once_with(rpath) + expected_rcu = bucket_layout.caching_tier == CachingTier.RAPID_CACHE_ULTRA + assert mock_init_aaow.await_args.kwargs["storage_class"] == ( + "RAPID" if expected_rcu else None + ) + mock_close_aaow.assert_awaited_once_with(writer, finalize_on_close=expected_rcu) + + @pytest.mark.asyncio + @pytest.mark.parametrize("finalize_on_close", [None, True]) + @pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) + async def test_pipe_file_failure_closes_without_finalizing( + self, bucket_layout, finalize_on_close + ): + """A failed pipe_file closes the writer without finalizing the partial object and still updates the dircache.""" + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + fs._grpc_client = mock.MagicMock() + path = f"{TEST_ZONAL_BUCKET}/dir/file.txt" + + writer = mock.MagicMock() + writer.append = mock.AsyncMock(side_effect=ConnectionError("append failed")) + + with ( + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", + new_callable=mock.AsyncMock, + return_value=writer, + ), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.close_aaow", + new_callable=mock.AsyncMock, + ) as mock_close_aaow, + mock.patch.object( + fs, "_write_file_cache_update", new_callable=mock.AsyncMock + ) as mock_update, + ): + with pytest.raises(ConnectionError, match="append failed"): + await fs._pipe_file( + path, b"some-data", finalize_on_close=finalize_on_close + ) + + mock_close_aaow.assert_awaited_once_with(writer, finalize_on_close=False) + mock_update.assert_awaited_once_with(path) + + @pytest.mark.asyncio + @pytest.mark.parametrize("finalize_on_close", [None, True]) + @pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) + async def test_put_file_failure_closes_without_finalizing( + self, bucket_layout, finalize_on_close + ): + """A failed put_file closes the writer without finalizing the partial object and still updates the dircache.""" + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + fs._grpc_client = mock.MagicMock() + rpath = f"{TEST_ZONAL_BUCKET}/dir/file.txt" + + writer = mock.MagicMock() + writer.append_from_file = mock.AsyncMock( + side_effect=ConnectionError("append failed") + ) + + with tmpfile() as lpath: + with open(lpath, "wb") as f: + f.write(b"some-data") + + with ( + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", + new_callable=mock.AsyncMock, + return_value=writer, + ), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.close_aaow", + new_callable=mock.AsyncMock, + ) as mock_close_aaow, + mock.patch.object( + fs, "_write_file_cache_update", new_callable=mock.AsyncMock + ) as mock_update, + ): + with pytest.raises(ConnectionError, match="append failed"): + await fs._put_file( + lpath, rpath, finalize_on_close=finalize_on_close + ) + + mock_close_aaow.assert_awaited_once_with(writer, finalize_on_close=False) + mock_update.assert_awaited_once_with(rpath) + + @pytest.mark.asyncio + @pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) + async def test_put_file_missing_lpath_does_not_open_writer( + self, bucket_layout, tmp_path + ): + """A missing local file fails before the writer replaces the object.""" + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + fs._grpc_client = mock.MagicMock() + + with ( + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", + new_callable=mock.AsyncMock, + ) as mock_init_aaow, + ): + with pytest.raises(FileNotFoundError): + await fs._put_file( + str(tmp_path / "missing"), f"{TEST_ZONAL_BUCKET}/dir/file.txt" + ) + + mock_init_aaow.assert_not_awaited() @pytest.fixture @@ -208,6 +357,7 @@ def test_zonal_file_write_success(extended_gcsfs, zonal_write_mocks, file_path): assert extended_gcsfs.cat(file_path) == data1 + data2 +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_open_write_mode(extended_gcsfs, zonal_write_mocks, file_path): """Test that opening a ZonalFile in write mode initializes the writer.""" bucket, key, _ = extended_gcsfs.split_path(file_path) @@ -222,11 +372,13 @@ def test_zonal_file_open_write_mode(extended_gcsfs, zonal_write_mocks, file_path None, _DEFAULT_FLUSH_INTERVAL_BYTES, if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) else: assert extended_gcsfs.exists(file_path) +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_open_write_mode_with_custom_flush_interval_bytes( extended_gcsfs, zonal_write_mocks, file_path ): @@ -249,11 +401,13 @@ def test_zonal_file_open_write_mode_with_custom_flush_interval_bytes( None, custom_flush_interval_bytes, if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) else: assert extended_gcsfs.exists(file_path) +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_open_create_mode(extended_gcsfs, zonal_write_mocks, file_path): """mode="xb" opens the writer with the exclusive-create precondition.""" bucket, key, _ = extended_gcsfs.split_path(file_path) @@ -268,11 +422,47 @@ def test_zonal_file_open_create_mode(extended_gcsfs, zonal_write_mocks, file_pat None, _DEFAULT_FLUSH_INTERVAL_BYTES, if_generation_match=0, + storage_class=zonal_write_mocks["storage_class"], ) else: assert extended_gcsfs.exists(file_path) +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) +def test_zonal_file_failed_writer_open_is_not_retried_on_close( + extended_gcsfs, zonal_write_mocks, file_path +): + """A writer that fails to open raises once; close() does not try to open it again.""" + if not zonal_write_mocks: + pytest.skip("Unit test for writer open failures with mocks.") + zonal_write_mocks["init_aaow"].side_effect = FileExistsError(file_path) + + with pytest.raises(FileExistsError): + with extended_gcsfs.open(file_path, "xb") as f: + f.write(test_data) + + zonal_write_mocks["init_aaow"].assert_awaited_once() + + +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) +def test_zonal_file_flush_failure_invalidates_cache( + extended_gcsfs, zonal_write_mocks, file_path +): + """A failed closing flush still drops the cached parent listing of the replaced object.""" + if not zonal_write_mocks: + pytest.skip("Unit test for closing-flush failures with mocks.") + path = f"{file_path}/obj" + extended_gcsfs.dircache[file_path] = [{"name": path, "type": "file", "size": 1}] + zonal_write_mocks["aaow"].flush.side_effect = NotFound("flush failed") + + with pytest.raises(NotFound, match="flush failed"): + with extended_gcsfs.open(path, "wb") as f: + f.write(test_data) + + assert file_path not in extended_gcsfs.dircache + + +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_open_write_mode_with_if_generation_match( extended_gcsfs, zonal_write_mocks, file_path ): @@ -290,6 +480,7 @@ def test_zonal_file_open_write_mode_with_if_generation_match( None, _DEFAULT_FLUSH_INTERVAL_BYTES, if_generation_match=7, + storage_class=zonal_write_mocks["storage_class"], ) else: # file_path does not exist, so generation 7 cannot match; SDK releases @@ -317,6 +508,7 @@ def test_zonal_file_rejects_if_generation_match_before_opening_pool( mock_sync.assert_not_called() +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_open_append_mode(extended_gcsfs, zonal_write_mocks, file_path): """Test that opening a ZonalFile in append mode initializes the writer with generation.""" bucket, key, _ = extended_gcsfs.split_path(file_path) @@ -334,11 +526,13 @@ def test_zonal_file_open_append_mode(extended_gcsfs, zonal_write_mocks, file_pat "12345", _DEFAULT_FLUSH_INTERVAL_BYTES, if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) else: assert extended_gcsfs.cat(file_path) == b"data" +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_open_append_mode_nonexistent_file( extended_gcsfs, zonal_write_mocks, file_path ): @@ -366,6 +560,7 @@ def test_zonal_file_open_append_mode_nonexistent_file( None, _DEFAULT_FLUSH_INTERVAL_BYTES, if_generation_match=None, + storage_class=zonal_write_mocks["storage_class"], ) # _info is called to get the generation, but it fails extended_gcsfs._info.assert_awaited_once() @@ -373,6 +568,59 @@ def test_zonal_file_open_append_mode_nonexistent_file( assert extended_gcsfs.cat(file_path) == test_data +@pytest.mark.parametrize("mode", ["wb", "xb", "ab"]) +@pytest.mark.parametrize( + "zonal_write_mocks", + [layouts.STANDARD_RCU, layouts.HIERARCHICAL_RCU], + indirect=True, +) +def test_open_write_modes_use_zonal_file_on_rcu( + extended_gcsfs, zonal_write_mocks, file_path, mode +): + """Every write mode on an RCU bucket opens a ZonalFile that writes RAPID and finalizes on close by default.""" + with extended_gcsfs.open(file_path, mode) as f: + assert isinstance(f, ZonalFile) + assert f.finalize_on_close is True + f.write(test_data) + + assert zonal_write_mocks["init_aaow"].await_args.kwargs["storage_class"] == "RAPID" + zonal_write_mocks["aaow"].close.assert_awaited_once_with(finalize_on_close=True) + + +@pytest.mark.parametrize( + "kwargs, warns", + [ + ({"metadata": {"k": "v"}}, True), + ({"fixed_key_metadata": {"cache_control": "no-cache"}}, True), + ({"kms_key_name": "projects/p/locations/l/keyRings/r/cryptoKeys/k"}, True), + ({"content_type": "text/plain"}, True), + ({"content_type": "application/octet-stream"}, False), + ({}, False), + ], + ids=[ + "metadata", + "fixed_key_metadata", + "kms_key_name", + "content_type", + "default_content_type", + "no_params", + ], +) +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) +def test_zonal_file_write_mode_warns_on_ignored_upload_params( + extended_gcsfs, zonal_write_mocks, file_path, caplog, kwargs, warns +): + """Write modes warn when an upload parameter the appendable writer ignores is set.""" + with caplog.at_level(logging.WARNING, logger="gcsfs.zonal_file"): + with extended_gcsfs.open(file_path, "wb", finalize_on_close=True, **kwargs): + pass + warned = any( + "These parameters will be ignored" in r.message and r.levelname == "WARNING" + for r in caplog.records + ) + assert warned is warns + + def test_zonal_file_flush(extended_gcsfs, zonal_write_mocks, file_path): """Test that flush calls the underlying writer's flush method.""" with extended_gcsfs.open(file_path, "wb") as f: @@ -382,6 +630,7 @@ def test_zonal_file_flush(extended_gcsfs, zonal_write_mocks, file_path): zonal_write_mocks["aaow"].flush.assert_awaited() +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_lazy_init_creates_empty_file_on_close( extended_gcsfs, zonal_write_mocks, file_path ): @@ -398,10 +647,15 @@ def test_zonal_file_lazy_init_creates_empty_file_on_close( if zonal_write_mocks: zonal_write_mocks["init_aaow"].assert_awaited_once() + assert ( + zonal_write_mocks["init_aaow"].await_args.kwargs["storage_class"] + == zonal_write_mocks["storage_class"] + ) else: assert extended_gcsfs.exists(file_path) +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) def test_zonal_file_commit(extended_gcsfs, zonal_write_mocks, file_path): """Test that commit finalizes the write, sets finalized to True and does not finalize on close.""" with extended_gcsfs.open(file_path, "wb", finalize_on_close=True) as f: @@ -431,13 +685,225 @@ def test_zonal_file_finalize_on_close_true( def test_zonal_file_finalize_on_close_default_false( extended_gcsfs, zonal_write_mocks, file_path ): - """Test that finalize_on_close is False by default.""" + """Test that finalize_on_close is False by default for Zonal buckets.""" with extended_gcsfs.open(file_path, "wb") as f: assert f.finalize_on_close is False if zonal_write_mocks: zonal_write_mocks["aaow"].close.assert_awaited_with(finalize_on_close=False) +@pytest.mark.parametrize( + "zonal_write_mocks", + [layouts.STANDARD_RCU, layouts.HIERARCHICAL_RCU], + indirect=True, +) +def test_zonal_file_finalize_on_close_rcu_default_and_overrides( + extended_gcsfs, zonal_write_mocks, file_path +): + """On RCU buckets finalize_on_close defaults to True, but explicit False overrides it.""" + if not zonal_write_mocks: + pytest.skip("Unit test for finalize_on_close resolution with mocks.") + + aaow = zonal_write_mocks["aaow"] + + # 1. Default via open() is True + with extended_gcsfs.open(file_path, "wb") as f: + assert f.finalize_on_close is True + aaow.close.assert_awaited_once_with(finalize_on_close=True) + + # 2. Direct ZonalFile instantiation also defaults to True on RCU + aaow.close.reset_mock() + with ZonalFile(extended_gcsfs, file_path, mode="wb") as f_direct: + assert f_direct.finalize_on_close is True + aaow.close.assert_awaited_once_with(finalize_on_close=True) + + # 3. Per-call finalize_on_close=False overrides RCU default + aaow.close.reset_mock() + with extended_gcsfs.open(file_path, "wb", finalize_on_close=False) as f_false: + assert f_false.finalize_on_close is False + aaow.close.assert_awaited_once_with(finalize_on_close=False) + + # 4. Filesystem-level finalize_on_close=False overrides RCU default + orig = getattr(extended_gcsfs, "finalize_on_close", None) + try: + extended_gcsfs.finalize_on_close = False + aaow.close.reset_mock() + with extended_gcsfs.open(file_path, "wb") as f_fs_false: + assert f_fs_false.finalize_on_close is False + aaow.close.assert_awaited_once_with(finalize_on_close=False) + finally: + extended_gcsfs.finalize_on_close = orig + + +@pytest.mark.parametrize("call_value", [None, True, False]) +@pytest.mark.parametrize("fs_value", [None, True, False]) +@pytest.mark.parametrize("zonal_write_mocks", RAPID_WRITE_LAYOUTS, indirect=True) +def test_zonal_file_finalize_on_close_precedence( + extended_gcsfs, zonal_write_mocks, file_path, fs_value, call_value +): + """A per-call finalize_on_close wins over the filesystem value, which wins over the bucket default.""" + if not zonal_write_mocks: + pytest.skip("Unit test for finalize_on_close resolution with mocks.") + if call_value is not None: + expected = call_value + elif fs_value is not None: + expected = fs_value + else: + # Only RCU layouts finalize by default; they are also the only ones + # for which zonal_write_mocks expects storage_class "RAPID". + expected = zonal_write_mocks["storage_class"] == "RAPID" + + orig = getattr(extended_gcsfs, "finalize_on_close", None) + try: + extended_gcsfs.finalize_on_close = fs_value + with extended_gcsfs.open(file_path, "wb", finalize_on_close=call_value) as f: + assert f.finalize_on_close is expected + finally: + extended_gcsfs.finalize_on_close = orig + zonal_write_mocks["aaow"].close.assert_awaited_once_with(finalize_on_close=expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) +async def test_simple_upload_finalize_on_close_defaults_by_bucket_layout(bucket_layout): + """simple_upload defaults finalize_on_close to True on RCU and False on Zonal, with override support.""" + from gcsfs.extended_gcsfs import simple_upload + + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + fs._grpc_client = mock.MagicMock() + writer = mock.MagicMock() + writer.append = mock.AsyncMock() + expected_rcu = bucket_layout.caching_tier == CachingTier.RAPID_CACHE_ULTRA + + with ( + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", + new_callable=mock.AsyncMock, + return_value=writer, + ), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.close_aaow", + new_callable=mock.AsyncMock, + ) as mock_close_aaow, + ): + await simple_upload(fs, TEST_ZONAL_BUCKET, "k", b"data") + mock_close_aaow.assert_awaited_once_with(writer, finalize_on_close=expected_rcu) + + mock_close_aaow.reset_mock() + await simple_upload( + fs, TEST_ZONAL_BUCKET, "k", b"data", finalize_on_close=not expected_rcu + ) + mock_close_aaow.assert_awaited_once_with( + writer, finalize_on_close=not expected_rcu + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("finalize_on_close", [None, True]) +@pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) +async def test_simple_upload_failure_closes_without_finalizing( + bucket_layout, finalize_on_close +): + """A failed simple_upload closes the writer without finalizing the partial object.""" + from gcsfs.extended_gcsfs import simple_upload + + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + fs._grpc_client = mock.MagicMock() + writer = mock.MagicMock() + writer.append = mock.AsyncMock(side_effect=ConnectionError("append failed")) + + with ( + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", + new_callable=mock.AsyncMock, + return_value=writer, + ), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.close_aaow", + new_callable=mock.AsyncMock, + ) as mock_close_aaow, + ): + with pytest.raises(ConnectionError, match="append failed"): + await simple_upload( + fs, + TEST_ZONAL_BUCKET, + "k", + b"data", + finalize_on_close=finalize_on_close, + ) + + mock_close_aaow.assert_awaited_once_with(writer, finalize_on_close=False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("helper", ["_put_file", "_pipe_file", "simple_upload"]) +@pytest.mark.parametrize("bucket_layout", RAPID_WRITE_LAYOUTS) +async def test_write_helpers_propagate_finalize_failure( + bucket_layout, helper, tmp_path +): + """With the real close_aaow, a failed finalize propagates and the writer is closed once. + + _put_file and _pipe_file still update the dircache; simple_upload leaves it + to its callers. + """ + from gcsfs.extended_gcsfs import simple_upload + + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + fs._grpc_client = mock.MagicMock() + path = f"{TEST_ZONAL_BUCKET}/k" + writer = mock.MagicMock(bucket_name=TEST_ZONAL_BUCKET, object_name="k") + writer.append = mock.AsyncMock() + writer.append_from_file = mock.AsyncMock() + writer.close = mock.AsyncMock(side_effect=NotFound("finalize failed")) + lpath = tmp_path / "data" + lpath.write_bytes(b"data") + + with ( + mock.patch.object( + fs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch.object(fs, "_get_grpc_client", new_callable=mock.AsyncMock), + mock.patch( + "gcsfs.extended_gcsfs.zb_hns_utils.init_aaow", + new_callable=mock.AsyncMock, + return_value=writer, + ), + mock.patch.object( + fs, "_write_file_cache_update", new_callable=mock.AsyncMock + ) as mock_update, + ): + with pytest.raises(NotFound, match="finalize failed"): + if helper == "_put_file": + await fs._put_file(str(lpath), path, finalize_on_close=True) + elif helper == "_pipe_file": + await fs._pipe_file(path, b"data", finalize_on_close=True) + else: + await simple_upload( + fs, TEST_ZONAL_BUCKET, "k", b"data", finalize_on_close=True + ) + + writer.close.assert_awaited_once_with(finalize_on_close=True) + if helper != "simple_upload": + mock_update.assert_awaited_once_with(path) + + def test_zonal_file_flush_after_finalize_logs_warning( extended_gcsfs, zonal_write_mocks, file_path ): @@ -486,7 +952,7 @@ def test_zonal_file_discard(extended_gcsfs, zonal_write_mocks, file_path): f.discard() mock_logger.warning.assert_called_once() assert ( - "Discard is not applicable for Zonal Buckets" + "Discard is not applicable for Zonal and RCU buckets" in mock_logger.warning.call_args[0][0] ) @@ -511,7 +977,13 @@ def test_zonal_file_overwrite(extended_gcsfs, zonal_write_mocks, file_path): ) as f: f.write(b"Sample text data.") - if not zonal_write_mocks: + if zonal_write_mocks: + # Each open finalizes its own writer on close. + assert ( + zonal_write_mocks["aaow"].close.await_args_list + == [mock.call(finalize_on_close=True)] * 2 + ) + else: assert extended_gcsfs.cat(file_path) == b"Sample text data." @@ -1118,6 +1590,27 @@ async def failing_close_aaow(aaow, finalize_on_close=False): zf._close_impl() +def test_close_impl_flush_failure_still_closes_aaow(mock_gcsfs): + """If the closing flush fails, the writer is closed without finalizing and the flush error is raised.""" + mock_gcsfs.loop = fsspec.asyn.get_loop() + zf = ZonalFile( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="wb", + finalize_on_close=True, + ) + aaow = mock.Mock(_is_stream_open=True) + aaow.flush = mock.AsyncMock(side_effect=RuntimeError("flush failed")) + aaow.close = mock.AsyncMock() + zf.aaow = aaow + + with pytest.raises(RuntimeError, match="flush failed"): + zf.close() + + aaow.close.assert_awaited_once_with(finalize_on_close=False) + assert zf.closed + + def test_zonal_file_fetch_range_unhandled_runtime_error(mock_gcsfs): """Tests that a RuntimeError not containing 'not satisfiable' is re-raised.""" diff --git a/gcsfs/tests/test_zonal_unit.py b/gcsfs/tests/test_zonal_unit.py index 364ce4dcd..437a08c97 100644 --- a/gcsfs/tests/test_zonal_unit.py +++ b/gcsfs/tests/test_zonal_unit.py @@ -14,9 +14,11 @@ ) from google.cloud.storage.exceptions import DataCorruption +import gcsfs.core from gcsfs.checkers import ConsistencyChecker, MD5Checker, SizeChecker from gcsfs.core import GCSFile from gcsfs.extended_gcsfs import ( + BucketLayout, ExtendedGcsFileSystem, _close_pool, initiate_upload, @@ -49,6 +51,12 @@ ), ] +GRPC_READ_LAYOUTS = [ + layouts.ZONAL_HIERARCHICAL, + layouts.STANDARD_RCU, + layouts.HIERARCHICAL_RCU, +] + @pytest.mark.parametrize("bucket_type_val", layouts.ALL_BUCKET_LAYOUTS) def test_open_uses_default_blocksize_and_consistency_from_fs( @@ -61,7 +69,7 @@ def test_open_uses_default_blocksize_and_consistency_from_fs( with gcs_bucket_mocks(csv_data, bucket_type_val=bucket_type_val): with extended_gcsfs.open(csv_file_path, "rb") as f: assert isinstance(f, GCSFile) - assert isinstance(f, ZonalFile) is bucket_type_val.is_zonal + assert isinstance(f, ZonalFile) is bucket_type_val.uses_grpc_data_path assert f.blocksize == extended_gcsfs.default_block_size assert type(f.checker) is ConsistencyChecker @@ -92,6 +100,21 @@ def test_open_uses_correct_blocksize_and_consistency_for_all_bucket_types( assert isinstance(f.checker, SizeChecker) +@pytest.mark.parametrize( + "bucket_layout", + [layouts.STANDARD_RCU, layouts.HIERARCHICAL_RCU], +) +@pytest.mark.parametrize("mode", ["r+", "rb+"]) +def test_open_plus_mode_does_not_check_out_mrd_pool( + extended_gcsfs, gcs_bucket_mocks, bucket_layout, mode +): + """Read-write modes are rejected before an MRD pool is taken from the cache.""" + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: + with pytest.raises(NotImplementedError): + extended_gcsfs.open(file_path, mode) + mocks["pool_cache_get"].assert_not_called() + + @pytest.mark.parametrize( "start, end, exp_offset, exp_length", [ @@ -176,17 +199,18 @@ def test_process_limits_when_file_size_passed(extended_gcsfs): mock_info.assert_not_called() +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) @pytest.mark.parametrize( "exception_to_raise", [ValueError, DataCorruption, Exception], ) -def test_mrd_exception_handling(extended_gcsfs, gcs_bucket_mocks, exception_to_raise): +def test_mrd_exception_handling( + extended_gcsfs, gcs_bucket_mocks, exception_to_raise, bucket_layout +): """ Tests that _cat_file correctly propagates exceptions from mrd.download_ranges. """ - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: # Configure the mock to raise a specified exception if exception_to_raise is DataCorruption: # The first argument is 'response', the message is in '*args' @@ -204,14 +228,15 @@ def test_mrd_exception_handling(extended_gcsfs, gcs_bucket_mocks, exception_to_r mocks["downloader"].download_ranges.call_count = 2 -def test_mrd_created_once_for_zonal_file(extended_gcsfs, gcs_bucket_mocks): +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +def test_mrd_created_once_for_zonal_file( + extended_gcsfs, gcs_bucket_mocks, bucket_layout +): """ Tests that the AsyncMultiRangeDownloader (MRD) is created only once when a ZonalFile is opened, and not for each subsequent read operation. """ - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: with extended_gcsfs.open(file_path, "rb") as f: # The MRDPoolCache.get should be called upon opening the file. mocks["pool_cache_get"].assert_called_once() @@ -225,15 +250,14 @@ def test_mrd_created_once_for_zonal_file(extended_gcsfs, gcs_bucket_mocks): mocks["pool_cache_get"].assert_called_once() +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) def test_zonal_file_warning_on_missing_persisted_size( - extended_gcsfs, gcs_bucket_mocks, caplog + extended_gcsfs, gcs_bucket_mocks, caplog, bucket_layout ): """ Tests that a warning is logged when MRD has no 'persisted_size' attribute when opening ZonalFile. """ - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: # Force persisted_size to None to trigger the warning mocks["pool"].persisted_size = None with caplog.at_level(logging.WARNING, logger="gcsfs"): @@ -243,13 +267,14 @@ def test_zonal_file_warning_on_missing_persisted_size( @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_file_warning_on_missing_persisted_size( - extended_gcsfs, gcs_bucket_mocks, caplog + extended_gcsfs, gcs_bucket_mocks, caplog, bucket_layout ): """ Tests that a warning is logged in cat_file when MRD has no 'persisted_size' attribute. """ - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): # 'persisted_size' is set to None in the mock downloader with ( caplog.at_level(logging.WARNING, logger="gcsfs"), @@ -264,11 +289,14 @@ async def test_cat_file_warning_on_missing_persisted_size( @pytest.mark.asyncio -async def test_cat_file_passes_cache_type(extended_gcsfs, gcs_bucket_mocks): +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +async def test_cat_file_passes_cache_type( + extended_gcsfs, gcs_bucket_mocks, bucket_layout +): """ Tests that cache_type and cache_source are propagated to _mrd_pool_cache.get. """ - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): with mock.patch.object( extended_gcsfs._mrd_pool_cache, "get", new_callable=mock.AsyncMock ) as mock_get: @@ -302,12 +330,15 @@ async def test_cat_file_passes_cache_type(extended_gcsfs, gcs_bucket_mocks): @pytest.mark.asyncio -async def test_cat_file_zonal_default_concurrency(extended_gcsfs, gcs_bucket_mocks): - """Tests that cat_file defaults to pool_size=1 on zonal buckets.""" +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +async def test_cat_file_zonal_default_concurrency( + extended_gcsfs, gcs_bucket_mocks, bucket_layout +): + """Tests that cat_file defaults to pool_size=1 on zonal/RCU buckets.""" # Arrange from gcsfs.zb_hns_utils import MRDPool - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): with mock.patch.object( extended_gcsfs._mrd_pool_cache, "get", new_callable=mock.AsyncMock ) as mock_get: @@ -328,12 +359,15 @@ async def test_cat_file_zonal_default_concurrency(extended_gcsfs, gcs_bucket_moc @pytest.mark.asyncio -async def test_cat_file_zonal_explicit_concurrency(extended_gcsfs, gcs_bucket_mocks): - """Tests that cat_file propagates explicit concurrency to pool_size on zonal buckets.""" +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +async def test_cat_file_zonal_explicit_concurrency( + extended_gcsfs, gcs_bucket_mocks, bucket_layout +): + """Tests that cat_file propagates explicit concurrency to pool_size on zonal/RCU buckets.""" # Arrange from gcsfs.zb_hns_utils import MRDPool - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): with mock.patch.object( extended_gcsfs._mrd_pool_cache, "get", new_callable=mock.AsyncMock ) as mock_get: @@ -355,10 +389,13 @@ async def test_cat_file_zonal_explicit_concurrency(extended_gcsfs, gcs_bucket_mo assert mock_get.call_args.kwargs["pool_size"] == 3 -def test_zonal_prefetcher_default_concurrency(extended_gcsfs, gcs_bucket_mocks): +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +def test_zonal_prefetcher_default_concurrency( + extended_gcsfs, gcs_bucket_mocks, bucket_layout +): """Tests that ZonalFile streaming prefetcher defaults to concurrency=4 and pool_size=4.""" # Arrange - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): # Act with extended_gcsfs.open(file_path, "rb") as f: # Assert @@ -390,14 +427,13 @@ def test_split_ranges_across_streams_edge_cases(): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_uses_one_pool_and_splits_across_mrds( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): starts = [0, 10, 20, 30, 40, 50, file_size - 5, 5] ends = [5, 15, 25, 35, 45, 55, file_size + 100, 5] - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: res = await extended_gcsfs._cat_ranges( [file_path] * len(starts), starts, ends, batch_size=3 ) @@ -415,14 +451,13 @@ async def test_cat_ranges_zonal_uses_one_pool_and_splits_across_mrds( @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_multi_file_shares_stream_budget( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): other = f"{TEST_ZONAL_BUCKET}/test/accounts.2.json" paths = [file_path, other] * 4 - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: res = await extended_gcsfs._cat_ranges(paths, 0, 10, batch_size=4) assert res == [json_data[:10]] * len(paths) @@ -455,16 +490,15 @@ def test_allocate_streams_stays_within_budget(): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_many_objects_bounded_by_batch_size( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): paths = [f"{TEST_ZONAL_BUCKET}/obj{i}" for i in range(10)] in_flight = 0 peak = 0 - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: pool = mocks["pool"] async def tracking_get(*args, **kwargs): @@ -492,12 +526,11 @@ async def tracking_close(): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_returns_group_exceptions( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["downloader"].download_ranges.side_effect = ValueError("boom") res = await extended_gcsfs._cat_ranges([file_path] * 2, [0, 10], [5, 15]) assert all(isinstance(r, ValueError) for r in res) @@ -505,10 +538,20 @@ async def test_cat_ranges_zonal_returns_group_exceptions( @pytest.mark.asyncio +@pytest.mark.parametrize( + "bucket_layout", + [ + layouts.STANDARD, + layouts.STANDARD_HIERARCHICAL, + layouts.STANDARD_RAPID_CACHE, + layouts.HIERARCHICAL_RAPID_CACHE, + BucketLayout.UNKNOWN, + ], +) async def test_cat_ranges_non_zonal_delegates_to_fsspec( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.STANDARD): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): with mock.patch( "fsspec.asyn.AsyncFileSystem._cat_ranges", new_callable=mock.AsyncMock ) as mock_super: @@ -519,22 +562,22 @@ async def test_cat_ranges_non_zonal_delegates_to_fsspec( @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_with_max_gap_delegates_to_fsspec( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: with pytest.raises(NotImplementedError): await extended_gcsfs._cat_ranges([file_path], [0], [1], max_gap=0) mocks["pool_cache_get"].assert_not_called() @pytest.mark.asyncio -async def test_cat_ranges_zonal_on_error_raise(extended_gcsfs, gcs_bucket_mocks): - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +async def test_cat_ranges_zonal_on_error_raise( + extended_gcsfs, gcs_bucket_mocks, bucket_layout +): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["downloader"].download_ranges.side_effect = ValueError("boom") with pytest.raises(ValueError, match="boom"): await extended_gcsfs._cat_ranges( @@ -543,12 +586,11 @@ async def test_cat_ranges_zonal_on_error_raise(extended_gcsfs, gcs_bucket_mocks) @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_respects_concurrency_kwarg( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: res = await extended_gcsfs._cat_ranges( [file_path] * 20, 0, 5, batch_size=64, concurrency=3 ) @@ -557,12 +599,11 @@ async def test_cat_ranges_zonal_respects_concurrency_kwarg( @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_batch_size_minus_one_is_unbounded( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: res = await extended_gcsfs._cat_ranges([file_path] * 20, 0, 5, batch_size=-1) assert res == [json_data[:5]] * 20 # fsspec's -1 means "no limit": only the per-object cap applies. @@ -575,18 +616,17 @@ async def test_cat_ranges_zonal_batch_size_minus_one_is_unbounded( @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_mixed_zonal_and_regional_paths( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): regional = f"{TEST_BUCKET}/regional.json" paths = [file_path, regional, file_path, regional] - async def is_zonal(bucket): - return bucket == TEST_ZONAL_BUCKET + async def lookup_layout(bucket): + return bucket_layout if bucket == TEST_ZONAL_BUCKET else layouts.STANDARD - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: pool = mocks["pool"] async def slow_pool_get(*args, **kwargs): @@ -596,7 +636,9 @@ async def slow_pool_get(*args, **kwargs): mocks["pool_cache_get"].side_effect = slow_pool_get err = RuntimeError("regional fail") with ( - mock.patch.object(extended_gcsfs, "_is_zonal_bucket", side_effect=is_zonal), + mock.patch.object( + extended_gcsfs, "_lookup_bucket_layout", side_effect=lookup_layout + ), mock.patch.object( extended_gcsfs, "_cat_file", @@ -620,17 +662,16 @@ async def slow_pool_get(*args, **kwargs): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_init_fallbacks_and_errors( - extended_gcsfs, gcs_bucket_mocks, caplog + extended_gcsfs, gcs_bucket_mocks, caplog, bucket_layout ): with pytest.raises(TypeError): await extended_gcsfs._cat_ranges(file_path, [0], [5]) with pytest.raises(ValueError): await extended_gcsfs._cat_ranges([file_path], [0, 1], [5]) - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["pool"].persisted_size = None mocks["downloader"].persisted_size = None @@ -705,8 +746,9 @@ async def counting_close_pool(pool): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_cancel_closes_pool_and_tasks( - extended_gcsfs, gcs_bucket_mocks, close_pool_cancels + extended_gcsfs, gcs_bucket_mocks, close_pool_cancels, bucket_layout ): started = asyncio.Event() calls = {"downloads": 0, "cancelled": 0, "close_completed": 0} @@ -728,9 +770,7 @@ async def slow_close(): await asyncio.sleep(0) calls["close_completed"] += 1 - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["downloader"].download_ranges.side_effect = blocking_download mocks["pool"].close.side_effect = slow_close @@ -748,8 +788,9 @@ async def slow_close(): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_cancel_while_closing_pool( - extended_gcsfs, gcs_bucket_mocks, close_pool_cancels + extended_gcsfs, gcs_bucket_mocks, close_pool_cancels, bucket_layout ): paths = [f"{TEST_ZONAL_BUCKET}/done", f"{TEST_ZONAL_BUCKET}/in_flight"] closing = asyncio.Event() @@ -769,9 +810,7 @@ async def slow_close(): await asyncio.sleep(0) close_completed += 1 - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["downloader"].download_ranges.side_effect = download_ranges mocks["pool"].close.side_effect = slow_close @@ -790,8 +829,9 @@ async def slow_close(): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_cancel_many_objects_closes_open_pools( - extended_gcsfs, gcs_bucket_mocks, close_pool_cancels + extended_gcsfs, gcs_bucket_mocks, close_pool_cancels, bucket_layout ): paths = [f"{TEST_ZONAL_BUCKET}/obj{i}" for i in range(10)] started = asyncio.Event() @@ -804,9 +844,7 @@ async def blocking_download(read_requests, metadata=None): started.set() await asyncio.Event().wait() - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["downloader"].download_ranges.side_effect = blocking_download leftover = await _cancel_once( @@ -821,8 +859,9 @@ async def blocking_download(read_requests, metadata=None): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_pool_close_error_raised_after_streams_finish( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): paths = [f"{TEST_ZONAL_BUCKET}/done", f"{TEST_ZONAL_BUCKET}/slow"] release = asyncio.Event() @@ -842,9 +881,7 @@ async def close(): close_failed.set() raise RuntimeError("close boom") - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["downloader"].download_ranges.side_effect = download_ranges mocks["pool"].close.side_effect = close @@ -865,8 +902,9 @@ async def close(): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_cancel_while_opening_pool( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): started = asyncio.Event() @@ -874,9 +912,7 @@ async def blocking_get(*args, **kwargs): started.set() await asyncio.Event().wait() - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["pool_cache_get"].side_effect = blocking_get leftover = await _cancel_once( @@ -891,14 +927,13 @@ async def blocking_get(*args, **kwargs): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_chunks_download_ranges_at_mrd_max_ranges( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): starts = list(range(0, 70, 10)) ends = [s + 5 for s in starts] - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: with mock.patch("gcsfs.extended_gcsfs.MRD_MAX_RANGES", 3): res = await extended_gcsfs._cat_ranges( [file_path] * len(starts), starts, ends, batch_size=1 @@ -909,16 +944,15 @@ async def test_cat_ranges_zonal_chunks_download_ranges_at_mrd_max_ranges( @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_short_read_is_an_error( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): async def short_read(read_requests, metadata=None): for offset, length, view in read_requests: view.write(json_data[offset : offset + length - 1]) - with gcs_bucket_mocks( - json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL - ) as mocks: + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: mocks["downloader"].download_ranges.side_effect = short_read res = await extended_gcsfs._cat_ranges( [file_path] * 2, [0, 10], [5, 15], batch_size=1 @@ -927,8 +961,9 @@ async def short_read(read_requests, metadata=None): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_cat_ranges_zonal_closes_all_buffers_when_close_raises( - extended_gcsfs, gcs_bucket_mocks + extended_gcsfs, gcs_bucket_mocks, bucket_layout ): created = [] real_buf = DirectMemmoveBuffer @@ -950,7 +985,7 @@ def close(): return buf with ( - gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL), + gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout), mock.patch("gcsfs.extended_gcsfs.DirectMemmoveBuffer", side_effect=make_buf), pytest.raises(RuntimeError, match="buf0 close failed"), ): @@ -1156,7 +1191,7 @@ async def test_simple_upload_delegates_to_core_for_non_zonal(async_gcs): data = b"test data" with ( - mock.patch.object(async_gcs, "_is_zonal_bucket", return_value=False), + mock.patch.object(async_gcs, "_uses_grpc_data_path", return_value=False), mock.patch( "gcsfs.core.simple_upload", new_callable=mock.AsyncMock ) as mock_core_simple, @@ -1201,7 +1236,7 @@ async def test_initiate_upload_delegates_to_core_for_non_zonal(async_gcs): key = "test-key" with ( - mock.patch.object(async_gcs, "_is_zonal_bucket", return_value=False), + mock.patch.object(async_gcs, "_uses_grpc_data_path", return_value=False), mock.patch( "gcsfs.core.initiate_upload", new_callable=mock.AsyncMock ) as mock_core_initiate, @@ -1283,13 +1318,14 @@ async def test_upload_chunk_delegates_to_core_for_http_url(async_gcs, http_locat @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_get_file_warning_on_missing_persisted_size( - async_gcs, gcs_bucket_mocks, caplog, tmp_path, file_path + async_gcs, gcs_bucket_mocks, caplog, tmp_path, file_path, bucket_layout ): """ Tests that a warning is logged in _get_file when MRD has no 'persisted_size' attribute. """ - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): lpath = tmp_path / "output.txt" with ( caplog.at_level(logging.WARNING, logger="gcsfs"), @@ -1304,15 +1340,16 @@ async def test_get_file_warning_on_missing_persisted_size( @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_get_file_incomplete_download( - async_gcs, gcs_bucket_mocks, tmp_path, file_path + async_gcs, gcs_bucket_mocks, tmp_path, file_path, bucket_layout ): """ Tests that _get_file raises ClientError and correctly removes the local file when the downloaded bytes are less than the expected size. """ - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): lpath = tmp_path / "output.txt" # Simulate download_range returning a shorter byte string than expected @@ -1349,15 +1386,16 @@ async def mock_download_range(*args, **kwargs): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) async def test_get_file_exception_cleanup( - async_gcs, gcs_bucket_mocks, tmp_path, file_path + async_gcs, gcs_bucket_mocks, tmp_path, file_path, bucket_layout ): """ Tests that _get_file correctly removes the local file when an exception occurs during download. """ - with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout): lpath = tmp_path / "output.txt" error_message = "Simulated network failure during download" with ( @@ -1740,6 +1778,7 @@ def _enable_crc32c_checker(monkeypatch): @pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) @pytest.mark.parametrize( "consistency,persisted_size,details_size,crc_payload", [ @@ -1764,7 +1803,13 @@ def _enable_crc32c_checker(monkeypatch): ], ) async def test_get_file_concurrent_downloads_persisted_size( - tmp_path, monkeypatch, consistency, persisted_size, details_size, crc_payload + tmp_path, + monkeypatch, + consistency, + persisted_size, + details_size, + crc_payload, + bucket_layout, ): """Test _get_file_concurrent downloads persisted_size bytes without calling _info.""" from types import SimpleNamespace @@ -1795,7 +1840,7 @@ async def fake_cat_file(rpath, start=None, end=None, **kwargs): # Act with ( mock.patch.object( - fs, "_is_zonal_bucket", new=mock.AsyncMock(return_value=True) + fs, "_lookup_bucket_layout", new=mock.AsyncMock(return_value=bucket_layout) ), mock.patch.object(fs, "_info", new=mock_info), mock.patch.object(fs, "_cat_file", side_effect=fake_cat_file), @@ -1816,7 +1861,10 @@ async def fake_cat_file(rpath, start=None, end=None, **kwargs): @pytest.mark.asyncio -async def test_get_file_concurrent_crc32c_mismatch_raises(tmp_path, monkeypatch): +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +async def test_get_file_concurrent_crc32c_mismatch_raises( + tmp_path, monkeypatch, bucket_layout +): """Test _get_file_concurrent raises ChecksumError when crc32c does not match.""" from types import SimpleNamespace @@ -1847,7 +1895,7 @@ async def fake_cat_file(rpath, start=None, end=None, **kwargs): # Act & Assert with ( mock.patch.object( - fs, "_is_zonal_bucket", new=mock.AsyncMock(return_value=True) + fs, "_lookup_bucket_layout", new=mock.AsyncMock(return_value=bucket_layout) ), mock.patch.object(fs, "_cat_file", side_effect=fake_cat_file), mock.patch("gcsfs.retry.asyncio.sleep", new=mock.AsyncMock()), @@ -1862,3 +1910,150 @@ async def fake_cat_file(rpath, start=None, end=None, **kwargs): ) assert pool.closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bucket_layout", GRPC_READ_LAYOUTS) +async def test_get_threshold_for_disk_reads_grpc_buckets(bucket_layout): + """Test _get_threshold_for_disk_reads returns 5 MiB for Zonal and RCU buckets.""" + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + with mock.patch.object( + fs, "_lookup_bucket_layout", new=mock.AsyncMock(return_value=bucket_layout) + ): + threshold = await fs._get_threshold_for_disk_reads("test-bucket") + assert threshold == 5 * 1024 * 1024 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "bucket_layout", + [ + layouts.STANDARD, + layouts.STANDARD_HIERARCHICAL, + layouts.STANDARD_RAPID_CACHE, + layouts.HIERARCHICAL_RAPID_CACHE, + BucketLayout.UNKNOWN, + ], +) +async def test_get_threshold_for_disk_reads_http_buckets_delegates_to_super( + bucket_layout, +): + """Test _get_threshold_for_disk_reads delegates to super for HTTP buckets.""" + fs = ExtendedGcsFileSystem(token="anon", skip_instance_cache=True) + with ( + mock.patch.object( + fs, "_lookup_bucket_layout", new=mock.AsyncMock(return_value=bucket_layout) + ), + mock.patch( + "gcsfs.core.GCSFileSystem._get_threshold_for_disk_reads", + new=mock.AsyncMock(return_value=12345), + ) as mock_super, + ): + threshold = await fs._get_threshold_for_disk_reads("test-bucket") + assert threshold == 12345 + mock_super.assert_awaited_once_with("test-bucket") + + +@pytest.mark.parametrize( + "bucket_layout", [layouts.STANDARD_RCU, layouts.HIERARCHICAL_RCU] +) +def test_open_rcu_bucket_read_uses_zonal_file( + extended_gcsfs, gcs_bucket_mocks, bucket_layout +): + """Test _open returns ZonalFile for 'rb' on RCU buckets.""" + with gcs_bucket_mocks(json_data, bucket_type_val=bucket_layout) as mocks: + with extended_gcsfs.open(file_path, "rb") as rf: + assert type(rf) is ZonalFile + assert rf.read(10) == json_data[:10] + mocks["pool_cache_get"].assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "bucket_layout", [layouts.STANDARD_RCU, layouts.HIERARCHICAL_RCU] +) +async def test_cp_file_and_merge_on_rcu_bucket_delegate_to_core( + async_gcs, bucket_layout +): + """cp_file and merge on RCU buckets keep using the HTTP implementation.""" + with ( + mock.patch.object( + async_gcs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch( + "gcsfs.core.GCSFileSystem._cp_file", new_callable=mock.AsyncMock + ) as mock_core_cp_file, + mock.patch( + "gcsfs.core.GCSFileSystem._merge", new_callable=mock.AsyncMock + ) as mock_core_merge, + ): + await async_gcs._cp_file(f"{TEST_BUCKET}/a", f"{TEST_BUCKET}/b") + await async_gcs._merge( + f"{TEST_BUCKET}/m", [f"{TEST_BUCKET}/a", f"{TEST_BUCKET}/b"] + ) + + mock_core_cp_file.assert_awaited_once() + mock_core_merge.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_anywhere_cache_bucket_writes_stay_on_http(extended_gcsfs, async_gcs): + """Anywhere Cache (RAPID_CACHE) buckets keep writing over HTTP.""" + layout = layouts.STANDARD_RAPID_CACHE + path = f"{TEST_BUCKET}/anywhere-cache-file" + with ( + mock.patch( + "gcsfs.extended_gcsfs.ExtendedGcsFileSystem._lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=layout, + ), + mock.patch( + "gcsfs.extended_gcsfs.ExtendedGcsFileSystem._sync_lookup_bucket_layout", + return_value=layout, + ), + mock.patch("gcsfs.extended_gcsfs.GCSFile") as mock_gcs_file, + mock.patch("gcsfs.extended_gcsfs.ZonalFile") as mock_zonal_file, + mock.patch( + "gcsfs.core.GCSFileSystem._pipe_file", new_callable=mock.AsyncMock + ) as mock_core_pipe_file, + mock.patch( + "gcsfs.zb_hns_utils.init_aaow", new_callable=mock.AsyncMock + ) as mock_init_aaow, + ): + extended_gcsfs._open(path, "wb") + await async_gcs._pipe_file(path, b"data") + + mock_gcs_file.assert_called_once() + mock_zonal_file.assert_not_called() + mock_core_pipe_file.assert_awaited_once() + mock_init_aaow.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "bucket_layout", [layouts.STANDARD_RCU, layouts.HIERARCHICAL_RCU] +) +async def test_core_upload_helpers_route_rcu_to_extended(async_gcs, bucket_layout): + """core.simple_upload and core.initiate_upload hand RCU buckets to gRPC.""" + with ( + mock.patch.object( + async_gcs, + "_lookup_bucket_layout", + new_callable=mock.AsyncMock, + return_value=bucket_layout, + ), + mock.patch( + "gcsfs.extended_gcsfs.simple_upload", new_callable=mock.AsyncMock + ) as mock_ext_simple_upload, + mock.patch( + "gcsfs.extended_gcsfs.initiate_upload", new_callable=mock.AsyncMock + ) as mock_ext_initiate_upload, + ): + await gcsfs.core.simple_upload(async_gcs, TEST_BUCKET, "key", b"data") + await gcsfs.core.initiate_upload(async_gcs, TEST_BUCKET, "key") + + mock_ext_simple_upload.assert_awaited_once() + mock_ext_initiate_upload.assert_awaited_once() diff --git a/gcsfs/zb_hns_utils.py b/gcsfs/zb_hns_utils.py index 08483ac70..083e80b47 100644 --- a/gcsfs/zb_hns_utils.py +++ b/gcsfs/zb_hns_utils.py @@ -169,6 +169,7 @@ async def init_aaow( generation=None, flush_interval_bytes=None, if_generation_match=None, + storage_class=None, ): """ Creates and opens the AsyncAppendableObjectWriter. @@ -179,6 +180,11 @@ async def init_aaow( A failed precondition raises ``FileExistsError`` for ``0`` and ``GenerationMismatchError`` otherwise; without ``if_generation_match`` the SDK's ``FailedPrecondition`` propagates unchanged. + + ``storage_class`` is passed to the writer only when set. RCU buckets + default new objects to ``STANDARD``, which appendable writes reject, so + writes to them must request ``"RAPID"``; Zonal buckets already default to + ``RAPID`` and leave it unset. """ writer_options = {} # Only pass flush_interval_bytes if the user explicitly provided a @@ -197,6 +203,8 @@ async def init_aaow( generation = 0 elif if_generation_match is not None: kwargs["if_generation_match"] = if_generation_match + if storage_class is not None: + kwargs["storage_class"] = storage_class try: writer = AsyncAppendableObjectWriter( client=grpc_client, @@ -210,7 +218,7 @@ async def init_aaow( if "if_generation_match" not in str(e): raise raise NotImplementedError( - "if_generation_match > 0 on zonal buckets requires a google-cloud-storage " + "if_generation_match > 0 on zonal and RCU buckets requires a google-cloud-storage " "release whose AsyncAppendableObjectWriter accepts if_generation_match" ) from e try: @@ -244,8 +252,10 @@ async def close_mrd(mrd): async def close_aaow(aaow, finalize_on_close=False): """ - Closes the AsyncAppendableObjectWriter gracefully. - Logs a warning if closing fails, instead of raising an exception. + Closes the AsyncAppendableObjectWriter. + Logs a warning if closing fails, and re-raises the error when finalizing, + so that a failed finalize is not reported as a successful write. Closes + that do not finalize, such as the abort paths, stay best-effort. """ if aaow: try: @@ -254,6 +264,8 @@ async def close_aaow(aaow, finalize_on_close=False): logger.warning( f"Error closing AsyncAppendableObjectWriter for {aaow.bucket_name}/{aaow.object_name}: {e}" ) + if finalize_on_close: + raise # Default timeout for synchronous teardowns when no explicit timeout is configured. diff --git a/gcsfs/zonal_file.py b/gcsfs/zonal_file.py index e4cfc128e..01e147d36 100644 --- a/gcsfs/zonal_file.py +++ b/gcsfs/zonal_file.py @@ -27,7 +27,7 @@ class ZonalFile(GCSFile): """ ZonalFile is subclass of GCSFile and handles data operations from - Zonal buckets only using a high-performance gRPC path. + Zonal and RCU buckets using a high-performance gRPC path. """ def __init__( @@ -49,20 +49,24 @@ def __init__( kms_key_name=None, if_generation_match=None, pool_size=zb_hns_utils.DEFAULT_CONCURRENCY, - finalize_on_close=False, + finalize_on_close=None, flush_interval_bytes=_DEFAULT_FLUSH_INTERVAL_BYTES, **kwargs, ): """ Initializes the ZonalFile object. - For Zonal buckets, `finalize_on_close` is set to `False` by default to optimize - for write throughput and keep the file appendable. This means that when exiting - a `with` block or closing, the file will not be automatically finalized. To - ensure the write is finalized, `.commit()` must be called explicitly or - `finalize_on_close` must be set to `True` when opening the file. - - For Zonal buckets, `flush_interval_bytes` controls the write buffer size before + For Zonal buckets, `finalize_on_close` defaults to `False` to optimize + for write throughput and keep the file appendable. For RCU buckets, + `finalize_on_close` defaults to `True`. When `finalize_on_close` is + `False`, exiting a `with` block or closing the file will not + automatically finalize it; to finalize the write, `.commit()` must be + called explicitly or `finalize_on_close` must be set to `True`. + On Zonal and RCU buckets an overwrite replaces the existing object as + soon as the writer opens, so a failed write leaves the partially + written object in its place, unfinalized. + + For Zonal and RCU buckets, `flush_interval_bytes` controls the write buffer size before persisting data to GCS (default: 16 MiB). This value must be a multiple of `_MAX_CHUNK_SIZE_BYTES` (2 MiB). Note that this higher default value may increase memory usage. @@ -75,6 +79,15 @@ def __init__( # MRD pool below has been opened; validate first so an invalid value # fails without a lookup. _resolve_open_generation_match(mode, if_generation_match) + if finalize_on_close is None: + from gcsfs.extended_gcsfs import ExtendedGcsFileSystem + + if isinstance(gcsfs, ExtendedGcsFileSystem): + finalize_on_close = gcsfs._resolve_finalize_on_close( + gcsfs._sync_lookup_bucket_layout(bucket) + ) + else: + finalize_on_close = False self.aaow = None self.finalize_on_close = finalize_on_close self.finalized = False @@ -107,10 +120,19 @@ def __init__( "This may result in incorrect behavior for unfinalized objects." ) elif "w" in self.mode or "a" in self.mode or "x" in self.mode: - pass + if ( + metadata + or fixed_key_metadata + or kms_key_name + or (content_type and content_type != "application/octet-stream") + ): + logger.warning( + "Zonal and RCU buckets do not support content_type, metadata, fixed_key_metadata, " + "or kms_key_name during upload. These parameters will be ignored." + ) else: raise NotImplementedError( - "Only read, write and append operations are currently supported for Zonal buckets." + "Only read, write and append operations are currently supported for Zonal and RCU buckets." ) super().__init__( @@ -161,9 +183,7 @@ async def _init_aaow( except FileNotFoundError: # if file doesn't exist, we don't need generation pass - await self.gcsfs._get_grpc_client() - return await zb_hns_utils.init_aaow( - self.gcsfs.grpc_client, + return await self.gcsfs._init_aaow( bucket_name, object_name, generation, @@ -173,14 +193,20 @@ async def _init_aaow( def _ensure_aaow(self): if self.aaow is None: - self.aaow = asyn.sync( - self.gcsfs.loop, - self._init_aaow, - self.bucket, - self.key, - self.generation, - self.flush_interval_bytes, - ) + try: + self.aaow = asyn.sync( + self.gcsfs.loop, + self._init_aaow, + self.bucket, + self.key, + self.generation, + self.flush_interval_bytes, + ) + except BaseException: + # Mark the file force-flushed so that close() does not try to + # open the writer again. + self.forced = True + raise def _fetch_range( self, @@ -330,14 +356,14 @@ def commit(self): self.finalize_on_close = False def discard(self): - """Discard is not applicable for Zonal Buckets. Log a warning instead.""" + """Discard is not applicable for Zonal and RCU buckets. Log a warning instead.""" logger.warning( - "Discard is not applicable for Zonal Buckets. \ + "Discard is not applicable for Zonal and RCU buckets. \ Data is uploaded via streaming and cannot be cancelled." ) def _initiate_upload(self): - """Initiates the upload for Zonal buckets using gRPC.""" + """Initiates the upload for Zonal and RCU buckets using gRPC.""" from gcsfs.extended_gcsfs import initiate_upload self.location = asyn.sync( @@ -356,7 +382,7 @@ def _initiate_upload(self): ) def _simple_upload(self): - """Performs a simple upload for Zonal buckets using gRPC.""" + """Performs a simple upload for Zonal and RCU buckets using gRPC.""" from gcsfs.extended_gcsfs import simple_upload self.buffer.seek(0) @@ -385,10 +411,23 @@ def _upload_chunk(self, final=False): ) def _close_impl(self): - super()._close_impl() + errors = [] + + # A failed closing flush must not skip the teardown below, and the + # partially written object must not be finalized. + try: + super()._close_impl() + except Exception as e: + errors.append(e) + if not self.finalized: + self.finalize_on_close = False + if self.aaow is not None: + # fsspec skips its cache invalidation when the closing flush + # fails, but the open writer has already replaced the object. + self.fs.invalidate_cache(self.path) + self.fs.invalidate_cache(self.fs._parent(self.path)) timeout = self.timeout or DEFAULT_TEARDOWN_TIMEOUT_SECONDS - errors = [] # Teardown the read-side MRD pool if initialized. if hasattr(self, "mrd_pool") and self.mrd_pool: diff --git a/pyproject.toml b/pyproject.toml index beeb954a3..c5d0c7488 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,8 +29,8 @@ dependencies = [ "fsspec>=2026.9.0", "google-auth>=1.2", "google-auth-oauthlib", - "google-cloud-storage>=3.14.1", - "google-cloud-storage-control", + "google-cloud-storage>=3.15.0", + "google-cloud-storage-control>=1.14.0", "requests", ]