adds support from sqllite and mysql in sqlalchemy

rudolfix · rudolfix · commit bd25c3624ab2 · 2025-11-20T23:25:00.000+01:00
diff --git a/.github/workflows/test_destinations_remote.yml b/.github/workflows/test_destinations_remote.yml
@@ -112,7 +112,7 @@ jobs:
             - name: mssql
               destinations: "[\"mssql\"]"
               filesystem_drivers: "[\"memory\"]"
-              extras: "--extra mssql --extra s3 --extra gs --extra az --extra parquet --extra adbc"
+              extras: "--extra mssql --extra s3 --extra gs --extra az --extra parquet --group adbc"
               pre_install_commands: "sudo ACCEPT_EULA=Y apt-get install --yes msodbcsql18"
               post_install_commands: "uv run dbc install mssql"
               always_run_all_tests: true
diff --git a/dlt/destinations/_adbc_jobs.py b/dlt/destinations/_adbc_jobs.py
@@ -8,6 +8,7 @@
 from dlt.common.destination.capabilities import LoaderFileFormatSelector
 from dlt.common.schema.typing import TTableSchema
 from dlt.common.typing import TLoaderFileFormat
+from dlt.common.utils import without_none
 from dlt.destinations.job_client_impl import SqlJobClientBase
 
 if TYPE_CHECKING:
@@ -20,11 +21,33 @@ class AdbcParquetCopyJob(RunnableLoadJob, ABC):
     def __init__(self, file_path: str) -> None:
         super().__init__(file_path)
         self._job_client: SqlJobClientBase = None
+        # override default schema handling
+        self._connect_catalog_name: str = None
+        self._connect_schema_name: str = None
 
     @abstractmethod
     def _connect(self) -> Connection:
         pass
 
+    def _set_catalog_and_schema(self) -> Tuple[str, str]:
+        catalog_name = self._connect_catalog_name
+        if catalog_name is None:
+            catalog_name = self._job_client.sql_client.catalog_name(quote=False)
+        elif catalog_name == "":
+            # empty string disables catalog
+            catalog_name = None
+
+        schema_name = self._connect_schema_name
+        if schema_name is None:
+            schema_name = self._job_client.sql_client.escape_column_name(
+                self._job_client.sql_client.dataset_name, quote=False, casefold=True
+            )
+        elif schema_name == "":
+            # empty string disables schema
+            schema_name = None
+
+        return catalog_name, schema_name
+
     def run(self) -> None:
         from dlt.common.libs.pyarrow import pq_stream_with_new_columns
         from dlt.common.libs.pyarrow import pyarrow
@@ -36,15 +59,15 @@ def _iter_batches(file_path: str) -> Iterator[pyarrow.RecordBatch]:
         with self._connect() as conn, conn.cursor() as cur:
             import time
 
+            catalog_name, schema_name = self._set_catalog_and_schema()
+            kwargs = dict(catalog_name=catalog_name, db_schema_name=schema_name)
+
             t_ = time.time()
             rows = cur.adbc_ingest(
                 self.load_table_name,
                 _iter_batches(self._file_path),
                 mode="append",
-                catalog_name=self._job_client.sql_client.catalog_name(quote=False),
-                db_schema_name=self._job_client.sql_client.fully_qualified_dataset_name(
-                    quote=False
-                ),
+                **without_none(kwargs),  # type: ignore[arg-type]
             )
             conn.commit()
             logger.warning(
diff --git a/dlt/destinations/impl/mssql/mssql.py b/dlt/destinations/impl/mssql/mssql.py
@@ -94,7 +94,7 @@ def _connect(self) -> "Connection":
         self._config = self._job_client.config  # type: ignore[assignment]
         conn_dsn = self.odbc_to_go_mssql_dsn(self._config.credentials.get_odbc_dsn_dict())
         conn_str = ";".join([f"{k}={v}" for k, v in conn_dsn.items()])
-        logger.warning(f"ADBC connecting to {conn_str}")
+        logger.info(f"ADBC connect to {conn_str}")
         return dbapi.connect(driver="mssql", db_kwargs={"uri": conn_str})
 
     @staticmethod
diff --git a/dlt/destinations/impl/sqlalchemy/load_jobs.py b/dlt/destinations/impl/sqlalchemy/load_jobs.py
@@ -1,17 +1,21 @@
+from __future__ import annotations
+
 from typing import IO, Any, Dict, Iterator, List, Sequence, TYPE_CHECKING, Optional
 import math
 
 import sqlalchemy as sa
 
+from dlt.common import logger
 from dlt.common.destination.client import (
     RunnableLoadJob,
     HasFollowupJobs,
     PreparedTableSchema,
 )
 from dlt.common.storages import FileStorage
 from dlt.common.json import json, PY_DATETIME_DECODERS
-from dlt.destinations.sql_jobs import SqlFollowupJob
 
+from dlt.destinations._adbc_jobs import AdbcParquetCopyJob
+from dlt.destinations.sql_jobs import SqlFollowupJob
 from dlt.destinations.impl.sqlalchemy.db_api_client import SqlalchemyClient
 from dlt.destinations.impl.sqlalchemy.merge_job import SqlalchemyMergeFollowupJob
 
@@ -74,6 +78,90 @@ def run(self) -> None:
                 _sql_client.execute_sql(table.insert(), chunk)
 
 
+class SqlalchemyParquetADBCJob(AdbcParquetCopyJob):
+    """ADBC Parquet copy job for SQLAlchemy (sqlite, mysql) with query param handling."""
+
+    def __init__(self, file_path: str, table: sa.Table) -> None:
+        super().__init__(file_path)
+        self._job_client: "SqlalchemyJobClient" = None
+        self.table = table
+
+    if TYPE_CHECKING:
+        from adbc_driver_manager.dbapi import Connection
+
+    def _connect(self) -> Connection:
+        from adbc_driver_manager import dbapi
+
+        engine = self._job_client.config.credentials.engine
+        dialect = engine.dialect.name.lower()
+        url = engine.url
+
+        query = dict(url.query or {})
+
+        if dialect == "sqlite":
+            # disable schema and catalog when ingest
+            self._connect_schema_name = ""
+            self._connect_catalog_name = ""
+
+            # attach directly to dataset sqlite file as "main"
+            if self._job_client.sql_client.dataset_name == "main":
+                db_path = url.database
+            else:
+                db_path = self._job_client.sql_client._sqlite_dataset_filename(
+                    self._job_client.sql_client.dataset_name
+                )
+            conn_str = f"file:{db_path}"
+
+            if query:
+                qs = "&".join(f"{k}={v}" for k, v in query.items())
+                conn_str = f"{conn_str}?{qs}"
+
+            logger.info(f"ADBC connect to {conn_str}")
+            return dbapi.connect(driver="sqlite", db_kwargs={"uri": conn_str})
+
+        elif dialect == "mysql":
+            # disable schema and catalog when ingest
+            self._connect_schema_name = ""
+            self._connect_catalog_name = ""
+
+            # mysql: convert SSL params into go-mysql ADBC parameters
+            mapped = {}
+            for k, v in query.items():
+                lk = k.lower()
+                if lk == "ssl_ca":
+                    mapped["tls-ca"] = v
+                elif lk == "ssl_cert":
+                    mapped["tls-cert"] = v
+                elif lk == "ssl_key":
+                    mapped["tls-key"] = v
+                elif lk == "ssl_mode":
+                    mapped["tls"] = v
+                else:
+                    mapped[k] = v
+
+            username = url.username or ""
+            password = url.password or ""
+            auth = f"{username}:{password}@" if username or password else ""
+
+            host = url.host or "localhost"
+            port = url.port or 3306
+            # dataset name is schema name is database name. each database is a schema in mysql
+            database = self._job_client.sql_client.dataset_name  # url.database or ""
+
+            base = f"{auth}tcp({host}:{port})/{database}"
+            if mapped:
+                qs = "&".join(f"{k}={v}" for k, v in mapped.items())
+                conn_str = f"{base}?{qs}"
+            else:
+                conn_str = base
+
+            logger.info(f"ADBC connect to {conn_str}")
+            return dbapi.connect(driver="mysql", db_kwargs={"uri": conn_str})
+
+        else:
+            raise NotImplementedError(f"ADBC not supported for sqlalchemy dialect {dialect}")
+
+
 class SqlalchemyParquetInsertJob(SqlalchemyJsonLInsertJob):
     def _iter_data_item_chunks(self) -> Iterator[Sequence[Dict[str, Any]]]:
         from dlt.common.libs.pyarrow import ParquetFile
diff --git a/dlt/destinations/impl/sqlalchemy/sqlalchemy_job_client.py b/dlt/destinations/impl/sqlalchemy/sqlalchemy_job_client.py
@@ -14,7 +14,6 @@
     PreparedTableSchema,
     FollowupJobRequest,
 )
-from dlt.destinations.job_client_impl import SqlJobClientWithStagingDataset, SqlLoadJob
 from dlt.common.destination.capabilities import DestinationCapabilitiesContext
 from dlt.common.schema import Schema, TTableSchema, TColumnSchema, TSchemaTables
 from dlt.common.schema.typing import (
@@ -30,11 +29,15 @@
     get_columns_names_with_prop,
 )
 from dlt.common.storages.load_storage import ParsedLoadJobFileName
+
+from dlt.destinations.job_client_impl import SqlJobClientWithStagingDataset
+from dlt.destinations._adbc_jobs import has_driver as adbc_has_driver
 from dlt.destinations.exceptions import DatabaseUndefinedRelation
 from dlt.destinations.impl.sqlalchemy.db_api_client import SqlalchemyClient
 from dlt.destinations.impl.sqlalchemy.configuration import SqlalchemyClientConfiguration
 from dlt.destinations.impl.sqlalchemy.load_jobs import (
     SqlalchemyJsonLInsertJob,
+    SqlalchemyParquetADBCJob,
     SqlalchemyParquetInsertJob,
     SqlalchemyReplaceJob,
     SqlalchemyMergeFollowupJob,
@@ -138,7 +141,11 @@ def create_load_job(
             return SqlalchemyJsonLInsertJob(file_path, table_obj)
         elif parsed_file.file_format == "parquet":
             table_obj = self._to_table_object(table)
-            return SqlalchemyParquetInsertJob(file_path, table_obj)
+            # if driver for a given dialect is installed
+            if adbc_has_driver(self.config.credentials.engine.dialect.name):
+                return SqlalchemyParquetADBCJob(file_path, table_obj)
+            else:
+                return SqlalchemyParquetInsertJob(file_path, table_obj)
         return None
 
     def complete_load(self, load_id: str) -> None:
diff --git a/tests/load/pipeline/test_adbc_loading.py b/tests/load/pipeline/test_adbc_loading.py
@@ -1,6 +1,7 @@
 import pytest
 
 import dlt
+from dlt.common import Decimal
 
 from tests.cases import table_update_and_row
 from tests.load.pipeline.utils import get_load_package_jobs
@@ -10,46 +11,61 @@
 )
 
 
-# def test_adbc_detection() -> None:
-#     from adbc_driver_manager import dbapi, ProgrammingError
-#     import adbc_driver_manager as dm
+@pytest.mark.parametrize(
+    "destination_config",
+    destinations_configs(default_sql_configs=True, subset=["postgres", "mssql", "sqlalchemy"]),
+    ids=lambda x: x.name,
+)
+def test_adbc_detection(destination_config: DestinationTestConfiguration) -> None:
+    from dlt.destinations._adbc_jobs import has_driver
+
+    driver = destination_config.destination_name or destination_config.destination_type
+    if driver == "postgres":
+        driver = "postgresql"
+    elif driver == "sqlalchemy_sqlite":
+        driver = "sqlite"
+    elif driver == "sqlalchemy_mysql":
+        driver = "mysql"
 
-#     try:
-#         db = dm.AdbcDatabase(driver="mssqll")
-#         db.close()
-#     # try:
-#     #     dbapi.connect(driver="postgresql", db_kwargs={"uri": "server"})
-#     except ProgrammingError as pr_ex:
-#         print(str(pr_ex))
-#         print(pr_ex.sqlstate)
+    assert has_driver(driver)[0] is True
 
 
 @pytest.mark.parametrize(
     "destination_config",
-    destinations_configs(default_sql_configs=True, subset=["postgres", "mssql"]),
+    destinations_configs(default_sql_configs=True, subset=["postgres", "mssql", "sqlalchemy"]),
     ids=lambda x: x.name,
 )
 def test_adbc_parquet_loading(destination_config: DestinationTestConfiguration) -> None:
-    column_schemas, data_types = table_update_and_row()
+    # if destination_config.destination_name == "sqlalchemy_sqlite":
+    #     pytest.skip("skip generic ADBC test for sqlite because just a few data types are supported")
+    column_schemas, data_ = table_update_and_row()
 
     pipeline = destination_config.setup_pipeline("pipeline_adbc", dev_mode=True)
 
-    # postgres
-    del column_schemas["col6_precision"]  # adbc cannot process decimal(6,2)
-    # mssql
-    del column_schemas["col7_precision"]  # adbc cannot process fixed binary
+    if destination_config.destination_type in ("postgres", "mssql"):
+        del column_schemas["col11_precision"]  # TIME(3) not supported
+        if destination_config.destination_type == "postgres":
+            del column_schemas["col6_precision"]  # adbc cannot process decimal(6,2)
+        else:
+            del column_schemas["col7_precision"]  # adbc cannot process fixed binary
 
-    # both
-    del column_schemas["col11_precision"]  # TIME(3) not supported
+    if destination_config.destination_name == "sqlalchemy_sqlite":
+        for k, v in column_schemas.items():
+            # decimals not supported
+            if v["data_type"] in ("decimal", "wei", "time"):
+                data_[k] = str(data_[k])
+                column_schemas[k]["data_type"] = "text"
 
     @dlt.resource(file_format="parquet", columns=column_schemas, max_table_nesting=0)
     def complex_resource():
-        yield data_types
+        yield data_
 
     info = pipeline.run(complex_resource())
     jobs = get_load_package_jobs(
         info.load_packages[0], "completed_jobs", "complex_resource", ".parquet"
     )
     # there must be a parquet job or adbc is not installed so we fall back to other job type
     assert len(jobs) == 1
-    print(pipeline.dataset().table("complex_resource").fetchall())
+    # make sure we can read data back. TODO: verify data types
+    rows = pipeline.dataset().table("complex_resource").fetchall()
+    assert len(rows) == 1
