I/O: Apache Iceberg (suggestions by CodeRabbit)

Author: amotl

cratedb_toolkit/cluster/core.py

@@ -20,7 +20,6 @@
     DatabaseAddressMissingError,
     OperationFailed,
 )
-from cratedb_toolkit.io.iceberg import from_iceberg, to_iceberg
 from cratedb_toolkit.io.ingestr.api import ingestr_copy, ingestr_select
 from cratedb_toolkit.model import ClusterAddressOptions, DatabaseAddress, InputOutputResource, TableAddress
 from cratedb_toolkit.util.client import jwt_token_patch
@@ -617,7 +616,10 @@ def load_table(
                     self._load_table_result = False
 
         elif source_url_obj.scheme.startswith("iceberg") or source_url_obj.scheme.endswith("iceberg"):
-            return from_iceberg(str(source_url_obj), target_url)
+            from cratedb_toolkit.io.iceberg import from_iceberg
+
+            if from_iceberg(str(source_url_obj), target_url):
+                self._load_table_result = True
 
         elif ingestr_select(source_url):
             if ingestr_copy(source_url, self.address, progress=True):
@@ -648,7 +650,9 @@ def save_table(
         target_url_obj = URL(target.url)
 
         if target_url_obj.scheme.startswith("iceberg") or target_url_obj.scheme.endswith("iceberg"):
-            return to_iceberg(source_url, target.url)
+            from cratedb_toolkit.io.iceberg import to_iceberg
+
+            to_iceberg(source_url, target.url)
         else:
             raise NotImplementedError(f"Exporting resource not implemented yet: {target_url_obj}")
 

cratedb_toolkit/io/cli.py

@@ -116,7 +116,7 @@ def save_table(
     source = TableAddress(schema=schema, table=table)
     target = InputOutputResource(url=url, format=format_, compression=compression)
 
-    # Dispatch "load table" operation.
+    # Dispatch "save table" operation.
     cluster = DatabaseCluster.create(
         cluster_id=cluster_id,
         cluster_name=cluster_name,

cratedb_toolkit/io/iceberg.py

@@ -1,3 +1,10 @@
+"""
+Apache Iceberg integration for CrateDB Toolkit.
+
+This module provides functionality to transfer data between Iceberg tables
+and CrateDB databases, supporting both import and export operations.
+"""
+
 import dataclasses
 import logging
 import tempfile
@@ -44,7 +51,7 @@ def from_url(cls, url: str):
             if iceberg_url.scheme.endswith("+iceberg"):
                 iceberg_url.scheme = iceberg_url.scheme.replace("+iceberg", "")
             u2 = copy(iceberg_url)
-            u2._query = ""
+            u2.query_params.clear()
             location = str(u2)
         return cls(
             url=iceberg_url,
@@ -55,6 +62,11 @@ def from_url(cls, url: str):
         )
 
     def load_catalog(self) -> Catalog:
+        """
+        Load the Iceberg catalog with appropriate configuration.
+        """
+        # TODO: Consider accepting catalog configuration as parameters
+        #       to support different catalog types (Hive, REST, etc.).
         return load_catalog(self.catalog, **self.catalog_properties)
 
     @property
@@ -69,29 +81,34 @@ def catalog_properties(self):
 
     @property
     def storage_options(self):
-        return {
+        opts = {
             "s3.endpoint": self.url.query_params.get("s3.endpoint"),
             "s3.region": self.url.query_params.get("s3.region"),
             "s3.access-key-id": self.url.query_params.get("s3.access-key-id"),
             "s3.secret-access-key": self.url.query_params.get("s3.secret-access-key"),
         }
+        return {k: v for k, v in opts.items() if v is not None}
 
     @property
     def identifier(self):
+        """
+        Return the catalog-table identifier tuple.
+        """
         return (self.namespace, self.table)
 
     def load_table(self) -> pl.LazyFrame:
         """
-        Load a table from a catalog, or by scanning the filesystem.
+        Load the Iceberg table as a Polars LazyFrame.
+
+        Either load a table from a catalog, or by scanning the filesystem.
         """
         if self.catalog is not None:
             catalog = self.load_catalog()
             return catalog.load_table(self.identifier).to_polars()
-        else:
-            return pl.scan_iceberg(self.location, storage_options=self.storage_options)
+        return pl.scan_iceberg(self.location, storage_options=self.storage_options)
 
 
-def from_iceberg(source_url, cratedb_url, progress: bool = False):
+def from_iceberg(source_url, target_url, progress: bool = False):
     """
     Scan an Iceberg table from local filesystem or object store, and load into CrateDB.
     https://docs.pola.rs/api/python/stable/reference/api/polars.scan_iceberg.html
@@ -121,11 +138,11 @@ def from_iceberg(source_url, cratedb_url, progress: bool = False):
     # Display parameters.
     logger.info(f"Iceberg address: Path: {iceberg_address.location}")
 
-    cratedb_address = DatabaseAddress.from_string(cratedb_url)
+    cratedb_address = DatabaseAddress.from_string(target_url)
     cratedb_url, cratedb_table = cratedb_address.decode()
     if cratedb_table.table is None:
         raise ValueError("Table name is missing. Please adjust CrateDB database URL.")
-    logger.info(f"Target address: {cratedb_address}")
+    logger.info("Target address: %s", cratedb_address)
 
     # Invoke copy operation.
     logger.info("Running Iceberg copy")
@@ -139,7 +156,7 @@ def from_iceberg(source_url, cratedb_url, progress: bool = False):
     # https://github.com/pola-rs/polars/issues/7852
     # Note: This code also uses the most efficient `insert_bulk` method with CrateDB.
     # https://cratedb.com/docs/sqlalchemy-cratedb/dataframe.html#efficient-insert-operations-with-pandas
-    table.collect(streaming=True).to_pandas().to_sql(
+    table.collect(engine="streaming").to_pandas().to_sql(
         name=cratedb_table.table,
         schema=cratedb_table.schema,
         con=engine,
@@ -153,8 +170,10 @@ def from_iceberg(source_url, cratedb_url, progress: bool = False):
     # table.to_polars().collect(streaming=True) \  # noqa: ERA001
     #     .write_database(table_name=table_address.fullname, connection=engine, if_table_exists="replace")
 
+    return True
 
-def to_iceberg(cratedb_url, target_url, progress: bool = False):
+
+def to_iceberg(source_url, target_url, progress: bool = False):
     """
     Synopsis
     --------
@@ -167,7 +186,7 @@ def to_iceberg(cratedb_url, target_url, progress: bool = False):
         "s3+iceberg://bucket1/?catalog=default&namespace=demo&table=taxi-tiny&s3.access-key-id=<your_access_key_id>&s3.secret-access-key=<your_secret_access_key>&s3.endpoint=<endpoint_url>&s3.region=<s3-region>"
     """  # noqa:E501
 
-    cratedb_address = DatabaseAddress.from_string(cratedb_url)
+    cratedb_address = DatabaseAddress.from_string(source_url)
     cratedb_url, cratedb_table = cratedb_address.decode()
     if cratedb_table.table is None:
         raise ValueError("Table name is missing. Please adjust CrateDB database URL.")
@@ -197,3 +216,5 @@ def to_iceberg(cratedb_url, target_url, progress: bool = False):
             catalog_properties=catalog_properties,
             append=False,  # TODO: Make available via parameter.
         )
+
+    return True

tests/io/test_iceberg.py

@@ -1,28 +1,32 @@
 from pathlib import Path
 
 import pandas as pd
-import polars as pl
 import pytest
 from click.testing import CliRunner
 from pueblo.testing.dataframe import DataFrameFactory
-from pyiceberg.catalog import load_catalog
 
 from cratedb_toolkit.cli import cli
 
-if not hasattr(pd.DataFrame, "to_iceberg2"):
+polars = pytest.importorskip("polars", reason="Skipping Iceberg tests because 'polars' package is not installed")
+
+import polars as pl  # noqa: E402
+
+if not hasattr(pd.DataFrame, "to_iceberg"):
     raise pytest.skip("Older pandas releases do not support Apache Iceberg", allow_module_level=True)
 
 
 @pytest.fixture
 def example_iceberg(tmp_path) -> Path:
+    from pyiceberg.catalog import load_catalog  # noqa: E402
+
     catalog_properties = {
         "uri": f"sqlite:///{tmp_path}/pyiceberg_catalog.db",
         "warehouse": str(tmp_path),
     }
     catalog = load_catalog("default", **catalog_properties)
     catalog.create_namespace_if_not_exists("demo")
 
-    dff = DataFrameFactory(rows=42)
+    dff = DataFrameFactory()
     df = dff.make_mixed()
     df.to_iceberg(
         "demo.mixed",
@@ -41,10 +45,13 @@ def find_iceberg_data_metadata_location(table_path: Path) -> Path:
     Resolve path to metadata.json file in Iceberg table.
     This path is needed for `polars.scan_iceberg()`.
     """
-    return sorted((table_path / "metadata").glob("*.json"))[-1]
+    files = sorted((table_path / "metadata").glob("*.json"))
+    if not files:
+        raise FileNotFoundError(f"No Iceberg metadata JSON found under {table_path / 'metadata'}")
+    return files[-1]
 
 
-def test_load_iceberg_table(caplog, cratedb, example_iceberg):
+def test_load_iceberg_table(cratedb, example_iceberg):
     """
     Verify loading data from an Iceberg table into CrateDB.
     """
@@ -69,7 +76,7 @@ def test_load_iceberg_table(caplog, cratedb, example_iceberg):
     assert db.count_records("testdrive.demo") == 5, "Table `testdrive.demo` does not include expected amount of records"
 
 
-def test_save_iceberg_table(caplog, cratedb, tmp_path):
+def test_save_iceberg_table(cratedb, tmp_path):
     """
     Verify saving data from CrateDB into an Iceberg table.
     """
@@ -90,7 +97,7 @@ def test_save_iceberg_table(caplog, cratedb, tmp_path):
     # Verify data in Iceberg table.
     metadata_location = find_iceberg_data_metadata_location(tmp_path / "sys" / "summits")
     table = pl.scan_iceberg(str(metadata_location))
-    assert table.collect_schema().names() == [
+    assert sorted(table.collect_schema().names()) == [
         "classification",
         "coordinates",
         "country",
@@ -101,4 +108,4 @@ def test_save_iceberg_table(caplog, cratedb, tmp_path):
         "range",
         "region",
     ]
-    assert table.collect().height == 1605, "Iceberg table does not include expected amount of records"
+    assert table.collect().height >= 1600, "Iceberg table does not include expected amount of records"