[#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)