tabsdata-conn-common-sql 2.0.0__cp312-abi3-win_amd64.whl

tabsdata_conn_common_sql-2.0.0.data/purelib/tabsdatak/conn/common/sql/_engine/constants.py

@@ -0,0 +1,34 @@
+#
+# Copyright 2026 Tabsdata Inc.
+#
+
+from tabsdatak.conn.common import _config
+
+SQL_KEY_PREFIX = _config.TABSDATA_CFG_PREFIX + "sql."
+
+# --- Read --------------------------------------------------------
+
+READ_ENGINE_KEY = SQL_KEY_PREFIX + "read_engine"
+POLARS_SQL_ALCHEMY_READ_ENGINE = "polarsSqlAlchemyRead"
+PANDAS_SQL_ALCHEMY_READ_ENGINE = "pandasSqlAlchemyRead"
+
+AVAIL_READ_ENGINES_KEY = SQL_KEY_PREFIX + "engine.avail_read_engines"
+
+READ_CHUNK_SIZE_KEY = SQL_KEY_PREFIX + "chunk_size"
+READ_CHUNK_SIZE_DEFAULT = 100_000
+
+POLARS_SQL_ALCHEMY_READ_CHUNK_SIZE_DEFAULT = READ_CHUNK_SIZE_DEFAULT
+PANDAS_SQL_ALCHEMY_READ_CHUNK_SIZE_DEFAULT = READ_CHUNK_SIZE_DEFAULT
+
+
+# --- Write -------------------------------------------------------
+
+WRITE_ENGINE_KEY = SQL_KEY_PREFIX + "write_engine"
+POLARS_SQL_ALCHEMY_WRITE_ENGINE = "polarsSqlAlchemyWrite"
+
+AVAIL_WRITE_ENGINES_KEY = SQL_KEY_PREFIX + "engine.avail_write_engines"
+
+WRITE_CHUNK_SIZE_KEY = SQL_KEY_PREFIX + "chunk_size"
+WRITE_CHUNK_SIZE_DEFAULT = 100_000
+
+POLARS_SQL_ALCHEMY_WRITE_CHUNK_SIZE_DEFAULT = WRITE_CHUNK_SIZE_DEFAULT

tabsdata_conn_common_sql-2.0.0.data/purelib/tabsdatak/conn/common/sql/_engine/factory.py

@@ -0,0 +1,60 @@
+#
+# Copyright 2026 Tabsdata Inc.
+#
+
+from typing import Any, Mapping
+
+# noinspection PyProtectedMember
+import tabsdatak.conn.common.sql._engine.constants as const
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common import _config
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common.sql._engine.reader import SqlReadEngine
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common.sql._engine.sqlalchemy_reader import (
+    PandasReadEngine,
+    PolarsReadEngine,
+)
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common.sql._engine.sqlalchemy_writer import PolarsWriteEngine
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common.sql._engine.writer import SqlWriteEngine
+from tabsdatak.conn.common.sql.error import SqlCommonErrorCode
+from tabsdatak.conn.common.types import CfgSpec
+
+
+def get_read_engine(
+    src_cfg: CfgSpec,
+    defaults: Mapping[str, Any],
+) -> SqlReadEngine:
+    engine = _config.get_cfg(
+        src_cfg,
+        const.READ_ENGINE_KEY,
+        defaults[const.READ_ENGINE_KEY],
+        defaults[const.AVAIL_READ_ENGINES_KEY],
+    )
+    if engine == const.PANDAS_SQL_ALCHEMY_READ_ENGINE:
+        return PandasReadEngine()
+    if engine == const.POLARS_SQL_ALCHEMY_READ_ENGINE:
+        return PolarsReadEngine()
+    raise SqlCommonErrorCode.SQL_COMMON_1.exception(engine=engine)
+
+
+def get_write_engine(
+    dest_cfg: CfgSpec,
+    defaults: Mapping[str, Any],
+) -> SqlWriteEngine:
+    engine = _config.get_cfg(
+        dest_cfg,
+        const.WRITE_ENGINE_KEY,
+        defaults[const.WRITE_ENGINE_KEY],
+        defaults[const.AVAIL_WRITE_ENGINES_KEY],
+    )
+    if engine == const.POLARS_SQL_ALCHEMY_WRITE_ENGINE:
+        return PolarsWriteEngine()
+    raise SqlCommonErrorCode.SQL_COMMON_1.exception(engine=engine)

tabsdata_conn_common_sql-2.0.0.data/purelib/tabsdatak/conn/common/sql/_engine/initial_values.py

@@ -0,0 +1,37 @@
+#
+# Copyright 2026 Tabsdata Inc.
+#
+
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+from tabsdatak.spi import SrcPluginCtx
+
+
+def resolve_initial_values(
+    ctx: SrcPluginCtx,
+    seed: Mapping[str, Any] | None,
+) -> dict[str, Any] | None:
+    """Look up the active bind dict for this run.
+
+    Walks ``seed`` (the user's first-run ``Src.initial_values``) bind
+    name by bind name. For each, prefers ``ctx.get_attr(name)`` --
+    that's where the previous run's persisted value lands -- and
+    falls back to the seed when the ctx scratch has no entry for
+    that name.
+
+    Returns ``None`` when ``seed`` is None (no placeholders in any
+    query), signalling no substitution is needed.
+
+    Uses ``is not None`` rather than truthiness so a legitimately
+    falsy value stored last run (``0``, ``False``, ``""``) isn't
+    silently shadowed by the seed.
+    """
+    if seed is None:
+        return None
+    out: dict[str, Any] = {}
+    for name, value in seed.items():
+        from_ctx = ctx.get_attr(name)
+        out[name] = from_ctx if from_ctx is not None else value
+    return out

tabsdata_conn_common_sql-2.0.0.data/purelib/tabsdatak/conn/common/sql/_engine/reader.py

@@ -0,0 +1,35 @@
+#
+# Copyright 2026 Tabsdata Inc.
+#
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Mapping
+
+from sqlalchemy import Engine
+
+from tabsdatak.conn.common.sql.types import SchemaOverrideSpec
+from tabsdatak.conn.common.types import CfgSpec
+from tabsdatak.spi import Conn, TableFileSpec
+
+
+@dataclass(kw_only=True)
+class SqlReadCfg[C: Conn](ABC):
+    conn: C
+    queries: list[str]
+    transactional: bool
+    schema_overrides: SchemaOverrideSpec | None
+    initial_values: Mapping[str, Any] | None
+    cfg: CfgSpec | None
+
+    def create_sql_alchemy_engine(self) -> Engine:
+        """
+        Implemented by subclasses that use SqlAlchemy.
+        """
+        raise NotImplementedError()
+
+
+class SqlReadEngine(ABC):
+    @abstractmethod
+    def read(self, cfg: SqlReadCfg, work_dir: Path) -> list[list[TableFileSpec]]: ...

tabsdata_conn_common_sql-2.0.0.data/purelib/tabsdatak/conn/common/sql/_engine/sqlalchemy_reader.py

@@ -0,0 +1,199 @@
+#
+# Copyright 2026 Tabsdata Inc.
+#
+
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Callable, Mapping
+
+from sqlalchemy import Connection
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.orm import Session
+from tenacity import (
+    Retrying,
+    before_sleep_log,
+    retry_if_exception_type,
+    stop_after_delay,
+    wait_fixed,
+)
+
+# noinspection PyProtectedMember
+import tabsdatak.conn.common.sql._engine.constants as const
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common import _config
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common.sql._engine.reader import SqlReadCfg, SqlReadEngine
+
+# noinspection PyProtectedMember
+from tabsdatak.conn.common.sql._engine.type_convertion import resolve_schema_overrides
+from tabsdatak.spi import TableFileSpec
+
+# --- Retry ----------------------------------------------------------
+
+logger = logging.getLogger(__name__)
+
+RETRY_TIMEOUT_S = 5 * 60
+RETRY_WAIT_S = 5
+
+
+def _retry() -> Retrying:
+    return Retrying(
+        stop=stop_after_delay(RETRY_TIMEOUT_S),
+        wait=wait_fixed(RETRY_WAIT_S),
+        retry=retry_if_exception_type(OperationalError),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        reraise=True,
+    )
+
+
+ReadLoader = Callable[..., list[Path]]
+
+
+class SqlAlchemyReadEngine(SqlReadEngine, ABC):
+    @abstractmethod
+    def loader(self) -> ReadLoader: ...
+
+    def read(self, src: SqlReadCfg, work_dir: Path) -> list[list[TableFileSpec]]:
+        read_loader = self.loader()
+
+        # Empty dict -> None (per the public contract; saves the
+        # loaders from having to special-case the empty case).
+        parameters: Mapping[str, Any] | None = src.initial_values or None
+
+        def _run() -> list[list[TableFileSpec]]:
+            logger.debug("sql-reader: establishing engine")
+            engine = src.create_sql_alchemy_engine()
+            try:
+                results: list[list[TableFileSpec]] = []
+                if src.transactional:
+                    with Session(engine) as session, session.begin():
+                        conn = session.connection()
+                        self._run_queries(
+                            read_loader, conn, src, work_dir, parameters, results
+                        )
+                else:
+                    with engine.connect() as conn:
+                        self._run_queries(
+                            read_loader, conn, src, work_dir, parameters, results
+                        )
+                return results
+            finally:
+                logger.debug("sql-reader: disposing engine")
+                engine.dispose()
+
+        return _retry()(_run)
+
+    @staticmethod
+    def _run_queries(
+        read_loader: Any,
+        conn: Any,
+        src: SqlReadCfg,
+        work_dir: Path,
+        parameters: Mapping[str, Any] | None,
+        results: list[list[TableFileSpec]],
+    ):
+        chunk_size = _config.get_cfg(
+            src.cfg,
+            const.READ_CHUNK_SIZE_KEY,
+            const.READ_CHUNK_SIZE_DEFAULT,
+        )
+        total = len(src.queries)
+        logger.info("sql-reader: reading %d quer(ies)", total)
+        for i, query in enumerate(src.queries):
+            logger.info("sql-reader: reading query %d/%d", i + 1, total)
+            overrides = (
+                resolve_schema_overrides(src.schema_overrides[i])
+                if src.schema_overrides is not None
+                else None
+            )
+            files = read_loader(
+                connection=conn,
+                query=query,
+                work_dir=work_dir,
+                slot=i,
+                chunk_size=chunk_size,
+                schema_overrides=overrides,
+                parameters=parameters,
+            )
+            logger.debug(
+                "sql-reader: query %d/%d produced %d chunk(s)", i + 1, total, len(files)
+            )
+            results.append([Path(p) for p in files])
+        logger.info("sql-reader: finished reading %d quer(ies)", total)
+
+
+def _pandas_read(
+    *,
+    connection: Connection,
+    query: str,
+    work_dir: Path,
+    slot: int,
+    chunk_size: int,
+    schema_overrides: Mapping[str, Any] | None,
+    parameters: Mapping[str, Any] | None = None,
+) -> list[Path]:
+    import pandas as pd
+    from sqlalchemy import text
+
+    files: list[Path] = []
+    iterator = pd.read_sql_query(
+        text(query),
+        con=connection,
+        chunksize=chunk_size,
+        dtype=dict(schema_overrides) if schema_overrides else None,
+        params=dict(parameters) if parameters else None,
+    )
+    for chunk_idx, frame in enumerate(iterator):
+        out = work_dir / f"slot{slot}_chunk{chunk_idx}.parquet"
+        frame.to_parquet(out, engine="pyarrow")
+        files.append(out)
+
+    return files
+
+
+class PandasReadEngine(SqlAlchemyReadEngine):
+    def loader(self) -> ReadLoader:
+        return _pandas_read
+
+
+def _polars_read(
+    *,
+    connection: Any,
+    query: str,
+    work_dir: Path,
+    slot: int,
+    chunk_size: int,
+    schema_overrides: Mapping[str, Any] | None,
+    parameters: Mapping[str, Any] | None = None,
+) -> list[Path]:
+    import polars as pl
+    from sqlalchemy import text
+
+    execute_options: dict[str, Any] = {}
+    if parameters:
+        execute_options["parameters"] = dict(parameters)
+
+    files: list[Path] = []
+    streamed_conn = connection.execution_options(stream_results=True)
+    batches = pl.read_database(
+        text(query),
+        connection=streamed_conn,
+        iter_batches=True,
+        batch_size=chunk_size,
+        schema_overrides=schema_overrides,
+        execute_options=execute_options or None,
+    )
+    for chunk_idx, batch in enumerate(batches):
+        out = work_dir / f"slot{slot}_chunk{chunk_idx}.parquet"
+        batch.write_parquet(out, use_pyarrow=True)
+        files.append(out)
+
+    return files
+
+
+class PolarsReadEngine(SqlAlchemyReadEngine):
+    def loader(self) -> ReadLoader:
+        return _polars_read

tabsdata_conn_common_sql-2.0.0.data/purelib/tabsdatak/conn/common/sql/_engine/sqlalchemy_utils.py

@@ -0,0 +1,63 @@
+#
+# Copyright 2026 Tabsdata Inc.
+#
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Mapping
+
+import sqlalchemy
+from sqlalchemy.engine import URL, Engine
+
+logger = logging.getLogger(__name__)
+
+
+def create_engine(url: str | URL, **kwargs: Any) -> Engine:
+    """Build a SQLAlchemy engine through :func:`sqlalchemy.create_engine`,
+    logging how the URL resolved to a driver and a dialect implementation.
+
+    Every connector creates its engines here so that the resolution is
+    traceable from a function's log alone. The URL is rendered with the
+    password hidden, so the caller may pass one with credentials in it.
+    """
+    engine = sqlalchemy.create_engine(url, **kwargs)
+    logger.info(
+        "SQLAlchemy engine: URI: %s - URL driver name: %s - Engine driver: %s"
+        " - Dialect: %s",
+        engine.url.render_as_string(hide_password=True),
+        engine.url.drivername,
+        engine.driver,
+        type(engine.dialect).__module__,
+    )
+    return engine
+
+
+def normalize_scheme(uri: str, *, aliases: Mapping[str, str]) -> str:
+    """Replace the URI's scheme with its canonical form, if it appears
+    in ``aliases``. ``aliases`` maps the user-typed scheme to the
+    canonical one (e.g. ``{"mariadb": "mysql"}``).
+
+    The match is exact on the bare scheme (``mariadb://``) and on the
+    driver-qualified prefix (``mariadb+...``). The rest of the URI is
+    untouched.
+    """
+    for alias, canonical in aliases.items():
+        if uri.startswith(f"{alias}://") or uri.startswith(f"{alias}+"):
+            return uri.replace(alias, canonical, 1)
+    return uri
+
+
+def add_driver_to_uri(uri: str, *, drivers: Mapping[str, str]) -> str:
+    """Inject a default DBAPI driver into a bare ``<scheme>://`` URI
+    if one isn't already present. ``drivers`` maps the scheme to the
+    driver to inject (e.g. ``{"mysql": "mysqlconnector"}``).
+
+    If the URI already carries a ``+driver`` qualifier the function
+    is a no-op.
+    """
+    for scheme, driver in drivers.items():
+        bare_prefix = f"{scheme}://"
+        if uri.startswith(bare_prefix):
+            return uri.replace(bare_prefix, f"{scheme}+{driver}://", 1)
+    return uri