sql2deltalake 0.1.0rc1__cp310-abi3-win_amd64.whl

sql2deltalake/__init__.py

@@ -0,0 +1,43 @@
+"""sql2deltalake: incremental loads from SQL Server and Postgres into append-only SCD2 Delta tables.
+
+The Spark data source lives in ``sql2deltalake.spark``."""
+
+from .write_init import (
+    make_writer,
+    DBDeltaPathConfigs,
+    WriteConfig,
+    DEFAULT_DATA_TYPE_MAP,
+    WriteConfigAndInfos,
+)
+from .reader.reader import DataSourceReader
+from .destination.destination import Destination
+from .load_lock import LoadAlreadyRunning
+from pathlib import Path
+from typing import Union
+
+
+def write_db_to_delta(
+    source: Union[DataSourceReader, str],
+    table_or_query: tuple[str, str],
+    destination: Union[Destination, Path],
+    write_config: Union[WriteConfig, None] = None,
+):
+    return make_writer(
+        source=source,
+        table_or_query=table_or_query,
+        destination=destination,
+        write_config=write_config,
+    ).execute()
+
+
+__all__ = [
+    "DEFAULT_DATA_TYPE_MAP",
+    "DataSourceReader",
+    "DBDeltaPathConfigs",
+    "Destination",
+    "LoadAlreadyRunning",
+    "WriteConfig",
+    "WriteConfigAndInfos",
+    "make_writer",
+    "write_db_to_delta",
+]

sql2deltalake/_core.pyd

@@ -0,0 +1,4 @@
+[diffend] Oversized file quarantined before diffing.
+name: sql2deltalake/_core.pyd
+size: 44678656 bytes
+sha256: 1e2694007749412cd8f9bd969180ac606f7a1f8138bfc967b0715bad487d6af4

sql2deltalake/_core.pyi

@@ -0,0 +1,72 @@
+"""Type stubs of the Rust extension module (crates/sql2deltalake-py)."""
+
+from typing import Any, Mapping, Optional, Sequence, Union
+
+__version__: str
+
+class ArrowStream:
+    """An Arrow C stream (`__arrow_c_stream__`), consumable once."""
+
+    def __arrow_c_stream__(self, requested_schema: Optional[object] = None) -> object: ...
+    def __arrow_c_schema__(self) -> object: ...
+
+class SqlBuilder:
+    def __init__(
+        self,
+        dialect_name: str,
+        source_ref: Union[Sequence[str], str],
+        type_map_spec: Mapping[str, tuple[str, bool]],
+        no_trim: bool,
+        columns: Sequence[Mapping[str, Any]],
+    ) -> None: ...
+    def full_load(self, valid_from: Optional[str] = None) -> str: ...
+    def pk_ts(self, columns: Sequence[int]) -> str: ...
+    def updates(self, criteria: Sequence[str], valid_from: Optional[str] = None) -> str: ...
+    def criterion(self, column: int, op: str, value: Any) -> str: ...
+    def source_ts_cnt(self, column: Optional[int] = None) -> str: ...
+    def rows_by_pk_json(
+        self, pks: Sequence[int], json: str, valid_from: Optional[str] = None
+    ) -> str: ...
+
+def load_timestamp_literal(dialect_name: str, micros: int) -> str: ...
+def compat_name(name: str) -> str: ...
+def normalize(
+    source: object,
+    time_scales: Optional[Mapping[int, int]] = None,
+    target: Optional[object] = None,
+) -> ArrowStream: ...
+def read_sql(
+    connection_string: str,
+    sql: str,
+    batch_size: int = 65536,
+    utf8_strings: bool = False,
+    partition_on: Optional[str] = None,
+    partitions: int = 1,
+) -> ArrowStream:
+    """Only in builds with the cargo feature `arrowsql-core` (the default)."""
+    ...
+def pk_diff(
+    snapshot: object,
+    previous: object,
+    keys: Sequence[str],
+    delta: str,
+    watermark: object,
+) -> dict[str, Any]: ...
+def count_new_keys(rows: object, previous: object, keys: Sequence[str]) -> int: ...
+def merge_latest(
+    previous: Optional[object],
+    loaded: object,
+    keys: Sequence[str],
+    delta: str,
+    order: str = "__timestamp",
+    deleted: str = "__is_deleted",
+) -> dict[str, Any]: ...
+def replay_state(
+    log: object,
+    keys: Sequence[str],
+    delta: str,
+    seq: str = "__state_seq",
+    deleted: str = "__is_deleted",
+) -> ArrowStream: ...
+def run_load(host: object, load_spec: Mapping[str, Any]) -> dict[str, Any]: ...
+def restore_last_pk(host: object, load_spec: Mapping[str, Any]) -> bool: ...

sql2deltalake/_native.py

@@ -0,0 +1,327 @@
+"""The boundary to the Rust core (`sql2deltalake._core`): sqlglot objects and user callables are
+turned into plain values here, once per load, so the core needs no Python while it runs."""
+
+from typing import TYPE_CHECKING, Any, Mapping, Union
+
+import sqlglot.expressions as ex
+
+from .sql_schema import is_string_type
+
+if TYPE_CHECKING:
+    from .metadata import InformationSchemaColInfo
+    from .write_init import WriteConfig, WriteConfigAndInfos
+
+
+def column_spec(
+    c: "InformationSchemaColInfo", write_config: "WriteConfig"
+) -> dict[str, Any]:
+    dialect = write_config.dialect
+    return {
+        "name": c.column_name,
+        "target": write_config.get_target_name(c),
+        "type_key": ex.DataType(this=c.data_type.this).sql(dialect).lower(),
+        "type_sql": c.data_type.sql(dialect),
+        "orig_type": c.data_type_str or "",
+        "is_string": bool(is_string_type(c.data_type)),
+        "is_numeric": c.data_type.this
+        in [ex.DataType.Type.BOOLEAN, *ex.DataType.NUMERIC_TYPES],
+        "is_identity": bool(c.is_identity),
+    }
+
+
+def type_map_spec(write_config: "WriteConfig") -> dict[str, tuple[str, bool]]:
+    from .write_init import DEFAULT_DATA_TYPE_MAP
+
+    tm: Mapping[str, Union[ex.DataType, ex.DataType.Type]] = (
+        DEFAULT_DATA_TYPE_MAP.get(write_config.dialect, {})
+        if write_config.data_type_map == "default"
+        else write_config.data_type_map
+    )
+    res: dict[str, tuple[str, bool]] = {}
+    for key, mapped in tm.items():
+        if not isinstance(mapped, (ex.DataType, ex.DataType.Type, str)):
+            continue  # never looked up, e.g. a dialect-keyed map merged in by mistake
+        to = ex.cast(ex.column("x"), mapped).to  # what sqlglot renders for CAST(x AS mapped)
+        res[key] = (to.sql(write_config.dialect), bool(is_string_type(mapped)))
+    return res
+
+
+def source_spec(infos: "WriteConfigAndInfos") -> Union[list[str], str]:
+    """The table name parts, or the query rendered in the source dialect."""
+    toq = infos.table_or_query
+    if isinstance(toq, ex.Query):
+        return toq.sql(infos.write_config.dialect)
+    if isinstance(toq, str):
+        return [toq]
+    return list(toq)
+
+
+def fail_point(name: str) -> "Union[bool, None]":
+    """Test hook, called by the load at a few named steps (`load_updates`, `additional_updates`,
+    `restore_last_pk`, `write_delta2`, `local_pk_count`). Tests patch it to inject a failure or to
+    change the source mid-load; returning False from `restore_last_pk` skips the restore."""
+    return None
+
+
+class PanicError(RuntimeError):
+    """A Rust panic during a load (pyo3's ``PanicException``, a ``BaseException``; arrowsql
+    raises one when the server kills the connection), turned into an ``Exception`` once the load
+    restored its state, so ``except Exception`` in callers sees it."""
+
+
+def is_panic(e: BaseException) -> bool:
+    """pyo3's ``pyo3_runtime.PanicException``, of any extension module (it cannot be imported)."""
+    return not isinstance(e, Exception) and type(e).__name__ == "PanicException"
+
+
+def as_load_error(e: BaseException) -> BaseException:
+    """``e``, or a :class:`PanicError` caused by it if it is a panic."""
+    if is_panic(e):
+        err = PanicError(f"the load panicked: {e}")
+        err.__cause__ = e
+        if getattr(e, "__notes__", None):  # e.g. a restore that failed too
+            err.__notes__ = list(e.__notes__)  # type: ignore[attr-defined]
+        return err
+    return e
+
+
+def enabled() -> bool:
+    import os
+
+    return os.getenv("SQL2DELTA_NATIVE", "1") not in ("0", "false", "False", "")
+
+
+def supports(infos: "WriteConfigAndInfos") -> bool:
+    """The Rust load handles readers with local Delta tables; for now SQL Server through arrowsql."""
+    from .reader.arrowsql_reader import ArrowSQLReader
+
+    return (
+        enabled()
+        and isinstance(infos.source, ArrowSQLReader)
+        and infos.write_config.dialect in ("tsql", "mssql")
+    )
+
+
+def _overrides(reader, name: str) -> bool:
+    """True if a subclass replaced the reader's own implementation of `name`."""
+    from .reader.adbc_reader import ADBCReader
+    from .reader.arrowsql_reader import ArrowSQLReader
+
+    impl = getattr(type(reader), name)
+    return all(
+        impl is not getattr(base, name, None) for base in (ArrowSQLReader, ADBCReader)
+    )
+
+
+class Host:
+    """What the Rust load calls back into: the source, the Delta tables below the destination and
+    the logger. Data crosses as Arrow C streams."""
+
+    def __init__(self, infos: "WriteConfigAndInfos"):
+        self.infos = infos
+        self.reader = infos.source
+        self.dest = infos.destination
+        self.logger = infos.logger
+        self._copies = _overrides(self.reader, "source_write_sql_to_delta")
+
+    # source
+    def source_query(self, sql: str):
+        reader = self.reader
+        if hasattr(reader, "_source_arrow") and not _overrides(reader, "source_sql_to_py"):
+            return reader._source_arrow(sql, utf8_strings=False)  # type: ignore
+        import pyarrow as pa
+
+        return pa.Table.from_pylist(reader.source_sql_to_py(sql))
+
+    def source_connection(self):
+        """(connection string, utf8_strings, partitions) if Rust reads the source itself through
+        arrowsql's core, None if the reads go through the reader's Python methods."""
+        from .reader.arrowsql_reader import ArrowSQLSource, core_reads
+
+        reader = self.reader
+        if (
+            self._copies
+            or not isinstance(reader, ArrowSQLSource)
+            or _overrides(reader, "source_sql_to_py")
+            or _overrides(reader, "_source_arrow")
+            or not core_reads(reader.connection_string)
+        ):
+            return None
+        return (reader.connection_string, reader.utf8_strings, reader.partitions)
+
+    def source_stream(self, sql: str, partition_on=None):
+        if self._copies or not hasattr(self.reader, "_source_arrow"):
+            return None
+        return self.reader._source_arrow(  # type: ignore
+            sql,
+            utf8_strings=getattr(self.reader, "utf8_strings", False),
+            partition_on=partition_on,
+        )
+
+    def source_copy(self, sql: str, rel: str, mode, drift):
+        self.reader.source_write_sql_to_delta(
+            sql, self.dest / rel, mode, allow_schema_drift=drift
+        )
+
+    # target
+    def kernel_root(self):
+        """The target folder as a local path, whose Delta tables Rust reads itself (delta-kernel-rs)."""
+        from .destination.file_system import FileSystemDestination
+
+        if type(self.dest) is FileSystemDestination:
+            return str(self.dest.path.absolute())
+        return None
+
+    def _dt(self, rel: str, version=None):
+        from .reader.local_reader import _delta_table
+
+        dt = _delta_table(self.dest / rel)
+        if version is not None:
+            dt.load_as_version(version)
+        return dt
+
+    def table_exists(self, rel: str) -> bool:
+        return self.reader.local_delta_table_exists(self.dest / rel)
+
+    def path_exists(self, rel: str) -> bool:
+        return (self.dest / rel).exists()
+
+    def version(self, rel: str) -> int:
+        return self._dt(rel).version()
+
+    def table_id(self, rel: str) -> "Union[str, None]":
+        return str(self._dt(rel).metadata().id)  # a UUID string (typed int by deltalake)
+
+    def read(self, rel: str, version, columns):
+        import pyarrow as pa
+
+        ds = self._dt(rel, version).to_pyarrow_dataset()
+        if columns is not None:
+            names = {n.lower(): n for n in ds.schema.names}
+            columns = [c if c in ds.schema.names else names.get(c.lower(), c) for c in columns]
+        return ds.scanner(columns=columns).to_reader()
+
+    def column_names(self, rel: str) -> "list[str]":
+        return [f.name for f in self._dt(rel).schema().fields]
+
+    def write(self, rel: str, stream, mode, drift):
+        import pyarrow as pa
+
+        self.reader._write_reader_to_delta(  # type: ignore
+            pa.RecordBatchReader.from_stream(stream),
+            self.dest / rel,
+            mode,
+            drift,
+            all_nullable=True,
+        )
+
+    def restore(self, rel: str, version: int):
+        self.reader.get_local_delta_ops(self.dest / rel).restore(version)
+
+    def remove(self, rel: str, recursive: bool):
+        (self.dest / rel).remove(recursive)
+
+    def vacuum(self, rel: str):
+        self.reader.get_local_delta_ops(self.dest / rel).vacuum()
+
+    def mkdir(self, rel: str):
+        (self.dest / rel).mkdir()
+
+    def acquire_lock(self, rel: str) -> str:
+        from .load_lock import acquire
+
+        return acquire(self.dest / rel, self.logger)
+
+    def release_lock(self, rel: str, token: str):
+        from .load_lock import release
+
+        release(self.dest / rel, token, self.logger)
+
+    def log(self, level: str, msg: str, kwargs: dict):
+        getattr(self.logger, level)(msg, **kwargs)
+
+    def fail_point(self, name: str):
+        return fail_point(name)
+
+
+def load_spec(infos: "WriteConfigAndInfos") -> dict[str, Any]:
+    from .db_to_delta import (
+        _shortcut_allowed,
+        _source_schema_fingerprint,
+        _watermark_parts,
+        _watermark_sql,
+    )
+
+    wc = infos.write_config
+    cols = list(infos.col_infos)
+
+    def _index(c) -> int:
+        for i, x in enumerate(cols):
+            if x is c:
+                return i
+        for i, x in enumerate(cols):
+            if x.column_name == c.column_name:
+                return i
+        raise ValueError(f"column {c.column_name} not in the columns of the load")
+
+    return {
+        "dialect": wc.dialect,
+        "source": source_spec(infos),
+        "type_map": type_map_spec(wc),
+        "no_trim": wc.no_trim,
+        "columns": [column_spec(c, wc) for c in cols],
+        "pks": [_index(c) for c in infos.pk_cols],
+        "delta_col": _index(infos.delta_col) if infos.delta_col is not None else None,
+        "load_mode": wc.load_mode,
+        "allow_schema_drift": wc.allow_schema_drift,
+        "no_complex_entries_load": wc.no_complex_entries_load,
+        "schema_fingerprint": _source_schema_fingerprint(infos),
+        "shortcut_allowed": _shortcut_allowed(infos),
+        "watermark_sql": _watermark_sql(infos),
+        "watermark_parts": _watermark_parts(infos),
+    }
+
+
+def _result(res: dict):
+    from .load_result import (
+        AppendOnlyLoadResult,
+        DeltaLoadResult,
+        FullLoadResult,
+        NoLoadResult,
+    )
+
+    t = res["type"]
+    if t == "full":
+        return FullLoadResult()
+    if t == "none":
+        return NoLoadResult()
+    if t == "append_only":
+        return AppendOnlyLoadResult()
+    r = DeltaLoadResult()
+    r.starting_local_state = res["starting_local_state"]
+    r.starting_source_state = res["starting_source_state"]
+    if "end_source_state" in res:
+        r.end_source_state = res["end_source_state"]
+    r.dirty = res["dirty"]
+    return r
+
+
+def execute(infos: "WriteConfigAndInfos"):
+    """Runs the load in Rust (`_core.run_load`); the GIL is released while it runs."""
+    from . import _core
+
+    try:
+        return _result(_core.run_load(Host(infos), load_spec(infos)))
+    except BaseException as e:
+        err = as_load_error(e)
+        if err is e:
+            raise
+        raise err from e
+    finally:
+        infos.logger.flush()
+
+
+def restore_last_pk(infos: "WriteConfigAndInfos") -> bool:
+    from . import _core
+
+    return _core.restore_last_pk(Host(infos), load_spec(infos))

sql2deltalake/consistency.py

@@ -0,0 +1,56 @@
+from typing import Any
+from .write_init import WriteConfigAndInfos, DBDeltaPathConfigs
+from .write_utils.restore_pk import create_last_pk_version_view
+import sqlglot as sg
+import sqlglot.expressions as ex
+
+
+class InconsistentPrimaryKeyError(Exception):
+    """Exception raised when primary keys are inconsistent."""
+
+    def __init__(self, message, invalid_data: list[dict[str, Any]]):
+        self.message = message
+        self.invalid_data = invalid_data
+        super().__init__(self.message)
+
+
+def check_latest_pk(infos: WriteConfigAndInfos, raise_if_not_consistent=True):
+    if not infos.delta_col or not infos.pk_cols:
+        raise ValueError("Primary keys and delta column must be defined")
+    postfix = "_" + str(abs(hash(str(infos.destination))))
+    lpk_view = "lastest_pk_" + postfix
+    infos.source.local_register_update_view(
+        infos.destination / "delta_load" / DBDeltaPathConfigs.LATEST_PK_VERSION,
+        lpk_view,
+    )
+    _, view_name, success = create_last_pk_version_view(
+        infos, view_prefix="v_" + postfix
+    )
+    d_cols = list(infos.pk_cols) + [infos.delta_col]
+    col_names = [infos.write_config.get_target_name(p) for p in d_cols]
+    assert success
+    assert view_name is not None
+    query1 = sg.except_(
+        sg.from_(ex.table_(lpk_view, alias="lpk"))
+        .select(*[ex.column(c, "lpk", quoted=True) for c in col_names])
+        .select(ex.convert("added in persisted data"), append=True),
+        sg.from_(ex.table_(view_name, alias="rs"))
+        .select(*[ex.column(c, "rs", quoted=True) for c in col_names])
+        .select(ex.convert("added in persisted data"), append=True),
+    )
+    query2 = sg.except_(
+        sg.from_(ex.table_(view_name, alias="rs"))
+        .select(*[ex.column(c, "rs", quoted=True) for c in col_names])
+        .select(ex.convert("missing in persisted data"), append=True),
+        sg.from_(ex.table_(lpk_view, alias="lpk"))
+        .select(*[ex.column(c, "lpk", quoted=True) for c in col_names])
+        .select(ex.convert("missing in persisted data"), append=True),
+    )
+    result = infos.source.local_execute_sql_to_py(
+        ex.union(query1, query2, distinct=False)
+    )
+    if result:
+        print(result)
+        if raise_if_not_consistent:
+            raise InconsistentPrimaryKeyError("Primary keys are not consistent", result)
+    return result