Feat: Introduce ExtendedGcsFileSystem for Zonal Bucket gRPC Read Path

.github/workflows/ci.yml

@@ -35,13 +35,22 @@ jobs:
       - name: install
         run: |
           pip install -e .
-      - name: Run tests
+      - name: Run Standard Tests
         run: |
           export GOOGLE_APPLICATION_CREDENTIALS=$(pwd)/gcsfs/tests/fake-secret.json
           pytest -vv -s \
           --log-format="%(asctime)s %(levelname)s %(message)s" \
           --log-date-format="%H:%M:%S" \
-          gcsfs/
+          gcsfs/ \
+          --ignore=gcsfs/tests/test_extended_gcsfs.py
+      - name: Run Extended Tests
+        run: |
+          export GOOGLE_APPLICATION_CREDENTIALS=$(pwd)/gcsfs/tests/fake-secret.json
+          export GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT="true"
+          pytest -vv -s \
+            --log-format="%(asctime)s %(levelname)s %(message)s" \
+            --log-date-format="%H:%M:%S" \
+            gcsfs/tests/test_extended_gcsfs.py
 
   lint:
     name: lint

.isort.cfg

@@ -1,2 +1,3 @@
 [settings]
+profile = black
 known_third_party = aiohttp,click,decorator,fsspec,fuse,google,google_auth_oauthlib,pytest,requests,setuptools

environment_gcsfs.yaml

@@ -13,9 +13,11 @@ dependencies:
   - google-auth-oauthlib
   - google-cloud-core
   - google-cloud-storage
+  - grpcio
   - pytest
   - pytest-timeout
   - pytest-asyncio
+  - pytest-subtests
   - requests
   - ujson
   - pip:

gcsfs/__init__.py

@@ -1,10 +1,27 @@
+import logging
+import os
+
 from ._version import get_versions
 
+logger = logging.getLogger(__name__)
 __version__ = get_versions()["version"]
 del get_versions
 from .core import GCSFileSystem
 from .mapping import GCSMap
 
+if os.getenv("GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT", "false").lower() in ("true", "1"):
+    try:
+        from .extended_gcsfs import ExtendedGcsFileSystem as GCSFileSystem
+
+        logger.info(
+            "gcsfs experimental features enabled via GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT."
+        )
+    except ImportError as e:
+        logger.warning(
+            f"GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT is set, but failed to import experimental features: {e}"
+        )
+        # Fallback to core GCSFileSystem, do not register here
+
 __all__ = ["GCSFileSystem", "GCSMap"]
 
 from . import _version

gcsfs/extended_gcsfs.py

@@ -0,0 +1,228 @@
+import logging
+from enum import Enum
+
+from fsspec import asyn
+from google.api_core import exceptions as api_exceptions
+from google.api_core import gapic_v1
+from google.api_core.client_info import ClientInfo
+from google.cloud import storage_control_v2
+from google.cloud.storage._experimental.asyncio.async_grpc_client import AsyncGrpcClient
+
+from gcsfs import __version__ as version
+from gcsfs import zb_hns_utils
+from gcsfs.core import GCSFile, GCSFileSystem
+from gcsfs.zonal_file import ZonalFile
+
+logger = logging.getLogger("gcsfs")
+
+USER_AGENT = "python-gcsfs"
+
+
+class BucketType(Enum):
+    ZONAL_HIERARCHICAL = "ZONAL_HIERARCHICAL"
+    HIERARCHICAL = "HIERARCHICAL"
+    NON_HIERARCHICAL = "NON_HIERARCHICAL"
+    UNKNOWN = "UNKNOWN"
+
+
+gcs_file_types = {
+    BucketType.ZONAL_HIERARCHICAL: ZonalFile,
+    BucketType.NON_HIERARCHICAL: GCSFile,
+    BucketType.HIERARCHICAL: GCSFile,
+    BucketType.UNKNOWN: GCSFile,
+}
+
+
+class ExtendedGcsFileSystem(GCSFileSystem):
+    """
+    This class will be used when GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT env variable is set to true.
+    ExtendedGcsFileSystem is a subclass of GCSFileSystem that adds new logic for bucket types
+    including zonal and hierarchical. For buckets without special properties, it forwards requests
+    to the parent class GCSFileSystem for default processing.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.grpc_client = None
+        self.storage_control_client = None
+        # initializing grpc and storage control client for Hierarchical and
+        # zonal bucket operations
+        self.grpc_client = asyn.sync(self.loop, self._create_grpc_client)
+        self._storage_control_client = asyn.sync(
+            self.loop, self._create_control_plane_client
+        )
+        self._storage_layout_cache = {}
+
+    async def _create_grpc_client(self):
+        if self.grpc_client is None:
+            return AsyncGrpcClient(
+                client_info=ClientInfo(user_agent=f"{USER_AGENT}/{version}"),
+            ).grpc_client
+        else:
+            return self.grpc_client
+
+    async def _create_control_plane_client(self):
+        # Initialize the storage control plane client for bucket
+        # metadata operations
+        client_info = gapic_v1.client_info.ClientInfo(
+            user_agent=f"{USER_AGENT}/{version}"
+        )
+        return storage_control_v2.StorageControlAsyncClient(
+            credentials=self.credentials.credentials, client_info=client_info
+        )
+
+    async def _get_bucket_type(self, bucket):
+        if bucket in self._storage_layout_cache:
+            return self._storage_layout_cache[bucket]
+        try:
+
+            # Bucket name details
+            bucket_name_value = f"projects/_/buckets/{bucket}/storageLayout"
+            # Make the request to get bucket type
+            response = await self._storage_control_client.get_storage_layout(
+                name=bucket_name_value
+            )
+
+            if response.location_type == "zone":
+                return BucketType.ZONAL_HIERARCHICAL
+            else:
+                # This should be updated to include HNS in the future
+                return BucketType.NON_HIERARCHICAL
+        except api_exceptions.NotFound:
+            print(f"Error: Bucket {bucket} not found or you lack permissions.")
+            return BucketType.UNKNOWN
+        except Exception as e:
+            logger.error(
+                f"Could not determine bucket type for bucket name {bucket}: {e}"
+            )
+            # Default to UNKNOWN
+            return BucketType.UNKNOWN
+
+    _sync_get_bucket_type = asyn.sync_wrapper(_get_bucket_type)
+
+    def _open(
+        self,
+        path,
+        mode="rb",
+        block_size=None,
+        cache_options=None,
+        acl=None,
+        consistency=None,
+        metadata=None,
+        autocommit=True,
+        fixed_key_metadata=None,
+        generation=None,
+        **kwargs,
+    ):
+        """
+        Open a file.
+        """
+        bucket, _, _ = self.split_path(path)
+        bucket_type = self._sync_get_bucket_type(bucket)
+        self._storage_layout_cache[bucket] = bucket_type
+        return gcs_file_types[bucket_type](
+            self,
+            path,
+            mode,
+            block_size,
+            cache_options=cache_options,
+            consistency=consistency,
+            metadata=metadata,
+            acl=acl,
+            autocommit=autocommit,
+            fixed_key_metadata=fixed_key_metadata,
+            generation=generation,
+            **kwargs,
+        )
+
+    # Replacement method for _process_limits to support new params (offset and length) for MRD.
+    async def _process_limits_to_offset_and_length(self, path, start, end):
+        """
+        Calculates the read offset and length from start and end parameters.
+
+        Args:
+            path (str): The path to the file.
+            start (int | None): The starting byte position.
+            end (int | None): The ending byte position.
+
+        Returns:
+            tuple: A tuple containing (offset, length).
+
+        Raises:
+            ValueError: If the calculated range is invalid.
+        """
+        size = None
+
+        if start is None:
+            offset = 0
+        elif start < 0:
+            size = size or (await self._info(path))["size"]
+            offset = size + start
+        else:
+            offset = start
+
+        if end is None:
+            size = size or (await self._info(path))["size"]
+            effective_end = size
+        elif end < 0:
+            size = size or (await self._info(path))["size"]
+            effective_end = size + end
+        else:
+            effective_end = end
+
+        size = size or (await self._info(path))["size"]
+        if offset < 0:
+            raise ValueError(f"Calculated start offset ({offset}) cannot be negative.")
+        if effective_end < offset:
+            raise ValueError(
+                f"Calculated end position ({effective_end}) cannot be before start offset ({offset})."
+            )
+        elif effective_end == offset:
+            length = 0  # Handle zero-length slice
+        elif effective_end > size:
+            length = max(0, size - offset)  # Clamp and ensure non-negative
+        else:
+            length = effective_end - offset  # Normal case
+
+        return offset, length
+
+    sync_process_limits_to_offset_and_length = asyn.sync_wrapper(
+        _process_limits_to_offset_and_length
+    )
+
+    async def _is_zonal_bucket(self, bucket):
+        bucket_type = await self._get_bucket_type(bucket)
+        self._storage_layout_cache[bucket] = bucket_type
+        return bucket_type == BucketType.ZONAL_HIERARCHICAL
+
+    async def _cat_file(self, path, start=None, end=None, **kwargs):
+        """
+        Fetch a file's contents as bytes.
+        """
+        mrd = kwargs.pop("mrd", None)
+        mrd_created = False
+
+        # A new MRD is required when read is done directly by the
+        # GCSFilesystem class without creating a GCSFile object first.
+        if mrd is None:
+            bucket, object_name, generation = self.split_path(path)
+            # Fall back to default implementation if not a zonal bucket
+            if not await self._is_zonal_bucket(bucket):
+                return await super()._cat_file(path, start=start, end=end, **kwargs)
+
+            mrd = await zb_hns_utils.create_mrd(
+                self.grpc_client, bucket, object_name, generation
+            )
+            mrd_created = True
+
+        offset, length = await self._process_limits_to_offset_and_length(
+            path, start, end
+        )
+        try:
+            return await zb_hns_utils.download_range(
+                offset=offset, length=length, mrd=mrd
+            )
+        finally:
+            # Explicit cleanup if we created the MRD and it has a close method
+            if mrd_created:
+                await mrd.close()

gcsfs/tests/conftest.py

@@ -1,7 +1,10 @@
+import logging
 import os
 import shlex
 import subprocess
 import time
+from contextlib import nullcontext
+from unittest.mock import patch
 
 import fsspec
 import pytest
@@ -91,10 +94,9 @@ def docker_gcs():
 def gcs_factory(docker_gcs):
     params["endpoint_url"] = docker_gcs
 
-    def factory(default_location=None):
+    def factory(**kwargs):
         GCSFileSystem.clear_instance_cache()
-        params["default_location"] = default_location
-        return fsspec.filesystem("gcs", **params)
+        return fsspec.filesystem("gcs", **params, **kwargs)
 
     return factory
 
@@ -125,6 +127,57 @@ def gcs(gcs_factory, populate=True):
             pass
 
 
+def _cleanup_gcs(gcs, is_real_gcs):
+    """Only remove the bucket/contents if we are NOT using the real GCS, logging a warning on failure."""
+    del os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"]
+    if is_real_gcs:
+        return
+    try:
+        gcs.rm(gcs.find(TEST_BUCKET), recursive=True)
+        gcs.rm(TEST_BUCKET)
+    except Exception as e:
+        logging.warning(f"Failed to clean up GCS bucket {TEST_BUCKET}: {e}")
+
+
+@pytest.fixture
+def extended_gcsfs(gcs_factory, populate=True):
+    # Check if we are running against a real GCS endpoint
+    is_real_gcs = (
+        os.environ.get("STORAGE_EMULATOR_HOST") == "https://storage.googleapis.com"
+    )
+
+    # Mock authentication if not using a real GCS endpoint,
+    # since grpc client in extended_gcsfs does not work with anon access
+    mock_authentication_manager = (
+        patch("google.auth.default", return_value=(None, "fake-project"))
+        if not is_real_gcs
+        else nullcontext()
+    )
+
+    with mock_authentication_manager:
+        os.environ["GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"] = "true"
+        extended_gcsfs = gcs_factory()
+        try:
+            # Only create/delete/populate the bucket if we are NOT using the real GCS endpoint
+            if not is_real_gcs:
+                try:
+                    extended_gcsfs.rm(TEST_BUCKET, recursive=True)
+                except FileNotFoundError:
+                    pass
+                try:
+                    extended_gcsfs.mkdir(TEST_BUCKET)
+                except Exception:
+                    pass
+                if populate:
+                    extended_gcsfs.pipe(
+                        {TEST_BUCKET + "/" + k: v for k, v in allfiles.items()}
+                    )
+            extended_gcsfs.invalidate_cache()
+            yield extended_gcsfs
+        finally:
+            _cleanup_gcs(extended_gcsfs, is_real_gcs)
+
+
 @pytest.fixture
 def gcs_versioned(gcs_factory):
     gcs = gcs_factory()