Support loading custom catalog impl

Author: Jack Ye

pyiceberg/catalog/__init__.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import importlib
 import logging
 import re
 import uuid
@@ -77,6 +78,7 @@
 
 TOKEN = "token"
 TYPE = "type"
+PY_CATALOG_IMPL = "py-catalog-impl"
 ICEBERG = "iceberg"
 TABLE_TYPE = "table_type"
 WAREHOUSE_LOCATION = "warehouse"
@@ -230,6 +232,13 @@ def load_catalog(name: Optional[str] = None, **properties: Optional[str]) -> Cat
     env = _ENV_CONFIG.get_catalog_config(name)
     conf: RecursiveDict = merge_config(env or {}, cast(RecursiveDict, properties))
 
+    if catalog_impl := properties.get(PY_CATALOG_IMPL):
+        if catalog := _import_catalog(catalog_impl, properties):
+            logger.info("Loaded Catalog: %s", catalog_impl)
+            return catalog
+        else:
+            raise ValueError(f"Could not initialize Catalog: {catalog_impl}")
+
     catalog_type: Optional[CatalogType]
     provided_catalog_type = conf.get(TYPE)
 
@@ -283,6 +292,20 @@ def delete_data_files(io: FileIO, manifests_to_delete: List[ManifestFile]) -> No
                 deleted_files[path] = True
 
 
+def _import_catalog(catalog_impl: str, properties: Properties) -> Optional[Catalog]:
+    try:
+        path_parts = catalog_impl.split(".")
+        if len(path_parts) < 2:
+            raise ValueError(f"py-catalog-impl should be full path (module.CustomCatalog), got: {catalog_impl}")
+        module_name, class_name = ".".join(path_parts[:-1]), path_parts[-1]
+        module = importlib.import_module(module_name)
+        class_ = getattr(module, class_name)
+        return class_(properties)
+    except ModuleNotFoundError:
+        logger.warning("Could not initialize Catalog: %s", catalog_impl)
+        return None
+
+
 @dataclass
 class PropertiesUpdateSummary:
     removed: List[str]

tests/catalog/test_base.py

@@ -32,11 +32,7 @@
 from pydantic_core import ValidationError
 from pytest_lazyfixture import lazy_fixture
 
-from pyiceberg.catalog import (
-    Catalog,
-    MetastoreCatalog,
-    PropertiesUpdateSummary,
-)
+from pyiceberg.catalog import PY_CATALOG_IMPL, Catalog, MetastoreCatalog, PropertiesUpdateSummary, load_catalog
 from pyiceberg.exceptions import (
     NamespaceAlreadyExistsError,
     NamespaceNotEmptyError,
@@ -295,6 +291,13 @@ def given_catalog_has_a_table(
     )
 
 
+def test_load_catalog_does_not_exist() -> None:
+    with pytest.raises(ValueError) as exc_info:
+        load_catalog("catalog", **{PY_CATALOG_IMPL: "pyiceberg.does.not.exist.Catalog"})
+
+    assert "Could not initialize Catalog: pyiceberg.does.not.exist.Catalog" in str(exc_info.value)
+
+
 def test_namespace_from_tuple() -> None:
     # Given
     identifier = ("com", "organization", "department", "my_table")

tests/catalog/test_glue.py

@@ -22,6 +22,7 @@
 import pytest
 from moto import mock_aws
 
+from pyiceberg.catalog import PY_CATALOG_IMPL, load_catalog
 from pyiceberg.catalog.glue import GlueCatalog
 from pyiceberg.exceptions import (
     NamespaceAlreadyExistsError,
@@ -46,6 +47,10 @@
 )
 
 
+def test_load_catalog_from_impl() -> None:
+    assert isinstance(load_catalog("catalog", **{PY_CATALOG_IMPL: "pyiceberg.catalog.glue.GlueCatalog"}), GlueCatalog)
+
+
 @mock_aws
 def test_create_table_with_database_location(
     _glue: boto3.client,