replace InMemoryCatalog by a subclass of SqlCatalog with SQLite in-memory

Author: hussein-awala

tests/catalog/test_base.py

@@ -17,43 +17,31 @@
 # pylint:disable=redefined-outer-name
 
 
-import uuid
 from pathlib import PosixPath
-from typing import (
-    Dict,
-    List,
-    Optional,
-    Set,
-    Tuple,
-    Union,
-)
+from typing import Optional, Union
 
 import pyarrow as pa
 import pytest
 from pydantic_core import ValidationError
 from pytest_lazyfixture import lazy_fixture
 
-from pyiceberg.catalog import Catalog, MetastoreCatalog, PropertiesUpdateSummary, load_catalog
+from pyiceberg.catalog import Catalog, load_catalog
+from pyiceberg.catalog.sql import SqlCatalog
 from pyiceberg.exceptions import (
     NamespaceAlreadyExistsError,
     NamespaceNotEmptyError,
     NoSuchNamespaceError,
     NoSuchTableError,
     TableAlreadyExistsError,
 )
-from pyiceberg.io import WAREHOUSE, load_file_io
+from pyiceberg.io import WAREHOUSE
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table import (
     AddSchemaUpdate,
-    CommitTableResponse,
     SetCurrentSchemaUpdate,
     Table,
-    TableRequirement,
-    TableUpdate,
-    update_table_metadata,
 )
-from pyiceberg.table.metadata import new_table_metadata
 from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder
 from pyiceberg.transforms import IdentityTransform
 from pyiceberg.typedef import EMPTY_DICT, Identifier, Properties
@@ -62,21 +50,20 @@
 DEFAULT_WAREHOUSE_LOCATION = "file:///tmp/warehouse"
 
 
-class InMemoryCatalog(MetastoreCatalog):
+class InMemoryCatalog(SqlCatalog):
     """
-    An in-memory catalog implementation that uses in-memory data-structures to store the namespaces and tables.
+    An in-memory catalog implementation that uses SqlCatalog with SQLite in-memory database.
 
     This is useful for test, demo, and playground but not in production as data is not persisted.
     """
 
-    __tables: Dict[Identifier, Table]
-    __namespaces: Dict[Identifier, Properties]
-
-    def __init__(self, name: str, **properties: str) -> None:
-        super().__init__(name, **properties)
-        self.__tables = {}
-        self.__namespaces = {}
-        self._warehouse_location = properties.get(WAREHOUSE, DEFAULT_WAREHOUSE_LOCATION)
+    def __init__(self, name: str, **kwargs: str) -> None:
+        if WAREHOUSE not in kwargs:
+            kwargs[WAREHOUSE] = DEFAULT_WAREHOUSE_LOCATION
+        self._warehouse_location = kwargs[WAREHOUSE]
+        if "uri" not in kwargs:
+            kwargs["uri"] = "sqlite:///:memory:"
+        super().__init__(name=name, **kwargs)
 
     def create_table(
         self,
@@ -86,181 +73,24 @@ def create_table(
         partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC,
         sort_order: SortOrder = UNSORTED_SORT_ORDER,
         properties: Properties = EMPTY_DICT,
-        table_uuid: Optional[uuid.UUID] = None,
     ) -> Table:
-        schema: Schema = self._convert_schema_if_needed(schema)  # type: ignore
-
-        identifier = Catalog.identifier_to_tuple(identifier)
-        namespace = Catalog.namespace_from(identifier)
-
-        if identifier in self.__tables:
-            raise TableAlreadyExistsError(f"Table already exists: {identifier}")
-        else:
-            if namespace not in self.__namespaces:
-                self.__namespaces[namespace] = {}
-
-            if not location:
-                location = f'{self._warehouse_location}/{"/".join(identifier)}'
-            location = location.rstrip("/")
-
-            metadata_location = self._get_metadata_location(location=location)
-            metadata = new_table_metadata(
-                schema=schema,
-                partition_spec=partition_spec,
-                sort_order=sort_order,
-                location=location,
-                properties=properties,
-                table_uuid=table_uuid,
-            )
-            io = load_file_io({**self.properties, **properties}, location=location)
-            self._write_metadata(metadata, io, metadata_location)
-
-            table = Table(
-                identifier=identifier,
-                metadata=metadata,
-                metadata_location=metadata_location,
-                io=io,
-                catalog=self,
-            )
-            self.__tables[identifier] = table
-            return table
-
-    def register_table(self, identifier: Union[str, Identifier], metadata_location: str) -> Table:
-        raise NotImplementedError
-
-    def commit_table(
-        self, table: Table, requirements: Tuple[TableRequirement, ...], updates: Tuple[TableUpdate, ...]
-    ) -> CommitTableResponse:
-        identifier_tuple = self._identifier_to_tuple_without_catalog(table.identifier)
-        current_table = self.load_table(identifier_tuple)
-        base_metadata = current_table.metadata
-
-        for requirement in requirements:
-            requirement.validate(base_metadata)
-
-        updated_metadata = update_table_metadata(base_metadata, updates)
-        if updated_metadata == base_metadata:
-            # no changes, do nothing
-            return CommitTableResponse(metadata=base_metadata, metadata_location=current_table.metadata_location)
-
-        # write new metadata
-        new_metadata_version = self._parse_metadata_version(current_table.metadata_location) + 1
-        new_metadata_location = self._get_metadata_location(current_table.metadata.location, new_metadata_version)
-        self._write_metadata(updated_metadata, current_table.io, new_metadata_location)
-
-        # update table state
-        current_table.metadata = updated_metadata
-
-        return CommitTableResponse(metadata=updated_metadata, metadata_location=new_metadata_location)
-
-    def load_table(self, identifier: Union[str, Identifier]) -> Table:
-        identifier_tuple = self._identifier_to_tuple_without_catalog(identifier)
-        try:
-            return self.__tables[identifier_tuple]
-        except KeyError as error:
-            raise NoSuchTableError(f"Table does not exist: {identifier_tuple}") from error
-
-    def drop_table(self, identifier: Union[str, Identifier]) -> None:
-        identifier_tuple = self._identifier_to_tuple_without_catalog(identifier)
-        try:
-            self.__tables.pop(identifier_tuple)
-        except KeyError as error:
-            raise NoSuchTableError(f"Table does not exist: {identifier_tuple}") from error
-
-    def purge_table(self, identifier: Union[str, Identifier]) -> None:
-        self.drop_table(identifier)
-
-    def rename_table(self, from_identifier: Union[str, Identifier], to_identifier: Union[str, Identifier]) -> Table:
-        identifier_tuple = self._identifier_to_tuple_without_catalog(from_identifier)
-        try:
-            table = self.__tables.pop(identifier_tuple)
-        except KeyError as error:
-            raise NoSuchTableError(f"Table does not exist: {identifier_tuple}") from error
-
-        to_identifier = Catalog.identifier_to_tuple(to_identifier)
-        to_namespace = Catalog.namespace_from(to_identifier)
-        if to_namespace not in self.__namespaces:
-            self.__namespaces[to_namespace] = {}
-
-        self.__tables[to_identifier] = Table(
-            identifier=to_identifier,
-            metadata=table.metadata,
-            metadata_location=table.metadata_location,
-            io=self._load_file_io(properties=table.metadata.properties, location=table.metadata_location),
-            catalog=self,
+        namespace_identifier = Catalog.namespace_from(identifier)
+        if not self._namespace_exists(namespace_identifier):
+            self.create_namespace(namespace_identifier)
+        return super().create_table(
+            identifier=identifier,
+            schema=schema,
+            location=location,
+            partition_spec=partition_spec,
+            sort_order=sort_order,
+            properties=properties,
         )
-        return self.__tables[to_identifier]
-
-    def create_namespace(self, namespace: Union[str, Identifier], properties: Properties = EMPTY_DICT) -> None:
-        namespace = Catalog.identifier_to_tuple(namespace)
-        if namespace in self.__namespaces:
-            raise NamespaceAlreadyExistsError(f"Namespace already exists: {namespace}")
-        else:
-            self.__namespaces[namespace] = properties if properties else {}
-
-    def drop_namespace(self, namespace: Union[str, Identifier]) -> None:
-        namespace = Catalog.identifier_to_tuple(namespace)
-        if [table_identifier for table_identifier in self.__tables.keys() if namespace == table_identifier[:-1]]:
-            raise NamespaceNotEmptyError(f"Namespace is not empty: {namespace}")
-        try:
-            self.__namespaces.pop(namespace)
-        except KeyError as error:
-            raise NoSuchNamespaceError(f"Namespace does not exist: {namespace}") from error
-
-    def list_tables(self, namespace: Optional[Union[str, Identifier]] = None) -> List[Identifier]:
-        if namespace:
-            namespace = Catalog.identifier_to_tuple(namespace)
-            list_tables = [table_identifier for table_identifier in self.__tables.keys() if namespace == table_identifier[:-1]]
-        else:
-            list_tables = list(self.__tables.keys())
-
-        return list_tables
-
-    def list_namespaces(self, namespace: Union[str, Identifier] = ()) -> List[Identifier]:
-        # Hierarchical namespace is not supported. Return an empty list
-        if namespace:
-            return []
-
-        return list(self.__namespaces.keys())
-
-    def load_namespace_properties(self, namespace: Union[str, Identifier]) -> Properties:
-        namespace = Catalog.identifier_to_tuple(namespace)
-        try:
-            return self.__namespaces[namespace]
-        except KeyError as error:
-            raise NoSuchNamespaceError(f"Namespace does not exist: {namespace}") from error
-
-    def update_namespace_properties(
-        self, namespace: Union[str, Identifier], removals: Optional[Set[str]] = None, updates: Properties = EMPTY_DICT
-    ) -> PropertiesUpdateSummary:
-        removed: Set[str] = set()
-        updated: Set[str] = set()
-
-        namespace = Catalog.identifier_to_tuple(namespace)
-        if namespace in self.__namespaces:
-            if removals:
-                for key in removals:
-                    if key in self.__namespaces[namespace]:
-                        del self.__namespaces[namespace][key]
-                        removed.add(key)
-            if updates:
-                for key, value in updates.items():
-                    self.__namespaces[namespace][key] = value
-                    updated.add(key)
-        else:
-            raise NoSuchNamespaceError(f"Namespace does not exist: {namespace}")
-
-        expected_to_change = removed.difference(removals or set())
-
-        return PropertiesUpdateSummary(
-            removed=list(removed or []), updated=list(updates.keys() if updates else []), missing=list(expected_to_change)
-        )
-
-    def list_views(self, namespace: Optional[Union[str, Identifier]] = None) -> List[Identifier]:
-        raise NotImplementedError
 
-    def drop_view(self, identifier: Union[str, Identifier]) -> None:
-        raise NotImplementedError
+    def rename_table(self, from_identifier: Union[str, Identifier], to_identifier: Union[str, Identifier]) -> Table:
+        namespace_identifier = Catalog.namespace_from(to_identifier)
+        if not self._namespace_exists(namespace_identifier):
+            self.create_namespace(namespace_identifier)
+        return super().rename_table(from_identifier, to_identifier)
 
 
 @pytest.fixture
@@ -278,11 +108,13 @@ def catalog(tmp_path: PosixPath) -> InMemoryCatalog:
 )
 TEST_TABLE_PARTITION_SPEC = PartitionSpec(PartitionField(name="x", transform=IdentityTransform(), source_id=1, field_id=1000))
 TEST_TABLE_PROPERTIES = {"key1": "value1", "key2": "value2"}
-NO_SUCH_TABLE_ERROR = "Table does not exist: \\('com', 'organization', 'department', 'my_table'\\)"
-TABLE_ALREADY_EXISTS_ERROR = "Table already exists: \\('com', 'organization', 'department', 'my_table'\\)"
-NAMESPACE_ALREADY_EXISTS_ERROR = "Namespace already exists: \\('com', 'organization', 'department'\\)"
-NO_SUCH_NAMESPACE_ERROR = "Namespace does not exist: \\('com', 'organization', 'department'\\)"
-NAMESPACE_NOT_EMPTY_ERROR = "Namespace is not empty: \\('com', 'organization', 'department'\\)"
+NO_SUCH_TABLE_ERROR = "Table does not exist: com.organization.department.my_table"
+TABLE_ALREADY_EXISTS_ERROR = "Table com.organization.department.my_table already exists"
+NAMESPACE_ALREADY_EXISTS_ERROR = "Namespace \\('com', 'organization', 'department'\\) already exists"
+# TODO: consolidate namespace error messages then remove this
+DROP_NOT_EXISTING_NAMESPACE_ERROR = "Namespace does not exist: \\('com', 'organization', 'department'\\)"
+NO_SUCH_NAMESPACE_ERROR = "Namespace com.organization.department does not exists"
+NAMESPACE_NOT_EMPTY_ERROR = "Namespace com.organization.department is not empty"
 
 
 def given_catalog_has_a_table(
@@ -591,7 +423,7 @@ def test_drop_namespace(catalog: InMemoryCatalog) -> None:
 
 
 def test_drop_namespace_raises_error_when_namespace_does_not_exist(catalog: InMemoryCatalog) -> None:
-    with pytest.raises(NoSuchNamespaceError, match=NO_SUCH_NAMESPACE_ERROR):
+    with pytest.raises(NoSuchNamespaceError, match=DROP_NOT_EXISTING_NAMESPACE_ERROR):
         catalog.drop_namespace(TEST_TABLE_NAMESPACE)
 
 
@@ -607,7 +439,7 @@ def test_list_tables(catalog: InMemoryCatalog) -> None:
     # Given
     given_catalog_has_a_table(catalog)
     # When
-    tables = catalog.list_tables()
+    tables = catalog.list_tables(namespace=TEST_TABLE_NAMESPACE)
     # Then
     assert tables
     assert TEST_TABLE_IDENTIFIER in tables
@@ -619,7 +451,7 @@ def test_list_tables_under_a_namespace(catalog: InMemoryCatalog) -> None:
     new_namespace = ("new", "namespace")
     catalog.create_namespace(new_namespace)
     # When
-    all_tables = catalog.list_tables()
+    all_tables = catalog.list_tables(namespace=TEST_TABLE_NAMESPACE)
     new_namespace_tables = catalog.list_tables(new_namespace)
     # Then
     assert all_tables
@@ -638,7 +470,9 @@ def test_update_namespace_metadata(catalog: InMemoryCatalog) -> None:
     # Then
     assert TEST_TABLE_NAMESPACE in catalog.list_namespaces()
     assert new_metadata.items() <= catalog.load_namespace_properties(TEST_TABLE_NAMESPACE).items()
-    assert summary == PropertiesUpdateSummary(removed=[], updated=["key3", "key4"], missing=[])
+    assert summary.removed == []
+    assert sorted(summary.updated) == ["key3", "key4"]
+    assert summary.missing == []
 
 
 def test_update_namespace_metadata_removals(catalog: InMemoryCatalog) -> None:
@@ -654,7 +488,9 @@ def test_update_namespace_metadata_removals(catalog: InMemoryCatalog) -> None:
     assert TEST_TABLE_NAMESPACE in catalog.list_namespaces()
     assert new_metadata.items() <= catalog.load_namespace_properties(TEST_TABLE_NAMESPACE).items()
     assert remove_metadata.isdisjoint(catalog.load_namespace_properties(TEST_TABLE_NAMESPACE).keys())
-    assert summary == PropertiesUpdateSummary(removed=["key1"], updated=["key3", "key4"], missing=[])
+    assert summary.removed == ["key1"]
+    assert sorted(summary.updated) == ["key3", "key4"]
+    assert summary.missing == []
 
 
 def test_update_namespace_metadata_raises_error_when_namespace_does_not_exist(catalog: InMemoryCatalog) -> None: