can override table location

Author: kevinjqliu

pyiceberg/catalog/memory.py

@@ -20,6 +20,7 @@
     NoSuchTableError,
     TableAlreadyExistsError,
 )
+from pyiceberg.io import WAREHOUSE
 from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table import (
@@ -41,12 +42,11 @@ class InMemoryCatalog(Catalog):
     __tables: Dict[Identifier, Table]
     __namespaces: Dict[Identifier, Properties]
 
-    def __init__(self, name: str, warehouse_location: Optional[str] = None, **properties: str) -> None:
+    def __init__(self, name: str, **properties: str) -> None:
         super().__init__(name, **properties)
-
-        self._warehouse_location = warehouse_location or DEFAULT_WAREHOUSE_LOCATION
         self.__tables = {}
         self.__namespaces = {}
+        self._warehouse_location = properties.get(WAREHOUSE, None) or DEFAULT_WAREHOUSE_LOCATION
 
     def create_table(
         self,
@@ -67,10 +67,9 @@ def create_table(
             if namespace not in self.__namespaces:
                 self.__namespaces[namespace] = {}
 
-            # if not location:
-            location = f'{self._warehouse_location}/{"/".join(identifier)}'
+            if not location:
+                location = f'{self._warehouse_location}/{"/".join(identifier)}'
 
-            # _get_default_warehouse_location
             metadata_location = f'{self._warehouse_location}/{"/".join(identifier)}/metadata/metadata.json'
 
             metadata = new_table_metadata(

tests/catalog/test_base.py

@@ -33,6 +33,7 @@
     NoSuchTableError,
     TableAlreadyExistsError,
 )
+from pyiceberg.io import WAREHOUSE
 from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table import (
@@ -49,9 +50,7 @@
 
 @pytest.fixture
 def catalog(tmp_path: PosixPath) -> InMemoryCatalog:
-    return InMemoryCatalog(
-        "test.in.memory.catalog", warehouse_location=tmp_path.absolute().as_posix(), **{"test.key": "test.value"}
-    )
+    return InMemoryCatalog("test.in.memory.catalog", **{WAREHOUSE: tmp_path.absolute().as_posix(), "test.key": "test.value"})
 
 
 TEST_TABLE_IDENTIFIER = ("com", "organization", "department", "my_table")
@@ -76,7 +75,6 @@ def given_catalog_has_a_table(catalog: InMemoryCatalog) -> Table:
     return catalog.create_table(
         identifier=TEST_TABLE_IDENTIFIER,
         schema=TEST_TABLE_SCHEMA,
-        location=TEST_TABLE_LOCATION,
         partition_spec=TEST_TABLE_PARTITION_SPEC,
         properties=TEST_TABLE_PROPERTIES,
     )
@@ -119,6 +117,16 @@ def test_name_from_str() -> None:
 
 
 def test_create_table(catalog: InMemoryCatalog) -> None:
+    table = catalog.create_table(
+        identifier=TEST_TABLE_IDENTIFIER,
+        schema=TEST_TABLE_SCHEMA,
+        partition_spec=TEST_TABLE_PARTITION_SPEC,
+        properties=TEST_TABLE_PROPERTIES,
+    )
+    assert catalog.load_table(TEST_TABLE_IDENTIFIER) == table
+
+
+def test_create_table_override(catalog: InMemoryCatalog) -> None:
     table = catalog.create_table(
         identifier=TEST_TABLE_IDENTIFIER,
         schema=TEST_TABLE_SCHEMA,

tests/cli/test_console.py

@@ -264,7 +264,6 @@ def test_location(catalog: InMemoryCatalog) -> None:
     catalog.create_table(
         identifier=TEST_TABLE_IDENTIFIER,
         schema=TEST_TABLE_SCHEMA,
-        location=TEST_TABLE_LOCATION,
         partition_spec=TEST_TABLE_PARTITION_SPEC,
     )
 
@@ -274,6 +273,20 @@ def test_location(catalog: InMemoryCatalog) -> None:
     assert result.output == f"""{DEFAULT_WAREHOUSE_LOCATION}/default/my_table\n"""
 
 
+def test_location_override(catalog: InMemoryCatalog) -> None:
+    catalog.create_table(
+        identifier=TEST_TABLE_IDENTIFIER,
+        schema=TEST_TABLE_SCHEMA,
+        location=TEST_TABLE_LOCATION,
+        partition_spec=TEST_TABLE_PARTITION_SPEC,
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(run, ["location", "default.my_table"])
+    assert result.exit_code == 0
+    assert result.output == f"""{TEST_TABLE_LOCATION}\n"""
+
+
 def test_location_does_not_exists(catalog: InMemoryCatalog) -> None:
     # pylint: disable=unused-argument
 
@@ -674,7 +687,6 @@ def test_json_location(catalog: InMemoryCatalog) -> None:
     catalog.create_table(
         identifier=TEST_TABLE_IDENTIFIER,
         schema=TEST_TABLE_SCHEMA,
-        location=TEST_TABLE_LOCATION,
         partition_spec=TEST_TABLE_PARTITION_SPEC,
     )
 
@@ -684,6 +696,20 @@ def test_json_location(catalog: InMemoryCatalog) -> None:
     assert result.output == f'"{DEFAULT_WAREHOUSE_LOCATION}/default/my_table"\n'
 
 
+def test_json_location_override(catalog: InMemoryCatalog) -> None:
+    catalog.create_table(
+        identifier=TEST_TABLE_IDENTIFIER,
+        schema=TEST_TABLE_SCHEMA,
+        location=TEST_TABLE_LOCATION,
+        partition_spec=TEST_TABLE_PARTITION_SPEC,
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(run, ["--output=json", "location", "default.my_table"])
+    assert result.exit_code == 0
+    assert result.output == f'"{TEST_TABLE_LOCATION}"\n'
+
+
 def test_json_location_does_not_exists(catalog: InMemoryCatalog) -> None:
     # pylint: disable=unused-argument