Tests: Add integration tests with CrateDB

pyproject.toml

@@ -85,7 +85,7 @@ dependencies = [
   "databricks-sqlalchemy<3",
   "dataclasses-json<0.7",
   "dlt<1.29",
-  "dlt-cratedb<0.2",
+  "dlt-cratedb @ git+https://github.com/zerotired/dlt-cratedb.git@fix-dynamic-install",
   "duckdb<1.6",
   "duckdb-engine<0.18",
   "elasticsearch<10",
@@ -126,6 +126,7 @@ dependencies = [
   "sqlalchemy-bigquery<2",
   "sqlalchemy-cratedb<1",
   "sqlalchemy-hana<5",
+  "sqlalchemy-postgresql-relaxed<0.2",
   "sqlalchemy-redshift<2",
   "sqlalchemy-spanner<2",
   "stripe<16",
@@ -177,7 +178,7 @@ optional-dependencies.test = [
   "pytest-asyncio<2",
   "pytest-cov<8",
   "pytest-xdist[psutil]<4",
-  "testcontainers[mysql,postgres]>=4.9.1,<4.15",
+  "testcontainers[cratedb,mysql,postgres] @ git+https://github.com/testcontainers/testcontainers-python.git@main",
   "types-requests<3",
   "verlib2",
 ]

tests/stream/test_kafka.py

@@ -52,18 +52,30 @@ def get_output_table():
 
     res = get_output_table()
     assert len(res) == 3
-    assert res[0] == ("message1",)
-    assert res[1] == ("message2",)
-    assert res[2] == ("message3",)
+    if dest_uri.startswith("cratedb://"):
+        messages_db = [res[0][0], res[1][0], res[2][0]]
+        assert "message1" in messages_db
+        assert "message2" in messages_db
+        assert "message3" in messages_db
+    else:
+        assert res[0] == ("message1",)
+        assert res[1] == ("message2",)
+        assert res[2] == ("message3",)
 
     # run again, nothing should be inserted into the output table
     run()
 
     res = get_output_table()
     assert len(res) == 3
-    assert res[0] == ("message1",)
-    assert res[1] == ("message2",)
-    assert res[2] == ("message3",)
+    if dest_uri.startswith("cratedb://"):
+        messages_db = [res[0][0], res[1][0], res[2][0]]
+        assert "message1" in messages_db
+        assert "message2" in messages_db
+        assert "message3" in messages_db
+    else:
+        assert res[0] == ("message1",)
+        assert res[1] == ("message2",)
+        assert res[2] == ("message3",)
 
     # add a new message
     producer.produce(topic, "message4".encode("utf-8"))
@@ -73,10 +85,17 @@ def get_output_table():
     run()
     res = get_output_table()
     assert len(res) == 4
-    assert res[0] == ("message1",)
-    assert res[1] == ("message2",)
-    assert res[2] == ("message3",)
-    assert res[3] == ("message4",)
+    if dest_uri.startswith("cratedb://"):
+        messages_db = [res[0][0], res[1][0], res[2][0], res[3][0]]
+        assert "message1" in messages_db
+        assert "message2" in messages_db
+        assert "message3" in messages_db
+        assert "message4" in messages_db
+    else:
+        assert res[0] == ("message1",)
+        assert res[1] == ("message2",)
+        assert res[2] == ("message3",)
+        assert res[3] == ("message4",)
 
 
 @pytest.mark.parametrize(

tests/stream/test_mqbridge_kafka.py

@@ -55,6 +55,12 @@ def test_mqbridge_kafka_to_db(kafka, dest, topic):
     with ThreadPoolExecutor() as executor:
         dest_uri = executor.submit(dest.start).result()
 
+    if dest_uri.startswith("cratedb://"):
+        pytest.skip(
+            "Fails on CrateDB with `DestinationSchemaTampered`, see "
+            "https://github.com/crate/dlt-cratedb/issues/14"
+        )
+
     _produce(kafka, topic, ROWS)
 
     def run():

tests/util/__init__.py

@@ -6,6 +6,7 @@
 from typer.testing import CliRunner
 
 from omniload.main import app
+from tests.util.db import dbquery
 
 # Sentinel for a required-but-positionally-late argument. `dest_uri` must stay
 # effectively required (omitting it is a caller bug), but `source_table` sits
@@ -189,6 +190,15 @@ def invoke_ingest_command(
             app,
             args,
         )
+
+        if result.exit_code == 0:
+            post_ingest(
+                source_uri=source_uri,
+                source_table=source_table,
+                dest_uri=dest_uri,
+                dest_table=dest_table,
+            )
+
         if result.exit_code != 0 and print_output:
             if result.exc_info is not None:
                 traceback.print_exception(*result.exc_info)
@@ -218,9 +228,28 @@ def __init__(self, exit_code, stdout, stderr, exc_info=None):
 
     result = Result(process.returncode, process.stdout, process.stderr)
 
+    if result.exit_code == 0:
+        post_ingest(
+            source_uri=source_uri,
+            source_table=source_table,
+            dest_uri=dest_uri,
+            dest_table=dest_table,
+        )
+
     if result.exit_code != 0 and print_output:
         print(result.stdout)
         print(result.stderr)
         # traceback.print_exception(result.exc_info)
 
     return result
+
+
+def post_ingest(
+    source_uri=None,
+    source_table=None,
+    dest_uri=None,
+    dest_table=None,
+):
+    # CrateDB needs an explicit flush to make data available for reads immediately.
+    if dest_uri and dest_uri.startswith("cratedb://"):
+        dbquery(dest_uri, f"REFRESH TABLE {dest_table}")

tests/util/container/model.py

@@ -7,9 +7,11 @@
 from typing import Callable, Optional
 
 import pytest
+from testcontainers.community.cratedb import CrateDBContainer
+from testcontainers.community.generic import SqlContainer
+from testcontainers.community.kafka import KafkaContainer
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.generic import DbContainer
-from testcontainers.kafka import KafkaContainer
 
 from tests.util.container.impl.couchbase import CouchbaseContainer
 
@@ -212,7 +214,10 @@ def get_connection_url(self):
             return self.container.get_bootstrap_server()
         elif isinstance(self.container, CouchbaseContainer):
             return self.container.get_connection_url() + self.connection_suffix
-        elif isinstance(self.container, DbContainer):
+        elif isinstance(self.container, CrateDBContainer):
+            uri = f"cratedb://{self.container.username}:{self.container.password}@{self.container.get_container_host_ip()}:{self.container.get_exposed_port(5432)}"
+            return uri + self.connection_suffix
+        elif isinstance(self.container, (DbContainer, SqlContainer)):
             return self.container.get_connection_url() + self.connection_suffix
         raise ValueError("Unable to get connection url")
 

tests/util/db.py

@@ -8,6 +8,12 @@ def dbquery(
 ) -> Union[Sequence[Union[sa.Row[Any], sa.RowMapping]], None]:
     """Query database using SQLAlchemy and optionally return results."""
 
+    # CrateDB needs a relaxed SQLAlchemy dialect for querying.
+    # It will not support advanced features of CrateDB,
+    # but that's okay in this case.
+    if uri.startswith("cratedb://"):
+        uri = uri.replace("cratedb://", "postgresql+psycopg_relaxed://")
+
     engine = sa.create_engine(uri, poolclass=sa.NullPool)
     response = None
     with engine.connect() as conn:

tests/warehouse/db/test_arrow.py

@@ -126,6 +126,12 @@ def run_command(df: pd.DataFrame, incremental_key: Optional[str] = None):
 
     dest_uri = dest.start()
 
+    if dest_uri.startswith("cratedb://"):
+        pytest.skip(
+            "Fails on CrateDB with `DestinationSchemaTampered`, see "
+            "https://github.com/crate/dlt-cratedb/issues/14"
+        )
+
     # let's start with a basic dataframe
     row_count = 1000
     df = pd.DataFrame(
@@ -265,6 +271,12 @@ def run_command(df: pd.DataFrame):
 
     dest_uri = dest.start()
 
+    if dest_uri.startswith("cratedb://"):
+        pytest.skip(
+            "Fails on CrateDB with `DestinationSchemaTampered`, see "
+            "https://github.com/crate/dlt-cratedb/issues/14"
+        )
+
     # let's start with a basic dataframe
     row_count = 1000
     df = pd.DataFrame({"id": range(row_count), "value": ["a"] * row_count})

tests/warehouse/db/test_cratedb.py

@@ -0,0 +1,15 @@
+import unittest
+
+import dlt
+from sqlalchemy.util import classproperty
+
+from omniload.target.cratedb import CrateDBDestination
+from tests.main.test_targets import GenericSqlDestinationFixture
+
+
+class CrateDBDestinationTest(unittest.TestCase, GenericSqlDestinationFixture):
+    destination = CrateDBDestination()
+
+    @classproperty
+    def expected_class(cls):
+        return dlt.destinations.cratedb  # ty: ignore[unresolved-attribute]

tests/warehouse/db/test_dynamodb.py

@@ -121,6 +121,12 @@ def smoke_test(dest_uri, dynamodb, schema: str):
             assert result[i][1] == pendulum.parse(dynamodb.data[i]["updated_at"])
 
     def append_test(dest_uri, dynamodb, schema: str):
+
+        if dest_uri.startswith("cratedb://"):
+            pytest.skip(
+                "Fails on CrateDB with twice the amount of expected results: AssertionError: assert 6 == 3"
+            )
+
         dest_table = f"{schema}.dynamodb_{get_random_string(5)}"
 
         # we run it twice to assert that the data in destination doesn't change
@@ -145,6 +151,13 @@ def append_test(dest_uri, dynamodb, schema: str):
 
     def incremental_test_factory(strategy):
         def incremental_test(dest_uri, dynamodb, schema: str):
+
+            if dest_uri.startswith("cratedb://"):
+                pytest.skip(
+                    "Fails on CrateDB with `DestinationSchemaTampered`, see "
+                    "https://github.com/crate/dlt-cratedb/issues/14"
+                )
+
             dest_table = f"{schema}.dynamodb_{get_random_string(5)}"
 
             result = invoke_ingest_command(

tests/warehouse/db/test_mongodb.py

@@ -362,6 +362,12 @@ def test_mongodb_source(mongodb_function, dest):
 
     dest_uri = dest.start()
 
+    if dest_uri.startswith("cratedb://"):
+        pytest.skip(
+            'Fails on CrateDB with `"_id" conflicts with system column`, '
+            "see https://github.com/crate/dlt-cratedb/issues/19"
+        )
+
     try:
         invoke_ingest_command(
             mongo.get_connection_url(),
@@ -775,5 +781,11 @@ def incremental_multiple_days(mongo, dest_uri: str):
 )
 def test_mongodb_custom_query(testcase, mongodb_function, dest):
     """Test MongoDB custom aggregation queries"""
-    testcase(mongodb_function, dest.start())
+    dest_uri = dest.start()
+    if dest_uri.startswith("cratedb://"):
+        pytest.skip(
+            'Fails on CrateDB with `"_id" conflicts with system column`, '
+            "see https://github.com/crate/dlt-cratedb/issues/19"
+        )
+    testcase(mongodb_function, dest_uri)
     dest.stop()

tests/warehouse/manager.py

@@ -1,7 +1,8 @@
-from testcontainers.clickhouse import ClickHouseContainer
-from testcontainers.mongodb import MongoDbContainer
-from testcontainers.mysql import MySqlContainer
-from testcontainers.postgres import PostgresContainer
+from testcontainers.community.clickhouse import ClickHouseContainer
+from testcontainers.community.cratedb import CrateDBContainer
+from testcontainers.community.mongodb import MongoDbContainer
+from testcontainers.community.mysql import MySqlContainer
+from testcontainers.community.postgres import PostgresContainer
 
 from tests.util.container.impl.clickhouse import ClickhouseService
 from tests.util.container.impl.duckdb import EphemeralDuckDb
@@ -11,6 +12,7 @@
 
 CLICKHOUSE_IMAGE = "docker.io/clickhouse/clickhouse-server:26.5"
 COUCHBASE_IMAGE = "docker.io/couchbase:7.6.9"
+CRATEDB_IMAGE = "docker.io/crate/crate:nightly"
 FLOCI_IMAGE = "docker.io/floci/floci:1.5.25"
 KAFKA_IMAGE = "docker.io/confluentinc/cp-kafka:7.6.0"
 MONGODB_IMAGE = "docker.io/mongo:8.3"
@@ -22,6 +24,19 @@
     clickhouse=ClickhouseService(
         "clickhouse", lambda: ClickHouseContainer(CLICKHOUSE_IMAGE)
     ),
+    cratedb=DockerService(
+        "cratedb",
+        lambda: CrateDBContainer(
+            CRATEDB_IMAGE,
+            cmd_opts=[
+                # The test suite creates lots of tables which are currently not purged.
+                # This leads to the following errors on CrateDB when running the whole suite:
+                # > [...] this action would add [4] total shards, but this
+                # > cluster currently has [1000]/[1000] maximum shards open
+                ("cluster.max_shards_per_node", "2000"),
+            ],
+        ),
+    ),
     duckdb_source=EphemeralDuckDb(),
     duckdb_destination=EphemeralDuckDb(),
     mongodb=DockerService("mongodb", lambda: MongoDbContainer(MONGODB_IMAGE)),

tests/warehouse/model.py

@@ -9,6 +9,7 @@ class ServiceRegistry:
     """Registry for all services."""
 
     clickhouse: Optional[AbstractService]
+    cratedb: Optional[AbstractService]
     duckdb_source: Optional[AbstractService]
     duckdb_destination: Optional[AbstractService]
     mongodb: Optional[AbstractService]