I/O: Adapter for PostgreSQL full-load using ingestr

Author: amotl

.dlt/config.toml

@@ -0,0 +1,3 @@
+[runtime]
+
+dlthub_telemetry=false

.github/workflows/postgresql.yml

@@ -0,0 +1,118 @@
+---
+name: "Tests: PostgreSQL"
+
+on:
+  pull_request:
+    paths:
+    - '.github/workflows/postgresql.yml'
+    - 'cratedb_toolkit/io/ingestr/**'
+    - 'tests/io/ingestr/*postgresql*'
+    - 'pyproject.toml'
+  push:
+    branches: [ main ]
+    paths:
+    - '.github/workflows/postgresql.yml'
+    - 'cratedb_toolkit/io/ingestr/**'
+    - 'tests/io/ingestr/*postgresql*'
+    - 'pyproject.toml'
+
+  # Allow job to be triggered manually.
+  workflow_dispatch:
+
+  # Run the job each night after CrateDB nightly has been published.
+  schedule:
+    - cron: '0 3 * * *'
+
+# Cancel in-progress jobs when pushing to the same branch.
+concurrency:
+  cancel-in-progress: true
+  group: ${{ github.workflow }}-${{ github.ref }}
+
+jobs:
+
+  tests:
+
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+        cratedb-version: [
+          "nightly",
+        ]
+        postgresql-version: [
+          "latest",
+        ]
+        python-version: [
+          "3.10",
+          "3.13",
+        ]
+
+    env:
+      OS: ${{ matrix.os }}
+      PYTHON: ${{ matrix.python-version }}
+      CRATEDB_VERSION: ${{ matrix.cratedb-version }}
+      POSTGRESQL_VERSION: ${{ matrix.postgresql-version }}
+      UV_SYSTEM_PYTHON: true
+      TC_KEEPALIVE: true  # Do not tear down Testcontainers
+      INGESTR_DISABLE_TELEMETRY: true
+
+    services:
+      cratedb:
+        image: crate/crate:${{ matrix.cratedb-version }}
+        ports:
+          - 4200:4200
+          - 5432:5432
+        env:
+          CRATE_HEAP_SIZE: 4g
+      postgresql:
+        image: postgres:${{ matrix.postgresql-version }}
+        ports:
+          - 5433:5432
+        env:
+          POSTGRES_HOST_AUTH_METHOD: trust
+
+    name: "
+    Python ${{ matrix.python-version }},
+    PostgreSQL ${{ matrix.postgresql-version }},
+    OS ${{ matrix.os }}
+    "
+    steps:
+
+    - name: Acquire sources
+      uses: actions/checkout@v4
+
+    - name: Install Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        cache-dependency-glob: |
+          pyproject.toml
+        cache-suffix: ${{ matrix.python-version }}
+        enable-cache: true
+        version: "latest"
+
+    - name: Set up project
+      run: |
+
+        # Install package in editable mode.
+        uv pip install --editable='.[io-ingestr,test,develop]'
+
+    - name: Run linter and software tests
+      run: |
+        pytest -m postgresql
+
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v5
+      env:
+        CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+      with:
+        files: ./coverage.xml
+        flags: postgresql
+        env_vars: OS,PYTHON
+        name: codecov-umbrella
+        fail_ci_if_error: true

CHANGES.md

@@ -16,6 +16,7 @@
 - CDC: Enabled loading DMS events from Kinesis streams and stream-dump files
 - CDC: Added subcommand `ctk dms table-mappings`
 - Dependencies: Updated to `commons-codec>=0.0.23`
+- I/O: Adapter for PostgreSQL full-load using ingestr
 
 ## 2025/05/13 v0.0.35
 - Added lost `pytest` dependencies to `cratedb-toolkit[testing]`

cratedb_toolkit/cluster/core.py

@@ -20,6 +20,7 @@
     DatabaseAddressMissingError,
     OperationFailed,
 )
+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
 from cratedb_toolkit.util.data import asbool
@@ -537,10 +538,16 @@ def load_table(
         """
         source_url = source.url
         target_url = self.address.dburi
-
         source_url_obj = URL(source.url)
 
-        if source_url_obj.scheme.startswith("dynamodb"):
+        if ingestr_select(source_url):
+            if ingestr_copy(source_url, self.address, progress=True):
+                self._load_table_result = True
+            else:
+                logger.error("Data loading failed or incomplete")
+                self._load_table_result = False
+
+        elif source_url_obj.scheme.startswith("dynamodb"):
             from cratedb_toolkit.io.dynamodb.api import dynamodb_copy
 
             if dynamodb_copy(str(source_url_obj), target_url, progress=True):

cratedb_toolkit/io/core.py

@@ -103,7 +103,7 @@ class BulkProcessor:
     https://cratedb.com/docs/crate/reference/en/latest/interfaces/http.html#bulk-operations
     """
 
-    connection: sa.Connection
+    connection: sa.engine.Connection
     data: t.Iterable[t.List[t.Dict[str, t.Any]]]
     batch_to_operation: t.Callable[[t.List[t.Dict[str, t.Any]]], SQLOperation]
     progress_bar: t.Union[tqdm, None] = None

cratedb_toolkit/io/ingestr/__init__.py

@@ -0,0 +1,106 @@
+import logging
+
+from yarl import URL
+
+from cratedb_toolkit.io.ingestr.boot import import_ingestr
+from cratedb_toolkit.model import DatabaseAddress
+
+logger = logging.getLogger(__name__)
+
+
+ingestr_available, ingestr, ConfigFieldMissingException = import_ingestr()
+
+
+def ingestr_select(source_url: str) -> bool:
+    """
+    Whether to select `ingestr` for this data source.
+    """
+    if not ingestr_available:
+        return False
+    try:
+        factory = ingestr.src.factory.SourceDestinationFactory(source_url, "csv:////tmp/foobar.csv")
+        factory.get_source()
+        scheme = ingestr.src.factory.parse_scheme_from_uri(source_url)
+        logger.info(f"Selecting ingestr for source scheme: {scheme}")
+        return True
+    except (ImportError, ValueError, AttributeError) as ex:
+        if "Unsupported source scheme" in str(ex):
+            logger.debug(f"Failed to select ingestr for source url '{source_url}': {ex}")
+        else:
+            logger.exception(f"Unexpected error with ingestr for source url: {source_url}")
+        return False
+    except Exception:
+        logger.exception(f"Unexpected error with ingestr for source url: {source_url}")
+        return False
+
+
+def ingestr_copy(source_url: str, target_address: DatabaseAddress, progress: bool = False):
+    """
+    Invoke data transfer to CrateDB from any source provided by `ingestr`.
+
+    https://cratedb-toolkit.readthedocs.io/io/ingestr/
+
+    Synopsis:
+
+        ctk load table \
+            "frankfurter://?base=EUR&table=latest" \
+            --cluster-url="crate://crate:na@localhost:4200/testdrive/exchange_latest"
+
+        ctk load table \
+            "frankfurter://?base=EUR&table=currencies" \
+            --cluster-url="crate://crate:na@localhost:4200/testdrive/exchange_currencies"
+
+        ctk load table \
+            "postgresql://pguser:[email protected]:5432/postgres?table=public.diamonds" \
+            --cluster-url="crate://crate:na@localhost:4200/testdrive/ibis_diamonds"
+    """
+
+    # Sanity checks.
+    if not ingestr_available:
+        raise ModuleNotFoundError("ingestr subsystem not installed")
+
+    # Compute source and target URLs and table names.
+    # Table names use dotted notation `<schema>.<table>`.
+
+    source_url_obj = URL(source_url)
+    source_table = source_url_obj.query.get("table")
+    source_fragment = source_url_obj.fragment
+    source_url_obj = source_url_obj.without_query_params("table").with_fragment("")
+
+    target_uri, target_table_address = target_address.decode()
+    target_table = target_table_address.fullname
+    target_url = target_address.to_ingestr_url()
+
+    if not source_table:
+        raise ValueError("Source table is required")
+    if not target_table:
+        target_table = source_table
+
+    if source_fragment:
+        source_table += f"#{source_fragment}"
+
+    logger.info("Invoking ingestr")
+    logger.info(f"Source URL: {source_url_obj}")
+    logger.info(f"Target URL: {target_url}")
+    logger.info(f"Source Table: {source_table}")
+    logger.info(f"Target Table: {target_table}")
+
+    try:
+        ingestr.main.ingest(
+            source_uri=str(source_url_obj),
+            dest_uri=str(target_url),
+            source_table=source_table,
+            dest_table=target_table,
+            yes=True,
+        )
+        return True
+    except ConfigFieldMissingException:
+        logger.error(
+            "A configuration field is missing. Please ensure all required credentials are provided. "
+            "For example, if your account does not use a password, use a dummy password `na` like "
+            "`export CRATEDB_CLUSTER_URL=crate://crate:na@localhost:4200/testdrive`"
+        )
+        raise
+    except Exception as ex:
+        logger.exception(f"Failed to ingest data: {ex}")
+        return False

cratedb_toolkit/io/ingestr/api.py

@@ -0,0 +1,40 @@
+import logging
+from unittest import mock
+
+from boltons.urlutils import URL
+
+logger = logging.getLogger(__name__)
+
+
+def import_ingestr():
+    """Import ingestr with CrateDB destination adapter."""
+    try:
+        with mock.patch("ingestr.src.telemetry.event.track"), mock.patch(
+            "dlt.common.runtime.telemetry._TELEMETRY_STARTED", True
+        ):
+            import dlt_cratedb  # noqa: F401
+            import ingestr.main
+            import ingestr.src.factory
+            from dlt.common.configuration import ConfigFieldMissingException
+            from ingestr.src.destinations import GenericSqlDestination
+
+            class CrateDBDestination(GenericSqlDestination):
+                def dlt_dest(self, uri: str, **kwargs):
+                    uri = self._replace_url(uri)
+                    import dlt_cratedb.impl.cratedb.factory
+
+                    return dlt_cratedb.impl.cratedb.factory.cratedb(credentials=uri, **kwargs)
+
+                @staticmethod
+                def _replace_url(uri: str) -> str:
+                    url_obj = URL(uri)
+                    if url_obj.scheme == "cratedb":
+                        url_obj.scheme = "postgres"
+                    return str(url_obj)
+
+            ingestr.src.factory.SourceDestinationFactory.destinations["cratedb"] = CrateDBDestination
+
+            return True, ingestr, ConfigFieldMissingException
+    except ImportError:
+        logger.error("Could not import ingestr")
+        return False, None, None

cratedb_toolkit/io/ingestr/boot.py

@@ -130,6 +130,25 @@ def httpuri(self) -> str:
             uri.scheme = "https"
         return str(uri)
 
+    def to_postgresql_url(self, port: int = 5432) -> URL:
+        """
+        Return the `postgresql://` variant of the database URI.
+        """
+        uri = deepcopy(self.uri)
+        uri.scheme = "postgresql"
+        if not uri.host:
+            uri.host = "localhost"
+        uri.port = port
+        return uri
+
+    def to_ingestr_url(self, port: int = 5432) -> URL:
+        """
+        Return the `cratedb://` variant of the database URI, suitable for `ingestr`.
+        """
+        uri = deepcopy(self.to_postgresql_url(port))
+        uri.scheme = "cratedb"
+        return uri
+
     @property
     def verify_ssl(self) -> bool:
         return self.uri.query_params.get("sslmode", "disable") not in ["disable", "require"]

cratedb_toolkit/model.py

@@ -54,7 +54,7 @@ class CrateDBContainer(DockerSkippingContainer, KeepaliveContainer, DbContainer)
     """
 
     CRATEDB_USER = os.environ.get("CRATEDB_USER", "crate")
-    CRATEDB_PASSWORD = os.environ.get("CRATEDB_PASSWORD", "")
+    CRATEDB_PASSWORD = os.environ.get("CRATEDB_PASSWORD", "crate")
     CRATEDB_DB = os.environ.get("CRATEDB_DB", "doc")
     KEEPALIVE = asbool(os.environ.get("CRATEDB_KEEPALIVE", os.environ.get("TC_KEEPALIVE", False)))
     CMD_OPTS = {