[#5139] feat(python-client): Support GCS fileset in the python GVFS client

bundles/gcp-bundle/src/main/java/org/apache/gravitino/gcs/fs/GCSFileSystemProvider.java

@@ -25,8 +25,12 @@
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class GCSFileSystemProvider implements FileSystemProvider {
+  private static final Logger LOGGER = LoggerFactory.getLogger(GCSFileSystemProvider.class);
+
   @Override
   public FileSystem getFileSystem(Path path, Map<String, String> config) throws IOException {
     Configuration configuration = new Configuration();
@@ -35,6 +39,7 @@ public FileSystem getFileSystem(Path path, Map<String, String> config) throws IO
           configuration.set(k.replace("gravitino.bypass.", ""), v);
         });
 
+    LOGGER.info("Creating GCS file system with config: {}", config);
     return GoogleHadoopFileSystem.newInstance(path.toUri(), configuration);
   }
 

catalogs/catalog-hadoop/src/test/java/org/apache/gravitino/catalog/hadoop/integration/test/HadoopGCSCatalogIT.java

@@ -44,7 +44,7 @@ public class HadoopGCSCatalogIT extends HadoopCatalogIT {
 
   @Override
   public void startIntegrationTest() throws Exception {
-    // Do nothing.
+    // Just overwrite super, do nothing.
   }
 
   @BeforeAll

clients/client-python/gravitino/filesystem/gvfs.py

@@ -14,19 +14,20 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
+import os
 from enum import Enum
 from pathlib import PurePosixPath
 from typing import Dict, Tuple
 import re
+import importlib
 import fsspec
 
 from cachetools import TTLCache, LRUCache
 from fsspec import AbstractFileSystem
 from fsspec.implementations.local import LocalFileSystem
 from fsspec.implementations.arrow import ArrowFSWrapper
 from fsspec.utils import infer_storage_options
-from pyarrow.fs import HadoopFileSystem
+
 from readerwriterlock import rwlock
 from gravitino.audit.caller_context import CallerContext, CallerContextHolder
 from gravitino.audit.fileset_audit_constants import FilesetAuditConstants
@@ -47,6 +48,7 @@
 class StorageType(Enum):
     HDFS = "hdfs"
     LOCAL = "file"
+    GCS = "gs"
 
 
 class FilesetContextPair:
@@ -66,7 +68,7 @@ def filesystem(self):
 
 
 class GravitinoVirtualFileSystem(fsspec.AbstractFileSystem):
-    """This is a virtual file system which users can access `fileset` and
+    """This is a virtual file system that users can access `fileset` and
     other resources.
 
     It obtains the actual storage location corresponding to the resource from the
@@ -149,6 +151,7 @@ def __init__(
         self._cache_lock = rwlock.RWLockFair()
         self._catalog_cache = LRUCache(maxsize=100)
         self._catalog_cache_lock = rwlock.RWLockFair()
+        self._options = options
 
         super().__init__(**kwargs)
 
@@ -309,7 +312,9 @@ def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs):
         )
         dst_actual_path = dst_context_pair.actual_file_location()
 
-        if storage_type == StorageType.HDFS:
+        # convert the following to in
+
+        if storage_type in [StorageType.HDFS, StorageType.GCS]:
             src_context_pair.filesystem().mv(
                 self._strip_storage_protocol(storage_type, src_actual_path),
                 self._strip_storage_protocol(storage_type, dst_actual_path),
@@ -540,7 +545,11 @@ def _convert_actual_path(
         :param virtual_location: Virtual location
         :return A virtual path
         """
-        if storage_location.startswith(f"{StorageType.HDFS.value}://"):
+
+        # If the storage path starts with hdfs, gcs, we should use the path as the prefix.
+        if storage_location.startswith(
+            f"{StorageType.HDFS.value}://"
+        ) or storage_location.startswith(f"{StorageType.GCS.value}://"):
             actual_prefix = infer_storage_options(storage_location)["path"]
         elif storage_location.startswith(f"{StorageType.LOCAL.value}:/"):
             actual_prefix = storage_location[len(f"{StorageType.LOCAL.value}:") :]
@@ -681,6 +690,8 @@ def _recognize_storage_type(path: str):
             return StorageType.HDFS
         if path.startswith(f"{StorageType.LOCAL.value}:/"):
             return StorageType.LOCAL
+        if path.startswith(f"{StorageType.GCS.value}://"):
+            return StorageType.GCS
         raise GravitinoRuntimeException(
             f"Storage type doesn't support now. Path:{path}"
         )
@@ -705,10 +716,11 @@ def _strip_storage_protocol(storage_type: StorageType, path: str):
         :param path: The path
         :return: The stripped path
         """
-        if storage_type == StorageType.HDFS:
+        if storage_type in (StorageType.HDFS, StorageType.GCS):
             return path
         if storage_type == StorageType.LOCAL:
             return path[len(f"{StorageType.LOCAL.value}:") :]
+
         raise GravitinoRuntimeException(
             f"Storage type:{storage_type} doesn't support now."
         )
@@ -774,9 +786,12 @@ def _get_filesystem(self, actual_file_location: str):
             if cache_value is not None:
                 return cache_value
             if storage_type == StorageType.HDFS:
-                fs = ArrowFSWrapper(HadoopFileSystem.from_uri(actual_file_location))
+                fs_class = importlib.import_module("pyarrow.fs").HadoopFileSystem
+                fs = ArrowFSWrapper(fs_class.from_uri(actual_file_location))
             elif storage_type == StorageType.LOCAL:
                 fs = LocalFileSystem()
+            elif storage_type == StorageType.GCS:
+                fs = ArrowFSWrapper(self._get_gcs_filesystem())
             else:
                 raise GravitinoRuntimeException(
                     f"Storage type: `{storage_type}` doesn't support now."
@@ -786,5 +801,23 @@ def _get_filesystem(self, actual_file_location: str):
         finally:
             write_lock.release()
 
+    def _get_gcs_filesystem(self):
+        # get All keys from the options that start with 'gravitino.bypass.gcs.' and remove the prefix
+        gcs_options = {
+            key[len(GVFSConfig.GVFS_FILESYSTEM_BY_PASS_GCS) :]: value
+            for key, value in self._options.items()
+            if key.startswith(GVFSConfig.GVFS_FILESYSTEM_BY_PASS_GCS)
+        }
+
+        # get 'service-account-key' from gcs_options, if the key is not found, throw an exception
+        service_account_key_path = gcs_options.get(GVFSConfig.GVFS_FILESYSTEM_KEY_FILE)
+        if service_account_key_path is None:
+            raise GravitinoRuntimeException(
+                "Service account key is not found in the options."
+            )
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = service_account_key_path
+
+        return importlib.import_module("pyarrow.fs").GcsFileSystem()
+
 
 fsspec.register_implementation(PROTOCOL_NAME, GravitinoVirtualFileSystem)

clients/client-python/gravitino/filesystem/gvfs_config.py

@@ -31,3 +31,7 @@ class GVFSConfig:
     OAUTH2_CREDENTIAL = "oauth2_credential"
     OAUTH2_PATH = "oauth2_path"
     OAUTH2_SCOPE = "oauth2_scope"
+
+    GVFS_FILESYSTEM_BY_PASS = "gravitino.bypass"
+    GVFS_FILESYSTEM_BY_PASS_GCS = "gravitino.bypass.gcs."
+    GVFS_FILESYSTEM_KEY_FILE = "service-account-key-path"

clients/client-python/requirements.txt

@@ -22,4 +22,5 @@ dataclasses-json==0.6.6
 readerwriterlock==1.0.9
 fsspec==2024.3.1
 pyarrow==15.0.2
-cachetools==5.3.3
+cachetools==5.3.3
+google-auth==2.35.0

clients/client-python/tests/integration/integration_test_env.py

@@ -80,6 +80,12 @@ def setUpClass(cls):
             )
             sys.exit(0)
 
+        # remove data dir under gravitino_home
+        data_dir = os.path.join(cls.gravitino_home, "data")
+        if os.path.exists(data_dir):
+            logger.info("Remove Gravitino data directory: %s", data_dir)
+            subprocess.run(["rm", "-rf", data_dir], check=False)
+
         logger.info("Starting integration test environment...")
 
         # Start Gravitino Server
@@ -141,6 +147,12 @@ def restart_server(cls):
                 "project root directory."
             )
 
+        # remove data dir under gravitino_home
+        data_dir = os.path.join(gravitino_home, "data")
+        if os.path.exists(data_dir):
+            logger.info("Remove Gravitino data directory: %s", data_dir)
+            subprocess.run(["rm", "-rf", data_dir], check=False)
+
         # Restart Gravitino Server
         env_vars = os.environ.copy()
         env_vars["HADOOP_USER_NAME"] = "anonymous"

clients/client-python/tests/integration/test_catalog.py

@@ -39,7 +39,8 @@
 class TestCatalog(IntegrationTestEnv):
     metalake_name: str = "TestSchema_metalake" + str(randint(1, 10000))
 
-    catalog_name: str = "testCatalog"
+    catalog_name: str = "testCatalog" + str(randint(1, 10000))
+    catalog_name_bak = catalog_name
     catalog_comment: str = "catalogComment"
     catalog_location_prop: str = "location"  # Fileset Catalog must set `location`
     catalog_provider: str = "hadoop"
@@ -81,21 +82,25 @@ def clean_test_data(self):
         )
         try:
             logger.info(
-                "Drop catalog %s[%s]",
+                "TestCatalog: drop catalog %s[%s]",
                 self.catalog_ident,
                 self.gravitino_client.drop_catalog(name=self.catalog_name, force=True),
             )
         except GravitinoRuntimeException:
-            logger.warning("Failed to drop catalog %s", self.catalog_name)
+            logger.warning("TestCatalog: failed to drop catalog %s", self.catalog_name)
 
         try:
             logger.info(
-                "Drop metalake %s[%s]",
+                "TestCatalog: drop metalake %s[%s]",
                 self.metalake_name,
                 self.gravitino_admin_client.drop_metalake(self.metalake_name),
             )
         except GravitinoRuntimeException:
-            logger.warning("Failed to drop metalake %s", self.metalake_name)
+            logger.warning(
+                "TestCatalog: failed to drop metalake %s", self.metalake_name
+            )
+
+        self.catalog_name = self.catalog_name_bak
 
     def test_list_catalogs(self):
         self.create_catalog(self.catalog_name)