feat(rcu): route Rapid Cache Ultra writes to the gRPC appendable writer

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,15 +3117,15 @@ 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
     from .extended_gcsfs import initiate_upload as ext_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,

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,17 +611,19 @@ 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),
     ):
         mocks = {
             "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",

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

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(

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):