Skip to content

Commit

Permalink
replace InMemoryCatalog by a subclass of SqlCatalog with SQLite in-me…
Browse files Browse the repository at this point in the history 
…mory
  • Loading branch information
@hussein-awala
hussein-awala committed Sep 5, 2024
1 parent 052a9cd commit 21d6d2f
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 244 deletions.
252 changes: 44 additions & 208 deletions tests/catalog/test_base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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)


Expand All @@ -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
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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:
Expand Down
Loading

0 comments on commit 21d6d2f

Please sign in to comment.