I/O: Apache Iceberg (save)

Author: amotl

cratedb_toolkit/cli.py

@@ -10,7 +10,8 @@
 from .dms.cli import cli as dms_cli
 from .docs.cli import cli as docs_cli
 from .info.cli import cli as info_cli
-from .io.cli import cli as io_cli
+from .io.cli import cli_load as io_cli_load
+from .io.cli import cli_save as io_cli_save
 from .query.cli import cli as query_cli
 from .settings.cli import cli as settings_cli
 from .shell.cli import cli as shell_cli
@@ -32,7 +33,8 @@ def cli(ctx: click.Context, verbose: bool, debug: bool):
 cli.add_command(cloud_cli, name="cluster")
 cli.add_command(docs_cli, name="docs")
 cli.add_command(dms_cli, name="dms")
-cli.add_command(io_cli, name="load")
+cli.add_command(io_cli_load, name="load")
+cli.add_command(io_cli_save, name="save")
 cli.add_command(query_cli, name="query")
 cli.add_command(rockset_cli, name="rockset")
 cli.add_command(shell_cli, name="shell")

cratedb_toolkit/cluster/core.py

@@ -20,7 +20,7 @@
     DatabaseAddressMissingError,
     OperationFailed,
 )
-from cratedb_toolkit.io.iceberg import from_iceberg
+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
@@ -396,6 +396,14 @@ def load_table(
 
         return self
 
+    def save_table(
+        self, source: TableAddress, target: InputOutputResource, transformation: t.Union[Path, None] = None
+    ) -> "ManagedCluster":
+        """
+        Export data from a database table on a managed CrateDB Server.
+        """
+        raise NotImplementedError("Not implemented for CrateDB Cloud yet")
+
     @property
     def adapter(self) -> DatabaseAdapter:
         """
@@ -623,6 +631,29 @@ def load_table(
 
         return self
 
+    def save_table(
+        self, source: TableAddress, target: InputOutputResource, transformation: t.Union[Path, None] = None
+    ) -> "StandaloneCluster":
+        """
+        Export data from a database table on a standalone CrateDB Server.
+
+        Synopsis
+        --------
+        export CRATEDB_CLUSTER_URL=crate://crate@localhost:4200/testdrive/demo
+
+        ctk load table influxdb2://example:token@localhost:8086/testdrive/demo
+        ctk load table mongodb://localhost:27017/testdrive/demo
+        """
+        source_url = self.address.dburi
+        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)
+        else:
+            raise NotImplementedError(f"Exporting resource not implemented yet: {target_url_obj}")
+
+        return self
+
 
 class DatabaseCluster:
     """

cratedb_toolkit/io/cli.py

@@ -18,14 +18,14 @@
 @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
 @click.version_option()
 @click.pass_context
-def cli(ctx: click.Context, verbose: bool, debug: bool):
+def cli_load(ctx: click.Context, verbose: bool, debug: bool):
     """
     Load data into CrateDB.
     """
     return boot_click(ctx, verbose, debug)
 
 
-@make_command(cli, name="table")
+@make_command(cli_load, name="table")
 @click.argument("url")
 @option_cluster_id
 @option_cluster_name
@@ -67,3 +67,59 @@ def load_table(
         cluster_url=cluster_url,
     )
     cluster.load_table(source=source, target=target, transformation=transformation)
+
+
+@click.group(cls=ClickAliasedGroup)  # type: ignore[arg-type]
+@click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
+@click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
+@click.version_option()
+@click.pass_context
+def cli_save(ctx: click.Context, verbose: bool, debug: bool):
+    """
+    Export data from CrateDB.
+    """
+    return boot_click(ctx, verbose, debug)
+
+
+@make_command(cli_save, name="table")
+@click.argument("url")
+@option_cluster_id
+@option_cluster_name
+@option_cluster_url
+@click.option("--schema", envvar="CRATEDB_SCHEMA", type=str, required=False, help="Schema where to import the data")
+@click.option("--table", envvar="CRATEDB_TABLE", type=str, required=False, help="Table where to import the data")
+@click.option("--format", "format_", type=str, required=False, help="File format of the import resource")
+@click.option("--compression", type=str, required=False, help="Compression format of the import resource")
+@click.option("--transformation", type=Path, required=False, help="Path to Zyp transformation file")
+@click.pass_context
+def save_table(
+    ctx: click.Context,
+    url: str,
+    cluster_id: str,
+    cluster_name: str,
+    cluster_url: str,
+    schema: str,
+    table: str,
+    format_: str,
+    compression: str,
+    transformation: t.Union[Path, None],
+):
+    """
+    Export data from CrateDB and CrateDB Cloud clusters.
+    """
+
+    # When `--transformation` is given, but empty, fix it.
+    if transformation is not None and transformation.name == "":
+        transformation = None
+
+    # Encapsulate source and target parameters.
+    source = TableAddress(schema=schema, table=table)
+    target = InputOutputResource(url=url, format=format_, compression=compression)
+
+    # Dispatch "load table" operation.
+    cluster = DatabaseCluster.create(
+        cluster_id=cluster_id,
+        cluster_name=cluster_name,
+        cluster_url=cluster_url,
+    )
+    cluster.save_table(source=source, target=target, transformation=transformation)

cratedb_toolkit/io/iceberg.py

@@ -1,6 +1,7 @@
 import dataclasses
 import logging
 
+import pandas as pd
 import polars as pl
 import sqlalchemy as sa
 from boltons.urlutils import URL
@@ -105,3 +106,40 @@ def from_iceberg(source_url, cratedb_url, progress: bool = False):
     # Note: This was much slower.
     # table.to_polars().collect(streaming=True) \  # noqa: ERA001
     #     .write_database(table_name=table_address.fullname, connection=engine, if_table_exists="replace")
+
+
+def to_iceberg(cratedb_url, target_url, progress: bool = False):
+    """
+    Synopsis
+    --------
+    ctk save table \
+        --cluster-url="crate://crate@localhost:4200/demo/taxi_dataset" \
+        "file+iceberg://./var/lib/iceberg/?catalog=default&namespace=demo&table=taxi_dataset"
+    """
+
+    cratedb_address = DatabaseAddress.from_string(cratedb_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"Source address: {cratedb_address}")
+
+    iceberg_address = IcebergAddress.from_url(target_url)
+    logger.info(
+        f"Iceberg address: Path: {iceberg_address.path}, "
+        f"catalog: {iceberg_address.catalog}, namespace: {iceberg_address.namespace}, table: {iceberg_address.table}"
+    )
+
+    # Prepare namespace.
+    catalog = iceberg_address.load_catalog()
+    catalog.create_namespace_if_not_exists(iceberg_address.namespace)
+    catalog.close()
+
+    # Invoke copy operation.
+    logger.info("Running Iceberg copy")
+    engine = sa.create_engine(str(cratedb_url))
+    with engine.connect() as connection:
+        pd.read_sql_table(table_name=cratedb_table.table, schema=cratedb_table.schema, con=connection).to_iceberg(
+            iceberg_address.namespace + "." + iceberg_address.table,
+            catalog_name=iceberg_address.catalog,
+            catalog_properties=iceberg_address.catalog_properties,
+        )

pyproject.toml

@@ -176,6 +176,7 @@ optional-dependencies.io = [
   "sqlalchemy>=2",
 ]
 optional-dependencies.io-base = [
+  "connectorx<0.5",
   "cr8",
   "dask[dataframe]>=2020",
   "fsspec[http,s3]",