From 5ad1483bc7d65db04004414272cdd567a46b08a8 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 13 Oct 2026 23:28:59 +0000 Subject: [PATCH 01/17] Implement DuckDB Support --- pyproject.toml | 5 +- src/orm_loader/backends/__init__.py | 2 + src/orm_loader/backends/base.py | 72 +++++- src/orm_loader/backends/duckdb.py | 281 ++++++++++++++++++++++++ src/orm_loader/backends/resolve.py | 2 + src/orm_loader/config.py | 10 + src/orm_loader/tables/loadable_table.py | 11 + src/orm_loader/tables/typing.py | 1 + tests/backends/test_base_backend.py | 30 +-- tests/backends/test_duckdb_backend.py | 195 ++++++++++++++++ tests/backends/test_shared_backend.py | 36 ++- tests/backends/test_sqlite_backend.py | 100 ++++----- tests/conftest.py | 41 +++- tests/loaders/test_dedupe.py | 8 +- tests/loaders/test_loader_e2e.py | 257 ++++++++++++---------- tests/loaders/test_pandas_loader.py | 32 +-- tests/loaders/test_parquet_loader.py | 24 +- tests/loaders/test_split_connection.py | 26 +-- tests/models.py | 26 ++- 19 files changed, 908 insertions(+), 251 deletions(-) create mode 100644 src/orm_loader/backends/duckdb.py create mode 100644 tests/backends/test_duckdb_backend.py diff --git a/pyproject.toml b/pyproject.toml index 4ca8cba..58a4c22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,8 +49,11 @@ packages = ["src/orm_loader"] postgres = [ "psycopg[binary]>=3.2", ] +duckdb = [ + "duckdb-sqlalchemy>=1.5", +] dev = [ - "oa-configurator[postgres]>=2.0.0,<3", + "oa-configurator[postgres,duckdb]>=2.0.0,<3", "pytest>=9.0.3", "ty>=0.0.59", "ruff>=0.14.11", diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index 2c9c799..56cc81c 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -1,3 +1,4 @@ +from .duckdb import DuckDBBackend from .postgres import PostgresBackend from .resolve import resolve_backend from .sqlite import SQLiteBackend @@ -20,6 +21,7 @@ "ConcurrentRefreshNotEligibleError", "DatabaseBackend", "Dialect", + "DuckDBBackend", "MaterializationError", "MaterializationFailure", "MaterializationOperation", diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 016bb2c..0baf6c5 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -35,6 +35,22 @@ class BackendCapabilities: supports_materialized_views: bool = False +class DeferredForeignKeyViolationError(RuntimeError): + """A foreign key this backend can't enforce at the DB level is violated. + """ + + def __init__(self, constraint: sa.ForeignKeyConstraint, orphan_keys: list[tuple[Any, ...]]) -> None: + self.constraint = constraint + self.orphan_keys = orphan_keys + local_cols = ", ".join(c.name for c in constraint.columns) + referred_table = next(iter(constraint.elements)).column.table.name + super().__init__( + f"{constraint.table.name}.({local_cols}) references {referred_table!r} via a " + f"constraint this backend can't enforce at the DB level; {len(orphan_keys)} " + f"row(s) violate it (showing up to 5): {orphan_keys}" + ) + + STAGING_SCHEMA: str = "staging" @@ -269,6 +285,52 @@ def _insertable_column_names( """ return [c.name for c in table_cls.__table__.columns if c.computed is None] + def deferred_fk_constraints( + self, + table_cls: Type["CSVTableProtocol"], + ) -> tuple[sa.ForeignKeyConstraint, ...]: + """FK constraints on *table_cls* that this backend never enforces at the DB level. + + Default: none. Only a backend whose dialect can't execute a deferred + ``ALTER TABLE ADD CONSTRAINT`` (DuckDB, for a cyclic FK graph) overrides this. + """ + return () + + def validate_deferred_fks( + self, + table_cls: Type["CSVTableProtocol"], + session: so.Session, + ) -> None: + """Raise if any row in *table_cls* violates one of its deferred_fk_constraints(). + + An anti-join check recovering referential-integrity validation this + backend's DB can't give for those constraints. No-op when there are + none (every backend but DuckDB, today). + + Raises + ------ + DeferredForeignKeyViolationError + If any row's FK column is non-null with no matching row in the + referenced table. + """ + for fk in self.deferred_fk_constraints(table_cls): + self._validate_one_deferred_fk(fk, session) + + @staticmethod + def _validate_one_deferred_fk(fk: sa.ForeignKeyConstraint, session: so.Session) -> None: + local_cols = list(fk.columns) + remote_cols = [element.column for element in fk.elements] + referred_table = remote_cols[0].table + matches = sa.and_(*(local_cols[i] == remote_cols[i] for i in range(len(local_cols)))) + not_null = sa.and_(*(col.isnot(None) for col in local_cols)) + orphan_exists = sa.exists(sa.select(sa.literal(1)).select_from(referred_table).where(matches)) + + orphans = session.execute( + sa.select(*local_cols).where(not_null, ~orphan_exists).limit(5) + ).all() + if orphans: + raise DeferredForeignKeyViolationError(fk, [tuple(row) for row in orphans]) + @abstractmethod def create_staging_table( self, @@ -361,9 +423,15 @@ def merge_context( table_cls: Type["CSVTableProtocol"], session: so.Session, ) -> AbstractContextManager[None]: - """Disable backend FK checks for a table's merge and restore them afterward.""" + """Disable backend FK checks for a table's merge and restore them afterward. + + A backend without supports_fk_toggle has no state to disable in the + first place, so this skips the attempt entirely rather than raising; + such a backend's own merge methods must avoid operations (like + deleting a referenced row) that would need the toggle to be safe. + """ return self.bulk_load_context( - session, mapper=table_cls, disable_fk=True, no_autoflush=False + session, mapper=table_cls, disable_fk=self.capabilities.supports_fk_toggle, no_autoflush=False ) @contextmanager diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py new file mode 100644 index 0000000..164a9e1 --- /dev/null +++ b/src/orm_loader/backends/duckdb.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Type + +import duckdb_sqlalchemy +import sqlalchemy as sa +import sqlalchemy.orm as so +from sqlalchemy.sql.compiler import IdentifierPreparer + +from .base import BackendCapabilities, DatabaseBackend, Dialect + +if TYPE_CHECKING: + from ..tables.typing import CSVTableProtocol + + +class DuckDBBackend(DatabaseBackend): + @staticmethod + def staging_name_for_table(tablename: str) -> str: + return f"_staging_{tablename}" + + @property + def name(self) -> str: + return Dialect.DUCKDB + + @property + def dialect(self) -> Dialect: + return Dialect.DUCKDB + + @property + def identifier_preparer(self) -> IdentifierPreparer: + return duckdb_sqlalchemy.Dialect().identifier_preparer + + @property + def capabilities(self) -> BackendCapabilities: + return BackendCapabilities( + supports_fast_load=False, + supports_unlogged_staging=False, + supports_fk_toggle=False, + supports_materialized_views=False, + ) + + def deferred_fk_constraints( + self, + table_cls: Type["CSVTableProtocol"], + ) -> tuple[sa.ForeignKeyConstraint, ...]: + """FK constraints on *table_cls* excluded from its CREATE TABLE and never + added afterward: DuckDB can't execute the ALTER TABLE ADD CONSTRAINT that + would otherwise add them (oa-configurator's _skip_deferred_fk_on_duckdb no-ops it). + + sort_tables_and_constraints needs every table in the schema to tell + which FKs are genuinely deferred (cyclic), so this runs against + table_cls's whole shared metadata, not table_cls alone. + """ + metadata = table_cls.__table__.metadata + entries = sa.schema.sort_tables_and_constraints(metadata.tables.values()) + deferred = next(fks for tbl, fks in entries if tbl is None) + return tuple(fk for fk in deferred if fk.table is table_cls.__table__) + + def create_staging_table( + self, + table_cls: type["CSVTableProtocol"], + session: so.Session, + ) -> None: + bind_arguments = {"mapper": table_cls} + session.execute( + sa.text(f'DROP TABLE IF EXISTS {self.qualified_staging_name(table_cls.__tablename__)};'), + bind_arguments=bind_arguments, + ) + + staging_name = self.staging_name_for_table(table_cls.__tablename__) + metadata = sa.MetaData() + staging_columns = [ + sa.Column(col.name, col.type, nullable=True) + for col in table_cls.__table__.columns + ] + staging_table = sa.Table(staging_name, metadata, *staging_columns, schema=self.staging_schema) + metadata.create_all( + bind=session.connection(bind_arguments=bind_arguments), tables=[staging_table] + ) + session.commit() + + def drop_staging_table( + self, + table_cls: type["CSVTableProtocol"], + session: so.Session, + ) -> None: + session.execute( + sa.text(f'DROP TABLE IF EXISTS {self.qualified_staging_name(table_cls.__tablename__)}'), + bind_arguments={"mapper": table_cls}, + ) + + @staticmethod + def _normalize_fk_check_state(previous_state: str | int) -> str | int: + raise NotImplementedError( + "DuckDB has no foreign-key-toggle mechanism; there is no state to normalize." + ) + + def disable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: + raise NotImplementedError( + "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." + ) + + def enable_fk_check(self, session: so.Session, *, mapper: Any | None = None) -> str | int: + raise NotImplementedError( + "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." + ) + + def restore_fk_check( + self, + session: so.Session, + previous_state: str | int, + *, + mapper: Any | None = None, + ) -> None: + raise NotImplementedError( + "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." + ) + + @staticmethod + def _staging_rowid() -> sa.ColumnElement[int]: + """DuckDB's own rowid: dense per table, but not guaranteed to start at 0 or 1.""" + return sa.literal_column("rowid") + + def _staging_rowid_bounds(self, staging: sa.Table, session: so.Session) -> tuple[int, int] | None: + """Lowest and highest rowid in *staging*, or None when empty. + + DuckDB's rowid isn't guaranteed to start at any fixed value (verified + live), unlike SQLite's, so pagination bounds against the real + min/max instead of assuming a 0-based range. + """ + lowest, highest = session.execute( + sa.select(sa.func.min(self._staging_rowid()), sa.func.max(self._staging_rowid())) + .select_from(staging) + ).one() + if lowest is None: + return None + return lowest, highest + + def merge_replace( + self, + table_cls: type["CSVTableProtocol"], + session: so.Session, + pk_cols: list[str], + *, + merge_batch_size: int | None = None, + ) -> None: + """Update matching target rows from staging in place, then drop them from staging. + + DuckDB rejects deleting a still-referenced row outright, with no + FK-toggle or deferred-constraint escape hatch, so replace is an + UPDATE rather than orm-loader's usual delete-then-reinsert. Matched + rows are removed from staging so the follow-up merge_insert() call + only inserts rows that weren't already handled here. + """ + self._validate_merge_batch_size(merge_batch_size) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) + update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] + pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) + + def _batch_keys(extra: sa.ColumnElement[bool] | None = None) -> sa.Select[Any]: + keys = sa.select(*(staging.c[c] for c in pk_cols)) + if extra is not None: + keys = keys.where(extra) + return keys + + def _pk_in(keys: sa.Select[Any]) -> sa.ColumnElement[bool]: + if len(pk_cols) == 1: + return target.c[pk_cols[0]].in_(keys) + return sa.tuple_(*(target.c[c] for c in pk_cols)).in_(keys) + + def _update(extra: sa.ColumnElement[bool] | None = None) -> None: + stmt = sa.update(target).values({c: staging.c[c] for c in update_cols}).where(pk_join) + if extra is not None: + stmt = stmt.where(_pk_in(_batch_keys(extra))) + session.execute(stmt) + + def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> None: + matched = sa.select(*(target.c[c] for c in pk_cols)).where(pk_join) + if len(pk_cols) == 1: + condition: sa.ColumnElement[bool] = staging.c[pk_cols[0]].in_(matched) + else: + condition = sa.tuple_(*(staging.c[c] for c in pk_cols)).in_(matched) + stmt = sa.delete(staging).where(condition) + if extra is not None: + stmt = stmt.where(extra) + session.execute(stmt) + + if not update_cols: + # No non-pk columns to update: matching rows are already identical, just drop them. + _drop_matched_from_staging() + return + + if merge_batch_size is None: + _update() + _drop_matched_from_staging() + return + + bounds = self._staging_rowid_bounds(staging, session) + if bounds is None: + return + start, maximum = bounds + start -= 1 + rowid = self._staging_rowid() + while start < maximum: + end = start + merge_batch_size + batch_range = sa.and_(rowid > start, rowid <= end) + _update(batch_range) + _drop_matched_from_staging(batch_range) + session.commit() + start = end + + def merge_upsert( + self, + table_cls: type["CSVTableProtocol"], + session: so.Session, + pk_cols: list[str], + *, + merge_batch_size: int | None = None, + ) -> None: + self._validate_merge_batch_size(merge_batch_size) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) + insertable_cols = self._insertable_column_names(table_cls) + + def _upsert(select_: Any) -> sa.Insert: + return ( + duckdb_sqlalchemy.insert(target) + .from_select(insertable_cols, select_) + .on_conflict_do_nothing(index_elements=pk_cols) + ) + + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + + if merge_batch_size is None: + session.execute(_upsert(non_paginated_select)) + return + + bounds = self._staging_rowid_bounds(staging, session) + if bounds is None: + return + start, maximum = bounds + start -= 1 + rowid = self._staging_rowid() + while start < maximum: + end = start + merge_batch_size + batch_select = non_paginated_select.where(rowid > start, rowid <= end) + session.execute(_upsert(batch_select)) + start = end + + def merge_insert( + self, + table_cls: type["CSVTableProtocol"], + session: so.Session, + *, + merge_batch_size: int | None = None, + ) -> None: + self._validate_merge_batch_size(merge_batch_size) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) + insertable_cols = self._insertable_column_names(table_cls) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + + def _insert(select_: Any) -> sa.Insert: + return sa.insert(target).from_select(insertable_cols, select_) + + if merge_batch_size is None: + session.execute(_insert(non_paginated_select)) + return + + bounds = self._staging_rowid_bounds(staging, session) + if bounds is None: + return + start, maximum = bounds + start -= 1 + rowid = self._staging_rowid() + while start < maximum: + end = start + merge_batch_size + batch_select = non_paginated_select.where(rowid > start, rowid <= end) + session.execute(_insert(batch_select)) + start = end diff --git a/src/orm_loader/backends/resolve.py b/src/orm_loader/backends/resolve.py index c2fee8e..0b69066 100644 --- a/src/orm_loader/backends/resolve.py +++ b/src/orm_loader/backends/resolve.py @@ -6,12 +6,14 @@ from oa_configurator import Bindable, UnregisteredSchemaTagError, physical_schema_of from .base import DatabaseBackend, Dialect +from .duckdb import DuckDBBackend from .postgres import PostgresBackend from .sqlite import SQLiteBackend _BACKEND_TYPES: dict[Dialect, type[DatabaseBackend]] = { Dialect.POSTGRESQL: PostgresBackend, Dialect.SQLITE: SQLiteBackend, + Dialect.DUCKDB: DuckDBBackend, } diff --git a/src/orm_loader/config.py b/src/orm_loader/config.py index dc1e524..5017ab8 100644 --- a/src/orm_loader/config.py +++ b/src/orm_loader/config.py @@ -28,6 +28,9 @@ class exists to register orm-loader in the oa-configurator ecosystem, against SQLite specifically, regardless of what ``test_orm_db_pg`` happens to be configured to. Left unconfigured by design in every environment. + test_orm_db_duckdb : str, optional + Same shape as ``test_orm_db_sqlite``, for tests that must always + run against DuckDB specifically. Notes ----- @@ -49,3 +52,10 @@ class exists to register orm-loader in the oa-configurator ecosystem, "(isolated_test_database(..., dialect='sqlite') provisions one automatically)." ), ) + test_orm_db_duckdb: Annotated[str | None, RefTo(CDMDatabaseConfig, is_test=True)] = Field( + default=None, + description=( + "Disposable DuckDB test database; left unconfigured by design " + "(isolated_test_database(..., dialect='duckdb') provisions one automatically)." + ), + ) diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 1f29b6d..efe1142 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -373,6 +373,7 @@ def load_csv( index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, + validate_deferred_fks: bool = False, ) -> int: """ @@ -413,6 +414,11 @@ def load_csv( means no schema qualification (backend-default behavior). Threaded through every internal step of the load lifecycle so they all resolve the same backend/schema. + validate_deferred_fks + After merging, run an anti-join check for any FK this backend's + DB can't enforce itself (see DatabaseBackend.validate_deferred_fks()). + Off by default: a full-table scan per such FK, on top of the load. + No-op for a backend with none (every backend but DuckDB, today). Returns ------- @@ -473,6 +479,11 @@ def load_csv( staging_schema_tag=staging_schema_tag, ) + if validate_deferred_fks: + logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) + backend.validate_deferred_fks(cls, session) + cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) logger.info(f"Table `{cls.__tablename__}`: Successfully finished ingestion. Total rows: {total}") diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index b1a3b14..ec3a8fe 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -87,6 +87,7 @@ def load_csv( index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, + validate_deferred_fks: bool = False, ) -> int: ... @classmethod diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 34154a2..2add029 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -223,28 +223,28 @@ def test_insertable_column_names_exclude_computed_columns(): assert backend._insertable_column_names(_ComputedTableCls) == ["id", "name"] -def test_bulk_load_context_toggles_fk_and_restores(session): +def test_bulk_load_context_toggles_fk_and_restores(sqlite_session): backend = FakeBackend() - with backend.bulk_load_context(session): + with backend.bulk_load_context(sqlite_session): pass assert backend.calls == [ - ("disable_fk_check", session), + ("disable_fk_check", sqlite_session), ("restore_fk_check", "enabled"), ] -def test_bulk_load_context_without_fk_toggle(session): +def test_bulk_load_context_without_fk_toggle(sqlite_session): backend = FakeBackend() - with backend.bulk_load_context(session, disable_fk=False): + with backend.bulk_load_context(sqlite_session, disable_fk=False): pass assert backend.calls == [] -def test_bulk_load_context_raises_when_capability_missing(session): +def test_bulk_load_context_raises_when_capability_missing(sqlite_session): class NoFKBackend(FakeBackend): @property def capabilities(self) -> BackendCapabilities: @@ -253,19 +253,19 @@ def capabilities(self) -> BackendCapabilities: backend = NoFKBackend() with pytest.raises(NotImplementedError, match="does not support foreign key toggling"): - with backend.bulk_load_context(session): + with backend.bulk_load_context(sqlite_session): pass -def test_bulk_load_context_rolls_back_and_restores(session): +def test_bulk_load_context_rolls_back_and_restores(sqlite_session): backend = FakeBackend() with pytest.raises(RuntimeError, match="boom"): - with backend.bulk_load_context(session): + with backend.bulk_load_context(sqlite_session): raise RuntimeError("boom") assert backend.calls == [ - ("disable_fk_check", session), + ("disable_fk_check", sqlite_session), ("restore_fk_check", "enabled"), ] @@ -280,17 +280,17 @@ def test_backends_package_exports(): def test_resolve_backend_for_sqlite_engine_and_session(): - engine = sa.create_engine("sqlite:///:memory:", future=True) - session = so.Session(engine) + sqlite_engine = sa.create_engine("sqlite:///:memory:", future=True) + sqlite_session = so.Session(sqlite_engine) try: - engine_backend = resolve_backend(engine) - session_backend = resolve_backend(session) + engine_backend = resolve_backend(sqlite_engine) + session_backend = resolve_backend(sqlite_session) assert engine_backend.name == "sqlite" assert session_backend.name == "sqlite" finally: - session.close() + sqlite_session.close() def test_resolve_backend_raises_for_unknown_dialect(): diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py new file mode 100644 index 0000000..2cba7b3 --- /dev/null +++ b/tests/backends/test_duckdb_backend.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Type, cast + +import pytest +import sqlalchemy as sa +import sqlalchemy.orm as so + +from orm_loader.backends import Dialect, DuckDBBackend +from orm_loader.backends.base import DeferredForeignKeyViolationError +from tests.models import ComputedColumnTable, CycleATable, CycleBTable + +if TYPE_CHECKING: + from orm_loader.tables.typing import CSVTableProtocol + +_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) +_CycleATableCls = cast("Type[CSVTableProtocol]", CycleATable) +_CycleBTableCls = cast("Type[CSVTableProtocol]", CycleBTable) +_TARGET_TABLE = ComputedColumnTable.__tablename__ +_STAGING_TABLE = f"_staging_{_TARGET_TABLE}" + + +def test_duckdb_backend_identity_and_capabilities(): + backend = DuckDBBackend() + + assert backend.name == "duckdb" + assert backend.dialect == Dialect.DUCKDB + assert backend.supports_dialect(Dialect.DUCKDB) is True + assert backend.capabilities.supports_fast_load is False + assert backend.capabilities.supports_unlogged_staging is False + assert backend.capabilities.supports_fk_toggle is False + assert backend.capabilities.supports_materialized_views is False + assert backend.resolve_index_strategy("auto") == "drop_rebuild" + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +@pytest.mark.parametrize("merge_batch_size", [0, -1]) +@pytest.mark.parametrize("merge_strategy", ["replace", "upsert", "insert"]) +def test_duckdb_backend_rejects_non_positive_merge_batch_size(merge_batch_size, merge_strategy, duckdb_session): + backend = DuckDBBackend() + args = [_ComputedTableCls, duckdb_session] + if merge_strategy in {"replace", "upsert"}: + args.append(["id"]) + merge = getattr(backend, f"merge_{merge_strategy}") + + with pytest.raises(ValueError, match="merge_batch_size must be a positive integer or None"): + merge(*args, merge_batch_size=merge_batch_size) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_create_staging_table(duckdb_session): + backend = DuckDBBackend() + + backend.create_staging_table(_ComputedTableCls, duckdb_session) + inspector = sa.inspect(duckdb_session.get_bind()) + assert inspector.has_table(_STAGING_TABLE) is True + cols = inspector.get_columns(_STAGING_TABLE) + assert [c["name"] for c in cols] == ["id", "name", "slug"] + assert all(c["nullable"] is True for c in cols) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_drop_staging_table(duckdb_session): + backend = DuckDBBackend() + + backend.create_staging_table(_ComputedTableCls, duckdb_session) + backend.drop_staging_table(_ComputedTableCls, duckdb_session) + + inspector = sa.inspect(duckdb_session.get_bind()) + assert inspector.has_table(_STAGING_TABLE) is False + + +def test_duckdb_backend_fk_toggle_methods_raise(): + backend = DuckDBBackend() + + with pytest.raises(NotImplementedError, match="no foreign-key-toggle mechanism"): + backend.disable_fk_check(cast("so.Session", None)) + with pytest.raises(NotImplementedError, match="no foreign-key-toggle mechanism"): + backend.enable_fk_check(cast("so.Session", None)) + with pytest.raises(NotImplementedError, match="no foreign-key-toggle mechanism"): + backend.restore_fk_check(cast("so.Session", None), "ON") + with pytest.raises(NotImplementedError, match="no state to normalize"): + DuckDBBackend._normalize_fk_check_state("ON") + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_context_runs_without_toggling_fk(duckdb_session): + """merge_context() must not require supports_fk_toggle: DuckDB has none.""" + backend = DuckDBBackend() + backend.create_staging_table(_ComputedTableCls, duckdb_session) + + with backend.merge_context(_ComputedTableCls, duckdb_session): + pass + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_materialized_view_methods_raise(duckdb_session): + backend = DuckDBBackend() + selectable = sa.select(sa.literal(1).label("n")) + engine = duckdb_session.get_bind() + + with pytest.raises(NotImplementedError, match="does not support materialized views"): + backend.create_materialized_view(engine, "mv_test", selectable) + with pytest.raises(NotImplementedError, match="does not support materialized views"): + backend.refresh_materialized_view(engine, "mv_test") + with pytest.raises(NotImplementedError, match="does not support materialized views"): + backend.drop_materialized_view(engine, "mv_test") + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_deferred_fk_constraints_finds_the_cycle(duckdb_session): + backend = DuckDBBackend() + + a_fks = backend.deferred_fk_constraints(_CycleATableCls) + b_fks = backend.deferred_fk_constraints(_CycleBTableCls) + + assert [c.name for c in a_fks[0].columns] == ["b_id"] + assert [c.name for c in b_fks[0].columns] == ["a_id"] + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_validate_deferred_fks_passes_for_consistent_data(duckdb_session): + backend = DuckDBBackend() + duckdb_session.execute(sa.insert(CycleBTable), [{"id": 1, "a_id": None}]) + duckdb_session.execute(sa.insert(CycleATable), [{"id": 1, "b_id": 1}]) + + backend.validate_deferred_fks(_CycleATableCls, duckdb_session) + backend.validate_deferred_fks(_CycleBTableCls, duckdb_session) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_validate_deferred_fks_catches_an_orphan(duckdb_session): + """DuckDB's engine never created this FK as a real constraint (it's part + of a cycle, so oa-configurator's shim no-opped the ALTER TABLE ADD + CONSTRAINT), so an orphaned reference inserts without error; the + anti-join check is the only thing that can catch it.""" + backend = DuckDBBackend() + duckdb_session.execute(sa.insert(CycleATable), [{"id": 1, "b_id": 999}]) + + with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): + backend.validate_deferred_fks(_CycleATableCls, duckdb_session) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_replace_updates_in_place(duckdb_session): + """merge_replace cannot delete a still-referenced row (no FK-toggle escape + hatch), so it updates matching rows in place and leaves non-matching staging + rows for the follow-up merge_insert() call, matching merge_from_staging()'s + real two-step orchestration rather than merge_replace alone. + """ + backend = DuckDBBackend() + backend.create_staging_table(_ComputedTableCls, duckdb_session) + staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + + duckdb_session.execute( + sa.insert(ComputedColumnTable), + [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}], + ) + duckdb_session.execute( + sa.insert(staging), [{"id": 1, "name": "alpha-staged"}, {"id": 3, "name": "gamma"}] + ) + + backend.merge_replace(_ComputedTableCls, duckdb_session, ["id"]) + backend.merge_insert(_ComputedTableCls, duckdb_session) + + rows = {r.id: r.name for r in duckdb_session.execute(sa.select(ComputedColumnTable)).scalars().all()} + assert rows == {1: "alpha-staged", 2: "beta", 3: "gamma"} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_replace_paginated_updates_in_place(duckdb_session): + backend = DuckDBBackend() + backend.create_staging_table(_ComputedTableCls, duckdb_session) + staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + + duckdb_session.execute( + sa.insert(ComputedColumnTable), [{"id": i, "name": f"orig{i}"} for i in range(10)] + ) + duckdb_session.execute(sa.insert(staging), [{"id": i, "name": f"staged{i}"} for i in range(10)]) + + backend.merge_replace(_ComputedTableCls, duckdb_session, ["id"], merge_batch_size=3) + backend.merge_insert(_ComputedTableCls, duckdb_session, merge_batch_size=3) + + rows = {r.id: r.name for r in duckdb_session.execute(sa.select(ComputedColumnTable)).scalars().all()} + assert rows == {i: f"staged{i}" for i in range(10)} diff --git a/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py index 8c85250..b85da3f 100644 --- a/tests/backends/test_shared_backend.py +++ b/tests/backends/test_shared_backend.py @@ -5,6 +5,7 @@ import pytest import sqlalchemy as sa +from oa_configurator import Dialect from oa_configurator.testing import DIALECT_PARAMS from orm_loader.backends import STAGING_SCHEMA, DatabaseBackend from orm_loader.backends.resolve import resolve_backend @@ -18,26 +19,37 @@ _ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) _CompositeTableCls = cast("Type[CSVTableProtocol]", CompositeTable) +_SESSION_FIXTURE_BY_DIALECT = { + Dialect.POSTGRESQL: "pg_session", + Dialect.SQLITE: "sqlite_session", + Dialect.DUCKDB: "duckdb_session", +} +assert set(_SESSION_FIXTURE_BY_DIALECT) == {param.values[0] for param in DIALECT_PARAMS}, ( + "_SESSION_FIXTURE_BY_DIALECT is missing an entry for a dialect DIALECT_PARAMS now covers." +) + +# merge_replace deletes on sqlite/postgres but updates in place on DuckDB +# (see DuckDBBackend.merge_replace); these tests check the delete-only behavior. +_REPLACE_DELETE_SEMANTICS_PARAMS = tuple(p for p in DIALECT_PARAMS if p.values[0] != Dialect.DUCKDB) + @pytest.fixture(params=DIALECT_PARAMS) def merge_backend(request: pytest.FixtureRequest) -> tuple[DatabaseBackend, "so.Session"]: - """Same merge-method contract exercised against both real backends. - Only the postgresql param ever requests pg_session, so the sqlite - param never needs a database. + """Same merge-method contract exercised against every real backend. DIALECT_PARAMS carries each dialect's own mark plus `forked` directly on the param value, so this still works correctly even though - request.getfixturevalue("pg_session") is a dynamic, runtime lookup - invisible to pytest's collection-time fixturenames computation (the - usual pg_db-in-fixturenames auto-detection can't see it). + request.getfixturevalue(...) is a dynamic, runtime lookup invisible to + pytest's collection-time fixturenames computation (the usual + pg_db-in-fixturenames auto-detection can't see it). """ - if request.param == "postgresql": - session = request.getfixturevalue("pg_session") - return resolve_backend(session, staging_schema_tag=STAGING_SCHEMA), session - session = request.getfixturevalue("session") - return resolve_backend(session), session + session_fixture = _SESSION_FIXTURE_BY_DIALECT[request.param] + session = request.getfixturevalue(session_fixture) + staging_schema_tag = STAGING_SCHEMA if request.param != Dialect.SQLITE else None + return resolve_backend(session, staging_schema_tag=staging_schema_tag), session +@pytest.mark.parametrize("merge_backend", _REPLACE_DELETE_SEMANTICS_PARAMS, indirect=True) def test_merge_replace_single_pk(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend backend.create_staging_table(_ComputedTableCls, session) @@ -55,6 +67,7 @@ def test_merge_replace_single_pk(merge_backend: tuple[DatabaseBackend, "so.Sessi assert remaining == [2] +@pytest.mark.parametrize("merge_backend", _REPLACE_DELETE_SEMANTICS_PARAMS, indirect=True) def test_merge_replace_composite_pk(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend backend.create_staging_table(_CompositeTableCls, session) @@ -133,6 +146,7 @@ def test_postgres_merge_insert_paginates_rownums_with_initial_gap(pg_session) -> assert ids == list(range(10)) +@pytest.mark.parametrize("merge_backend", _REPLACE_DELETE_SEMANTICS_PARAMS, indirect=True) def test_merge_replace_paginated_path(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: backend, session = merge_backend backend.create_staging_table(_ComputedTableCls, session) diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 94f5f7b..7daefb8 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -62,9 +62,9 @@ def test_sqlite_backend_identity_and_capabilities(): @pytest.mark.parametrize("merge_strategy", ["replace", "upsert", "insert"]) def test_sqlite_backend_rejects_non_positive_merge_batch_size(merge_batch_size, merge_strategy): backend = SQLiteBackend() - session = _sess(_FakeSession()) + sqlite_session = _sess(_FakeSession()) merge = getattr(backend, f"merge_{merge_strategy}") - args = [_ComputedTableCls, session] + args = [_ComputedTableCls, sqlite_session] if merge_strategy in {"replace", "upsert"}: args.append(["id"]) @@ -72,11 +72,11 @@ def test_sqlite_backend_rejects_non_positive_merge_batch_size(merge_batch_size, merge(*args, merge_batch_size=merge_batch_size) -def test_sqlite_backend_create_staging_table(session, engine): +def test_sqlite_backend_create_staging_table(sqlite_session, sqlite_engine): backend = SQLiteBackend() - backend.create_staging_table(_ComputedTableCls, session) - inspector = sa.inspect(engine) + backend.create_staging_table(_ComputedTableCls, sqlite_session) + inspector = sa.inspect(sqlite_engine) assert inspector.has_table(_STAGING_TABLE) is True cols = inspector.get_columns(_STAGING_TABLE) assert [c["name"] for c in cols] == ["id", "name", "slug"] @@ -85,22 +85,22 @@ def test_sqlite_backend_create_staging_table(session, engine): def test_sqlite_backend_drop_staging_table(): backend = SQLiteBackend() - session = _FakeSession() + sqlite_session = _FakeSession() - backend.drop_staging_table(_ComputedTableCls, _sess(session)) + backend.drop_staging_table(_ComputedTableCls, _sess(sqlite_session)) - assert session.statements == [f'DROP TABLE IF EXISTS "{_STAGING_TABLE}"'] + assert sqlite_session.statements == [f'DROP TABLE IF EXISTS "{_STAGING_TABLE}"'] def test_sqlite_backend_disable_fk_reads_then_sets(): backend = SQLiteBackend() - session = _FakeSession(scalar_result=1) + sqlite_session = _FakeSession(scalar_result=1) - previous = backend.disable_fk_check(_sess(session)) + previous = backend.disable_fk_check(_sess(sqlite_session)) assert previous == 1 - assert session.statements == [ + assert sqlite_session.statements == [ "PRAGMA foreign_keys", # read current state "PRAGMA foreign_keys = OFF", # set to OFF ] @@ -108,12 +108,12 @@ def test_sqlite_backend_disable_fk_reads_then_sets(): def test_sqlite_backend_enable_fk_reads_then_sets(): backend = SQLiteBackend() - session = _FakeSession(scalar_result=0) + sqlite_session = _FakeSession(scalar_result=0) - previous = backend.enable_fk_check(_sess(session)) + previous = backend.enable_fk_check(_sess(sqlite_session)) assert previous == 0 - assert session.statements == [ + assert sqlite_session.statements == [ "PRAGMA foreign_keys", # read current state "PRAGMA foreign_keys = ON", # set to ON ] @@ -121,12 +121,12 @@ def test_sqlite_backend_enable_fk_reads_then_sets(): def test_sqlite_backend_restore_fk_normalises_int_and_emits(): backend = SQLiteBackend() - session = _FakeSession() + sqlite_session = _FakeSession() - backend.restore_fk_check(_sess(session), 1) - backend.restore_fk_check(_sess(session), 0) + backend.restore_fk_check(_sess(sqlite_session), 1) + backend.restore_fk_check(_sess(sqlite_session), 0) - assert session.statements == [ + assert sqlite_session.statements == [ "PRAGMA foreign_keys = ON", "PRAGMA foreign_keys = OFF", ] @@ -159,49 +159,49 @@ def test_sqlite_backend_normalize_fk_check_state(): raise AssertionError("Expected ValueError for unrecognised string") -def test_sqlite_backend_materialized_view_methods_raise(engine): +def test_sqlite_backend_materialized_view_methods_raise(sqlite_engine): backend = SQLiteBackend() selectable = sa.select(sa.literal(1).label("n")) try: - backend.create_materialized_view(engine, "mv_test", selectable) + backend.create_materialized_view(sqlite_engine, "mv_test", selectable) except NotImplementedError as exc: assert "does not support materialized views" in str(exc) else: raise AssertionError("Expected create_materialized_view() to raise NotImplementedError") try: - backend.refresh_materialized_view(engine, "mv_test") + backend.refresh_materialized_view(sqlite_engine, "mv_test") except NotImplementedError as exc: assert "does not support materialized views" in str(exc) else: raise AssertionError("Expected refresh_materialized_view() to raise NotImplementedError") -def test_sqlite_backend_drop_materialized_view_raises(engine): +def test_sqlite_backend_drop_materialized_view_raises(sqlite_engine): backend = SQLiteBackend() with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.drop_materialized_view(engine, "mv_test") + backend.drop_materialized_view(sqlite_engine, "mv_test") -def test_sqlite_backend_create_materialized_view_index_raises(engine): +def test_sqlite_backend_create_materialized_view_index_raises(sqlite_engine): from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex backend = SQLiteBackend() index = MaterializedViewIndex(name="mv_test_uq", columns=("id",), unique=True) with pytest.raises(NotImplementedError, match="does not support materialized views"): - backend.create_materialized_view_index(engine, "mv_test", index) + backend.create_materialized_view_index(sqlite_engine, "mv_test", index) def test_sqlite_backend_configures_bulk_load_pragmas(tmp_path: Path): backend = SQLiteBackend() db_path = tmp_path / "test.db" - engine = sa.create_engine(f"sqlite:///{db_path}", future=True) - backend.install_engine_hooks(engine) + sqlite_engine = sa.create_engine(f"sqlite:///{db_path}", future=True) + backend.install_engine_hooks(sqlite_engine) - with engine.connect() as conn: + with sqlite_engine.connect() as conn: busy_timeout = conn.execute(sa.text("PRAGMA busy_timeout")).scalar_one() journal_mode = conn.execute(sa.text("PRAGMA journal_mode")).scalar_one() foreign_keys = conn.execute(sa.text("PRAGMA foreign_keys")).scalar_one() @@ -214,14 +214,14 @@ def test_sqlite_backend_configures_bulk_load_pragmas(tmp_path: Path): def test_sqlite_backend_restore_journal_mode(tmp_path: Path): backend = SQLiteBackend() db_path = tmp_path / "journal.db" - engine = sa.create_engine(f"sqlite:///{db_path}", future=True) - backend.install_engine_hooks(engine) + sqlite_engine = sa.create_engine(f"sqlite:///{db_path}", future=True) + backend.install_engine_hooks(sqlite_engine) - with engine.begin() as conn: + with sqlite_engine.begin() as conn: conn.execute(sa.text("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)")) conn.execute(sa.text("INSERT INTO t (name) VALUES ('x')")) - engine.dispose() + sqlite_engine.dispose() backend.restore_journal_mode(db_path) with sqlite3.connect(db_path.resolve()) as conn: @@ -232,11 +232,11 @@ def test_sqlite_backend_restore_journal_mode(tmp_path: Path): def test_attach_sqlite_bulk_load_pragmas_installs_backend_hook(tmp_path: Path): db_path = tmp_path / "attached.db" - engine = sa.create_engine(f"sqlite:///{db_path}", future=True) + sqlite_engine = sa.create_engine(f"sqlite:///{db_path}", future=True) - attach_sqlite_bulk_load_pragmas(engine, busy_timeout_ms=45000) + attach_sqlite_bulk_load_pragmas(sqlite_engine, busy_timeout_ms=45000) - with engine.connect() as conn: + with sqlite_engine.connect() as conn: busy_timeout = conn.execute(sa.text("PRAGMA busy_timeout")).scalar_one() journal_mode = conn.execute(sa.text("PRAGMA journal_mode")).scalar_one() foreign_keys = conn.execute(sa.text("PRAGMA foreign_keys")).scalar_one() @@ -257,10 +257,10 @@ def test_sqlite_backend_rejects_invalid_journal_mode(): def test_sqlite_backend_disable_fk_raises_when_pragma_returns_non_int(): backend = SQLiteBackend() - session = _FakeSession(scalar_result="not_an_int") + sqlite_session = _FakeSession(scalar_result="not_an_int") try: - backend.disable_fk_check(_sess(session)) + backend.disable_fk_check(_sess(sqlite_session)) except RuntimeError as exc: assert "Expected SQLite FK state to be an int" in str(exc) else: @@ -269,10 +269,10 @@ def test_sqlite_backend_disable_fk_raises_when_pragma_returns_non_int(): def test_sqlite_backend_enable_fk_raises_when_pragma_returns_non_int(): backend = SQLiteBackend() - session = _FakeSession(scalar_result="not_an_int") + sqlite_session = _FakeSession(scalar_result="not_an_int") try: - backend.enable_fk_check(_sess(session)) + backend.enable_fk_check(_sess(sqlite_session)) except RuntimeError as exc: assert "Expected SQLite FK state to be an int" in str(exc) else: @@ -281,25 +281,25 @@ def test_sqlite_backend_enable_fk_raises_when_pragma_returns_non_int(): def test_sqlite_backend_restore_fk_accepts_string_values(): backend = SQLiteBackend() - session = _FakeSession() + sqlite_session = _FakeSession() - backend.restore_fk_check(_sess(session), "ON") - backend.restore_fk_check(_sess(session), "OFF") + backend.restore_fk_check(_sess(sqlite_session), "ON") + backend.restore_fk_check(_sess(sqlite_session), "OFF") - assert session.statements == [ + assert sqlite_session.statements == [ "PRAGMA foreign_keys = ON", "PRAGMA foreign_keys = OFF", ] -def test_sqlite_backend_fk_toggle_round_trip(session): +def test_sqlite_backend_fk_toggle_round_trip(sqlite_session): backend = SQLiteBackend() - session.execute(sa.text("PRAGMA foreign_keys = ON")) - assert session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 1 + sqlite_session.execute(sa.text("PRAGMA foreign_keys = ON")) + assert sqlite_session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 1 - previous = backend.disable_fk_check(session) - assert session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 0 + previous = backend.disable_fk_check(sqlite_session) + assert sqlite_session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 0 - backend.restore_fk_check(session, previous) - assert session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 1 + backend.restore_fk_check(sqlite_session, previous) + assert sqlite_session.execute(sa.text("PRAGMA foreign_keys")).scalar() == 1 diff --git a/tests/conftest.py b/tests/conftest.py index af3c066..489b672 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,18 +17,18 @@ @pytest.fixture -def engine(): +def sqlite_engine(): with isolated_test_database( OrmLoaderConfig, "test_orm_db_sqlite", dialect="sqlite", ) as db: - engine = db.connection.engine - Base.metadata.create_all(engine) - yield engine + sqlite_engine = db.connection.engine + Base.metadata.create_all(sqlite_engine) + yield sqlite_engine @pytest.fixture -def session(engine): - with so.Session(engine) as s: +def sqlite_session(sqlite_engine): + with so.Session(sqlite_engine) as s: yield s @@ -58,6 +58,31 @@ def pg_session(pg_db): return pg_db.session +# --------------------------------------------------------------------------- +# DuckDB fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def duckdb_db(request): + """Isolated DuckDB test database, a fresh disposable tempfile per test. + Schema-claimed the same way as pg_db: DuckDB genuinely supports schemas, + unlike SQLite.""" + with isolated_test_database( + OrmLoaderConfig, "test_orm_db_duckdb", dialect="duckdb", request=request, + schema_claims=[staging_schema_claim()], + ) as db: + yield db + + +@pytest.fixture +def duckdb_session(duckdb_db): + """The standard fixture for tests needing real tables ready to query: + creates the staging schema and Base.metadata, then returns duckdb_db.session.""" + conn = duckdb_db.connection + Base.metadata.create_all(conn) + return duckdb_db.session + + @contextmanager def schema_scoped_session( pg_db, table: sa.Table, *, prefix: str, split_roles: Iterable[Role] = () @@ -69,5 +94,5 @@ def schema_scoped_session( ) as scoped: with scoped.engine.begin() as conn: table.create(conn, checkfirst=True) - with so.Session(scoped.engine) as session: - yield scoped, session + with so.Session(scoped.engine) as sqlite_session: + yield scoped, sqlite_session diff --git a/tests/loaders/test_dedupe.py b/tests/loaders/test_dedupe.py index 34a6c31..af46839 100644 --- a/tests/loaders/test_dedupe.py +++ b/tests/loaders/test_dedupe.py @@ -38,8 +38,8 @@ def test_arrow_drop_duplicates_simple(): -def test_internal_deduplication(session, engine, tmp_path): - Base.metadata.create_all(engine) +def test_internal_deduplication(sqlite_session, sqlite_engine, tmp_path): + Base.metadata.create_all(sqlite_engine) csv = tmp_path / "dedup_table.csv" pd.DataFrame( @@ -51,11 +51,11 @@ def test_internal_deduplication(session, engine, tmp_path): ).to_csv(csv, index=False) inserted = _DedupTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), dedupe=True, ) - session.commit() + sqlite_session.commit() assert inserted == 2 \ No newline at end of file diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 6275c1a..16790bd 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -10,6 +10,7 @@ from oa_configurator import Role as SchemaRole from orm_loader.backends import resolve_backend +from orm_loader.backends.base import DeferredForeignKeyViolationError from orm_loader.helpers import IngestError from orm_loader.loaders.data.converters import _COLUMN_CAST_RULES, register_column_cast_rule from orm_loader.loaders.data_classes import _clean_nulls @@ -19,6 +20,7 @@ from tests.models import ( Base, CompositeTable, + CycleATable, EnumTable, Flag, ImpliedEnumTable, @@ -36,6 +38,7 @@ _EnumTable = cast(Type[CSVTableProtocol], EnumTable) _ImpliedEnumTable = cast(Type[CSVTableProtocol], ImpliedEnumTable) _VocabSchemaTable = cast(Type[CSVTableProtocol], VocabSchemaTable) +_CycleATable = cast(Type[CSVTableProtocol], CycleATable) @pytest.fixture(autouse=True) @@ -45,7 +48,7 @@ def _clear_column_cast_rules(): _COLUMN_CAST_RULES.clear() -def test_initial_csv_load(session, tmp_path): +def test_initial_csv_load(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -59,16 +62,16 @@ def test_initial_csv_load(session, tmp_path): loader = PandasLoader() inserted = _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, ) - session.commit() + sqlite_session.commit() assert inserted == 3 - rows = session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() + rows = sqlite_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() assert [(r.id, r.name) for r in rows] == [ (1, "alpha"), @@ -77,7 +80,7 @@ def test_initial_csv_load(session, tmp_path): ] -def test_initial_csv_load_for_a_non_primary_schema_table(session, tmp_path): +def test_initial_csv_load_for_a_non_primary_schema_table(sqlite_session, tmp_path): """SQLite has no real schema concept, so every Role folds to None on this connection (see oa_configurator's SQLiteTestStrategy). A VOCAB-tagged table's load path must not error out just because @@ -91,20 +94,20 @@ def test_initial_csv_load_for_a_non_primary_schema_table(session, tmp_path): ).to_csv(csv_path, index=False, sep="\t") inserted = _VocabSchemaTable.load_csv( - session, csv_path, dedupe=False, loader=PandasLoader() + sqlite_session, csv_path, dedupe=False, loader=PandasLoader() ) - session.commit() + sqlite_session.commit() assert inserted == 2 - rows = session.execute( + rows = sqlite_session.execute( sa.select(VocabSchemaTable).order_by(VocabSchemaTable.id) ).scalars().all() assert [(r.id, r.name) for r in rows] == [(1, "alpha"), (2, "beta")] -def test_replace_merge_strategy(session, tmp_path): +def test_replace_merge_strategy(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" # Initial load @@ -119,12 +122,12 @@ def test_replace_merge_strategy(session, tmp_path): loader = PandasLoader() _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, ) - session.commit() + sqlite_session.commit() # Replace rows 2 and 3 pd.DataFrame( @@ -135,17 +138,17 @@ def test_replace_merge_strategy(session, tmp_path): ).to_csv(csv_path, index=False, sep="\t") replaced = _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, merge_strategy="replace", ) - session.commit() + sqlite_session.commit() assert replaced == 2 - rows = session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() + rows = sqlite_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() assert [(r.id, r.name) for r in rows] == [ (1, "alpha"), @@ -154,27 +157,27 @@ def test_replace_merge_strategy(session, tmp_path): ] -def test_empty_csv_is_noop(session, tmp_path): +def test_empty_csv_is_noop(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" csv_path.touch() loader = PandasLoader() inserted = _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, ) - session.commit() + sqlite_session.commit() assert inserted == 0 - rows = session.execute(sa.select(SimpleTable)).scalars().all() + rows = sqlite_session.execute(sa.select(SimpleTable)).scalars().all() assert rows == [] -def test_required_column_violation_drops_rows(session, tmp_path): +def test_required_column_violation_drops_rows(sqlite_session, tmp_path): csv = tmp_path / "required_table.csv" pd.DataFrame( [ @@ -184,17 +187,17 @@ def test_required_column_violation_drops_rows(session, tmp_path): ).to_csv(csv, index=False) inserted = _RequiredTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), dedupe=False, ) - session.commit() + sqlite_session.commit() assert inserted == 1 -def test_required_column_entirely_missing_raises(session, tmp_path): +def test_required_column_entirely_missing_raises(sqlite_session, tmp_path): # Distinct from test_required_column_violation_drops_rows above: here # the required column isn't present with a bad value, it's absent from # the source file entirely. No merge strategy supports a partial-column @@ -205,14 +208,14 @@ def test_required_column_entirely_missing_raises(session, tmp_path): with pytest.raises(IngestError, match="required_table.*name"): _RequiredTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), dedupe=False, ) -def test_enum_column_cast_rule_resolves_known_and_drops_unknown(session, tmp_path): +def test_enum_column_cast_rule_resolves_known_and_drops_unknown(sqlite_session, tmp_path): # sa.Enum is itself a subclass of sa.String, so without a registered # column-specific rule this would just pass every raw value through # untouched via the existing String rule -- no membership check at all. @@ -227,15 +230,15 @@ def test_enum_column_cast_rule_resolves_known_and_drops_unknown(session, tmp_pat ] ).to_csv(csv, index=False) - total = _EnumTable.load_csv(session, csv, loader=PandasLoader(), dedupe=False) - session.commit() + total = _EnumTable.load_csv(sqlite_session, csv, loader=PandasLoader(), dedupe=False) + sqlite_session.commit() assert total == 3 - rows = dict(session.execute(sa.select(EnumTable.id, EnumTable.role)).all()) + rows = dict(sqlite_session.execute(sa.select(EnumTable.id, EnumTable.role)).all()) assert rows == {1: Role.FIRST_AUTHOR, 2: None, 3: Role.LAST_AUTHOR} -def test_implied_enum_on_plain_string_column(session, tmp_path): +def test_implied_enum_on_plain_string_column(sqlite_session, tmp_path): # The OMOP CDM concept.standard_concept/invalid_reason shape: a plain # String(1) column with no type-level enum signal at all -- sa.Enum-keyed # dispatch could never reach this, only a per-column rule can. This column @@ -261,15 +264,15 @@ def _to_flag_value(raw): ] ).to_csv(csv, index=False) - total = _ImpliedEnumTable.load_csv(session, csv, loader=PandasLoader(), dedupe=False) - session.commit() + total = _ImpliedEnumTable.load_csv(sqlite_session, csv, loader=PandasLoader(), dedupe=False) + sqlite_session.commit() assert total == 3 - rows = dict(session.execute(sa.select(ImpliedEnumTable.id, ImpliedEnumTable.flag)).all()) + rows = dict(sqlite_session.execute(sa.select(ImpliedEnumTable.id, ImpliedEnumTable.flag)).all()) assert rows == {1: "S", 2: None, 3: "C"} -def test_composite_pk_dedup(session, tmp_path): +def test_composite_pk_dedup(sqlite_session, tmp_path): csv = tmp_path / "composite_table.csv" pd.DataFrame( [ @@ -280,12 +283,12 @@ def test_composite_pk_dedup(session, tmp_path): ).to_csv(csv, index=False) inserted = _CompositeTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), dedupe=True, ) - session.commit() + sqlite_session.commit() assert inserted == 2 @@ -313,7 +316,7 @@ def test_composite_pk_dedup(session, tmp_path): ), ], ) -def test_merge_strategies(session, tmp_path, merge_strategy, expected_rows, expected_inserted): +def test_merge_strategies(sqlite_session, tmp_path, merge_strategy, expected_rows, expected_inserted): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -325,8 +328,8 @@ def test_merge_strategies(session, tmp_path, merge_strategy, expected_rows, expe ).to_csv(csv_path, index=False, sep="\t") loader = PandasLoader() - SimpleTable.load_csv(session, csv_path, dedupe=False, loader=loader) - session.commit() + SimpleTable.load_csv(sqlite_session, csv_path, dedupe=False, loader=loader) + sqlite_session.commit() pd.DataFrame( [ @@ -336,18 +339,18 @@ def test_merge_strategies(session, tmp_path, merge_strategy, expected_rows, expe ).to_csv(csv_path, index=False, sep="\t") inserted = SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, merge_strategy=merge_strategy, ) - session.commit() + sqlite_session.commit() assert inserted == expected_inserted rows = ( - session.execute(sa.select(SimpleTable).order_by(SimpleTable.id, SimpleTable.name)) + sqlite_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id, SimpleTable.name)) .scalars() .all() ) @@ -355,7 +358,7 @@ def test_merge_strategies(session, tmp_path, merge_strategy, expected_rows, expe assert [(r.id, r.name) for r in rows] == expected_rows -def test_insert_if_empty_merge_strategy(session, tmp_path): +def test_insert_if_empty_merge_strategy(sqlite_session, tmp_path): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -367,17 +370,17 @@ def test_insert_if_empty_merge_strategy(session, tmp_path): loader = PandasLoader() inserted = _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, merge_strategy="insert_if_empty", ) - session.commit() + sqlite_session.commit() assert inserted == 2 - rows = session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() + rows = sqlite_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() assert [(r.id, r.name) for r in rows] == [ (1, "alpha"), @@ -385,33 +388,33 @@ def test_insert_if_empty_merge_strategy(session, tmp_path): ] -def test_insert_if_empty_raises_on_non_empty_target(session, engine, tmp_path): +def test_insert_if_empty_raises_on_non_empty_target(sqlite_session, sqlite_engine, tmp_path): csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") loader = PandasLoader() - _SimpleTable.load_csv(session, csv_path, dedupe=False, loader=loader) - session.commit() + _SimpleTable.load_csv(sqlite_session, csv_path, dedupe=False, loader=loader) + sqlite_session.commit() pd.DataFrame([{"id": 2, "name": "beta"}]).to_csv(csv_path, index=False, sep="\t") with pytest.raises(ValueError, match="is not empty; cannot use merge strategy 'insert_if_empty'"): _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, merge_strategy="insert_if_empty", ) - inspector = sa.inspect(engine) - staging_name = resolve_backend(engine).staging_name_for_table(SimpleTable.__tablename__) + inspector = sa.inspect(sqlite_engine) + staging_name = resolve_backend(sqlite_engine).staging_name_for_table(SimpleTable.__tablename__) assert not inspector.has_table(staging_name) @pytest.mark.parametrize("merge_strategy", ["replace", "upsert"]) -def test_empty_target_routes_merge_to_insert_if_empty(session, tmp_path, caplog, merge_strategy): +def test_empty_target_routes_merge_to_insert_if_empty(sqlite_session, tmp_path, caplog, merge_strategy): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -424,17 +427,17 @@ def test_empty_target_routes_merge_to_insert_if_empty(session, tmp_path, caplog, caplog.set_level(logging.INFO, logger="orm_loader.tables.loadable_table") inserted = _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=PandasLoader(), merge_strategy=merge_strategy, ) - session.commit() + sqlite_session.commit() assert inserted == 2 - rows = session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() + rows = sqlite_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() assert [(r.id, r.name) for r in rows] == [ (1, "alpha"), (2, "beta"), @@ -454,25 +457,25 @@ def test_empty_target_routes_merge_to_insert_if_empty(session, tmp_path, caplog, assert not any("Merge upsert phase starting." in message for message in messages) -def test_staging_table_is_created_and_dropped(session, engine, tmp_path): +def test_staging_table_is_created_and_dropped(sqlite_session, sqlite_engine, tmp_path): csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False) _SimpleTable.load_csv( - session, + sqlite_session, csv_path, loader=PandasLoader(), dedupe=False, ) - session.commit() + sqlite_session.commit() - inspector = sa.inspect(engine) - staging_name = resolve_backend(engine).staging_name_for_table(SimpleTable.__tablename__) + inspector = sa.inspect(sqlite_engine) + staging_name = resolve_backend(sqlite_engine).staging_name_for_table(SimpleTable.__tablename__) assert not inspector.has_table(staging_name) -def test_composite_pk_replace_merge(session, tmp_path): +def test_composite_pk_replace_merge(sqlite_session, tmp_path): csv = tmp_path / "composite_table.csv" pd.DataFrame( @@ -482,8 +485,8 @@ def test_composite_pk_replace_merge(session, tmp_path): ] ).to_csv(csv, index=False) - CompositeTable.load_csv(session, csv, loader=PandasLoader()) - session.commit() + CompositeTable.load_csv(sqlite_session, csv, loader=PandasLoader()) + sqlite_session.commit() pd.DataFrame( [ @@ -492,15 +495,15 @@ def test_composite_pk_replace_merge(session, tmp_path): ).to_csv(csv, index=False) CompositeTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), merge_strategy="replace", ) - session.commit() + sqlite_session.commit() rows = ( - session.execute(sa.select(CompositeTable).order_by(CompositeTable.a, CompositeTable.b)) + sqlite_session.execute(sa.select(CompositeTable).order_by(CompositeTable.a, CompositeTable.b)) .scalars() .all() ) @@ -511,13 +514,13 @@ def test_composite_pk_replace_merge(session, tmp_path): ] -def test_filename_must_match_tablename(session, tmp_path): +def test_filename_must_match_tablename(sqlite_session, tmp_path): csv = tmp_path / "wrong_name.csv" pd.DataFrame([{"id": 1, "name": "x"}]).to_csv(csv, index=False) with pytest.raises(ValueError, match="does not match table"): SimpleTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), ) @@ -537,7 +540,7 @@ def test_clean_nulls_passthrough(): assert _clean_nulls("S") == "S" -def test_nullable_column_with_nan_does_not_crash(session, engine, tmp_path): +def test_nullable_column_with_nan_does_not_crash(sqlite_session, sqlite_engine, tmp_path): class NullableTable(Base, CSVLoadableTableInterface): __tablename__ = "nullable_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -545,7 +548,7 @@ class NullableTable(Base, CSVLoadableTableInterface): id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) flag: so.Mapped[str | None] = so.mapped_column(sa.String, nullable=True) - Base.metadata.create_all(engine) + Base.metadata.create_all(sqlite_engine) _NullableTable = cast(Type[CSVTableProtocol], NullableTable) csv = tmp_path / "nullable_table.csv" @@ -557,16 +560,16 @@ class NullableTable(Base, CSVLoadableTableInterface): ).to_csv(csv, index=False) inserted = _NullableTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), dedupe=False, ) - session.commit() + sqlite_session.commit() assert inserted == 2 - rows = session.execute(sa.select(NullableTable).order_by(NullableTable.id)).scalars().all() + rows = sqlite_session.execute(sa.select(NullableTable).order_by(NullableTable.id)).scalars().all() assert [(r.id, r.flag) for r in rows] == [ (1, "S"), @@ -574,7 +577,7 @@ class NullableTable(Base, CSVLoadableTableInterface): ] -def test_embedded_newline_in_field_is_preserved(session, engine, tmp_path): +def test_embedded_newline_in_field_is_preserved(sqlite_session, sqlite_engine, tmp_path): class TextTable(Base, CSVLoadableTableInterface): __tablename__ = "text_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -582,7 +585,7 @@ class TextTable(Base, CSVLoadableTableInterface): id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) name: so.Mapped[str] = so.mapped_column(sa.String) - Base.metadata.create_all(engine) + Base.metadata.create_all(sqlite_engine) _TextTable = cast(Type[CSVTableProtocol], TextTable) csv = tmp_path / "text_table.csv" @@ -591,18 +594,18 @@ class TextTable(Base, CSVLoadableTableInterface): csv.write_text('id\tname\n1\t"hello\nworld"\n') _TextTable.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), dedupe=False, ) - session.commit() + sqlite_session.commit() - rows = session.execute(sa.select(TextTable)).scalars().all() + rows = sqlite_session.execute(sa.select(TextTable)).scalars().all() assert rows[0].name == "hello\nworld" -def test_embedded_tab_in_field(session, engine, tmp_path): +def test_embedded_tab_in_field(sqlite_session, sqlite_engine, tmp_path): class TextTable2(Base, CSVLoadableTableInterface): __tablename__ = "tab_table" __table_args__ = {"schema": SchemaRole.PRIMARY.value} @@ -610,28 +613,28 @@ class TextTable2(Base, CSVLoadableTableInterface): id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) name: so.Mapped[str] = so.mapped_column(sa.String) - Base.metadata.create_all(engine) + Base.metadata.create_all(sqlite_engine) _TextTable2 = cast(Type[CSVTableProtocol], TextTable2) csv = tmp_path / "tab_table.csv" csv.write_text('id\tname\n1\t"foo\tbar"\n') _TextTable2.load_csv( - session, + sqlite_session, csv, loader=PandasLoader(), dedupe=False, ) - session.commit() + sqlite_session.commit() - rows = session.execute(sa.select(TextTable2)).scalars().all() + rows = sqlite_session.execute(sa.select(TextTable2)).scalars().all() assert rows[0].name == "foo\tbar" # --- index_strategy tests --- -def _make_ddl_tracker(engine): +def _make_ddl_tracker(sqlite_engine): """Return a list that is populated with DROP/CREATE INDEX statements as they execute.""" ddl_log: list[str] = [] @@ -641,62 +644,62 @@ def _capture(*args): if upper.startswith("DROP INDEX") or upper.startswith("CREATE INDEX"): ddl_log.append(statement.strip()) - sae.listen(engine, "before_cursor_execute", _capture) + sae.listen(sqlite_engine, "before_cursor_execute", _capture) return ddl_log -def test_auto_strategy_keeps_indices_on_sqlite(session, engine, tmp_path): +def test_auto_strategy_keeps_indices_on_sqlite(sqlite_session, sqlite_engine, tmp_path): """On SQLite, 'auto' resolves to 'keep' — no index DDL should be emitted.""" - ddl_log = _make_ddl_tracker(engine) + ddl_log = _make_ddl_tracker(sqlite_engine) csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]).to_csv( csv_path, index=False, sep="\t" ) - _SimpleTable.load_csv(session, csv_path, loader=PandasLoader(), index_strategy="auto") - session.commit() + _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="auto") + sqlite_session.commit() assert not any("DROP INDEX" in s.upper() for s in ddl_log) assert not any("CREATE INDEX" in s.upper() for s in ddl_log) - inspector = sa.inspect(engine) + inspector = sa.inspect(sqlite_engine) inspector.clear_cache() assert "ix_test_table_name" in {idx["name"] for idx in inspector.get_indexes("test_table")} -def test_explicit_keep_preserves_indices(session, engine, tmp_path): +def test_explicit_keep_preserves_indices(sqlite_session, sqlite_engine, tmp_path): """Explicit 'keep' emits no index DDL regardless of dialect.""" - ddl_log = _make_ddl_tracker(engine) + ddl_log = _make_ddl_tracker(sqlite_engine) csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") - _SimpleTable.load_csv(session, csv_path, loader=PandasLoader(), index_strategy="keep") - session.commit() + _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="keep") + sqlite_session.commit() assert not any("DROP INDEX" in s.upper() for s in ddl_log) - inspector = sa.inspect(engine) + inspector = sa.inspect(sqlite_engine) inspector.clear_cache() assert "ix_test_table_name" in {idx["name"] for idx in inspector.get_indexes("test_table")} -def test_explicit_drop_rebuild_on_sqlite_restores_index(session, engine, tmp_path): +def test_explicit_drop_rebuild_on_sqlite_restores_index(sqlite_session, sqlite_engine, tmp_path): """Explicit 'drop_rebuild' drops then restores the index even on SQLite.""" - ddl_log = _make_ddl_tracker(engine) + ddl_log = _make_ddl_tracker(sqlite_engine) csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]).to_csv( csv_path, index=False, sep="\t" ) - _SimpleTable.load_csv(session, csv_path, loader=PandasLoader(), index_strategy="drop_rebuild") - session.commit() + _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="drop_rebuild") + sqlite_session.commit() assert any("DROP INDEX" in s.upper() for s in ddl_log) assert any("CREATE INDEX" in s.upper() for s in ddl_log) - inspector = sa.inspect(engine) + inspector = sa.inspect(sqlite_engine) inspector.clear_cache() assert "ix_test_table_name" in {idx["name"] for idx in inspector.get_indexes("test_table")} -def test_drop_rebuild_logging_shows_merge_phases(session, tmp_path, caplog): +def test_drop_rebuild_logging_shows_merge_phases(sqlite_session, tmp_path, caplog): csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -707,8 +710,8 @@ def test_drop_rebuild_logging_shows_merge_phases(session, tmp_path, caplog): ).to_csv(csv_path, index=False, sep="\t") loader = PandasLoader() - _SimpleTable.load_csv(session, csv_path, dedupe=False, loader=loader) - session.commit() + _SimpleTable.load_csv(sqlite_session, csv_path, dedupe=False, loader=loader) + sqlite_session.commit() pd.DataFrame( [ @@ -720,14 +723,14 @@ def test_drop_rebuild_logging_shows_merge_phases(session, tmp_path, caplog): caplog.set_level(logging.INFO, logger="orm_loader.tables.loadable_table") _SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=loader, merge_strategy="replace", index_strategy="drop_rebuild", ) - session.commit() + sqlite_session.commit() messages = [record.getMessage() for record in caplog.records] @@ -758,13 +761,13 @@ def test_drop_rebuild_logging_shows_merge_phases(session, tmp_path, caplog): assert any("Index verification/rebuild completed in " in message for message in messages) -def test_invalid_index_strategy_raises(session, tmp_path): +def test_invalid_index_strategy_raises(sqlite_session, tmp_path): """An unrecognised strategy value raises ValueError before any DB work.""" csv_path = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv_path, index=False, sep="\t") with pytest.raises(ValueError, match="Unknown index_strategy"): - _SimpleTable.load_csv(session, csv_path, loader=PandasLoader(), index_strategy="not-valid") + _SimpleTable.load_csv(sqlite_session, csv_path, loader=PandasLoader(), index_strategy="not-valid") # from hypothesis import given, strategies as st @@ -791,17 +794,17 @@ def test_invalid_index_strategy_raises(session, tmp_path): # id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) # txt: so.Mapped[str] = so.mapped_column(sa.String) -# engine = sa.create_engine("sqlite:///:memory:", future=True) -# BaseLocal.metadata.create_all(engine) +# sqlite_engine = sa.create_engine("sqlite:///:memory:", future=True) +# BaseLocal.metadata.create_all(sqlite_engine) -# with Session(engine) as session: +# with Session(sqlite_engine) as sqlite_session: # csv = tmp_path / "fuzz_table.csv" # pd.DataFrame([{"id": 1, "txt": s}]).to_csv(csv, index=False, sep="\t") -# FuzzTable.load_csv(session, csv, loader=PandasLoader(), dedupe=False) -# session.commit() +# FuzzTable.load_csv(sqlite_session, csv, loader=PandasLoader(), dedupe=False) +# sqlite_session.commit() -# rows = session.execute(sa.select(FuzzTable)).scalars().all() +# rows = sqlite_session.execute(sa.select(FuzzTable)).scalars().all() # # What the loader *should* produce according to your spec # expected = perform_cast(s, sa.String(), on_error=None) @@ -813,9 +816,9 @@ def test_invalid_index_strategy_raises(session, tmp_path): # assert rows[0].txt.encode("utf-8", errors="replace") == s.encode("utf-8", errors="replace") -def test_sqlite_failed_batched_replace_keeps_original_target_rows(session, tmp_path): - session.add_all([SimpleTable(id=1, name="original-1"), SimpleTable(id=2, name="original-2")]) - session.commit() +def test_sqlite_failed_batched_replace_keeps_original_target_rows(sqlite_session, tmp_path): + sqlite_session.add_all([SimpleTable(id=1, name="original-1"), SimpleTable(id=2, name="original-2")]) + sqlite_session.commit() csv_path = tmp_path / "test_table.csv" pd.DataFrame( [ @@ -827,14 +830,32 @@ def test_sqlite_failed_batched_replace_keeps_original_target_rows(session, tmp_p with pytest.raises(sa.exc.IntegrityError): SimpleTable.load_csv( - session, + sqlite_session, csv_path, dedupe=False, loader=PandasLoader(), merge_strategy="replace", merge_batch_size=2, ) - session.rollback() + sqlite_session.rollback() - rows = session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() + rows = sqlite_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() assert [(row.id, row.name) for row in rows] == [(1, "original-1"), (2, "original-2")] + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_session, tmp_path): + """cycle_a_table.b_id is part of a genuine FK cycle with cycle_b_table, so + DuckDB never got a real constraint for it (oa-configurator's + _skip_deferred_fk_on_duckdb no-oped the ALTER TABLE ADD CONSTRAINT): a + load introducing an orphaned reference succeeds silently unless + validate_deferred_fks=True catches it. + """ + csv_path = tmp_path / "cycle_a_table.csv" + pd.DataFrame([{"id": 1, "b_id": 999}]).to_csv(csv_path, index=False, sep="\t") + + with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): + _CycleATable.load_csv( + duckdb_session, csv_path, dedupe=False, loader=PandasLoader(), validate_deferred_fks=True + ) diff --git a/tests/loaders/test_pandas_loader.py b/tests/loaders/test_pandas_loader.py index 3cc0e0a..b85e3c7 100644 --- a/tests/loaders/test_pandas_loader.py +++ b/tests/loaders/test_pandas_loader.py @@ -5,11 +5,11 @@ from orm_loader.loaders.loader_interface import PandasLoader from tests.models import PandasLoaderTable -def test_loader_context_fields(session): +def test_loader_context_fields(sqlite_session): ctx = LoaderContext( tableclass=PandasLoaderTable, - session=session, + session=sqlite_session, path=Path("file.csv"), staging_table=PandasLoaderTable.__table__, chunksize=100, @@ -21,7 +21,7 @@ def test_loader_context_fields(session): assert ctx.normalise is False assert ctx.dedupe is True -def test_pandas_loader_case_insensitive_headers(tmp_path, session): +def test_pandas_loader_case_insensitive_headers(tmp_path, sqlite_session): """ PandasLoader should accept mixed-case CSV headers and map them correctly to lowercase model columns. @@ -36,7 +36,7 @@ def test_pandas_loader_case_insensitive_headers(tmp_path, session): ctx = LoaderContext( tableclass=PandasLoaderTable, - session=session, + session=sqlite_session, path=csv, staging_table=PandasLoaderTable.__table__, chunksize=100, @@ -46,7 +46,7 @@ def test_pandas_loader_case_insensitive_headers(tmp_path, session): n = PandasLoader.orm_file_load(ctx) - rows = session.execute( + rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') ).all() @@ -54,7 +54,7 @@ def test_pandas_loader_case_insensitive_headers(tmp_path, session): assert rows == [(1, "alpha"), (2, "beta")] -def test_pandas_loader_literal_mode_preserves_quotes(tmp_path, session): +def test_pandas_loader_literal_mode_preserves_quotes(tmp_path, sqlite_session): """ The pandas fallback must honour quote_mode. Under 'literal' it uses QUOTE_NONE, so double-quotes are literal data and are preserved verbatim @@ -70,7 +70,7 @@ def test_pandas_loader_literal_mode_preserves_quotes(tmp_path, session): ctx = LoaderContext( tableclass=PandasLoaderTable, - session=session, + session=sqlite_session, path=csv, staging_table=PandasLoaderTable.__table__, chunksize=100, @@ -81,7 +81,7 @@ def test_pandas_loader_literal_mode_preserves_quotes(tmp_path, session): n = PandasLoader.orm_file_load(ctx) - rows = session.execute( + rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') ).all() @@ -89,7 +89,7 @@ def test_pandas_loader_literal_mode_preserves_quotes(tmp_path, session): assert rows == [(1, '"hello"'), (2, '"world"')] -def test_pandas_loader_csv_mode_strips_quotes(tmp_path, session): +def test_pandas_loader_csv_mode_strips_quotes(tmp_path, sqlite_session): """ Contrast with the literal case above: the same file under 'csv' mode uses QUOTE_MINIMAL, so the wrapping quotes are consumed. Proves quote_mode @@ -105,7 +105,7 @@ def test_pandas_loader_csv_mode_strips_quotes(tmp_path, session): ctx = LoaderContext( tableclass=PandasLoaderTable, - session=session, + session=sqlite_session, path=csv, staging_table=PandasLoaderTable.__table__, chunksize=100, @@ -116,7 +116,7 @@ def test_pandas_loader_csv_mode_strips_quotes(tmp_path, session): n = PandasLoader.orm_file_load(ctx) - rows = session.execute( + rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') ).all() @@ -124,7 +124,7 @@ def test_pandas_loader_csv_mode_strips_quotes(tmp_path, session): assert rows == [(1, "hello"), (2, "world")] -def test_pandas_loader_by_delimiter_tab_preserves_quotes(tmp_path, session): +def test_pandas_loader_by_delimiter_tab_preserves_quotes(tmp_path, sqlite_session): """ 'by_delimiter' derives the mode from the delimiter alone: a tab-delimited file resolves to 'literal', so embedded double-quotes are kept as data. @@ -139,7 +139,7 @@ def test_pandas_loader_by_delimiter_tab_preserves_quotes(tmp_path, session): ctx = LoaderContext( tableclass=PandasLoaderTable, - session=session, + session=sqlite_session, path=tsv, staging_table=PandasLoaderTable.__table__, chunksize=100, @@ -150,7 +150,7 @@ def test_pandas_loader_by_delimiter_tab_preserves_quotes(tmp_path, session): n = PandasLoader.orm_file_load(ctx) - rows = session.execute( + rows = sqlite_session.execute( sa.text('SELECT id, value FROM test_pandas_loader ORDER BY id') ).all() @@ -158,7 +158,7 @@ def test_pandas_loader_by_delimiter_tab_preserves_quotes(tmp_path, session): assert rows == [(1, '"hi"'), (2, "bye")] -def test_pandas_dedupe_internal(session, tmp_path): +def test_pandas_dedupe_internal(sqlite_session, tmp_path): df = pd.DataFrame({ "id": [1, 1, 2], "value": [10, 10, 20], @@ -166,7 +166,7 @@ def test_pandas_dedupe_internal(session, tmp_path): ctx = LoaderContext( tableclass=PandasLoaderTable, - session=session, + session=sqlite_session, path=Path("file.csv"), staging_table=PandasLoaderTable.__table__, chunksize=100, diff --git a/tests/loaders/test_parquet_loader.py b/tests/loaders/test_parquet_loader.py index 96b00fe..17d93ef 100644 --- a/tests/loaders/test_parquet_loader.py +++ b/tests/loaders/test_parquet_loader.py @@ -27,8 +27,8 @@ class ParquetTable(Base, CSVLoadableTableInterface): _ParquetTable = cast(Type[CSVTableProtocol], ParquetTable) -def test_parquet_loader(session, engine, tmp_path): - Base.metadata.create_all(engine) +def test_parquet_loader(sqlite_session, sqlite_engine, tmp_path): + Base.metadata.create_all(sqlite_engine) df = pd.DataFrame( [ @@ -41,12 +41,12 @@ def test_parquet_loader(session, engine, tmp_path): pq.write_table(table, path) inserted = _ParquetTable.load_csv( - session, + sqlite_session, path, loader=ParquetLoader(), dedupe=False, ) - session.commit() + sqlite_session.commit() assert inserted == 2 @@ -68,26 +68,26 @@ class NullSentinelTable(Base, CSVLoadableTableInterface): _SENTINELS = ["NULL", "null", " NULL ", "n/a", "N/A", "NaN", "none", "None", " ", ""] -def test_arrow_path_normalises_text_null_sentinels(session, engine, tmp_path): +def test_arrow_path_normalises_text_null_sentinels(sqlite_session, sqlite_engine, tmp_path): # cast_scalar normalises these before dispatching; without the same # treatment here the plain String branch wrote the literal text into # staging, so the same file loaded via the two loaders disagreed. - Base.metadata.create_all(engine) + Base.metadata.create_all(sqlite_engine) rows = [{"id": i, "name": s, "role": None} for i, s in enumerate(_SENTINELS, start=1)] rows.append({"id": 99, "name": " keep me ", "role": None}) path = tmp_path / "null_sentinel_table.parquet" pq.write_table(pa.Table.from_pandas(pd.DataFrame(rows)), path) - _NullSentinelTable.load_csv(session, path, loader=ParquetLoader(), dedupe=False) - session.commit() + _NullSentinelTable.load_csv(sqlite_session, path, loader=ParquetLoader(), dedupe=False) + sqlite_session.commit() - stored = dict(session.execute(sa.select(NullSentinelTable.id, NullSentinelTable.name)).all()) + stored = dict(sqlite_session.execute(sa.select(NullSentinelTable.id, NullSentinelTable.name)).all()) assert [stored[i] for i in range(1, len(_SENTINELS) + 1)] == [None] * len(_SENTINELS) assert stored[99] == "keep me" -def test_arrow_and_pandas_paths_agree_on_null_sentinels(session, engine, tmp_path): +def test_arrow_and_pandas_paths_agree_on_null_sentinels(sqlite_session, sqlite_engine, tmp_path): from orm_loader.loaders.data.converters import cast_arrow_column, cast_scalar col = NullSentinelTable.__table__.c.name @@ -97,7 +97,7 @@ def test_arrow_and_pandas_paths_agree_on_null_sentinels(session, engine, tmp_pat assert arrow == scalar, f"{raw!r}: arrow={arrow!r} scalar={scalar!r}" -def test_arrow_column_rule_does_not_count_nulls_as_cast_failures(session, engine, tmp_path): +def test_arrow_column_rule_does_not_count_nulls_as_cast_failures(sqlite_session, sqlite_engine, tmp_path): # The override branch resolved these to None either way, but recorded each # one as a failure -- warning about rows that are simply null. from enum import Enum @@ -128,7 +128,7 @@ class Role(str, Enum): _COLUMN_CAST_RULES.clear() -def test_arrow_path_nulls_float_nan(session, engine, tmp_path): +def test_arrow_path_nulls_float_nan(sqlite_session, sqlite_engine, tmp_path): from orm_loader.loaders.data.converters import cast_arrow_column class FloatTable(Base, CSVLoadableTableInterface): diff --git a/tests/loaders/test_split_connection.py b/tests/loaders/test_split_connection.py index 36c1fec..3f183df 100644 --- a/tests/loaders/test_split_connection.py +++ b/tests/loaders/test_split_connection.py @@ -47,9 +47,9 @@ def get_bind(self, mapper=None, clause=None, **kwargs): if mapper is None and clause is None: raise UnboundExecutionError("a mapper or clause is required") if mapper is not None: - engine = vocab_engine if mapper.__table__.schema == SchemaRole.VOCAB else primary_engine - self._last_mapper_engine = engine - return engine + sqlite_engine = vocab_engine if mapper.__table__.schema == SchemaRole.VOCAB else primary_engine + self._last_mapper_engine = sqlite_engine + return sqlite_engine table = getattr(clause, "table", None) if table is not None and table.schema in {SchemaRole.PRIMARY, SchemaRole.VOCAB}: return vocab_engine if table.schema == SchemaRole.VOCAB else primary_engine @@ -57,19 +57,19 @@ def get_bind(self, mapper=None, clause=None, **kwargs): csv_path = tmp_path / "test_vocab_role_table.csv" pd.DataFrame([{"id": 7, "name": "vocab-row"}]).to_csv(csv_path, index=False, sep="\t") - session = MapperRoutedSession() + sqlite_session = MapperRoutedSession() try: VocabSchemaTable.load_csv( - session, + sqlite_session, csv_path, merge_strategy="insert_if_empty", dedupe=False, loader=PandasLoader(), staging_schema_tag=staging_schema_claim().schema_tag, ) - session.commit() + sqlite_session.commit() finally: - session.close() + sqlite_session.close() with vocab_engine.connect() as connection: rows = connection.execute( @@ -81,9 +81,9 @@ def get_bind(self, mapper=None, clause=None, **kwargs): vocab_engine.dispose() -def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_path): - """``pg_db`` (real Postgres, primary connection) and ``session`` (real - SQLite, vocab connection, via the module-level ``engine``/``session`` +def test_load_csv_across_two_genuinely_separate_connections(pg_db, sqlite_session, tmp_path): + """``pg_db`` (real Postgres, primary connection) and ``sqlite_session`` (real + SQLite, vocab connection, via the module-level ``sqlite_engine``/``sqlite_session`` fixtures) are two entirely different engines against two entirely different database systems.""" with schema_scoped_session(pg_db, SimpleTable.__table__, prefix="loader_split") as (scoped, primary_session): @@ -100,8 +100,8 @@ def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_ SimpleTable.load_csv(primary_session, primary_csv, dedupe=False, loader=PandasLoader()) primary_session.commit() - VocabSchemaTable.load_csv(session, vocab_csv, dedupe=False, loader=PandasLoader()) - session.commit() + VocabSchemaTable.load_csv(sqlite_session, vocab_csv, dedupe=False, loader=PandasLoader()) + sqlite_session.commit() with scoped.engine.connect() as conn: primary_rows = conn.execute( @@ -114,7 +114,7 @@ def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_ assert primary_rows == [(1, "primary-alpha")] assert leaked_vocab_table_in_pg is None - vocab_rows = session.execute( + vocab_rows = sqlite_session.execute( sa.select(VocabSchemaTable).order_by(VocabSchemaTable.id) ).scalars().all() assert [(r.id, r.name) for r in vocab_rows] == [(1, "vocab-alpha")] diff --git a/tests/models.py b/tests/models.py index efc3b17..130d45e 100644 --- a/tests/models.py +++ b/tests/models.py @@ -54,6 +54,29 @@ class CompositeTable(Base, CSVLoadableTableInterface): value: so.Mapped[str] = so.mapped_column(sa.String) +class CycleATable(Base, CSVLoadableTableInterface): + """Mutually-referencing pair with CycleBTable, for deferred-FK tests: + DuckDB can't execute the ALTER TABLE ADD CONSTRAINT this cycle forces.""" + + __tablename__ = "cycle_a_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + b_id: so.Mapped[int | None] = so.mapped_column( + sa.Integer, sa.ForeignKey(f"{SchemaRole.PRIMARY.value}.cycle_b_table.id") + ) + + +class CycleBTable(Base, CSVLoadableTableInterface): + __tablename__ = "cycle_b_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + a_id: so.Mapped[int | None] = so.mapped_column( + sa.Integer, sa.ForeignKey(f"{SchemaRole.PRIMARY.value}.cycle_a_table.id") + ) + + class EnumTable(Base, CSVLoadableTableInterface): """A genuine sa.Enum column, unmodified -- register_column_cast_rule's enum_type matches a raw value (e.g. "first author") by .value but stores @@ -80,7 +103,8 @@ class ComputedColumnTable(Base, CSVLoadableTableInterface): id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) name: so.Mapped[str] = so.mapped_column(sa.String) - slug: so.Mapped[str] = so.mapped_column(sa.String, sa.Computed("lower(name)")) + # nullable: DuckDB doesn't support any constraint (including NOT NULL) on a generated column. + slug: so.Mapped[str | None] = so.mapped_column(sa.String, sa.Computed("lower(name)")) class ImpliedEnumTable(Base, CSVLoadableTableInterface): From 061bbe53fe845c73eaf3cbe9d21621954dceefd3 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 14 Oct 2026 23:14:44 +0000 Subject: [PATCH 02/17] Batched loading, UPDATE-in-place in merge replace instead of delete and reinsert, update CI, smaller helpers --- .github/workflows/ci.yml | 2 + src/orm_loader/backends/base.py | 1 + src/orm_loader/backends/duckdb.py | 136 +++++++++++++++--------- src/orm_loader/backends/postgres.py | 1 + src/orm_loader/backends/sqlite.py | 1 + src/orm_loader/tables/loadable_table.py | 21 ++-- tests/backends/test_duckdb_backend.py | 44 +++++++- tests/loaders/test_loader_e2e.py | 26 +++++ 8 files changed, 171 insertions(+), 61 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7e7a7f4..63b393a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -8,6 +8,8 @@ jobs: uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/label-gate.yml@main build-test-sqlite: uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test.yml@main + build-test-duckdb: + uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-duckdb.yml@main build-test-postgres: uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres-v2.yml@main with: diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 0baf6c5..974f662 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -33,6 +33,7 @@ class BackendCapabilities: supports_unlogged_staging: bool = False supports_fk_toggle: bool = False supports_materialized_views: bool = False + supports_replace_with_kept_indexes: bool = False class DeferredForeignKeyViolationError(RuntimeError): diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 164a9e1..7f0018a 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -1,10 +1,12 @@ from __future__ import annotations +from collections.abc import Iterator from typing import TYPE_CHECKING, Any, Type import duckdb_sqlalchemy import sqlalchemy as sa import sqlalchemy.orm as so +from oa_configurator import qualified from sqlalchemy.sql.compiler import IdentifierPreparer from .base import BackendCapabilities, DatabaseBackend, Dialect @@ -37,6 +39,7 @@ def capabilities(self) -> BackendCapabilities: supports_unlogged_staging=False, supports_fk_toggle=False, supports_materialized_views=False, + supports_replace_with_kept_indexes=False, ) def deferred_fk_constraints( @@ -50,22 +53,44 @@ def deferred_fk_constraints( sort_tables_and_constraints needs every table in the schema to tell which FKs are genuinely deferred (cyclic), so this runs against table_cls's whole shared metadata, not table_cls alone. + + Only reads the ``tbl is None`` bucket of sort_tables_and_constraints's + output. A standalone self-referential FK (not part of a cross-table + cycle) is still deferred by SQLAlchemy either way, but lands in its + *own* table's bucket unless the column declares ``use_alter=True`` -- + without that, this method can't see it, and it will insert without + error instead of being caught by validate_deferred_fks(). Model + authors: mark a self-referential FK ``use_alter=True`` explicitly if + it needs anti-join validation on DuckDB. """ metadata = table_cls.__table__.metadata entries = sa.schema.sort_tables_and_constraints(metadata.tables.values()) deferred = next(fks for tbl, fks in entries if tbl is None) return tuple(fk for fk in deferred if fk.table is table_cls.__table__) + def _staging_rownum_sequence_ref(self, tablename: str) -> str: + """Fully schema-qualified, quoted name of tablename's staging _rownum sequence. + + Notes + ----- + DuckDB can't ALTER TABLE ADD COLUMN ... GENERATED ALWAYS AS IDENTITY after creation. + Pagniation uses a sequence-backed DEFAULT instead to achieve the Postgres-like + Identity column: a stable, monotonic pagination key populated automatically on insert, independent of + DuckDB's own rowid (not guaranteed stable across a checkpoint). + """ + seq_name = f"{self.staging_name_for_table(tablename)}_rownum_seq" + return qualified(self.identifier_preparer, seq_name, physical_schema=self.staging_schema) + def create_staging_table( self, table_cls: type["CSVTableProtocol"], session: so.Session, ) -> None: bind_arguments = {"mapper": table_cls} - session.execute( - sa.text(f'DROP TABLE IF EXISTS {self.qualified_staging_name(table_cls.__tablename__)};'), - bind_arguments=bind_arguments, - ) + staging_ref = self.qualified_staging_name(table_cls.__tablename__) + seq_ref = self._staging_rownum_sequence_ref(table_cls.__tablename__) + session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};'), bind_arguments=bind_arguments) + session.execute(sa.text(f'DROP SEQUENCE IF EXISTS {seq_ref};'), bind_arguments=bind_arguments) staging_name = self.staging_name_for_table(table_cls.__tablename__) metadata = sa.MetaData() @@ -77,6 +102,12 @@ def create_staging_table( metadata.create_all( bind=session.connection(bind_arguments=bind_arguments), tables=[staging_table] ) + + session.execute(sa.text(f'CREATE SEQUENCE {seq_ref};'), bind_arguments=bind_arguments) + session.execute( + sa.text(f"ALTER TABLE {staging_ref} ADD COLUMN _rownum BIGINT DEFAULT nextval('{seq_ref}');"), + bind_arguments=bind_arguments, + ) session.commit() def drop_staging_table( @@ -84,9 +115,14 @@ def drop_staging_table( table_cls: type["CSVTableProtocol"], session: so.Session, ) -> None: + bind_arguments = {"mapper": table_cls} session.execute( sa.text(f'DROP TABLE IF EXISTS {self.qualified_staging_name(table_cls.__tablename__)}'), - bind_arguments={"mapper": table_cls}, + bind_arguments=bind_arguments, + ) + session.execute( + sa.text(f'DROP SEQUENCE IF EXISTS {self._staging_rownum_sequence_ref(table_cls.__tablename__)};'), + bind_arguments=bind_arguments, ) @staticmethod @@ -116,25 +152,39 @@ def restore_fk_check( "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." ) - @staticmethod - def _staging_rowid() -> sa.ColumnElement[int]: - """DuckDB's own rowid: dense per table, but not guaranteed to start at 0 or 1.""" - return sa.literal_column("rowid") + def _staging_rownum_index( + self, table_cls: type["CSVTableProtocol"], staging: sa.Table, session: so.Session + ) -> None: + """Index staging's _rownum column for O(N log N) windowed range scans.""" + staging_name = self.staging_name_for_table(table_cls.__tablename__) + idx = sa.Index(f"{staging_name}_rownum_idx", staging.c._rownum) + idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) + session.commit() - def _staging_rowid_bounds(self, staging: sa.Table, session: so.Session) -> tuple[int, int] | None: - """Lowest and highest rowid in *staging*, or None when empty. + @staticmethod + def _staging_rownum_bound(staging: sa.Table, session: so.Session) -> int: + """Highest _rownum in staging, as the pagination upper bound (0 when empty).""" + return session.execute( + sa.select(sa.func.coalesce(sa.func.max(staging.c._rownum), 0)).select_from(staging) + ).scalar_one() - DuckDB's rowid isn't guaranteed to start at any fixed value (verified - live), unlike SQLite's, so pagination bounds against the real - min/max instead of assuming a 0-based range. + def _staging_rownum_batches( + self, + table_cls: type["CSVTableProtocol"], + staging: sa.Table, + session: so.Session, + merge_batch_size: int, + ) -> Iterator[sa.ColumnElement[bool]]: + """Yield successive _rownum window predicates covering every row in + staging, each at most merge_batch_size wide. """ - lowest, highest = session.execute( - sa.select(sa.func.min(self._staging_rowid()), sa.func.max(self._staging_rowid())) - .select_from(staging) - ).one() - if lowest is None: - return None - return lowest, highest + self._staging_rownum_index(table_cls, staging, session) + bound = self._staging_rownum_bound(staging, session) + start = 0 + while start < bound: + end = start + merge_batch_size + yield sa.and_(staging.c._rownum > start, staging.c._rownum <= end) + start = end def merge_replace( self, @@ -151,6 +201,13 @@ def merge_replace( UPDATE rather than orm-loader's usual delete-then-reinsert. Matched rows are removed from staging so the follow-up merge_insert() call only inserts rows that weren't already handled here. + + Only correct when index_strategy drops indexes before this UPDATE: + updating an indexed column with the index still live raises + DuckDB's own (opaque) constraint error. capabilities. + supports_replace_with_kept_indexes is False for this reason; the + caller is expected to check it before resolving index_strategy to + "keep" for a "replace" merge. """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ @@ -196,19 +253,10 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N _drop_matched_from_staging() return - bounds = self._staging_rowid_bounds(staging, session) - if bounds is None: - return - start, maximum = bounds - start -= 1 - rowid = self._staging_rowid() - while start < maximum: - end = start + merge_batch_size - batch_range = sa.and_(rowid > start, rowid <= end) + for batch_range in self._staging_rownum_batches(table_cls, staging, session, merge_batch_size): _update(batch_range) _drop_matched_from_staging(batch_range) session.commit() - start = end def merge_upsert( self, @@ -236,17 +284,8 @@ def _upsert(select_: Any) -> sa.Insert: session.execute(_upsert(non_paginated_select)) return - bounds = self._staging_rowid_bounds(staging, session) - if bounds is None: - return - start, maximum = bounds - start -= 1 - rowid = self._staging_rowid() - while start < maximum: - end = start + merge_batch_size - batch_select = non_paginated_select.where(rowid > start, rowid <= end) - session.execute(_upsert(batch_select)) - start = end + for batch_range in self._staging_rownum_batches(table_cls, staging, session, merge_batch_size): + session.execute(_upsert(non_paginated_select.where(batch_range))) def merge_insert( self, @@ -268,14 +307,5 @@ def _insert(select_: Any) -> sa.Insert: session.execute(_insert(non_paginated_select)) return - bounds = self._staging_rowid_bounds(staging, session) - if bounds is None: - return - start, maximum = bounds - start -= 1 - rowid = self._staging_rowid() - while start < maximum: - end = start + merge_batch_size - batch_select = non_paginated_select.where(rowid > start, rowid <= end) - session.execute(_insert(batch_select)) - start = end + for batch_range in self._staging_rownum_batches(table_cls, staging, session, merge_batch_size): + session.execute(_insert(non_paginated_select.where(batch_range))) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 94eb39a..52cefa7 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -61,6 +61,7 @@ def capabilities(self) -> BackendCapabilities: supports_unlogged_staging=True, supports_fk_toggle=True, supports_materialized_views=True, + supports_replace_with_kept_indexes=True, ) def create_staging_table( diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 6cc2753..d22a5f9 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -97,6 +97,7 @@ def capabilities(self) -> BackendCapabilities: supports_unlogged_staging=False, supports_fk_toggle=True, supports_materialized_views=False, + supports_replace_with_kept_indexes=True, ) @property diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index efe1142..928715c 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -450,6 +450,14 @@ def load_csv( f"'insert_if_empty'" ) + if merge_strategy == "replace": + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) + if backend.resolve_index_strategy(index_strategy) == "keep": + backend._require_capability( + "supports_replace_with_kept_indexes", + "merge_strategy='replace' with index_strategy='keep'", + ) + loader_context = LoaderContext( tableclass=cls, session=session, @@ -479,12 +487,13 @@ def load_csv( staging_schema_tag=staging_schema_tag, ) - if validate_deferred_fks: - logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - backend.validate_deferred_fks(cls, session) - - cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) + try: + if validate_deferred_fks: + logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) + backend.validate_deferred_fks(cls, session) + finally: + cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) logger.info(f"Table `{cls.__tablename__}`: Successfully finished ingestion. Total rows: {total}") return total diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 2cba7b3..56b1b70 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -30,6 +30,7 @@ def test_duckdb_backend_identity_and_capabilities(): assert backend.capabilities.supports_unlogged_staging is False assert backend.capabilities.supports_fk_toggle is False assert backend.capabilities.supports_materialized_views is False + assert backend.capabilities.supports_replace_with_kept_indexes is False assert backend.resolve_index_strategy("auto") == "drop_rebuild" @@ -57,8 +58,8 @@ def test_duckdb_backend_create_staging_table(duckdb_session): inspector = sa.inspect(duckdb_session.get_bind()) assert inspector.has_table(_STAGING_TABLE) is True cols = inspector.get_columns(_STAGING_TABLE) - assert [c["name"] for c in cols] == ["id", "name", "slug"] - assert all(c["nullable"] is True for c in cols) + assert [c["name"] for c in cols] == ["id", "name", "slug", "_rownum"] + assert all(c["nullable"] is True for c in cols if c["name"] != "_rownum") @pytest.mark.duckdb @@ -73,6 +74,45 @@ def test_duckdb_backend_drop_staging_table(duckdb_session): assert inspector.has_table(_STAGING_TABLE) is False +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_drop_staging_table_also_drops_rownum_sequence(duckdb_session): + """The staging table's _rownum column is backed by a real sequence, so dropping + the staging table must drop that sequence too. + + Notes + ----- + DuckDB can't ALTER TABLE ADD a GENERATED ALWAYS AS IDENTITY column after creation + """ + backend = DuckDBBackend() + backend.create_staging_table(_ComputedTableCls, duckdb_session) + + sequences_while_present = duckdb_session.execute( + sa.text("SELECT sequence_name FROM duckdb_sequences()") + ).scalars().all() + assert any(f"{_STAGING_TABLE}_rownum_seq" in name for name in sequences_while_present) + + backend.drop_staging_table(_ComputedTableCls, duckdb_session) + + sequences_after_drop = duckdb_session.execute( + sa.text("SELECT sequence_name FROM duckdb_sequences()") + ).scalars().all() + assert not any(f"{_STAGING_TABLE}_rownum_seq" in name for name in sequences_after_drop) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_rownum_populates_monotonically_on_insert(duckdb_session): + backend = DuckDBBackend() + backend.create_staging_table(_ComputedTableCls, duckdb_session) + staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + + duckdb_session.execute(sa.insert(staging), [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}]) + + rownums = duckdb_session.execute(sa.select(staging.c._rownum).order_by(staging.c.id)).scalars().all() + assert rownums == [1, 2] + + def test_duckdb_backend_fk_toggle_methods_raise(): backend = DuckDBBackend() diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 16790bd..ec0e2cb 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -859,3 +859,29 @@ def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_s _CycleATable.load_csv( duckdb_session, csv_path, dedupe=False, loader=PandasLoader(), validate_deferred_fks=True ) + + staging_name = resolve_backend(duckdb_session, mapper=_CycleATable).staging_name_for_table( + _CycleATable.__tablename__ + ) + assert sa.inspect(duckdb_session.get_bind()).has_table(staging_name) is False + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_load_csv_rejects_replace_with_keep_index_strategy_on_duckdb_upfront(duckdb_session, tmp_path): + """merge_strategy='replace' does an UPDATE in place on DuckDB, which + breaks with a live index on an updated column. This combination must be + rejected before the staging load even starts, not surface as DuckDB's + own opaque constraint error partway through the merge.""" + csv_path = tmp_path / f"{SimpleTable.__tablename__}.csv" + pd.DataFrame([{"id": 1, "name": "a"}]).to_csv(csv_path, index=False, sep="\t") + + with pytest.raises(NotImplementedError, match="merge_strategy='replace' with index_strategy='keep'"): + SimpleTable.load_csv( + duckdb_session, + csv_path, + dedupe=False, + loader=PandasLoader(), + merge_strategy="replace", + index_strategy="keep", + ) From 0dcfe3126a6fdd0a1ff222cabeb07b4a443ed735 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 15 Oct 2026 02:26:26 +0000 Subject: [PATCH 03/17] Pagination centralised, core functionalities replace raw text --- src/orm_loader/backends/base.py | 62 ++++++++++- src/orm_loader/backends/duckdb.py | 88 +++++++--------- src/orm_loader/backends/postgres.py | 130 +++++++----------------- src/orm_loader/backends/sqlite.py | 92 +++++------------ tests/backends/test_base_backend.py | 5 + tests/backends/test_postgres_backend.py | 2 +- tests/backends/test_sqlite_backend.py | 9 +- 7 files changed, 169 insertions(+), 219 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 974f662..82085f6 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass -from collections.abc import Generator +from collections.abc import Generator, Iterator from functools import wraps from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Type, TypeVar, cast @@ -419,6 +419,66 @@ def merge_insert( ) -> None: """Insert all staging rows into the target table.""" + @abstractmethod + def _pagination_column( + self, + table_cls: Type["CSVTableProtocol"], + staging: sa.Table, + ) -> sa.ColumnElement[int]: + """Column this backend orders a staging-table scan by, for pagination.""" + + def _ensure_pagination_index( + self, + table_cls: Type["CSVTableProtocol"], + staging: sa.Table, + session: so.Session, + ) -> None: + """Make sure _pagination_column() is indexed, if it needs to be. + No-op by default. Overwrite in a backend that needs an index to paginate efficiently (e.g. DuckDB). + """ + return + + @property + def _commits_between_pagination_batches(self) -> bool: + """Whether _paginate_batches()'s caller should commit after each batch. + True by default, to bound how much uncommitted work a large load + accumulates. + """ + return True + + def _paginate_batches( + self, + table_cls: Type["CSVTableProtocol"], + staging: sa.Table, + session: so.Session, + merge_batch_size: int | None, + ) -> Iterator[sa.ColumnElement[bool] | None]: + """Yield one window predicate per batch for a full scan of staging. + + Yields + ------ + sa.ColumnElement[bool] or None + A predicate to filter the staging table to a single batch of rows. + None if merge_batch_size is None or the staging table is smaller than + merge_batch_size. + """ + if merge_batch_size is not None: + total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() + if merge_batch_size is None or total <= merge_batch_size: + yield None + return + + column = self._pagination_column(table_cls, staging) + self._ensure_pagination_index(table_cls, staging, session) + bound = session.execute( + sa.select(sa.func.coalesce(sa.func.max(column), 0)).select_from(staging) + ).scalar_one() + start = 0 + while start < bound: + end = start + merge_batch_size + yield sa.and_(column > start, column <= end) + start = end + def merge_context( self, table_cls: Type["CSVTableProtocol"], diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 7f0018a..5aa6a91 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -1,12 +1,12 @@ from __future__ import annotations -from collections.abc import Iterator from typing import TYPE_CHECKING, Any, Type import duckdb_sqlalchemy import sqlalchemy as sa import sqlalchemy.orm as so from oa_configurator import qualified +from sqlalchemy.schema import CreateSequence, DropSequence, DropTable from sqlalchemy.sql.compiler import IdentifierPreparer from .base import BackendCapabilities, DatabaseBackend, Dialect @@ -87,12 +87,17 @@ def create_staging_table( session: so.Session, ) -> None: bind_arguments = {"mapper": table_cls} + staging_name = self.staging_name_for_table(table_cls.__tablename__) staging_ref = self.qualified_staging_name(table_cls.__tablename__) seq_ref = self._staging_rownum_sequence_ref(table_cls.__tablename__) - session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};'), bind_arguments=bind_arguments) - session.execute(sa.text(f'DROP SEQUENCE IF EXISTS {seq_ref};'), bind_arguments=bind_arguments) + sequence = sa.Sequence(f"{staging_name}_rownum_seq", schema=self.staging_schema) + + session.execute( + DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), + bind_arguments=bind_arguments, + ) + session.execute(DropSequence(sequence, if_exists=True), bind_arguments=bind_arguments) - staging_name = self.staging_name_for_table(table_cls.__tablename__) metadata = sa.MetaData() staging_columns = [ sa.Column(col.name, col.type, nullable=True) @@ -103,7 +108,8 @@ def create_staging_table( bind=session.connection(bind_arguments=bind_arguments), tables=[staging_table] ) - session.execute(sa.text(f'CREATE SEQUENCE {seq_ref};'), bind_arguments=bind_arguments) + session.execute(CreateSequence(sequence), bind_arguments=bind_arguments) + # ALTER TABLE ADD COLUMN has no Core DDL construct, so this stays raw. session.execute( sa.text(f"ALTER TABLE {staging_ref} ADD COLUMN _rownum BIGINT DEFAULT nextval('{seq_ref}');"), bind_arguments=bind_arguments, @@ -116,12 +122,15 @@ def drop_staging_table( session: so.Session, ) -> None: bind_arguments = {"mapper": table_cls} + staging_name = self.staging_name_for_table(table_cls.__tablename__) session.execute( - sa.text(f'DROP TABLE IF EXISTS {self.qualified_staging_name(table_cls.__tablename__)}'), + DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), bind_arguments=bind_arguments, ) session.execute( - sa.text(f'DROP SEQUENCE IF EXISTS {self._staging_rownum_sequence_ref(table_cls.__tablename__)};'), + DropSequence( + sa.Sequence(f"{staging_name}_rownum_seq", schema=self.staging_schema), if_exists=True + ), bind_arguments=bind_arguments, ) @@ -152,40 +161,19 @@ def restore_fk_check( "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." ) - def _staging_rownum_index( + def _pagination_column( + self, table_cls: type["CSVTableProtocol"], staging: sa.Table + ) -> sa.ColumnElement[int]: + return staging.c._rownum + + def _ensure_pagination_index( self, table_cls: type["CSVTableProtocol"], staging: sa.Table, session: so.Session ) -> None: - """Index staging's _rownum column for O(N log N) windowed range scans.""" staging_name = self.staging_name_for_table(table_cls.__tablename__) idx = sa.Index(f"{staging_name}_rownum_idx", staging.c._rownum) idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) session.commit() - @staticmethod - def _staging_rownum_bound(staging: sa.Table, session: so.Session) -> int: - """Highest _rownum in staging, as the pagination upper bound (0 when empty).""" - return session.execute( - sa.select(sa.func.coalesce(sa.func.max(staging.c._rownum), 0)).select_from(staging) - ).scalar_one() - - def _staging_rownum_batches( - self, - table_cls: type["CSVTableProtocol"], - staging: sa.Table, - session: so.Session, - merge_batch_size: int, - ) -> Iterator[sa.ColumnElement[bool]]: - """Yield successive _rownum window predicates covering every row in - staging, each at most merge_batch_size wide. - """ - self._staging_rownum_index(table_cls, staging, session) - bound = self._staging_rownum_bound(staging, session) - start = 0 - while start < bound: - end = start + merge_batch_size - yield sa.and_(staging.c._rownum > start, staging.c._rownum <= end) - start = end - def merge_replace( self, table_cls: type["CSVTableProtocol"], @@ -248,15 +236,11 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N _drop_matched_from_staging() return - if merge_batch_size is None: - _update() - _drop_matched_from_staging() - return - - for batch_range in self._staging_rownum_batches(table_cls, staging, session, merge_batch_size): + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): _update(batch_range) _drop_matched_from_staging(batch_range) - session.commit() + if batch_range is not None and self._commits_between_pagination_batches: + session.commit() def merge_upsert( self, @@ -280,12 +264,11 @@ def _upsert(select_: Any) -> sa.Insert: non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) - if merge_batch_size is None: - session.execute(_upsert(non_paginated_select)) - return - - for batch_range in self._staging_rownum_batches(table_cls, staging, session, merge_batch_size): - session.execute(_upsert(non_paginated_select.where(batch_range))) + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) + session.execute(_upsert(select_)) + if batch_range is not None and self._commits_between_pagination_batches: + session.commit() def merge_insert( self, @@ -303,9 +286,8 @@ def merge_insert( def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) - if merge_batch_size is None: - session.execute(_insert(non_paginated_select)) - return - - for batch_range in self._staging_rownum_batches(table_cls, staging, session, merge_batch_size): - session.execute(_insert(non_paginated_select.where(batch_range))) + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) + session.execute(_insert(select_)) + if batch_range is not None and self._commits_between_pagination_batches: + session.commit() diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 52cefa7..516206b 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -14,6 +14,7 @@ ) from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError +from sqlalchemy.schema import DropTable from sqlalchemy.sql.compiler import IdentifierPreparer from ..mappers.materialised_view_errors import ( @@ -71,13 +72,17 @@ def create_staging_table( ) -> None: table = table_cls.__table__ preparer = self.identifier_preparer + staging_name = self.staging_name_for_table(table_cls.__tablename__) staging_ref = self.qualified_staging_name(table_cls.__tablename__) bind = session.get_bind(mapper=table_cls) source_ref = qualified( bind, table.name, physical_schema=physical_schema_of(bind, schema_tag=table.schema) ) bind_arguments = {"mapper": table_cls} - session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};'), bind_arguments=bind_arguments) + session.execute( + DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), + bind_arguments=bind_arguments, + ) session.execute( sa.text( f''' @@ -111,8 +116,9 @@ def drop_staging_table( table_cls: type["CSVTableProtocol"], session: so.Session, ) -> None: + staging_name = self.staging_name_for_table(table_cls.__tablename__) session.execute( - sa.text(f'DROP TABLE IF EXISTS {self.qualified_staging_name(table_cls.__tablename__)}'), + DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), bind_arguments={"mapper": table_cls}, ) @@ -182,7 +188,12 @@ def restore_fk_check( sa.text(f"SET session_replication_role = '{safe_state}'"), **execute_options ) - def _staging_rownum_index( + def _pagination_column( + self, table_cls: type["CSVTableProtocol"], staging: sa.Table + ) -> sa.ColumnElement[int]: + return staging.c._rownum + + def _ensure_pagination_index( self, table_cls: type["CSVTableProtocol"], staging: sa.Table, session: so.Session ) -> None: staging_name = self.staging_name_for_table(table_cls.__tablename__) @@ -190,29 +201,6 @@ def _staging_rownum_index( idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) session.commit() - @staticmethod - def _staging_rownum_bound(staging: sa.Table, session: so.Session) -> int: - """Highest ``_rownum`` in *staging*, as the pagination upper bound. - - ``_rownum`` is an identity column whose values need not start at 1, - so the row count is not a valid upper bound. - - Parameters - ---------- - staging : sqlalchemy.Table - Staging table being paginated. - session : sqlalchemy.orm.Session - Session to query on. - - Returns - ------- - int - 0 for an empty staging table. - """ - return session.execute( - sa.select(sa.func.coalesce(sa.func.max(staging.c._rownum), 0)) - ).scalar_one() - def merge_replace( self, table_cls: type["CSVTableProtocol"], @@ -233,30 +221,11 @@ def merge_replace( staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) - non_paginated_replace = sa.delete(target).where(pk_join) - - if merge_batch_size is None: - session.execute(non_paginated_replace) - return - - total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() - if total <= merge_batch_size: - session.execute(non_paginated_replace) - return - - self._staging_rownum_index(table_cls, staging, session) - - bound = self._staging_rownum_bound(staging, session) - start = 0 - while start < bound: - end = start + merge_batch_size - session.execute( - sa.delete(target).where( - pk_join, staging.c._rownum > start, staging.c._rownum <= end - ) - ) - session.commit() - start = end + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + where = pk_join if batch_range is None else sa.and_(pk_join, batch_range) + session.execute(sa.delete(target).where(where)) + if batch_range is not None and self._commits_between_pagination_batches: + session.commit() def merge_upsert( self, @@ -281,27 +250,11 @@ def _upsert(select_: Any) -> sa.Insert: non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) - if merge_batch_size is None: - session.execute(_upsert(non_paginated_select)) - return - - total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() - if total <= merge_batch_size: - session.execute(_upsert(non_paginated_select)) - return - - self._staging_rownum_index(table_cls, staging, session) - - bound = self._staging_rownum_bound(staging, session) - start = 0 - while start < bound: - end = start + merge_batch_size - batch_select = non_paginated_select.where( - staging.c._rownum > start, staging.c._rownum <= end - ) - session.execute(_upsert(batch_select)) - session.commit() - start = end + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) + session.execute(_upsert(select_)) + if batch_range is not None and self._commits_between_pagination_batches: + session.commit() def merge_insert( self, @@ -310,6 +263,11 @@ def merge_insert( *, merge_batch_size: int | None = None, ) -> None: + """Insert all staging rows into the target table. + + session_replication_role='replica' is session-level and persists + across commits, so FK checks stay disabled for every batch. + """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) @@ -319,31 +277,11 @@ def merge_insert( def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) - if merge_batch_size is None: - session.execute(_insert(non_paginated_select)) - return - - total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() - if total <= merge_batch_size: - session.execute(_insert(non_paginated_select)) - return - - # Paginated path: index _rownum for O(N log N) range scans then - # INSERT in batch-sized transactions to bound WAL per commit. - # session_replication_role='replica' is session-level and persists - # across commits, so FK checks stay disabled for all batches. - self._staging_rownum_index(table_cls, staging, session) - - bound = self._staging_rownum_bound(staging, session) - start = 0 - while start < bound: - end = start + merge_batch_size - batch_select = non_paginated_select.where( - staging.c._rownum > start, staging.c._rownum <= end - ) - session.execute(_insert(batch_select)) - session.commit() - start = end + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) + session.execute(_insert(select_)) + if batch_range is not None and self._commits_between_pagination_batches: + session.commit() @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view( diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index d22a5f9..182954a 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -10,6 +10,7 @@ from sqlalchemy import event, text from sqlalchemy.dialects import sqlite as sqlite_dialect from sqlalchemy.exc import IntegrityError +from sqlalchemy.schema import DropTable from sqlalchemy.sql.compiler import IdentifierPreparer from .base import BackendCapabilities, DatabaseBackend, Dialect @@ -111,7 +112,7 @@ def create_staging_table( ) -> None: staging_name = self.staging_name_for_table(table_cls.__tablename__) session.execute( - sa.text(f'DROP TABLE IF EXISTS {self.identifier_preparer.quote_identifier(staging_name)};'), + DropTable(sa.Table(staging_name, sa.MetaData()), if_exists=True), bind_arguments={"mapper": table_cls}, ) @@ -133,9 +134,9 @@ def drop_staging_table( table_cls: type["CSVTableProtocol"], session: so.Session, ) -> None: - staging_ref = self.identifier_preparer.quote_identifier(self.staging_name_for_table(table_cls.__tablename__)) + staging_name = self.staging_name_for_table(table_cls.__tablename__) session.execute( - sa.text(f'DROP TABLE IF EXISTS {staging_ref}'), + DropTable(sa.Table(staging_name, sa.MetaData()), if_exists=True), bind_arguments={"mapper": table_cls}, ) @@ -166,12 +167,20 @@ def restore_fk_check( execute_options = {"bind_arguments": {"mapper": mapper}} if mapper is not None else {} session.execute(text(f"PRAGMA foreign_keys = {safe_state}"), **execute_options) - @staticmethod - def _staging_rowid() -> sa.ColumnElement[int]: - """SQLite's implicit rowid: already gapless and indexed, so it needs - no added column or index the way Postgres's _rownum does.""" + def _pagination_column( + self, table_cls: type["CSVTableProtocol"], staging: sa.Table + ) -> sa.ColumnElement[int]: + """SQLite's implicit rowid is already gapless and indexed.""" return sa.literal_column("rowid") + @property + def _commits_between_pagination_batches(self) -> bool: + """False: keeps the whole merge in the caller's transaction, so a + later insert-phase failure can't leave a delete-then-reinsert merge + with the target emptied and nothing to replace it. + """ + return False + def merge_replace( self, table_cls: type["CSVTableProtocol"], @@ -180,16 +189,7 @@ def merge_replace( *, merge_batch_size: int | None = None, ) -> None: - """Delete every target row whose primary key appears in staging. - - Parameters - ---------- - merge_batch_size : int or None, optional - Rows per statement. Bounds statement size only; the batches are - not committed individually, so the whole merge stays in the - caller's transaction and a later failure in the insert phase - cannot leave the target emptied. - """ + """Delete every target row whose primary key appears in staging.""" self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) @@ -207,21 +207,8 @@ def _delete(extra: sa.ColumnElement[bool] | None = None) -> sa.Delete: sa.tuple_(*(target.c[c] for c in pk_cols)).in_(keys) ) - if merge_batch_size is None: - session.execute(_delete()) - return - - total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() - if total <= merge_batch_size: - session.execute(_delete()) - return - - rowid = self._staging_rowid() - start = 0 - while start < total: - end = start + merge_batch_size - session.execute(_delete(sa.and_(rowid > start, rowid <= end))) - start = end + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + session.execute(_delete(batch_range)) def merge_upsert( self, @@ -245,24 +232,10 @@ def _upsert(select_: Any) -> sa.Insert: non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) - if merge_batch_size is None: - # SQLite's grammar rejects INSERT...SELECT...ON CONFLICT with no - # WHERE on the SELECT (confirmed empirically); sa.true() supplies one. - session.execute(_upsert(non_paginated_select.where(sa.true()))) - return - - total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() - if total <= merge_batch_size: - session.execute(_upsert(non_paginated_select.where(sa.true()))) - return - - rowid = self._staging_rowid() - start = 0 - while start < total: - end = start + merge_batch_size - batch_select = non_paginated_select.where(rowid > start, rowid <= end) - session.execute(_upsert(batch_select)) - start = end + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + + select_ = non_paginated_select.where(batch_range if batch_range is not None else sa.true()) + session.execute(_upsert(select_)) def merge_insert( self, @@ -280,22 +253,9 @@ def merge_insert( def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) - if merge_batch_size is None: - session.execute(_insert(non_paginated_select)) - return - - total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() - if total <= merge_batch_size: - session.execute(_insert(non_paginated_select)) - return - - rowid = self._staging_rowid() - start = 0 - while start < total: - end = start + merge_batch_size - batch_select = non_paginated_select.where(rowid > start, rowid <= end) - session.execute(_insert(batch_select)) - start = end + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): + select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) + session.execute(_insert(select_)) def configure_dbapi_connection(self, dbapi_connection: sa.engine.interfaces.DBAPIConnection) -> None: if dbapi_connection.__class__.__module__.startswith("sqlite3"): diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 2add029..19d0ef7 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -107,6 +107,11 @@ def merge_insert( ) -> None: return None + def _pagination_column( + self, _table_cls: Type[CSVTableProtocol], _staging: sa.Table + ) -> sa.ColumnElement[int]: + return sa.literal_column("rowid") + @staticmethod def _normalize_fk_check_state(previous_state: str | int) -> str | int: return previous_state diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 774cc25..a720fc0 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -52,7 +52,7 @@ def get_execution_options(self) -> dict: def execute(self, statement, parameters=None, **kwargs): if hasattr(statement, "compile"): - sql = str(statement.compile(dialect=postgresql.dialect())) + sql = str(statement.compile(dialect=postgresql.dialect())).strip() else: sql = str(statement) self.statements.append(sql) diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 7daefb8..8d5301c 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -7,6 +7,7 @@ import pytest import sqlalchemy as sa import sqlalchemy.orm as so +from sqlalchemy.dialects import sqlite as sqlite_dialect from orm_loader.backends import Dialect, SQLiteBackend from orm_loader.helpers.sqlite import attach_sqlite_bulk_load_pragmas @@ -25,7 +26,11 @@ def __init__(self, scalar_result: int | str = 1) -> None: self.scalar_result = scalar_result def execute(self, statement, **kwargs): - self.statements.append(str(statement)) + if hasattr(statement, "compile"): + sql = str(statement.compile(dialect=sqlite_dialect.dialect())).strip() + else: + sql = str(statement) + self.statements.append(sql) class _Result: def __init__(self, value): @@ -89,7 +94,7 @@ def test_sqlite_backend_drop_staging_table(): backend.drop_staging_table(_ComputedTableCls, _sess(sqlite_session)) - assert sqlite_session.statements == [f'DROP TABLE IF EXISTS "{_STAGING_TABLE}"'] + assert sqlite_session.statements == [f"DROP TABLE IF EXISTS {_STAGING_TABLE}"] From 58004a6c39098d4832a88b8709a56363c5474a97 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 15 Oct 2026 02:47:10 +0000 Subject: [PATCH 04/17] Clearer docstrings on why some dialects commit during pagination --- src/orm_loader/backends/base.py | 7 ++----- src/orm_loader/backends/duckdb.py | 9 +++++++++ src/orm_loader/backends/postgres.py | 21 ++++++++++++++------- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 82085f6..606760a 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -439,12 +439,9 @@ def _ensure_pagination_index( return @property + @abstractmethod def _commits_between_pagination_batches(self) -> bool: - """Whether _paginate_batches()'s caller should commit after each batch. - True by default, to bound how much uncommitted work a large load - accumulates. - """ - return True + """Whether _paginate_batches()'s caller should commit after each batch.""" def _paginate_batches( self, diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 5aa6a91..5bca7ac 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -174,6 +174,15 @@ def _ensure_pagination_index( idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) session.commit() + @property + def _commits_between_pagination_batches(self) -> bool: + """True: bounds how much uncommitted work one load accumulates on a + large table. Safe here because merge_replace updates matching rows + in place rather than deleting and reinserting them, so a commit + between batches never leaves a row removed without its replacement. + """ + return True + def merge_replace( self, table_cls: type["CSVTableProtocol"], diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 516206b..91b8e88 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -201,6 +201,19 @@ def _ensure_pagination_index( idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) session.commit() + @property + def _commits_between_pagination_batches(self) -> bool: + """True: bounds the WAL growth, lock duration, and dead-tuple + accumulation of one uncommitted transaction over a large table. + + A crash between a delete batch and its insert batch leaves those + rows missing until the load runs again. The next run deletes them + from the target again (their primary keys are still in the freshly + reloaded staging table) before reinserting, so no manual cleanup + is needed, just a rerun. + """ + return True + def merge_replace( self, table_cls: type["CSVTableProtocol"], @@ -209,13 +222,7 @@ def merge_replace( *, merge_batch_size: int | None = None, ) -> None: - """Delete every target row whose primary key appears in staging. - - Notes - ----- - Each batch commits, so unlike SQLite's merge the delete phase is - already durable if a later insert phase fails. - """ + """Delete every target row whose primary key appears in staging.""" self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) From f4114be70bf223b4f956b878524c7b37bc847bcb Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 16 Oct 2026 00:04:20 +0000 Subject: [PATCH 05/17] Samll fixes, correct dialects, updated tests --- README.md | 2 +- docs/index.md | 2 +- src/orm_loader/backends/base.py | 30 ++++++++++++++++++--- src/orm_loader/backends/duckdb.py | 35 +++++++++++++++++++++--- src/orm_loader/backends/postgres.py | 7 +---- src/orm_loader/tables/loadable_table.py | 4 +-- src/orm_loader/tables/typing.py | 2 +- tests/backends/test_base_backend.py | 4 +++ tests/backends/test_duckdb_backend.py | 34 ++++++++++++++++++++++- tests/models.py | 34 +++++++++++++++++++++++ tests/test_optional_deps.py | 36 +++++++++++++++++++++++++ 11 files changed, 171 insertions(+), 19 deletions(-) create mode 100644 tests/test_optional_deps.py diff --git a/README.md b/README.md index 6c0b0b2..f6c3953 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ The library focuses on: * loader and validation infrastructure * operational helpers that work across supported backends -At the moment, the built-in backends are SQLite and PostgreSQL. +The built-in backends are SQLite, PostgreSQL, and DuckDB. See [oa-configurator's supported dialects](https://AustralianCancerDataNetwork.github.io/oa-configurator/config-reference/#supported-dialects) for the dialect list this is built against. ### What this library provides diff --git a/docs/index.md b/docs/index.md index 015be3c..f09f240 100644 --- a/docs/index.md +++ b/docs/index.md @@ -13,7 +13,7 @@ It focuses on: - model-agnostic validation scaffolding - operational helpers for supported backends -It currently ships with backend implementations for SQLite and PostgreSQL. +It ships with backend implementations for SQLite, PostgreSQL, and DuckDB. See [oa-configurator's supported dialects](https://AustralianCancerDataNetwork.github.io/oa-configurator/config-reference/#supported-dialects) for the dialect list this is built against. --- diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 606760a..182618c 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -319,12 +319,36 @@ def validate_deferred_fks( @staticmethod def _validate_one_deferred_fk(fk: sa.ForeignKeyConstraint, session: so.Session) -> None: + """Raise if any row of fk's local table has a non-null value with no match in the referred table. + + Parameters + ---------- + fk : sa.ForeignKeyConstraint + The deferred FK constraint to check. + session : so.Session + Session to run the anti-join query against. + + Raises + ------ + DeferredForeignKeyViolationError + If any row's FK column is non-null with no matching row in the + referenced table. + + Notes + ----- + The referred table is always aliased, even for an ordinary + cross-table FK where it changes nothing. A self-referential FK's + referred table is the same Table object as the row being checked, + so an unaliased subquery would compare each row to itself instead + of correlating to the outer row. + """ local_cols = list(fk.columns) remote_cols = [element.column for element in fk.elements] - referred_table = remote_cols[0].table - matches = sa.and_(*(local_cols[i] == remote_cols[i] for i in range(len(local_cols)))) + referred_alias = remote_cols[0].table.alias() + aliased_remote_cols = [referred_alias.c[col.name] for col in remote_cols] + matches = sa.and_(*(local_cols[i] == aliased_remote_cols[i] for i in range(len(local_cols)))) not_null = sa.and_(*(col.isnot(None) for col in local_cols)) - orphan_exists = sa.exists(sa.select(sa.literal(1)).select_from(referred_table).where(matches)) + orphan_exists = sa.exists(sa.select(sa.literal(1)).select_from(referred_alias).where(matches)) orphans = session.execute( sa.select(*local_cols).where(not_null, ~orphan_exists).limit(5) diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 5bca7ac..42bf9a1 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -2,7 +2,6 @@ from typing import TYPE_CHECKING, Any, Type -import duckdb_sqlalchemy import sqlalchemy as sa import sqlalchemy.orm as so from oa_configurator import qualified @@ -30,6 +29,8 @@ def dialect(self) -> Dialect: @property def identifier_preparer(self) -> IdentifierPreparer: + import duckdb_sqlalchemy + return duckdb_sqlalchemy.Dialect().identifier_preparer @property @@ -57,8 +58,8 @@ def deferred_fk_constraints( Only reads the ``tbl is None`` bucket of sort_tables_and_constraints's output. A standalone self-referential FK (not part of a cross-table cycle) is still deferred by SQLAlchemy either way, but lands in its - *own* table's bucket unless the column declares ``use_alter=True`` -- - without that, this method can't see it, and it will insert without + *own* table's bucket unless the column declares ``use_alter=True``. + Without that, this method can't see it, and it will insert without error instead of being caught by validate_deferred_fks(). Model authors: mark a self-referential FK ``use_alter=True`` explicitly if it needs anti-join validation on DuckDB. @@ -68,6 +69,16 @@ def deferred_fk_constraints( deferred = next(fks for tbl, fks in entries if tbl is None) return tuple(fk for fk in deferred if fk.table is table_cls.__table__) + @staticmethod + def _is_referenced_by_other_table(target: sa.Table) -> bool: + """True if some other table in target's own metadata has an FK pointing at target.""" + return any( + fk.referred_table is target + for other in target.metadata.tables.values() + if other is not target + for fk in other.foreign_key_constraints + ) + def _staging_rownum_sequence_ref(self, tablename: str) -> str: """Fully schema-qualified, quoted name of tablename's staging _rownum sequence. @@ -205,11 +216,27 @@ def merge_replace( supports_replace_with_kept_indexes is False for this reason; the caller is expected to check it before resolving index_strategy to "keep" for a "replace" merge. + + Raises + ------ + ValueError + If target is itself referenced by another table's FK and also + owns an FK of its own on a column being updated: DuckDB rejects + a SET on any FK-constrained column of a currently-referenced + row, dropped indexes or not. """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] + own_fk_cols = [c for c in update_cols if c in target.c and target.c[c].foreign_keys] + if own_fk_cols and self._is_referenced_by_other_table(target): + raise ValueError( + f"merge_strategy='replace' can't update {target.name}.{own_fk_cols} on DuckDB: " + f"{target.name} is referenced by another table's FK, and DuckDB rejects any SET " + "on an FK-constrained column of a currently-referenced row. Use " + "--merge-strategy upsert, or reload into an empty database." + ) pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) def _batch_keys(extra: sa.ColumnElement[bool] | None = None) -> sa.Select[Any]: @@ -259,6 +286,8 @@ def merge_upsert( *, merge_batch_size: int | None = None, ) -> None: + import duckdb_sqlalchemy + self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 91b8e88..cb537b8 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -351,12 +351,7 @@ def refresh_materialized_view( try: conn.execute(sa.text(f"REFRESH MATERIALIZED VIEW {concurrency}{safe_name};")) except OperationalError as error: - # Keep psycopg optional at import time; it is only needed when - # handling an actual PostgreSQL driver error. - try: - from psycopg.errors import ObjectNotInPrerequisiteState - except ImportError: - raise error from None + from psycopg.errors import ObjectNotInPrerequisiteState if not concurrently or not isinstance(error.orig, ObjectNotInPrerequisiteState): raise diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 928715c..47a1995 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -373,7 +373,7 @@ def load_csv( index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, - validate_deferred_fks: bool = False, + validate_deferred_fks: bool = True, ) -> int: """ @@ -417,8 +417,6 @@ def load_csv( validate_deferred_fks After merging, run an anti-join check for any FK this backend's DB can't enforce itself (see DatabaseBackend.validate_deferred_fks()). - Off by default: a full-table scan per such FK, on top of the load. - No-op for a backend with none (every backend but DuckDB, today). Returns ------- diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index ec3a8fe..93c6367 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -87,7 +87,7 @@ def load_csv( index_strategy: str = "auto", merge_batch_size: int | None = None, staging_schema_tag: str | None = None, - validate_deferred_fks: bool = False, + validate_deferred_fks: bool = True, ) -> int: ... @classmethod diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 19d0ef7..1a2db64 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -112,6 +112,10 @@ def _pagination_column( ) -> sa.ColumnElement[int]: return sa.literal_column("rowid") + @property + def _commits_between_pagination_batches(self) -> bool: + return True + @staticmethod def _normalize_fk_check_state(previous_state: str | int) -> str | int: return previous_state diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 56b1b70..85ea0a2 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -8,7 +8,14 @@ from orm_loader.backends import Dialect, DuckDBBackend from orm_loader.backends.base import DeferredForeignKeyViolationError -from tests.models import ComputedColumnTable, CycleATable, CycleBTable +from tests.models import ( + ComputedColumnTable, + CycleATable, + CycleBTable, + ReferencedParentTable, + ReferencingChildTable, + UnreferencedGrandparentTable, +) if TYPE_CHECKING: from orm_loader.tables.typing import CSVTableProtocol @@ -16,6 +23,7 @@ _ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) _CycleATableCls = cast("Type[CSVTableProtocol]", CycleATable) _CycleBTableCls = cast("Type[CSVTableProtocol]", CycleBTable) +_ReferencedParentTableCls = cast("Type[CSVTableProtocol]", ReferencedParentTable) _TARGET_TABLE = ComputedColumnTable.__tablename__ _STAGING_TABLE = f"_staging_{_TARGET_TABLE}" @@ -233,3 +241,27 @@ def test_duckdb_backend_merge_replace_paginated_updates_in_place(duckdb_session) rows = {r.id: r.name for r in duckdb_session.execute(sa.select(ComputedColumnTable)).scalars().all()} assert rows == {i: f"staged{i}" for i in range(10)} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_replace_rejects_referenced_table_with_own_fk(duckdb_session): + """DuckDB rejects a SET on any FK-constrained column of a row that's + still referenced elsewhere, regardless of index_strategy. merge_replace + must fail up front, before touching any data, rather than partway + through a batch.""" + backend = DuckDBBackend() + backend.create_staging_table(_ReferencedParentTableCls, duckdb_session) + staging = _ReferencedParentTableCls.get_staging_table( + duckdb_session, staging_schema_tag=backend.staging_schema_tag + ) + + duckdb_session.execute(sa.insert(UnreferencedGrandparentTable), [{"id": 1, "name": "g1"}]) + duckdb_session.execute( + sa.insert(ReferencedParentTable), [{"id": 1, "grandparent_id": 1, "name": "p1"}] + ) + duckdb_session.execute(sa.insert(ReferencingChildTable), [{"id": 1, "parent_id": 1}]) + duckdb_session.execute(sa.insert(staging), [{"id": 1, "grandparent_id": 1, "name": "p1-updated"}]) + + with pytest.raises(ValueError, match="referenced_parent_table"): + backend.merge_replace(_ReferencedParentTableCls, duckdb_session, ["id"]) diff --git a/tests/models.py b/tests/models.py index 130d45e..372acdb 100644 --- a/tests/models.py +++ b/tests/models.py @@ -120,6 +120,40 @@ class ImpliedEnumTable(Base, CSVLoadableTableInterface): flag: so.Mapped[str | None] = so.mapped_column(sa.String(1), nullable=True) +class ReferencedParentTable(Base, CSVLoadableTableInterface): + """Referenced by ReferencingChildTable *and* owns its own FK, for + merge_replace's up-front DuckDB guard: a table in this shape can't be + SET on its own FK column while still referenced, even with indexes + dropped.""" + + __tablename__ = "referenced_parent_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + grandparent_id: so.Mapped[int] = so.mapped_column( + sa.Integer, sa.ForeignKey(f"{SchemaRole.PRIMARY.value}.unreferenced_grandparent_table.id") + ) + name: so.Mapped[str] = so.mapped_column(sa.String) + + +class UnreferencedGrandparentTable(Base, CSVLoadableTableInterface): + __tablename__ = "unreferenced_grandparent_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + name: so.Mapped[str] = so.mapped_column(sa.String) + + +class ReferencingChildTable(Base, CSVLoadableTableInterface): + __tablename__ = "referencing_child_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + parent_id: so.Mapped[int] = so.mapped_column( + sa.Integer, sa.ForeignKey(f"{SchemaRole.PRIMARY.value}.referenced_parent_table.id") + ) + + class VocabSchemaTable(Base, CSVLoadableTableInterface): """A VOCAB-tagged table, so tests can prove the staging/index schema derivation actually resolves a non-primary schema tag correctly.""" diff --git a/tests/test_optional_deps.py b/tests/test_optional_deps.py new file mode 100644 index 0000000..cd4c06c --- /dev/null +++ b/tests/test_optional_deps.py @@ -0,0 +1,36 @@ +"""Tests that orm_loader stays importable without optional extras installed +""" + +from __future__ import annotations + +import subprocess +import sys + + +def _run_blocked(import_statement: str, *, blocked_module: str = "duckdb_sqlalchemy") -> subprocess.CompletedProcess[str]: + code = f"import sys; sys.modules['{blocked_module}'] = None; {import_statement}" + return subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + + +def test_backends_importable_without_duckdb_sqlalchemy(): + """orm_loader.backends imports cleanly when duckdb-sqlalchemy is not installed.""" + result = _run_blocked("import orm_loader.backends") + assert result.returncode == 0, result.stderr + + +def test_duckdb_backend_module_importable_without_duckdb_sqlalchemy(): + """orm_loader.backends.duckdb itself is importable without the extra.""" + result = _run_blocked("import orm_loader.backends.duckdb") + assert result.returncode == 0, result.stderr + + +def test_backends_importable_without_psycopg(): + """orm_loader.backends imports cleanly when psycopg is not installed.""" + result = _run_blocked("import orm_loader.backends", blocked_module="psycopg") + assert result.returncode == 0, result.stderr + + +def test_postgres_backend_module_importable_without_psycopg(): + """orm_loader.backends.postgres itself is importable without the driver.""" + result = _run_blocked("import orm_loader.backends.postgres", blocked_module="psycopg") + assert result.returncode == 0, result.stderr From ce5edac5ecd42d537b5d43e0ede52c5f7797a49c Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 16 Oct 2026 01:33:33 +0000 Subject: [PATCH 06/17] Proper cyclic table load commits --- src/orm_loader/backends/base.py | 13 ++++++++ src/orm_loader/backends/duckdb.py | 49 +++++++++++++++++++++++-------- 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 182618c..a8d3bfe 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -411,6 +411,19 @@ def restore_fk_check( ) -> None: """Restore FK checks to a previously returned backend-specific state.""" + def check_replace_feasible( + self, + table_cls: Type["CSVTableProtocol"], + pk_cols: list[str], + ) -> None: + """Raise if merge_strategy='replace' can't load table_cls on this backend. + + Default: always feasible. Only a backend that rejects certain + in-place updates on a currently-referenced row (DuckDB) overrides + this, so callers can preflight every table a replace load will + touch before loading any of them. + """ + @abstractmethod def merge_replace( self, diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 42bf9a1..9c32faa 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -194,6 +194,41 @@ def _commits_between_pagination_batches(self) -> bool: """ return True + def check_replace_feasible( + self, + table_cls: type["CSVTableProtocol"], + pk_cols: list[str], + ) -> None: + """Raise if merge_strategy='replace' can't update table_cls's own FK columns. + + Notes + ----- + Meant to be called before any replace load touches table_cls, so a table later + in load order failing this check doesn't leave earlier tables already replaced. + + Raises + ------ + ValueError + If target is itself referenced by another table's FK and also + owns an enforced FK of its own on a column being updated: DuckDB + rejects a SET on any FK-constrained column of a currently- + referenced row, dropped indexes or not. + """ + target = table_cls.__table__ + update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] + deferred_cols = {col.name for fk in self.deferred_fk_constraints(table_cls) for col in fk.columns} + enforced_fk_cols = [ + c for c in update_cols + if c in target.c and target.c[c].foreign_keys and c not in deferred_cols + ] + if enforced_fk_cols and self._is_referenced_by_other_table(target): + raise ValueError( + f"merge_strategy='replace' can't update {target.name}.{enforced_fk_cols} on DuckDB: " + f"{target.name} is referenced by another table's FK, and DuckDB rejects any SET " + "on an FK-constrained column of a currently-referenced row. Use " + "--merge-strategy upsert, or reload into an empty database." + ) + def merge_replace( self, table_cls: type["CSVTableProtocol"], @@ -220,23 +255,13 @@ def merge_replace( Raises ------ ValueError - If target is itself referenced by another table's FK and also - owns an FK of its own on a column being updated: DuckDB rejects - a SET on any FK-constrained column of a currently-referenced - row, dropped indexes or not. + See check_replace_feasible(). """ self._validate_merge_batch_size(merge_batch_size) + self.check_replace_feasible(table_cls, pk_cols) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] - own_fk_cols = [c for c in update_cols if c in target.c and target.c[c].foreign_keys] - if own_fk_cols and self._is_referenced_by_other_table(target): - raise ValueError( - f"merge_strategy='replace' can't update {target.name}.{own_fk_cols} on DuckDB: " - f"{target.name} is referenced by another table's FK, and DuckDB rejects any SET " - "on an FK-constrained column of a currently-referenced row. Use " - "--merge-strategy upsert, or reload into an empty database." - ) pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) def _batch_keys(extra: sa.ColumnElement[bool] | None = None) -> sa.Select[Any]: From d7cffa259f5f5e1156c9e506b3889585b18996f5 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 16 Oct 2026 02:45:30 +0000 Subject: [PATCH 07/17] Rectify replacement feasible for DuckDB, accurately revert on failed merge strategy --- src/orm_loader/backends/base.py | 4 +- src/orm_loader/backends/duckdb.py | 60 ++++++++++++++++--------- src/orm_loader/tables/loadable_table.py | 23 +++++----- tests/loaders/test_loader_e2e.py | 6 +++ 4 files changed, 59 insertions(+), 34 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index a8d3bfe..3873a5e 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -415,13 +415,15 @@ def check_replace_feasible( self, table_cls: Type["CSVTableProtocol"], pk_cols: list[str], + session: so.Session, ) -> None: """Raise if merge_strategy='replace' can't load table_cls on this backend. Default: always feasible. Only a backend that rejects certain in-place updates on a currently-referenced row (DuckDB) overrides this, so callers can preflight every table a replace load will - touch before loading any of them. + touch before loading any of them. session is only used as read-only + to check whether target or a referencing table already has rows """ @abstractmethod diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 9c32faa..b30ccde 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -4,7 +4,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import qualified +from oa_configurator import physical_schema_of, qualified from sqlalchemy.schema import CreateSequence, DropSequence, DropTable from sqlalchemy.sql.compiler import IdentifierPreparer @@ -70,13 +70,28 @@ def deferred_fk_constraints( return tuple(fk for fk in deferred if fk.table is table_cls.__table__) @staticmethod - def _is_referenced_by_other_table(target: sa.Table) -> bool: - """True if some other table in target's own metadata has an FK pointing at target.""" - return any( - fk.referred_table is target + def _tables_referencing(target: sa.Table) -> set[sa.Table]: + """Other tables in target's own metadata that have an FK pointing at target.""" + return { + other for other in target.metadata.tables.values() if other is not target for fk in other.foreign_key_constraints + if fk.referred_table is target + } + + def _is_referenced_by_existing_rows(self, target: sa.Table, session: so.Session) -> bool: + """True if some table referencing target has any rows. + + Table-level existence check. A referencing table may not exist yet, + since target's shared metadata spans every CDM table. + """ + bind = session.get_bind() + inspector = sa.inspect(bind) + return any( + inspector.has_table(other.name, schema=physical_schema_of(bind, schema_tag=other.schema)) + and session.execute(sa.select(sa.literal(1)).select_from(other).limit(1)).first() is not None + for other in self._tables_referencing(target) ) def _staging_rownum_sequence_ref(self, tablename: str) -> str: @@ -198,21 +213,19 @@ def check_replace_feasible( self, table_cls: type["CSVTableProtocol"], pk_cols: list[str], + session: so.Session, ) -> None: """Raise if merge_strategy='replace' can't update table_cls's own FK columns. - Notes - ----- - Meant to be called before any replace load touches table_cls, so a table later - in load order failing this check doesn't leave earlier tables already replaced. + Call before any replace load touches table_cls, so a later table + failing this check doesn't leave earlier tables already replaced. Raises ------ ValueError - If target is itself referenced by another table's FK and also - owns an enforced FK of its own on a column being updated: DuckDB - rejects a SET on any FK-constrained column of a currently- - referenced row, dropped indexes or not. + Target has rows, some table referencing it also has rows, and + target owns an enforced FK on a column being updated: DuckDB + rejects a SET on an FK-constrained column of a referenced row. """ target = table_cls.__table__ update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] @@ -221,13 +234,18 @@ def check_replace_feasible( c for c in update_cols if c in target.c and target.c[c].foreign_keys and c not in deferred_cols ] - if enforced_fk_cols and self._is_referenced_by_other_table(target): - raise ValueError( - f"merge_strategy='replace' can't update {target.name}.{enforced_fk_cols} on DuckDB: " - f"{target.name} is referenced by another table's FK, and DuckDB rejects any SET " - "on an FK-constrained column of a currently-referenced row. Use " - "--merge-strategy upsert, or reload into an empty database." - ) + if not enforced_fk_cols: + return + if not table_cls._target_has_rows(session=session): + return + if not self._is_referenced_by_existing_rows(target, session): + return + raise ValueError( + f"merge_strategy='replace' can't update {target.name}.{enforced_fk_cols} on DuckDB: " + f"{target.name} is referenced by another table's FK, and DuckDB rejects any SET " + "on an FK-constrained column of a currently-referenced row. Use " + "--merge-strategy upsert, or reload into an empty database." + ) def merge_replace( self, @@ -258,7 +276,7 @@ def merge_replace( See check_replace_feasible(). """ self._validate_merge_batch_size(merge_batch_size) - self.check_replace_feasible(table_cls, pk_cols) + self.check_replace_feasible(table_cls, pk_cols, session) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) update_cols = [c for c in self._insertable_column_names(table_cls) if c not in pk_cols] diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 47a1995..f0a9a70 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -477,19 +477,18 @@ def load_csv( # Merge staging to target (Wrapped in our index dropper!) logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") - with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): - cls.merge_from_staging( - session, - merge_strategy=merge_strategy, - merge_batch_size=merge_batch_size, - staging_schema_tag=staging_schema_tag, - ) - try: - if validate_deferred_fks: - logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) - backend.validate_deferred_fks(cls, session) + with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): + cls.merge_from_staging( + session, + merge_strategy=merge_strategy, + merge_batch_size=merge_batch_size, + staging_schema_tag=staging_schema_tag, + ) + if validate_deferred_fks: + logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) + backend.validate_deferred_fks(cls, session) finally: cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index ec0e2cb..7af044a 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -851,6 +851,12 @@ def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_s _skip_deferred_fk_on_duckdb no-oped the ALTER TABLE ADD CONSTRAINT): a load introducing an orphaned reference succeeds silently unless validate_deferred_fks=True catches it. + + Notes + ----- + Rollback on duckdb_session wipes out data that has already been committed. + Rollback-after-commit on this fixture's session is therefore unreliable for + asserting target-row survival. """ csv_path = tmp_path / "cycle_a_table.csv" pd.DataFrame([{"id": 1, "b_id": 999}]).to_csv(csv_path, index=False, sep="\t") From 17b2252300e581e4e4dfe2ee66983ca5d3a069d6 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 26 Oct 2026 04:57:19 +0000 Subject: [PATCH 08/17] Resolve schema_claims-only tags in create_tables()'s FK span check create_tables() only consulted resolved's static schema-translate map when deciding if a foreign key can span two tags, so a tag claimed only via create_engine(schema_claims=[...]) raised UnregisteredSchemaTagError even though the connection itself already resolves it. --- src/orm_loader/helpers/bootstrap.py | 14 +++++++++-- tests/helpers/test_bootstrap.py | 39 ++++++++++++++++++++++++++++- 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/src/orm_loader/helpers/bootstrap.py b/src/orm_loader/helpers/bootstrap.py index 77ccf14..c449f03 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -12,6 +12,7 @@ declared_schema_tags, is_ephemeral_url, open_connection, + physical_schema_of, referred_schema_tag, without_cross_engine_foreign_keys, ) @@ -91,9 +92,16 @@ def create_tables( "create_tables() requires all tables to share one database transaction; " f"got schema tags {sorted(schema_tags)}." ) + # Includes tags claimed only via schema_claims=[...] so + # foreign_key_can_span() can resolve them too. + physical_schemas = { + tag: physical_schema_of(connection, schema_tag=tag) for tag in claimed_schema_tags(connection) + } crossing = any( not resolved.foreign_key_can_span( - cast(str, table.schema), referred_schema_tag(fk, cast(str, table.schema)) + cast(str, table.schema), + referred_schema_tag(fk, cast(str, table.schema)), + physical_schemas=physical_schemas, ) for table in tables for fk in table.foreign_keys @@ -122,7 +130,9 @@ def create_tables( for metadata, related in related_by_metadata.items(): local_tables = [table for table in metadata.tables.values() if table in related] - copies = without_cross_engine_foreign_keys(local_tables, resolved=resolved) + copies = without_cross_engine_foreign_keys( + local_tables, resolved=resolved, physical_schemas=physical_schemas + ) requested = [table for table in copies if table.key in requested_by_metadata[metadata]] copies[0].metadata.create_all(connection, tables=requested, checkfirst=True) diff --git a/tests/helpers/test_bootstrap.py b/tests/helpers/test_bootstrap.py index 41224bb..c383411 100644 --- a/tests/helpers/test_bootstrap.py +++ b/tests/helpers/test_bootstrap.py @@ -2,7 +2,7 @@ import pytest import sqlalchemy as sa -from oa_configurator import CDMDatabaseConfig, ConnectionConfig, Resolver, StackConfig +from oa_configurator import CDMDatabaseConfig, ConnectionConfig, Resolver, SchemaClaim, StackConfig from orm_loader.helpers import bootstrap, create_db, create_tables @@ -160,3 +160,40 @@ def test_bootstrap_entry_points_require_resolved_database(entry): with pytest.raises(TypeError, match="requires a ResolvedDatabase"): entry(engine) engine.dispose() + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_create_tables_resolves_a_schema_claims_only_tag_on_duckdb(): + """A tag claimed only via create_engine(schema_claims=[...]), not present in + resolved's own config, must still resolve inside foreign_key_can_span() + instead of raising UnregisteredSchemaTagError.""" + stack = StackConfig.for_session( + connections={"main": ConnectionConfig(dialect="duckdb", database_name=":memory:")}, + databases={ + "test_cdm": CDMDatabaseConfig(connection="main", cdm_schema="cdm", vocab_schema="vocab") + }, + ) + resolved = Resolver(stack).resolve_database("test_cdm") + engine, _ = resolved.create_engines( + schema_claims=[SchemaClaim(schema_tag="custom", physical_schema="cdm")] + ) + + metadata = sa.MetaData() + primary = sa.Table( + "primary_table", metadata, sa.Column("id", sa.Integer, primary_key=True), schema="primary" + ) + custom = sa.Table( + "custom_table", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("primary_id", sa.Integer, sa.ForeignKey("primary.primary_table.id")), + schema="custom", + ) + + with engine.begin() as connection: + create_tables(connection, [primary, custom], resolved=resolved) + foreign_keys = sa.inspect(connection).get_foreign_keys("custom_table", schema="cdm") + + assert [fk["referred_table"] for fk in foreign_keys] == ["primary_table"] + engine.dispose() From 109a2a89733b63f5f055a9d2e6a3a86f6eabcc8b Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 26 Oct 2026 21:01:35 +0000 Subject: [PATCH 09/17] Update oa-configurator to 2.1.0 to receive DuckDB changes --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 58a4c22..01616fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ license = { text = "Apache-2.0" } requires-python = ">=3.12" dependencies = [ "chardet>=5.2.0", - "oa-configurator>=2.0.0,<3", + "oa-configurator>=2.1.0,<3", "pandas>=2.3.3", "pyarrow>=23.0.1", "sqlalchemy>=2.0.45", From 258f387cb907f8abaddd544aee26aeab601b5316 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 27 Oct 2026 04:25:36 +0000 Subject: [PATCH 10/17] Classify cross-schema and cyclic FKs via oa-configurator, validate before each batch commit - deferred_fk_constraints/check_replace_feasible now require resolved and delegate to unenforced_foreign_keys() instead of a local cyclic-only approximation, closing the cross-schema blind spot - validate deferred FKs before every intermediate batch commit, not just once after the whole merge, so an invalid batch is never committed in the first place --- src/orm_loader/backends/base.py | 72 +++++++- src/orm_loader/backends/duckdb.py | 129 ++++++++------ src/orm_loader/backends/postgres.py | 36 +++- src/orm_loader/backends/sqlite.py | 23 ++- src/orm_loader/tables/loadable_table.py | 63 ++++++- src/orm_loader/tables/typing.py | 8 + tests/backends/test_duckdb_backend.py | 227 ++++++++++++++++++++---- tests/loaders/test_loader_e2e.py | 9 +- tests/models.py | 16 ++ 9 files changed, 464 insertions(+), 119 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 3873a5e..e21d484 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass -from collections.abc import Generator, Iterator +from collections.abc import Generator, Iterator, Mapping from functools import wraps from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Type, TypeVar, cast @@ -12,7 +12,7 @@ from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer -from oa_configurator import Dialect, SchemaClaim, open_connection, qualified +from oa_configurator import Dialect, ResolvedDatabase, SchemaClaim, open_connection, qualified if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext @@ -116,7 +116,14 @@ class DatabaseBackend(ABC): without changing existing loader orchestration yet. """ - def __init__(self, *, staging_schema_tag: str | None = None, staging_schema: str | None = None) -> None: + def __init__( + self, + *, + staging_schema_tag: str | None = None, + staging_schema: str | None = None, + resolved: ResolvedDatabase | None = None, + physical_schemas: Mapping[str, str | None] | None = None, + ) -> None: """ Parameters ---------- @@ -132,9 +139,19 @@ def __init__(self, *, staging_schema_tag: str | None = None, staging_schema: str qualification — staging tables land in whatever schema the connection's search_path resolves to. SQLite has no schema concept and always resolves to ``None``. + resolved + Used by DuckDBBackend to classify a foreign key crossing a + schema or database boundary, via oa-configurator's + ``unenforced_foreign_keys()``. Ignored by every other backend. + physical_schemas + A schema tag's physical name beyond ``resolved``'s own + configured ones (e.g. one claimed via ``schema_claims``). + Forwarded alongside ``resolved``. """ self.staging_schema_tag = staging_schema_tag self.staging_schema = staging_schema + self.resolved = resolved + self.physical_schemas = physical_schemas @staticmethod def _validate_merge_batch_size(merge_batch_size: int | None) -> None: @@ -434,8 +451,15 @@ def merge_replace( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: - """Merge staging rows by replacing matching target rows first.""" + """Merge staging rows by replacing matching target rows first. + + validate_deferred_fks + Forwarded to :meth:`_commit_batch_if_needed` for each + intermediate commit, so a batch that violates a deferred FK is + caught before it commits instead of only after the whole merge. + """ @abstractmethod def merge_upsert( @@ -445,8 +469,13 @@ def merge_upsert( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: - """Merge staging rows using backend-specific upsert semantics.""" + """Merge staging rows using backend-specific upsert semantics. + + validate_deferred_fks + See :meth:`merge_replace`. + """ @abstractmethod def merge_insert( @@ -455,8 +484,13 @@ def merge_insert( session: so.Session, *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: - """Insert all staging rows into the target table.""" + """Insert all staging rows into the target table. + + validate_deferred_fks + See :meth:`merge_replace`. + """ @abstractmethod def _pagination_column( @@ -515,6 +549,32 @@ def _paginate_batches( yield sa.and_(column > start, column <= end) start = end + def _commit_batch_if_needed( + self, + table_cls: Type["CSVTableProtocol"], + session: so.Session, + batch_range: sa.ColumnElement[bool] | None, + *, + validate_deferred_fks: bool = True, + ) -> None: + """Validate deferred FKs and commit after one pagination batch, + if this backend commits between batches. + + Called from every merge_*() method after each batch of rows is processed. + Instead of committing directly, a batch that violates a deferred FK is caught + before the next batch starts. Without this, an earlier batch's commit cannot + be undone once a later batch's load_csv()-level validation fails. + + No-op when there is nothing to validate to not add cost. This is true for every + backend but DuckDB, and for DuckDB whenever table_cls has no deferred/ + unenforced FK at all. + """ + if batch_range is None or not self._commits_between_pagination_batches: + return + if validate_deferred_fks: + self.validate_deferred_fks(table_cls, session) + session.commit() + def merge_context( self, table_cls: Type["CSVTableProtocol"], diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index b30ccde..0efb47c 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -4,7 +4,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import physical_schema_of, qualified +from oa_configurator import physical_schema_of, qualified, unenforced_foreign_keys from sqlalchemy.schema import CreateSequence, DropSequence, DropTable from sqlalchemy.sql.compiler import IdentifierPreparer @@ -43,42 +43,60 @@ def capabilities(self) -> BackendCapabilities: supports_replace_with_kept_indexes=False, ) + def _unenforced_constraints(self, tables: tuple[sa.Table, ...]) -> frozenset[sa.ForeignKeyConstraint]: + """Every FK in *tables* DuckDB will not enforce: omitted (cross-database/ + cross-schema) or created but never added (cyclic/``use_alter=True``). + + Delegates to oa-configurator's :func:`unenforced_foreign_keys`, the + same classification DDL creation itself uses. + + Raises + ------ + ValueError + ``self.resolved`` was never given to ``resolve_backend()``. + """ + if self.resolved is None: + raise ValueError( + "DuckDBBackend needs resolved= to classify foreign-key " + "enforcement. Pass it to load_csv(...) or resolve_backend(...)." + ) + unenforced = unenforced_foreign_keys( + tables, resolved=self.resolved, physical_schemas=self.physical_schemas + ) + return frozenset(entry.constraint for entry in unenforced) + def deferred_fk_constraints( self, table_cls: Type["CSVTableProtocol"], ) -> tuple[sa.ForeignKeyConstraint, ...]: - """FK constraints on *table_cls* excluded from its CREATE TABLE and never - added afterward: DuckDB can't execute the ALTER TABLE ADD CONSTRAINT that - would otherwise add them (oa-configurator's _skip_deferred_fk_on_duckdb no-ops it). - - sort_tables_and_constraints needs every table in the schema to tell - which FKs are genuinely deferred (cyclic), so this runs against - table_cls's whole shared metadata, not table_cls alone. - - Only reads the ``tbl is None`` bucket of sort_tables_and_constraints's - output. A standalone self-referential FK (not part of a cross-table - cycle) is still deferred by SQLAlchemy either way, but lands in its - *own* table's bucket unless the column declares ``use_alter=True``. - Without that, this method can't see it, and it will insert without - error instead of being caught by validate_deferred_fks(). Model - authors: mark a self-referential FK ``use_alter=True`` explicitly if - it needs anti-join validation on DuckDB. + """FK constraints on *table_cls* this backend's database will not enforce: + - omitted (cross-database/cross-schema), or + - created but never added (cyclic/``use_alter=True``). + See :meth:`validate_deferred_fks`. """ + own_fks = tuple(table_cls.__table__.foreign_key_constraints) + if not own_fks: + return () metadata = table_cls.__table__.metadata - entries = sa.schema.sort_tables_and_constraints(metadata.tables.values()) - deferred = next(fks for tbl, fks in entries if tbl is None) - return tuple(fk for fk in deferred if fk.table is table_cls.__table__) + unenforced = self._unenforced_constraints(tuple(metadata.tables.values())) + return tuple(fk for fk in own_fks if fk in unenforced) - @staticmethod - def _tables_referencing(target: sa.Table) -> set[sa.Table]: - """Other tables in target's own metadata that have an FK pointing at target.""" - return { - other - for other in target.metadata.tables.values() - if other is not target - for fk in other.foreign_key_constraints + def _tables_referencing(self, target: sa.Table) -> set[sa.Table]: + """Other tables with an enforced FK pointing at target. + + Excludes an omitted/deferred FK: DuckDB has no real constraint + there, so it must not count toward check_replace_feasible()'s + preflight. + """ + tables = tuple(target.metadata.tables.values()) + incoming = tuple( + fk for other in tables if other is not target for fk in other.foreign_key_constraints if fk.referred_table is target - } + ) + if not incoming: + return set() + unenforced = self._unenforced_constraints(tables) + return {fk.table for fk in incoming if fk not in unenforced} def _is_referenced_by_existing_rows(self, target: sa.Table, session: so.Session) -> bool: """True if some table referencing target has any rows. @@ -96,13 +114,10 @@ def _is_referenced_by_existing_rows(self, target: sa.Table, session: so.Session) def _staging_rownum_sequence_ref(self, tablename: str) -> str: """Fully schema-qualified, quoted name of tablename's staging _rownum sequence. - - Notes - ----- - DuckDB can't ALTER TABLE ADD COLUMN ... GENERATED ALWAYS AS IDENTITY after creation. - Pagniation uses a sequence-backed DEFAULT instead to achieve the Postgres-like - Identity column: a stable, monotonic pagination key populated automatically on insert, independent of - DuckDB's own rowid (not guaranteed stable across a checkpoint). + + DuckDB can't ALTER TABLE ADD a GENERATED ALWAYS AS IDENTITY column + after creation, so pagination uses a sequence-backed DEFAULT + instead for a stable, monotonic key. """ seq_name = f"{self.staging_name_for_table(tablename)}_rownum_seq" return qualified(self.identifier_preparer, seq_name, physical_schema=self.staging_schema) @@ -202,10 +217,10 @@ def _ensure_pagination_index( @property def _commits_between_pagination_batches(self) -> bool: - """True: bounds how much uncommitted work one load accumulates on a - large table. Safe here because merge_replace updates matching rows - in place rather than deleting and reinserting them, so a commit - between batches never leaves a row removed without its replacement. + """True: bounds uncommitted work on a large table. Safe because + merge_replace updates rows in place rather than delete-then-reinsert, + so a commit between batches never leaves a row removed without its + replacement. """ return True @@ -254,21 +269,14 @@ def merge_replace( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: """Update matching target rows from staging in place, then drop them from staging. - DuckDB rejects deleting a still-referenced row outright, with no - FK-toggle or deferred-constraint escape hatch, so replace is an - UPDATE rather than orm-loader's usual delete-then-reinsert. Matched - rows are removed from staging so the follow-up merge_insert() call - only inserts rows that weren't already handled here. - - Only correct when index_strategy drops indexes before this UPDATE: - updating an indexed column with the index still live raises - DuckDB's own (opaque) constraint error. capabilities. - supports_replace_with_kept_indexes is False for this reason; the - caller is expected to check it before resolving index_strategy to - "keep" for a "replace" merge. + DuckDB can't delete a still-referenced row, so replace is an UPDATE + rather than delete-then-reinsert. Requires index_strategy to have + dropped indexes first; capabilities.supports_replace_with_kept_indexes + is False for this reason. Raises ------ @@ -318,8 +326,9 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): _update(batch_range) _drop_matched_from_staging(batch_range) - if batch_range is not None and self._commits_between_pagination_batches: - session.commit() + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks + ) def merge_upsert( self, @@ -328,6 +337,7 @@ def merge_upsert( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: import duckdb_sqlalchemy @@ -348,8 +358,9 @@ def _upsert(select_: Any) -> sa.Insert: for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_upsert(select_)) - if batch_range is not None and self._commits_between_pagination_batches: - session.commit() + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks + ) def merge_insert( self, @@ -357,6 +368,7 @@ def merge_insert( session: so.Session, *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ @@ -370,5 +382,6 @@ def _insert(select_: Any) -> sa.Insert: for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_insert(select_)) - if batch_range is not None and self._commits_between_pagination_batches: - session.commit() + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks + ) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index cb537b8..e8c2ea9 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from contextlib import contextmanager from typing import TYPE_CHECKING, Any @@ -26,6 +27,7 @@ from .base import BackendCapabilities, DatabaseBackend, requires_capability if TYPE_CHECKING: + from oa_configurator import ResolvedDatabase from sqlalchemy.engine import Connection, Engine from ..loaders.data_classes import LoaderContext @@ -36,8 +38,20 @@ class PostgresBackend(DatabaseBackend): - def __init__(self, *, staging_schema_tag: str | None = None, staging_schema: str | None = None) -> None: - super().__init__(staging_schema_tag=staging_schema_tag, staging_schema=staging_schema) + def __init__( + self, + *, + staging_schema_tag: str | None = None, + staging_schema: str | None = None, + resolved: ResolvedDatabase | None = None, + physical_schemas: Mapping[str, str | None] | None = None, + ) -> None: + super().__init__( + staging_schema_tag=staging_schema_tag, + staging_schema=staging_schema, + resolved=resolved, + physical_schemas=physical_schemas, + ) @staticmethod def staging_name_for_table(tablename: str) -> str: @@ -221,6 +235,7 @@ def merge_replace( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: """Delete every target row whose primary key appears in staging.""" self._validate_merge_batch_size(merge_batch_size) @@ -231,8 +246,9 @@ def merge_replace( for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): where = pk_join if batch_range is None else sa.and_(pk_join, batch_range) session.execute(sa.delete(target).where(where)) - if batch_range is not None and self._commits_between_pagination_batches: - session.commit() + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks + ) def merge_upsert( self, @@ -241,6 +257,7 @@ def merge_upsert( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ @@ -260,8 +277,9 @@ def _upsert(select_: Any) -> sa.Insert: for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_upsert(select_)) - if batch_range is not None and self._commits_between_pagination_batches: - session.commit() + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks + ) def merge_insert( self, @@ -269,6 +287,7 @@ def merge_insert( session: so.Session, *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: """Insert all staging rows into the target table. @@ -287,8 +306,9 @@ def _insert(select_: Any) -> sa.Insert: for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_insert(select_)) - if batch_range is not None and self._commits_between_pagination_batches: - session.commit() + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks + ) @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view( diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 182954a..4f8c7be 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping import logging import sqlite3 from pathlib import Path @@ -16,6 +17,7 @@ from .base import BackendCapabilities, DatabaseBackend, Dialect if TYPE_CHECKING: + from oa_configurator import ResolvedDatabase from sqlalchemy.engine import Connection, Engine from ..tables.typing import CSVTableProtocol @@ -37,6 +39,8 @@ def __init__( *, staging_schema_tag: str | None = None, staging_schema: str | None = None, + resolved: ResolvedDatabase | None = None, + physical_schemas: Mapping[str, str | None] | None = None, busy_timeout_ms: int = 60000, journal_mode: str = "WAL", defer_foreign_keys: bool = True, @@ -46,7 +50,12 @@ def __init__( "SQLite has no schema concept, so staging_schema must be None. " "physical_schema_of() already resolves it to None for this dialect." ) - super().__init__(staging_schema_tag=staging_schema_tag, staging_schema=staging_schema) + super().__init__( + staging_schema_tag=staging_schema_tag, + staging_schema=staging_schema, + resolved=resolved, + physical_schemas=physical_schemas, + ) self.busy_timeout_ms = busy_timeout_ms self.journal_mode = self._validate_journal_mode(journal_mode) self.defer_foreign_keys = defer_foreign_keys @@ -188,8 +197,16 @@ def merge_replace( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: - """Delete every target row whose primary key appears in staging.""" + """Delete every target row whose primary key appears in staging. + + Notes + ----- + validate_deferred_fks is accepted for interface parity only: + _commits_between_pagination_batches is always False here, so no + intermediate commit, and nothing to validate ahead of, ever happens. + """ self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) @@ -217,6 +234,7 @@ def merge_upsert( pk_cols: list[str], *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ @@ -243,6 +261,7 @@ def merge_insert( session: so.Session, *, merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, ) -> None: self._validate_merge_batch_size(merge_batch_size) target = table_cls.__table__ diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index f0a9a70..65a64a0 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -1,5 +1,6 @@ # pyright: reportPrivateUsage=false import logging +from collections.abc import Mapping from contextlib import contextmanager from pathlib import Path from time import perf_counter @@ -7,7 +8,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import Dialect, physical_schema_of +from oa_configurator import Dialect, ResolvedDatabase, physical_schema_of from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError from ..backends.resolve import resolve_backend @@ -374,6 +375,8 @@ def load_csv( merge_batch_size: int | None = None, staging_schema_tag: str | None = None, validate_deferred_fks: bool = True, + resolved: ResolvedDatabase | None = None, + physical_schemas: Mapping[str, str | None] | None = None, ) -> int: """ @@ -417,6 +420,20 @@ def load_csv( validate_deferred_fks After merging, run an anti-join check for any FK this backend's DB can't enforce itself (see DatabaseBackend.validate_deferred_fks()). + Also run before every intermediate commit a batched merge makes, + so a violation is caught before that batch commits rather than + only after the whole load, when an earlier batch's commit can + no longer be rolled back. + resolved + Topology this table's engine(s) were resolved from. Required on + DuckDB whenever table_cls has any foreign key at all, to + classify each as enforced or omitted via oa-configurator's + ``unenforced_foreign_keys()``; omitting it raises rather than + guessing. Ignored by every other backend, and by a DuckDB + table with no foreign keys. + physical_schemas + A schema tag's physical name beyond resolved's own configured + ones. Forwarded alongside resolved. Returns ------- @@ -484,10 +501,19 @@ def load_csv( merge_strategy=merge_strategy, merge_batch_size=merge_batch_size, staging_schema_tag=staging_schema_tag, + validate_deferred_fks=validate_deferred_fks, + resolved=resolved, + physical_schemas=physical_schemas, ) if validate_deferred_fks: logger.info(f"Table `{cls.__tablename__}`: Validating deferred foreign keys.") - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) + backend = resolve_backend( + session, + staging_schema_tag=staging_schema_tag, + mapper=cls, + resolved=resolved, + physical_schemas=physical_schemas, + ) backend.validate_deferred_fks(cls, session) finally: cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) @@ -519,6 +545,9 @@ def merge_from_staging( *, merge_batch_size: int | None = None, staging_schema_tag: str | None = None, + validate_deferred_fks: bool = True, + resolved: ResolvedDatabase | None = None, + physical_schemas: Mapping[str, str | None] | None = None, ): """ Merge data from the staging table into the target table. @@ -533,12 +562,24 @@ def merge_from_staging( staging_schema_tag schema_translate_map tag the staging table lives in. ``None`` means no schema qualification (backend-default behavior). + validate_deferred_fks + Forwarded to the backend's merge methods, which validate + deferred FKs before every intermediate commit a batched merge + makes. See ``load_csv()``. + resolved, physical_schemas + Forwarded to ``resolve_backend()``. See ``load_csv()``. """ target = cls.__tablename__ pk_cols = cls.pk_names() _require_bind(session, cls) - backend = resolve_backend(session, staging_schema_tag=staging_schema_tag, mapper=cls) + backend = resolve_backend( + session, + staging_schema_tag=staging_schema_tag, + mapper=cls, + resolved=resolved, + physical_schemas=physical_schemas, + ) target_empty_confirmed = False if merge_strategy in {"replace", "upsert"}: logger.info( @@ -561,14 +602,18 @@ def merge_from_staging( if merge_strategy == "replace": logger.info(f"Table `{target}`: Merge replace delete phase starting.") delete_started = perf_counter() - backend.merge_replace(cls, session, pk_cols, merge_batch_size=merge_batch_size) + backend.merge_replace( + cls, session, pk_cols, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + ) logger.info( f"Table `{target}`: Merge replace delete phase completed in " f"{_format_elapsed(perf_counter() - delete_started)}." ) logger.info(f"Table `{target}`: Merge insert phase starting.") insert_started = perf_counter() - backend.merge_insert(cls, session, merge_batch_size=merge_batch_size) + backend.merge_insert( + cls, session, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + ) logger.info( f"Table `{target}`: Merge insert phase completed in " f"{_format_elapsed(perf_counter() - insert_started)}." @@ -576,7 +621,9 @@ def merge_from_staging( elif merge_strategy == "upsert": logger.info(f"Table `{target}`: Merge upsert phase starting.") upsert_started = perf_counter() - backend.merge_upsert(cls, session, pk_cols, merge_batch_size=merge_batch_size) + backend.merge_upsert( + cls, session, pk_cols, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + ) logger.info( f"Table `{target}`: Merge upsert phase completed in " f"{_format_elapsed(perf_counter() - upsert_started)}." @@ -599,7 +646,9 @@ def merge_from_staging( logger.info(f"Table `{target}`: Merge insert-if-empty phase starting.") insert_started = perf_counter() - backend.merge_insert(cls, session, merge_batch_size=merge_batch_size) + backend.merge_insert( + cls, session, merge_batch_size=merge_batch_size, validate_deferred_fks=validate_deferred_fks + ) logger.info( f"Table `{target}`: Merge insert-if-empty phase completed in " f"{_format_elapsed(perf_counter() - insert_started)}." diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index 93c6367..0d74efd 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -1,9 +1,12 @@ from typing import Protocol, ClassVar, runtime_checkable, TYPE_CHECKING, Optional, Type, Dict, Any, Unpack, TypedDict import sqlalchemy.orm as so import sqlalchemy as sa +from collections.abc import Mapping from pathlib import Path from contextlib import AbstractContextManager if TYPE_CHECKING: + from oa_configurator import ResolvedDatabase + from ..loaders import LoaderContext, LoaderInterface class ToDictKwargs(TypedDict, total=False): @@ -88,6 +91,8 @@ def load_csv( merge_batch_size: int | None = None, staging_schema_tag: str | None = None, validate_deferred_fks: bool = True, + resolved: Optional["ResolvedDatabase"] = None, + physical_schemas: Optional[Mapping[str, str | None]] = None, ) -> int: ... @classmethod @@ -104,6 +109,9 @@ def merge_from_staging( *, merge_batch_size: int | None = None, staging_schema_tag: str | None = None, + validate_deferred_fks: bool = True, + resolved: Optional["ResolvedDatabase"] = None, + physical_schemas: Optional[Mapping[str, str | None]] = None, ) -> None: ... @classmethod diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 85ea0a2..9ca942e 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -1,20 +1,27 @@ from __future__ import annotations +import tempfile +from pathlib import Path from typing import TYPE_CHECKING, Type, cast import pytest import sqlalchemy as sa import sqlalchemy.orm as so +from oa_configurator import CDMDatabaseConfig, ConnectionConfig, Resolver, StackConfig -from orm_loader.backends import Dialect, DuckDBBackend +from orm_loader.backends import Dialect, DuckDBBackend, staging_schema_claim from orm_loader.backends.base import DeferredForeignKeyViolationError +from orm_loader.helpers.bootstrap import create_tables from tests.models import ( + Base, ComputedColumnTable, CycleATable, CycleBTable, ReferencedParentTable, ReferencingChildTable, UnreferencedGrandparentTable, + VocabReferencingTable, + VocabSchemaTable, ) if TYPE_CHECKING: @@ -24,10 +31,65 @@ _CycleATableCls = cast("Type[CSVTableProtocol]", CycleATable) _CycleBTableCls = cast("Type[CSVTableProtocol]", CycleBTable) _ReferencedParentTableCls = cast("Type[CSVTableProtocol]", ReferencedParentTable) +_VocabReferencingTableCls = cast("Type[CSVTableProtocol]", VocabReferencingTable) +_VocabSchemaTableCls = cast("Type[CSVTableProtocol]", VocabSchemaTable) _TARGET_TABLE = ComputedColumnTable.__tablename__ _STAGING_TABLE = f"_staging_{_TARGET_TABLE}" +@pytest.fixture +def duckdb_backend(duckdb_db): + """DuckDBBackend carrying duckdb_db's own resolved topology, needed by + any FK-classification call now that resolved is required unconditionally. + """ + return DuckDBBackend(resolved=duckdb_db.resolved) + + +@pytest.fixture +def duckdb_session_with_committed_schema(duckdb_db): + """Same database duckdb_session provides, with the schema DDL actually + committed, so a test can roll back its own later work mid-test without + also discarding the schema. + """ + conn = duckdb_db.connection + Base.metadata.create_all(conn) + conn.commit() + return duckdb_db.session + + +@pytest.fixture +def duckdb_split_schema_session(): + """A real DuckDB file with PRIMARY in ``clinical`` and VOCAB in + ``vocabulary``, so VocabReferencingTable's foreign key genuinely crosses + a physical schema and gets omitted from its DDL. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + stack = StackConfig.for_session( + connections={ + "db": ConnectionConfig( + dialect="duckdb", database_name=str(Path(tmp_dir) / "split.duckdb"), test_only=True + ) + }, + databases={ + "cdm": CDMDatabaseConfig( + connection="db", cdm_schema="clinical", vocab_schema="vocabulary", results_schema="results" + ) + }, + ) + resolved = Resolver(stack).resolve_database("cdm") + primary_engine, _ = resolved.create_engines(schema_claims=[staging_schema_claim()]) + try: + with primary_engine.connect() as conn: + create_tables( + conn, [VocabReferencingTable.__table__, VocabSchemaTable.__table__], resolved=resolved + ) + conn.commit() + with so.Session(bind=conn) as session: + yield resolved, session + finally: + primary_engine.dispose() + + def test_duckdb_backend_identity_and_capabilities(): backend = DuckDBBackend() @@ -85,12 +147,8 @@ def test_duckdb_backend_drop_staging_table(duckdb_session): @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_drop_staging_table_also_drops_rownum_sequence(duckdb_session): - """The staging table's _rownum column is backed by a real sequence, so dropping - the staging table must drop that sequence too. - - Notes - ----- - DuckDB can't ALTER TABLE ADD a GENERATED ALWAYS AS IDENTITY column after creation + """The staging table's _rownum column is backed by a real sequence, so + dropping the staging table must drop that sequence too. """ backend = DuckDBBackend() backend.create_staging_table(_ComputedTableCls, duckdb_session) @@ -162,11 +220,9 @@ def test_duckdb_backend_materialized_view_methods_raise(duckdb_session): @pytest.mark.duckdb @pytest.mark.db_dialect -def test_duckdb_backend_deferred_fk_constraints_finds_the_cycle(duckdb_session): - backend = DuckDBBackend() - - a_fks = backend.deferred_fk_constraints(_CycleATableCls) - b_fks = backend.deferred_fk_constraints(_CycleBTableCls) +def test_duckdb_backend_deferred_fk_constraints_finds_the_cycle(duckdb_session, duckdb_backend): + a_fks = duckdb_backend.deferred_fk_constraints(_CycleATableCls) + b_fks = duckdb_backend.deferred_fk_constraints(_CycleBTableCls) assert [c.name for c in a_fks[0].columns] == ["b_id"] assert [c.name for c in b_fks[0].columns] == ["a_id"] @@ -174,36 +230,31 @@ def test_duckdb_backend_deferred_fk_constraints_finds_the_cycle(duckdb_session): @pytest.mark.duckdb @pytest.mark.db_dialect -def test_duckdb_backend_validate_deferred_fks_passes_for_consistent_data(duckdb_session): - backend = DuckDBBackend() +def test_duckdb_backend_validate_deferred_fks_passes_for_consistent_data(duckdb_session, duckdb_backend): duckdb_session.execute(sa.insert(CycleBTable), [{"id": 1, "a_id": None}]) duckdb_session.execute(sa.insert(CycleATable), [{"id": 1, "b_id": 1}]) - backend.validate_deferred_fks(_CycleATableCls, duckdb_session) - backend.validate_deferred_fks(_CycleBTableCls, duckdb_session) + duckdb_backend.validate_deferred_fks(_CycleATableCls, duckdb_session) + duckdb_backend.validate_deferred_fks(_CycleBTableCls, duckdb_session) @pytest.mark.duckdb @pytest.mark.db_dialect -def test_duckdb_backend_validate_deferred_fks_catches_an_orphan(duckdb_session): - """DuckDB's engine never created this FK as a real constraint (it's part - of a cycle, so oa-configurator's shim no-opped the ALTER TABLE ADD - CONSTRAINT), so an orphaned reference inserts without error; the - anti-join check is the only thing that can catch it.""" - backend = DuckDBBackend() +def test_duckdb_backend_validate_deferred_fks_catches_an_orphan(duckdb_session, duckdb_backend): + """The cyclic FK was never created as a real constraint, so an orphaned + reference inserts without error; only the anti-join check catches it. + """ duckdb_session.execute(sa.insert(CycleATable), [{"id": 1, "b_id": 999}]) with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): - backend.validate_deferred_fks(_CycleATableCls, duckdb_session) + duckdb_backend.validate_deferred_fks(_CycleATableCls, duckdb_session) @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_merge_replace_updates_in_place(duckdb_session): - """merge_replace cannot delete a still-referenced row (no FK-toggle escape - hatch), so it updates matching rows in place and leaves non-matching staging - rows for the follow-up merge_insert() call, matching merge_from_staging()'s - real two-step orchestration rather than merge_replace alone. + """merge_replace updates matching rows in place and leaves non-matching + staging rows for the follow-up merge_insert() call. """ backend = DuckDBBackend() backend.create_staging_table(_ComputedTableCls, duckdb_session) @@ -245,15 +296,13 @@ def test_duckdb_backend_merge_replace_paginated_updates_in_place(duckdb_session) @pytest.mark.duckdb @pytest.mark.db_dialect -def test_duckdb_backend_merge_replace_rejects_referenced_table_with_own_fk(duckdb_session): - """DuckDB rejects a SET on any FK-constrained column of a row that's - still referenced elsewhere, regardless of index_strategy. merge_replace - must fail up front, before touching any data, rather than partway - through a batch.""" - backend = DuckDBBackend() - backend.create_staging_table(_ReferencedParentTableCls, duckdb_session) +def test_duckdb_backend_merge_replace_rejects_referenced_table_with_own_fk(duckdb_session, duckdb_backend): + """merge_replace must reject a still-referenced row up front, before + touching any data, rather than partway through a batch. + """ + duckdb_backend.create_staging_table(_ReferencedParentTableCls, duckdb_session) staging = _ReferencedParentTableCls.get_staging_table( - duckdb_session, staging_schema_tag=backend.staging_schema_tag + duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag ) duckdb_session.execute(sa.insert(UnreferencedGrandparentTable), [{"id": 1, "name": "g1"}]) @@ -264,4 +313,110 @@ def test_duckdb_backend_merge_replace_rejects_referenced_table_with_own_fk(duckd duckdb_session.execute(sa.insert(staging), [{"id": 1, "grandparent_id": 1, "name": "p1-updated"}]) with pytest.raises(ValueError, match="referenced_parent_table"): - backend.merge_replace(_ReferencedParentTableCls, duckdb_session, ["id"]) + duckdb_backend.merge_replace(_ReferencedParentTableCls, duckdb_session, ["id"]) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_batch_validation_prevents_partial_commit( + duckdb_session_with_committed_schema, duckdb_backend, +): + """merge_batch_size=1 commits each row separately. Row 3's own batch + (the orphan) must be caught and never committed, while rows 1 and 2, + already validated and committed in their own batches, survive. + """ + session = duckdb_session_with_committed_schema + duckdb_backend.create_staging_table(_CycleATableCls, session) + staging = _CycleATableCls.get_staging_table(session, staging_schema_tag=duckdb_backend.staging_schema_tag) + session.execute(sa.insert(staging), [{"id": 1, "b_id": None}, {"id": 2, "b_id": None}, {"id": 3, "b_id": 999}]) + + with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): + duckdb_backend.merge_insert(_CycleATableCls, session, merge_batch_size=1) + + session.rollback() + remaining = session.execute(sa.select(CycleATable.id).order_by(CycleATable.id)).scalars().all() + assert remaining == [1, 2] + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_batch_validation_is_a_no_op_without_deferred_fks(duckdb_session): + """A table with no deferred/unenforced FK must not pay for a validation + query on every batch commit. + """ + backend = DuckDBBackend() + backend.create_staging_table(_ComputedTableCls, duckdb_session) + staging = _ComputedTableCls.get_staging_table(duckdb_session, staging_schema_tag=backend.staging_schema_tag) + duckdb_session.execute(sa.insert(staging), [{"id": i, "name": f"n{i}"} for i in range(5)]) + + backend.merge_insert(_ComputedTableCls, duckdb_session, merge_batch_size=1) + + rows = duckdb_session.execute(sa.select(ComputedColumnTable.id)).scalars().all() + assert sorted(rows) == list(range(5)) + + +def test_duckdb_backend_deferred_fk_constraints_without_resolved_raises_on_cross_schema_fk(): + """`resolved` is required unconditionally, not just when a cross-schema + FK is actually present: a cheaper "cyclic-only" fallback could still + misclassify an unrelated cross-schema FK elsewhere in the metadata. + """ + backend = DuckDBBackend() + + with pytest.raises(ValueError, match="resolved"): + backend.deferred_fk_constraints(_VocabReferencingTableCls) + + +def test_duckdb_backend_deferred_fk_constraints_without_resolved_raises_even_for_a_cyclic_fk(): + """Same requirement for a purely cyclic FK that never crosses a schema + tag, since the method has no cheap way to tell that in advance. + """ + backend = DuckDBBackend() + + with pytest.raises(ValueError, match="resolved"): + backend.deferred_fk_constraints(_CycleATableCls) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_deferred_fk_constraints_classifies_cross_schema_fk_with_resolved( + duckdb_split_schema_session, +): + """With `resolved`, the omitted cross-schema key is classified the same + as a cyclic one: unenforced. + """ + resolved, _session = duckdb_split_schema_session + backend = DuckDBBackend(resolved=resolved) + + fks = backend.deferred_fk_constraints(_VocabReferencingTableCls) + + assert [c.name for c in fks[0].columns] == ["vocab_id"] + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_validate_deferred_fks_catches_a_cross_schema_orphan(duckdb_split_schema_session): + """The cross-schema key was never created, so an orphaned reference + inserts without error; only the anti-join check catches it. + """ + resolved, session = duckdb_split_schema_session + backend = DuckDBBackend(resolved=resolved) + session.execute(sa.insert(VocabReferencingTable), [{"id": 1, "vocab_id": 999}]) + + with pytest.raises(DeferredForeignKeyViolationError, match="vocab_referencing_table"): + backend.validate_deferred_fks(_VocabReferencingTableCls, session) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_tables_referencing_excludes_an_omitted_cross_schema_fk(duckdb_split_schema_session): + """VocabReferencingTable's foreign key into VocabSchemaTable was never + created, so check_replace_feasible() must not treat it as a real + referencing table. + """ + resolved, session = duckdb_split_schema_session + backend = DuckDBBackend(resolved=resolved) + session.execute(sa.insert(VocabSchemaTable), [{"id": 1, "name": "v1"}]) + session.execute(sa.insert(VocabReferencingTable), [{"id": 1, "vocab_id": 1}]) + + assert backend._tables_referencing(VocabSchemaTable.__table__) == set() + assert backend._is_referenced_by_existing_rows(VocabSchemaTable.__table__, session) is False diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 7af044a..5cf7bf2 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -845,7 +845,7 @@ def test_sqlite_failed_batched_replace_keeps_original_target_rows(sqlite_session @pytest.mark.duckdb @pytest.mark.db_dialect -def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_session, tmp_path): +def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_session, duckdb_db, tmp_path): """cycle_a_table.b_id is part of a genuine FK cycle with cycle_b_table, so DuckDB never got a real constraint for it (oa-configurator's _skip_deferred_fk_on_duckdb no-oped the ALTER TABLE ADD CONSTRAINT): a @@ -863,7 +863,12 @@ def test_validate_deferred_fks_catches_an_orphan_introduced_by_load_csv(duckdb_s with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): _CycleATable.load_csv( - duckdb_session, csv_path, dedupe=False, loader=PandasLoader(), validate_deferred_fks=True + duckdb_session, + csv_path, + dedupe=False, + loader=PandasLoader(), + validate_deferred_fks=True, + resolved=duckdb_db.resolved, ) staging_name = resolve_backend(duckdb_session, mapper=_CycleATable).staging_name_for_table( diff --git a/tests/models.py b/tests/models.py index 372acdb..b5697c1 100644 --- a/tests/models.py +++ b/tests/models.py @@ -166,3 +166,19 @@ class VocabSchemaTable(Base, CSVLoadableTableInterface): id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) name: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) + + +class VocabReferencingTable(Base, CSVLoadableTableInterface): + """PRIMARY-tagged table with a foreign key into VocabSchemaTable (VOCAB), + for cross-schema foreign-key classification tests: DuckDB cannot + enforce a key between two physical schemas, so oa-configurator omits + it from DDL, same practical outcome as a cyclic/deferred key but + reached a different way.""" + + __tablename__ = "vocab_referencing_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + vocab_id: so.Mapped[int | None] = so.mapped_column( + sa.Integer, sa.ForeignKey(f"{SchemaRole.VOCAB.value}.test_vocab_role_table.id") + ) From bcb7eade5028eccf0372d7ee378478c954c23a91 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 27 Oct 2026 21:56:04 +0000 Subject: [PATCH 11/17] Load self-referential FKs in dependency order, validate cross-database FKs for real - merge_insert/merge_upsert load a self-referential deferred FK wave by wave instead of by insertion order, so a child batched before its own parent no longer false-positives, while a genuine orphan is still caught before commit - DuckDBBackend.validate_deferred_fks checks a cross-database FK by querying each connection separately and comparing in Python, instead of skipping it --- src/orm_loader/backends/base.py | 96 +++++++++++++++++- src/orm_loader/backends/duckdb.py | 120 ++++++++++++++++++++-- tests/backends/test_duckdb_backend.py | 141 ++++++++++++++++++++++++++ tests/models.py | 21 +++- 4 files changed, 363 insertions(+), 15 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index e21d484..8888e25 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -1,19 +1,18 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Generator, Iterator, Mapping from contextlib import AbstractContextManager, contextmanager from dataclasses import dataclass -from collections.abc import Generator, Iterator, Mapping from functools import wraps from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Type, TypeVar, cast import sqlalchemy as sa import sqlalchemy.orm as so +from oa_configurator import Dialect, ResolvedDatabase, SchemaClaim, open_connection, qualified from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer -from oa_configurator import Dialect, ResolvedDatabase, SchemaClaim, open_connection, qualified - if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext from ..mappers.materialised_view_contracts import MaterializedViewIndex @@ -575,6 +574,97 @@ def _commit_batch_if_needed( self.validate_deferred_fks(table_cls, session) session.commit() + def _self_referential_deferred_fks( + self, table_cls: Type["CSVTableProtocol"] + ) -> tuple[sa.ForeignKeyConstraint, ...]: + """Deferred FKs on table_cls that point back at table_cls itself. + + A row's match for one of these can arrive in a later pagination + batch of the same load, since batching orders by insertion + position rather than by dependency. merge_insert()/merge_upsert() + must load table_cls in dependency order instead when any exist; + see :meth:`_merge_in_dependency_waves`. + """ + target = table_cls.__table__ + return tuple(fk for fk in self.deferred_fk_constraints(table_cls) if fk.referred_table is target) + + def _merge_in_dependency_waves( + self, + table_cls: Type["CSVTableProtocol"], + session: so.Session, + staging: sa.Table, + target: sa.Table, + self_referential_fks: tuple[sa.ForeignKeyConstraint, ...], + merge_batch_size: int | None, + validate_deferred_fks: bool, + insert_chunk: Callable[[sa.ColumnElement[bool]], Any], + ) -> None: + """Insert staging rows wave by wave for a self-referential deferred FK. + + Each wave is exactly the staging rows whose self-referential FK + value(s) are already satisfied, by being NULL or by matching a row + already committed to target (from an earlier wave, or already + there before this load). insert_chunk() inserts one wave (or a + merge_batch_size-sized slice of it); the matching staging rows are + then removed and the wave committed and validated, before the next + wave is computed against the now-larger target. + + Every wave is safe to validate and commit immediately. Nothing in it + can be a false positive, since its parent row is already committed. + A row that never becomes ready, once a pass makes no further progress, + is a genuine violation (an orphan, or a real cycle) and is raised directly + instead of being silently left in staging. + + Notes + ----- + One query round-trip per wave, so cost tracks how many generations + deep the self-reference goes, not how many rows there are. Fine for + a hierarchy that's a few levels deep with many rows per level (e.g. + a category tree). Slow for a long chain where each row's parent is + the previous row (i.e. one wave per row). + + Raises + ------ + DeferredForeignKeyViolationError + A row's self-referential FK value never matches any committed + row once no further wave makes progress. + """ + per_fk = [] + for fk in self_referential_fks: + local_cols: list[sa.ColumnElement[Any]] = [staging.c[col.name] for col in fk.columns] + remote_cols = [cast(sa.ColumnElement[Any], element.column) for element in fk.elements] + not_null = sa.and_(*(col.isnot(None) for col in local_cols)) + has_match = sa.exists( + sa.select(sa.literal(1)).select_from(target).where( + sa.and_(*(local_cols[i] == remote_cols[i] for i in range(len(remote_cols)))) + ) + ) + per_fk.append((fk, local_cols, not_null, has_match)) + ready = sa.and_(*(sa.or_(sa.not_(not_null), has_match) for _, _, not_null, has_match in per_fk)) + + order_column = self._pagination_column(table_cls, staging) + self._ensure_pagination_index(table_cls, staging, session) + + while True: + query = sa.select(order_column).where(ready).order_by(order_column) + if merge_batch_size is not None: + query = query.limit(merge_batch_size) + keys = session.execute(query).scalars().all() + if not keys: + break + chunk = order_column.in_(keys) + insert_chunk(chunk) + session.execute(sa.delete(staging).where(chunk)) + self._commit_batch_if_needed(table_cls, session, chunk, validate_deferred_fks=validate_deferred_fks) + + if session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one(): + for fk, local_cols, not_null, has_match in per_fk: + stuck = session.execute( + sa.select(*local_cols).where(not_null, sa.not_(has_match)).limit(5) + ).all() + if stuck: + raise DeferredForeignKeyViolationError(fk, [tuple(row) for row in stuck]) + def merge_context( self, table_cls: Type["CSVTableProtocol"], diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 0efb47c..282a069 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -4,11 +4,18 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import physical_schema_of, qualified, unenforced_foreign_keys +from oa_configurator import ( + ResolvedCDMDatabase, + Role, + UnenforcedForeignKey, + physical_schema_of, + qualified, + unenforced_foreign_keys, +) from sqlalchemy.schema import CreateSequence, DropSequence, DropTable from sqlalchemy.sql.compiler import IdentifierPreparer -from .base import BackendCapabilities, DatabaseBackend, Dialect +from .base import BackendCapabilities, DatabaseBackend, DeferredForeignKeyViolationError, Dialect if TYPE_CHECKING: from ..tables.typing import CSVTableProtocol @@ -43,9 +50,12 @@ def capabilities(self) -> BackendCapabilities: supports_replace_with_kept_indexes=False, ) - def _unenforced_constraints(self, tables: tuple[sa.Table, ...]) -> frozenset[sa.ForeignKeyConstraint]: - """Every FK in *tables* DuckDB will not enforce: omitted (cross-database/ - cross-schema) or created but never added (cyclic/``use_alter=True``). + def _unenforced_constraints( + self, tables: tuple[sa.Table, ...] + ) -> tuple[UnenforcedForeignKey, ...]: + """Every FK in *tables* DuckDB will not enforce, including reason: + - omitted (cross-database/cross-schema) or + - created but never added (cyclic/``use_alter=True``). Delegates to oa-configurator's :func:`unenforced_foreign_keys`, the same classification DDL creation itself uses. @@ -60,27 +70,97 @@ def _unenforced_constraints(self, tables: tuple[sa.Table, ...]) -> frozenset[sa. "DuckDBBackend needs resolved= to classify foreign-key " "enforcement. Pass it to load_csv(...) or resolve_backend(...)." ) - unenforced = unenforced_foreign_keys( + return unenforced_foreign_keys( tables, resolved=self.resolved, physical_schemas=self.physical_schemas ) - return frozenset(entry.constraint for entry in unenforced) def deferred_fk_constraints( self, table_cls: Type["CSVTableProtocol"], ) -> tuple[sa.ForeignKeyConstraint, ...]: """FK constraints on *table_cls* this backend's database will not enforce: - - omitted (cross-database/cross-schema), or - - created but never added (cyclic/``use_alter=True``). + - omitted (cross-database/cross-schema), or + - created but never added (cyclic/``use_alter=True``). See :meth:`validate_deferred_fks`. """ own_fks = tuple(table_cls.__table__.foreign_key_constraints) if not own_fks: return () metadata = table_cls.__table__.metadata - unenforced = self._unenforced_constraints(tuple(metadata.tables.values())) + unenforced = {entry.constraint for entry in self._unenforced_constraints(tuple(metadata.tables.values()))} return tuple(fk for fk in own_fks if fk in unenforced) + def validate_deferred_fks(self, table_cls: Type["CSVTableProtocol"], session: so.Session) -> None: + """Validate every deferred FK on table_cls, including a cross-database one. + + A cross-schema or cyclic deferred FK points to a table in the same + database, so one query can check both tables directly. A + cross-database FK points to a table in a different database: a + query touching both would raise CrossDatabaseStatementError no + matter what the data actually is, so _validate_cross_database_fk + runs one query per database instead. + """ + own_fks = tuple(table_cls.__table__.foreign_key_constraints) + if not own_fks: + return + metadata = table_cls.__table__.metadata + for entry in self._unenforced_constraints(tuple(metadata.tables.values())): + if entry.constraint not in own_fks: + continue + if entry.reason == "cross_database": + self._validate_cross_database_fk(entry.constraint, session) + else: + self._validate_one_deferred_fk(entry.constraint, session) + + def _validate_cross_database_fk(self, fk: sa.ForeignKeyConstraint, session: so.Session) -> None: + """Anti-join fk across two connections by comparing both sides in Python. + + Pulls fk's distinct non-null local values from *session*'s own + connection, then opens the referred table's own connection (via + self.resolved) to check which of them actually exist there. One + bulk query per side instead of one query spanning both. + + Raises + ------ + TypeError + self.resolved is not a ResolvedCDMDatabase: a cross-database FK + can only arise when PRIMARY/VOCAB are split across connections, + which only a CDM-aware resolved can describe. + DeferredForeignKeyViolationError + A local value has no match on the referred connection. + """ + if not isinstance(self.resolved, ResolvedCDMDatabase): + raise TypeError( + f"{fk.name} was classified cross_database, which requires split PRIMARY/VOCAB " + f"connections, but self.resolved is {type(self.resolved).__name__}, not " + "ResolvedCDMDatabase." + ) + local_cols = list(fk.columns) + remote_cols = [element.column for element in fk.elements] + not_null = sa.and_(*(col.isnot(None) for col in local_cols)) + local_values = session.execute(sa.select(*local_cols).where(not_null).distinct()).all() + if not local_values: + return + + remote_tag = remote_cols[0].table.schema + primary_engine, vocab_engine = self.resolved.create_engines(register_claims=False) + try: + remote_engine = vocab_engine if remote_tag == Role.VOCAB.value else primary_engine + with remote_engine.connect() as remote_conn: + if len(remote_cols) == 1: + condition = remote_cols[0].in_(value for (value,) in local_values) + else: + condition = sa.tuple_(*remote_cols).in_(tuple(value) for value in local_values) + existing = {tuple(row) for row in remote_conn.execute(sa.select(*remote_cols).where(condition))} + finally: + primary_engine.dispose() + if vocab_engine is not primary_engine: + vocab_engine.dispose() + + missing = [tuple(value) for value in local_values if tuple(value) not in existing] + if missing: + raise DeferredForeignKeyViolationError(fk, missing[:5]) + def _tables_referencing(self, target: sa.Table) -> set[sa.Table]: """Other tables with an enforced FK pointing at target. @@ -95,7 +175,7 @@ def _tables_referencing(self, target: sa.Table) -> set[sa.Table]: ) if not incoming: return set() - unenforced = self._unenforced_constraints(tables) + unenforced = {entry.constraint for entry in self._unenforced_constraints(tables)} return {fk.table for fk in incoming if fk not in unenforced} def _is_referenced_by_existing_rows(self, target: sa.Table, session: so.Session) -> bool: @@ -355,6 +435,15 @@ def _upsert(select_: Any) -> sa.Insert: non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + self_referential_fks = self._self_referential_deferred_fks(table_cls) + if self_referential_fks: + self._merge_in_dependency_waves( + table_cls, session, staging, target, self_referential_fks, + merge_batch_size, validate_deferred_fks, + insert_chunk=lambda chunk: session.execute(_upsert(non_paginated_select.where(chunk))), + ) + return + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_upsert(select_)) @@ -379,6 +468,15 @@ def merge_insert( def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) + self_referential_fks = self._self_referential_deferred_fks(table_cls) + if self_referential_fks: + self._merge_in_dependency_waves( + table_cls, session, staging, target, self_referential_fks, + merge_batch_size, validate_deferred_fks, + insert_chunk=lambda chunk: session.execute(_insert(non_paginated_select.where(chunk))), + ) + return + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_insert(select_)) diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 9ca942e..4b18653 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -19,6 +19,7 @@ CycleBTable, ReferencedParentTable, ReferencingChildTable, + SelfReferentialTable, UnreferencedGrandparentTable, VocabReferencingTable, VocabSchemaTable, @@ -31,6 +32,7 @@ _CycleATableCls = cast("Type[CSVTableProtocol]", CycleATable) _CycleBTableCls = cast("Type[CSVTableProtocol]", CycleBTable) _ReferencedParentTableCls = cast("Type[CSVTableProtocol]", ReferencedParentTable) +_SelfReferentialTableCls = cast("Type[CSVTableProtocol]", SelfReferentialTable) _VocabReferencingTableCls = cast("Type[CSVTableProtocol]", VocabReferencingTable) _VocabSchemaTableCls = cast("Type[CSVTableProtocol]", VocabSchemaTable) _TARGET_TABLE = ComputedColumnTable.__tablename__ @@ -90,6 +92,40 @@ def duckdb_split_schema_session(): primary_engine.dispose() +@pytest.fixture +def duckdb_split_database_session(): + """PRIMARY and VOCAB on two separate DuckDB files/connections, so + VocabReferencingTable's foreign key is classified "cross_database" + rather than "cross_schema". No single query can join the two. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + stack = StackConfig.for_session( + connections={ + "primary": ConnectionConfig( + dialect="duckdb", database_name=str(Path(tmp_dir) / "primary.duckdb"), test_only=True + ), + "vocab": ConnectionConfig( + dialect="duckdb", database_name=str(Path(tmp_dir) / "vocab.duckdb"), test_only=True + ), + }, + databases={"cdm": CDMDatabaseConfig(connection="primary", vocab_connection="vocab")}, + ) + resolved = Resolver(stack).resolve_database("cdm") + primary_engine, vocab_engine = resolved.create_engines(schema_claims=[staging_schema_claim()]) + try: + with vocab_engine.connect() as vocab_conn: + create_tables(vocab_conn, [VocabSchemaTable.__table__], resolved=resolved) + vocab_conn.commit() + with primary_engine.connect() as conn: + create_tables(conn, [VocabReferencingTable.__table__], resolved=resolved) + conn.commit() + with so.Session(bind=conn) as session: + yield resolved, session + finally: + primary_engine.dispose() + vocab_engine.dispose() + + def test_duckdb_backend_identity_and_capabilities(): backend = DuckDBBackend() @@ -338,6 +374,111 @@ def test_duckdb_backend_merge_insert_batch_validation_prevents_partial_commit( assert remaining == [1, 2] +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_loads_self_referential_rows_in_dependency_order( + duckdb_session, duckdb_backend, +): + """Staging rows arrive in reverse dependency order (each pointing at a + row not yet committed); merge_insert must still load every row, in + however many waves that takes, regardless of insertion order. + """ + duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) + staging = _SelfReferentialTableCls.get_staging_table( + duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + duckdb_session.execute( + sa.insert(staging), + [{"id": 3, "parent_id": 2}, {"id": 2, "parent_id": 1}, {"id": 1, "parent_id": None}], + ) + + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, merge_batch_size=1) + + rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} + assert rows == {1: None, 2: 1, 3: 2} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_upsert_loads_self_referential_rows_in_dependency_order( + duckdb_session, duckdb_backend, +): + duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) + staging = _SelfReferentialTableCls.get_staging_table( + duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + duckdb_session.execute(sa.insert(staging), [{"id": 2, "parent_id": 1}, {"id": 1, "parent_id": None}]) + + duckdb_backend.merge_upsert(_SelfReferentialTableCls, duckdb_session, ["id"], merge_batch_size=1) + + rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} + assert rows == {1: None, 2: 1} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_dependency_waves_catch_a_genuine_orphan( + duckdb_session_with_committed_schema, duckdb_backend, +): + """A row whose parent never arrives in any wave is a genuine violation. + A row that did become ready must still be committed. Only the stuck + row is left behind in staging and reported. + """ + session = duckdb_session_with_committed_schema + duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) + staging = _SelfReferentialTableCls.get_staging_table( + session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + session.execute(sa.insert(staging), [{"id": 1, "parent_id": None}, {"id": 2, "parent_id": 999}]) + + with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): + duckdb_backend.merge_insert(_SelfReferentialTableCls, session, merge_batch_size=1) + + assert session.execute(sa.select(SelfReferentialTable.id)).scalars().all() == [1] + assert session.execute(sa.select(staging.c.id)).scalars().all() == [2] + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_validate_deferred_fks_catches_a_cross_database_orphan(duckdb_split_database_session): + """A cross-database FK is genuinely checked, via a second connection, + not skipped. An orphan on the other database must still be caught. + """ + resolved, session = duckdb_split_database_session + backend = DuckDBBackend(resolved=resolved) + session.execute(sa.insert(VocabReferencingTable), [{"id": 1, "vocab_id": 999}]) + + with pytest.raises(DeferredForeignKeyViolationError, match="vocab_referencing_table"): + backend.validate_deferred_fks(_VocabReferencingTableCls, session) + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_validate_deferred_fks_passes_a_real_cross_database_reference( + duckdb_split_database_session, +): + """A cross-database FK pointing at a row that genuinely exists on the + other connection must pass, proving the check is real rather than a + pass-through. + """ + resolved, session = duckdb_split_database_session + backend = DuckDBBackend(resolved=resolved) + + primary_engine, vocab_engine = resolved.create_engines(register_claims=False) + try: + with vocab_engine.connect() as vocab_conn: + vocab_conn.execute(sa.insert(VocabSchemaTable), [{"id": 1, "name": "v1"}]) + vocab_conn.commit() + finally: + primary_engine.dispose() + if vocab_engine is not primary_engine: + vocab_engine.dispose() + + session.execute(sa.insert(VocabReferencingTable), [{"id": 1, "vocab_id": 1}]) + + backend.validate_deferred_fks(_VocabReferencingTableCls, session) + + @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_merge_insert_batch_validation_is_a_no_op_without_deferred_fks(duckdb_session): diff --git a/tests/models.py b/tests/models.py index b5697c1..7ad0adb 100644 --- a/tests/models.py +++ b/tests/models.py @@ -2,9 +2,10 @@ from enum import Enum import sqlalchemy as sa +import sqlalchemy.orm as so from oa_configurator import Role as SchemaRole from sqlalchemy.orm import declarative_base -import sqlalchemy.orm as so + from orm_loader.tables import CSVLoadableTableInterface Base = declarative_base() @@ -182,3 +183,21 @@ class VocabReferencingTable(Base, CSVLoadableTableInterface): vocab_id: so.Mapped[int | None] = so.mapped_column( sa.Integer, sa.ForeignKey(f"{SchemaRole.VOCAB.value}.test_vocab_role_table.id") ) + + +class SelfReferentialTable(Base, CSVLoadableTableInterface): + """A table referencing its own rows, for dependency-ordered batching + tests: a child row's parent may legitimately land in a later + pagination batch of the same load. use_alter=True since DuckDB can't + ALTER TABLE ADD this constraint, making it deferred/unenforced like a + cross-table cycle, rather than a plain inline (and thus DuckDB-enforced) + self-reference.""" + + __tablename__ = "self_referential_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + parent_id: so.Mapped[int | None] = so.mapped_column( + sa.Integer, + sa.ForeignKey(f"{SchemaRole.PRIMARY.value}.self_referential_table.id", use_alter=True), + ) From 36200a0ae44fd630153ff07ddd0a647639d3238b Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 27 Oct 2026 21:58:58 +0000 Subject: [PATCH 12/17] Load merge_replace's self-referential updates and inserts together, in dependency order Replace and insert ran as two sealed phases, so an existing row's self-referential FK could be re-pointed at a row only the later insert phase adds, validating it before that row exists. merge_replace now updates and inserts together through the same dependency-ordered wave loop, draining staging itself so the merge_insert() call that follows has nothing left to do. --- src/orm_loader/backends/duckdb.py | 34 ++++++++++++++++ tests/backends/test_duckdb_backend.py | 58 +++++++++++++++++++++++++++ 2 files changed, 92 insertions(+) diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 282a069..e89c040 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -358,6 +358,17 @@ def merge_replace( dropped indexes first; capabilities.supports_replace_with_kept_indexes is False for this reason. + If table_cls has a self-referential deferred FK, an existing row's + FK column can be re-pointed at a row that doesn't exist yet -- one + this same load will only add via the merge_insert() call that + normally follows. Processing replace and insert as two sealed + phases would validate that update before the row it points at + exists. Instead, both the update (for a row already in target) and + the insert (for one that isn't) are done together, wave by wave, + through merge_insert()'s dependency-ordered loop (see + _merge_in_dependency_waves()); staging ends up fully drained here, + so the merge_insert() call that follows has nothing left to do. + Raises ------ ValueError @@ -403,6 +414,29 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N _drop_matched_from_staging() return + self_referential_fks = self._self_referential_deferred_fks(table_cls) + if self_referential_fks: + insertable_cols = self._insertable_column_names(table_cls) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + + def _insert(select_: Any) -> sa.Insert: + return sa.insert(target).from_select(insertable_cols, select_) + + def _process_chunk(chunk: sa.ColumnElement[bool]) -> None: + # Rows in chunk already in target are updated in place; whatever + # staging rows chunk still matches afterward are new, and get + # inserted. _merge_in_dependency_waves() removes them from + # staging once this returns. + _update(chunk) + _drop_matched_from_staging(chunk) + session.execute(_insert(non_paginated_select.where(chunk))) + + self._merge_in_dependency_waves( + table_cls, session, staging, target, self_referential_fks, + merge_batch_size, validate_deferred_fks, insert_chunk=_process_chunk, + ) + return + for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): _update(batch_range) _drop_matched_from_staging(batch_range) diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 4b18653..4497442 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -415,6 +415,64 @@ def test_duckdb_backend_merge_upsert_loads_self_referential_rows_in_dependency_o assert rows == {1: None, 2: 1} +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_replace_updates_and_inserts_self_referential_rows_together( + duckdb_session, duckdb_backend, +): + """An existing row's self-referential FK is re-pointed at a row that + only this same load adds. Processed as two sealed phases (replace, + then insert), the update would be validated before its new target + exists. merge_replace must do both together, in dependency order, so + the update only lands once its target is actually there -- and drain + staging itself, leaving nothing for merge_insert() to do. + """ + duckdb_session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) + duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) + staging = _SelfReferentialTableCls.get_staging_table( + duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + duckdb_session.execute( + sa.insert(staging), [{"id": 1, "parent_id": 2}, {"id": 2, "parent_id": None}] + ) + + duckdb_backend.merge_replace(_SelfReferentialTableCls, duckdb_session, ["id"], merge_batch_size=1) + + rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} + assert rows == {1: 2, 2: None} + assert duckdb_session.execute(sa.select(staging.c.id)).scalars().all() == [] + + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session) + rows_after_insert = { + r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all() + } + assert rows_after_insert == rows + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_replace_dependency_waves_catch_a_genuine_orphan( + duckdb_session_with_committed_schema, duckdb_backend, +): + """A replace batch re-pointing an existing row at a genuine orphan (not + merely a not-yet-inserted new row) must still be caught, not mistaken + for a forward reference that will resolve later. + """ + session = duckdb_session_with_committed_schema + session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) + duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) + staging = _SelfReferentialTableCls.get_staging_table( + session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + session.execute(sa.insert(staging), [{"id": 1, "parent_id": 999}]) + + with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): + duckdb_backend.merge_replace(_SelfReferentialTableCls, session, ["id"], merge_batch_size=1) + + assert session.execute(sa.select(SelfReferentialTable.parent_id)).scalars().all() == [None] + assert session.execute(sa.select(staging.c.id)).scalars().all() == [1] + + @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_merge_insert_dependency_waves_catch_a_genuine_orphan( From b99acd09bc15ab4b442a2c6deaf7945883070e9d Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 27 Oct 2026 22:13:49 +0000 Subject: [PATCH 13/17] Trim yapping in merge_replace's self-referential docstring and comment --- src/orm_loader/backends/duckdb.py | 22 ++++++++-------------- tests/backends/test_duckdb_backend.py | 9 +++------ 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index e89c040..be1415b 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -358,16 +358,12 @@ def merge_replace( dropped indexes first; capabilities.supports_replace_with_kept_indexes is False for this reason. - If table_cls has a self-referential deferred FK, an existing row's - FK column can be re-pointed at a row that doesn't exist yet -- one - this same load will only add via the merge_insert() call that - normally follows. Processing replace and insert as two sealed - phases would validate that update before the row it points at - exists. Instead, both the update (for a row already in target) and - the insert (for one that isn't) are done together, wave by wave, - through merge_insert()'s dependency-ordered loop (see - _merge_in_dependency_waves()); staging ends up fully drained here, - so the merge_insert() call that follows has nothing left to do. + A self-referential FK can point an existing row at one this load + hasn't inserted yet. Doing replace then insert as separate phases + would validate that before the target exists, so both happen + together here, wave by wave (see _merge_in_dependency_waves()), + draining staging. The merge_insert() call that follows has + nothing left to do. Raises ------ @@ -423,10 +419,8 @@ def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) def _process_chunk(chunk: sa.ColumnElement[bool]) -> None: - # Rows in chunk already in target are updated in place; whatever - # staging rows chunk still matches afterward are new, and get - # inserted. _merge_in_dependency_waves() removes them from - # staging once this returns. + # chunk's rows already in target are updated; what's left in + # staging afterward is new and gets inserted here. _update(chunk) _drop_matched_from_staging(chunk) session.execute(_insert(non_paginated_select.where(chunk))) diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 4497442..9278dd0 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -420,12 +420,9 @@ def test_duckdb_backend_merge_upsert_loads_self_referential_rows_in_dependency_o def test_duckdb_backend_merge_replace_updates_and_inserts_self_referential_rows_together( duckdb_session, duckdb_backend, ): - """An existing row's self-referential FK is re-pointed at a row that - only this same load adds. Processed as two sealed phases (replace, - then insert), the update would be validated before its new target - exists. merge_replace must do both together, in dependency order, so - the update only lands once its target is actually there -- and drain - staging itself, leaving nothing for merge_insert() to do. + """An existing row is re-pointed at one only this load adds. Processed + as two sealed phases, the update would be validated before its target + exists; merge_replace must do both together, in dependency order. """ duckdb_session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) From bc9030102d7e8682dce1bdc8c429989e334ca732 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 27 Oct 2026 23:07:47 +0000 Subject: [PATCH 14/17] Share engine-dispose helper between bootstrap and cross-database FK validation Co-Authored-By: Claude Sonnet 5 --- src/orm_loader/backends/duckdb.py | 5 ++--- src/orm_loader/helpers/__init__.py | 2 ++ src/orm_loader/helpers/bootstrap.py | 5 ++--- src/orm_loader/helpers/engines.py | 8 ++++++++ 4 files changed, 14 insertions(+), 6 deletions(-) create mode 100644 src/orm_loader/helpers/engines.py diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index be1415b..68709c5 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -15,6 +15,7 @@ from sqlalchemy.schema import CreateSequence, DropSequence, DropTable from sqlalchemy.sql.compiler import IdentifierPreparer +from ..helpers.engines import dispose_engines from .base import BackendCapabilities, DatabaseBackend, DeferredForeignKeyViolationError, Dialect if TYPE_CHECKING: @@ -153,9 +154,7 @@ def _validate_cross_database_fk(self, fk: sa.ForeignKeyConstraint, session: so.S condition = sa.tuple_(*remote_cols).in_(tuple(value) for value in local_values) existing = {tuple(row) for row in remote_conn.execute(sa.select(*remote_cols).where(condition))} finally: - primary_engine.dispose() - if vocab_engine is not primary_engine: - vocab_engine.dispose() + dispose_engines(primary_engine, vocab_engine) missing = [tuple(value) for value in local_values if tuple(value) not in existing] if missing: diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index 18707e2..2c4ed42 100644 --- a/src/orm_loader/helpers/__init__.py +++ b/src/orm_loader/helpers/__init__.py @@ -1,5 +1,6 @@ from .errors import IngestError, ValidationError from .bootstrap import bootstrap, create_db, create_tables +from .engines import dispose_engines from .sqlite import ( attach_sqlite_bulk_load_pragmas, explain_sqlite_fk_error, @@ -16,6 +17,7 @@ "bootstrap", "create_db", "create_tables", + "dispose_engines", "attach_sqlite_bulk_load_pragmas", "explain_sqlite_fk_error", "restore_sqlite_journal_mode", diff --git a/src/orm_loader/helpers/bootstrap.py b/src/orm_loader/helpers/bootstrap.py index c449f03..b1b0566 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -17,6 +17,7 @@ without_cross_engine_foreign_keys, ) +from .engines import dispose_engines from .metadata import Base logger = logging.getLogger(__name__) @@ -209,9 +210,7 @@ def create_db(resolved: ResolvedDatabase, *, bindable: Bindable | None = None) - if owned: # _resolve_binds() only returns Engine instances on this path # (bindable is None): create_engines()/create_engine(). - cast(sa.engine.Engine, primary_bind).dispose() - if vocab_bind is not primary_bind: - cast(sa.engine.Engine, vocab_bind).dispose() + dispose_engines(cast(sa.engine.Engine, primary_bind), cast(sa.engine.Engine, vocab_bind)) def bootstrap( diff --git a/src/orm_loader/helpers/engines.py b/src/orm_loader/helpers/engines.py new file mode 100644 index 0000000..7a600d3 --- /dev/null +++ b/src/orm_loader/helpers/engines.py @@ -0,0 +1,8 @@ +import sqlalchemy as sa + + +def dispose_engines(primary: sa.engine.Engine, vocab: sa.engine.Engine) -> None: + """Dispose both engines, without disposing a colocated one twice.""" + primary.dispose() + if vocab is not primary: + vocab.dispose() From 3f3ae05630ec033741ff8342c94c3d70ea268585 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 28 Oct 2026 02:57:25 +0000 Subject: [PATCH 15/17] Replace dependency-wave merge scheduling with batch-ordered pagination DuckDBBackend now overrides _paginate_batches to yield self-referential batches in dependency order, falling back to the base class's ordinary pagination otherwise. This replaces _merge_in_dependency_waves, fixing unbatched loads committing per wave, valid self-links/cycles being rejected as orphans, validate_deferred_fks=False being ignored by an unconditional final check, and upsert's conflict handling being shadowed by premature readiness checks. Co-Authored-By: Claude Sonnet 5 --- src/orm_loader/backends/base.py | 83 +-------------- src/orm_loader/backends/duckdb.py | 147 ++++++++++++++++++++------ tests/backends/test_duckdb_backend.py | 127 ++++++++++++++++++++-- 3 files changed, 235 insertions(+), 122 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 8888e25..1977286 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -581,90 +581,13 @@ def _self_referential_deferred_fks( A row's match for one of these can arrive in a later pagination batch of the same load, since batching orders by insertion - position rather than by dependency. merge_insert()/merge_upsert() - must load table_cls in dependency order instead when any exist; - see :meth:`_merge_in_dependency_waves`. + position rather than by dependency. A backend with any of these + must override :meth:`_paginate_batches` to yield batches in + dependency order instead. """ target = table_cls.__table__ return tuple(fk for fk in self.deferred_fk_constraints(table_cls) if fk.referred_table is target) - def _merge_in_dependency_waves( - self, - table_cls: Type["CSVTableProtocol"], - session: so.Session, - staging: sa.Table, - target: sa.Table, - self_referential_fks: tuple[sa.ForeignKeyConstraint, ...], - merge_batch_size: int | None, - validate_deferred_fks: bool, - insert_chunk: Callable[[sa.ColumnElement[bool]], Any], - ) -> None: - """Insert staging rows wave by wave for a self-referential deferred FK. - - Each wave is exactly the staging rows whose self-referential FK - value(s) are already satisfied, by being NULL or by matching a row - already committed to target (from an earlier wave, or already - there before this load). insert_chunk() inserts one wave (or a - merge_batch_size-sized slice of it); the matching staging rows are - then removed and the wave committed and validated, before the next - wave is computed against the now-larger target. - - Every wave is safe to validate and commit immediately. Nothing in it - can be a false positive, since its parent row is already committed. - A row that never becomes ready, once a pass makes no further progress, - is a genuine violation (an orphan, or a real cycle) and is raised directly - instead of being silently left in staging. - - Notes - ----- - One query round-trip per wave, so cost tracks how many generations - deep the self-reference goes, not how many rows there are. Fine for - a hierarchy that's a few levels deep with many rows per level (e.g. - a category tree). Slow for a long chain where each row's parent is - the previous row (i.e. one wave per row). - - Raises - ------ - DeferredForeignKeyViolationError - A row's self-referential FK value never matches any committed - row once no further wave makes progress. - """ - per_fk = [] - for fk in self_referential_fks: - local_cols: list[sa.ColumnElement[Any]] = [staging.c[col.name] for col in fk.columns] - remote_cols = [cast(sa.ColumnElement[Any], element.column) for element in fk.elements] - not_null = sa.and_(*(col.isnot(None) for col in local_cols)) - has_match = sa.exists( - sa.select(sa.literal(1)).select_from(target).where( - sa.and_(*(local_cols[i] == remote_cols[i] for i in range(len(remote_cols)))) - ) - ) - per_fk.append((fk, local_cols, not_null, has_match)) - ready = sa.and_(*(sa.or_(sa.not_(not_null), has_match) for _, _, not_null, has_match in per_fk)) - - order_column = self._pagination_column(table_cls, staging) - self._ensure_pagination_index(table_cls, staging, session) - - while True: - query = sa.select(order_column).where(ready).order_by(order_column) - if merge_batch_size is not None: - query = query.limit(merge_batch_size) - keys = session.execute(query).scalars().all() - if not keys: - break - chunk = order_column.in_(keys) - insert_chunk(chunk) - session.execute(sa.delete(staging).where(chunk)) - self._commit_batch_if_needed(table_cls, session, chunk, validate_deferred_fks=validate_deferred_fks) - - if session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one(): - for fk, local_cols, not_null, has_match in per_fk: - stuck = session.execute( - sa.select(*local_cols).where(not_null, sa.not_(has_match)).limit(5) - ).all() - if stuck: - raise DeferredForeignKeyViolationError(fk, [tuple(row) for row in stuck]) - def merge_context( self, table_cls: Type["CSVTableProtocol"], diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 68709c5..62a7f49 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Type +from collections.abc import Iterator +from typing import TYPE_CHECKING, Any, Type, cast import sqlalchemy as sa import sqlalchemy.orm as so @@ -294,6 +295,84 @@ def _ensure_pagination_index( idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) session.commit() + def _paginate_batches( + self, + table_cls: type["CSVTableProtocol"], + staging: sa.Table, + session: so.Session, + merge_batch_size: int | None, + ) -> Iterator[sa.ColumnElement[bool] | None]: + """Yield batches in dependency order when a self-referential deferred FK exists. + + Falls back to the base class's ordinary position-based pagination + for every other table. See :meth:`_self_referential_batches`. + """ + self_referential_fks = self._self_referential_deferred_fks(table_cls) + if not self_referential_fks: + yield from super()._paginate_batches(table_cls, staging, session, merge_batch_size) + return + yield from self._self_referential_batches( + table_cls, staging, session, merge_batch_size, self_referential_fks + ) + + def _self_referential_batches( + self, + table_cls: type["CSVTableProtocol"], + staging: sa.Table, + session: so.Session, + merge_batch_size: int | None, + self_referential_fks: tuple[sa.ForeignKeyConstraint, ...], + ) -> Iterator[sa.ColumnElement[bool] | None]: + """Yield staging batches in the order a self-referential deferred FK needs. + + Each batch is the staging rows whose self-referential FK + value(s) are already satisfied (NULL, or matching a row already + committed to target). The caller commits and cleans up staging + for each batch before the next is requested; this generator + never deletes staging itself, so that cleanup stays in the same + transaction as the batch's own insert. + + Once no further batch can make progress, everything left in + staging is yielded as one final, unconditional batch, deferring + to the caller's insert/upsert statement and its own + validate_deferred_fks() check rather than raising here. + + When merge_batch_size is None, yields the entire staging set as + a single unfiltered batch with no readiness computation. + """ + if merge_batch_size is None: + yield None + return + + target = table_cls.__table__ + order_column = self._pagination_column(table_cls, staging) + self._ensure_pagination_index(table_cls, staging, session) + + per_fk = [] + for fk in self_referential_fks: + local_cols: list[sa.ColumnElement[Any]] = [staging.c[col.name] for col in fk.columns] + remote_cols = [cast(sa.ColumnElement[Any], element.column) for element in fk.elements] + not_null = sa.and_(*(col.isnot(None) for col in local_cols)) + has_match = sa.exists( + sa.select(sa.literal(1)).select_from(target).where( + sa.and_(*(local_cols[i] == remote_cols[i] for i in range(len(remote_cols)))) + ) + ) + per_fk.append((not_null, has_match)) + ready = sa.and_(*(sa.or_(sa.not_(not_null), has_match) for not_null, has_match in per_fk)) + + while True: + query = sa.select(order_column).where(ready).order_by(order_column).limit(merge_batch_size) + keys = session.execute(query).scalars().all() + if keys: + chunk: sa.ColumnElement[bool] = order_column.in_(keys) + else: + remaining = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() + if not remaining: + return + chunk = sa.true() + yield chunk + @property def _commits_between_pagination_batches(self) -> bool: """True: bounds uncommitted work on a large table. Safe because @@ -360,9 +439,9 @@ def merge_replace( A self-referential FK can point an existing row at one this load hasn't inserted yet. Doing replace then insert as separate phases would validate that before the target exists, so both happen - together here, wave by wave (see _merge_in_dependency_waves()), - draining staging. The merge_insert() call that follows has - nothing left to do. + together here, batch by batch (see _paginate_batches()), draining + staging. The merge_insert() call that follows has nothing left to + do. Raises ------ @@ -417,22 +496,27 @@ def _drop_matched_from_staging(extra: sa.ColumnElement[bool] | None = None) -> N def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) - def _process_chunk(chunk: sa.ColumnElement[bool]) -> None: - # chunk's rows already in target are updated; what's left in - # staging afterward is new and gets inserted here. - _update(chunk) - _drop_matched_from_staging(chunk) - session.execute(_insert(non_paginated_select.where(chunk))) - - self._merge_in_dependency_waves( - table_cls, session, staging, target, self_referential_fks, - merge_batch_size, validate_deferred_fks, insert_chunk=_process_chunk, - ) - return + def _process_batch(batch_range: sa.ColumnElement[bool] | None) -> None: + # batch_range's rows already in target are updated; what's + # left in staging afterward is new and gets inserted here. + _update(batch_range) + _drop_matched_from_staging(batch_range) + select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) + session.execute(_insert(select_)) + # Runs even when batch_range is None: merge_insert() runs + # next and needs staging fully drained. + delete_stmt = sa.delete(staging) + if batch_range is not None: + delete_stmt = delete_stmt.where(batch_range) + session.execute(delete_stmt) + else: + + def _process_batch(batch_range: sa.ColumnElement[bool] | None) -> None: + _update(batch_range) + _drop_matched_from_staging(batch_range) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): - _update(batch_range) - _drop_matched_from_staging(batch_range) + _process_batch(batch_range) self._commit_batch_if_needed( table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks ) @@ -461,19 +545,15 @@ def _upsert(select_: Any) -> sa.Insert: ) non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) - - self_referential_fks = self._self_referential_deferred_fks(table_cls) - if self_referential_fks: - self._merge_in_dependency_waves( - table_cls, session, staging, target, self_referential_fks, - merge_batch_size, validate_deferred_fks, - insert_chunk=lambda chunk: session.execute(_upsert(non_paginated_select.where(chunk))), - ) - return + self_referential = bool(self._self_referential_deferred_fks(table_cls)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_upsert(select_)) + if self_referential and batch_range is not None: + # Must be removed from staging in this same transaction, + # or the next readiness query would select it again. + session.execute(sa.delete(staging).where(batch_range)) self._commit_batch_if_needed( table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks ) @@ -495,18 +575,15 @@ def merge_insert( def _insert(select_: Any) -> sa.Insert: return sa.insert(target).from_select(insertable_cols, select_) - self_referential_fks = self._self_referential_deferred_fks(table_cls) - if self_referential_fks: - self._merge_in_dependency_waves( - table_cls, session, staging, target, self_referential_fks, - merge_batch_size, validate_deferred_fks, - insert_chunk=lambda chunk: session.execute(_insert(non_paginated_select.where(chunk))), - ) - return + self_referential = bool(self._self_referential_deferred_fks(table_cls)) for batch_range in self._paginate_batches(table_cls, staging, session, merge_batch_size): select_ = non_paginated_select if batch_range is None else non_paginated_select.where(batch_range) session.execute(_insert(select_)) + if self_referential and batch_range is not None: + # Must be removed from staging in this same transaction, + # or the next readiness query would select it again. + session.execute(sa.delete(staging).where(batch_range)) self._commit_batch_if_needed( table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks ) diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 9278dd0..3fa45e6 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -381,7 +381,7 @@ def test_duckdb_backend_merge_insert_loads_self_referential_rows_in_dependency_o ): """Staging rows arrive in reverse dependency order (each pointing at a row not yet committed); merge_insert must still load every row, in - however many waves that takes, regardless of insertion order. + however many batches that takes, regardless of insertion order. """ duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) staging = _SelfReferentialTableCls.get_staging_table( @@ -448,12 +448,13 @@ def test_duckdb_backend_merge_replace_updates_and_inserts_self_referential_rows_ @pytest.mark.duckdb @pytest.mark.db_dialect -def test_duckdb_backend_merge_replace_dependency_waves_catch_a_genuine_orphan( +def test_duckdb_backend_merge_replace_catches_a_genuine_orphan( duckdb_session_with_committed_schema, duckdb_backend, ): """A replace batch re-pointing an existing row at a genuine orphan (not merely a not-yet-inserted new row) must still be caught, not mistaken - for a forward reference that will resolve later. + for a forward reference that will resolve later. Rolling back must + restore the original row. """ session = duckdb_session_with_committed_schema session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) @@ -466,18 +467,19 @@ def test_duckdb_backend_merge_replace_dependency_waves_catch_a_genuine_orphan( with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): duckdb_backend.merge_replace(_SelfReferentialTableCls, session, ["id"], merge_batch_size=1) + session.rollback() assert session.execute(sa.select(SelfReferentialTable.parent_id)).scalars().all() == [None] assert session.execute(sa.select(staging.c.id)).scalars().all() == [1] @pytest.mark.duckdb @pytest.mark.db_dialect -def test_duckdb_backend_merge_insert_dependency_waves_catch_a_genuine_orphan( +def test_duckdb_backend_merge_insert_catches_a_genuine_orphan_in_its_own_batch( duckdb_session_with_committed_schema, duckdb_backend, ): - """A row whose parent never arrives in any wave is a genuine violation. - A row that did become ready must still be committed. Only the stuck - row is left behind in staging and reported. + """A row whose parent never arrives is a genuine violation, caught in + its own batch's commit. A row that became ready in an earlier batch + already committed and survives rolling back the failed one. """ session = duckdb_session_with_committed_schema duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) @@ -489,10 +491,121 @@ def test_duckdb_backend_merge_insert_dependency_waves_catch_a_genuine_orphan( with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): duckdb_backend.merge_insert(_SelfReferentialTableCls, session, merge_batch_size=1) + session.rollback() assert session.execute(sa.select(SelfReferentialTable.id)).scalars().all() == [1] assert session.execute(sa.select(staging.c.id)).scalars().all() == [2] +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_unbatched_has_no_intermediate_commit( + duckdb_session_with_committed_schema, duckdb_backend, +): + """Without merge_batch_size, nothing may commit until the caller's own + validation runs, even when a later row is a genuine orphan. Rolling + back after that validation must leave nothing behind. Comparable to + test_duckdb_backend_merge_insert_batch_validation_prevents_partial_commit, + but with no batching requested at all. + """ + session = duckdb_session_with_committed_schema + duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) + staging = _SelfReferentialTableCls.get_staging_table( + session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + session.execute(sa.insert(staging), [{"id": 1, "parent_id": None}, {"id": 2, "parent_id": 999}]) + + duckdb_backend.merge_insert(_SelfReferentialTableCls, session, merge_batch_size=None) + with pytest.raises(DeferredForeignKeyViolationError, match="self_referential_table"): + duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, session) + + session.rollback() + assert session.execute(sa.select(SelfReferentialTable.id)).scalars().all() == [] + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_accepts_a_self_link(duckdb_session, duckdb_backend): + """A row whose parent is itself is valid once it exists, not an orphan + merely because the target was empty when it was considered. + """ + duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) + staging = _SelfReferentialTableCls.get_staging_table( + duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + duckdb_session.execute(sa.insert(staging), [{"id": 1, "parent_id": 1}]) + + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, merge_batch_size=1) + duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, duckdb_session) + + rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} + assert rows == {1: 1} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_accepts_a_two_row_cycle(duckdb_session, duckdb_backend): + """Two rows referencing each other are valid once both exist. Neither + can ever become ready on its own, which is not a reason to reject them. + """ + duckdb_backend.create_staging_table(_SelfReferentialTableCls, duckdb_session) + staging = _SelfReferentialTableCls.get_staging_table( + duckdb_session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + duckdb_session.execute(sa.insert(staging), [{"id": 1, "parent_id": 2}, {"id": 2, "parent_id": 1}]) + + duckdb_backend.merge_insert(_SelfReferentialTableCls, duckdb_session, merge_batch_size=1) + duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, duckdb_session) + + rows = {r.id: r.parent_id for r in duckdb_session.execute(sa.select(SelfReferentialTable)).scalars().all()} + assert rows == {1: 2, 2: 1} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_upsert_ignores_a_duplicate_with_a_bad_reference( + duckdb_session_with_committed_schema, duckdb_backend, +): + """A staging row ON CONFLICT DO NOTHING will discard must not be + judged as if it were going to be stored. + """ + session = duckdb_session_with_committed_schema + session.execute(sa.insert(SelfReferentialTable), [{"id": 1, "parent_id": None}]) + duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) + staging = _SelfReferentialTableCls.get_staging_table( + session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + session.execute(sa.insert(staging), [{"id": 1, "parent_id": 999}]) + + duckdb_backend.merge_upsert(_SelfReferentialTableCls, session, ["id"], merge_batch_size=1) + duckdb_backend.validate_deferred_fks(_SelfReferentialTableCls, session) + + rows = {r.id: r.parent_id for r in session.execute(sa.select(SelfReferentialTable)).scalars().all()} + assert rows == {1: None} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_merge_insert_respects_validate_deferred_fks_false( + duckdb_session_with_committed_schema, duckdb_backend, +): + """Opting out of application validation must not be overridden by the + self-referential batching path. + """ + session = duckdb_session_with_committed_schema + duckdb_backend.create_staging_table(_SelfReferentialTableCls, session) + staging = _SelfReferentialTableCls.get_staging_table( + session, staging_schema_tag=duckdb_backend.staging_schema_tag + ) + session.execute(sa.insert(staging), [{"id": 1, "parent_id": 999}]) + + duckdb_backend.merge_insert( + _SelfReferentialTableCls, session, merge_batch_size=1, validate_deferred_fks=False + ) + + rows = {r.id: r.parent_id for r in session.execute(sa.select(SelfReferentialTable)).scalars().all()} + assert rows == {1: 999} + + @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_validate_deferred_fks_catches_a_cross_database_orphan(duckdb_split_database_session): From bf14559d9788f0dd90aad9ea1187620c380ec3ed Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 28 Oct 2026 03:02:50 +0000 Subject: [PATCH 16/17] Guard cross-database FK validation against an ephemeral remote connection _validate_cross_database_fk rebuilt the remote role's engine from config, which silently gets an empty, disconnected database for a DuckDB ':memory:' connection instead of reconnecting to real data. Raise clearly instead, matching bootstrap.py's existing is_ephemeral_url guard for the same failure mode. Co-Authored-By: Claude Sonnet 5 --- src/orm_loader/backends/duckdb.py | 14 +++++++++ tests/backends/test_duckdb_backend.py | 45 +++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 62a7f49..25928a4 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -9,6 +9,7 @@ ResolvedCDMDatabase, Role, UnenforcedForeignKey, + is_ephemeral_url, physical_schema_of, qualified, unenforced_foreign_keys, @@ -128,6 +129,10 @@ def _validate_cross_database_fk(self, fk: sa.ForeignKeyConstraint, session: so.S self.resolved is not a ResolvedCDMDatabase: a cross-database FK can only arise when PRIMARY/VOCAB are split across connections, which only a CDM-aware resolved can describe. + ValueError + The remote role's connection is ephemeral (DuckDB ``:memory:``): + a freshly built engine would get its own, empty database rather + than the one already holding the remote data. DeferredForeignKeyViolationError A local value has no match on the referred connection. """ @@ -145,6 +150,15 @@ def _validate_cross_database_fk(self, fk: sa.ForeignKeyConstraint, session: so.S return remote_tag = remote_cols[0].table.schema + remote_role = Role.VOCAB if remote_tag == Role.VOCAB.value else Role.PRIMARY + remote_connection = self.resolved.connection_for_role(remote_role) + if is_ephemeral_url(remote_connection.safe_url): + raise ValueError( + f"{fk.name} crosses to {remote_connection.name!r}, an ephemeral connection " + "(DuckDB ':memory:'): a freshly built engine would see an empty, disconnected " + "database instead of the data already loaded there. Cross-database FK " + "validation needs a real, file-backed or server-backed remote connection." + ) primary_engine, vocab_engine = self.resolved.create_engines(register_claims=False) try: remote_engine = vocab_engine if remote_tag == Role.VOCAB.value else primary_engine diff --git a/tests/backends/test_duckdb_backend.py b/tests/backends/test_duckdb_backend.py index 3fa45e6..7043220 100644 --- a/tests/backends/test_duckdb_backend.py +++ b/tests/backends/test_duckdb_backend.py @@ -126,6 +126,34 @@ def duckdb_split_database_session(): vocab_engine.dispose() +@pytest.fixture +def duckdb_split_database_session_with_ephemeral_vocab(): + """Same split topology as duckdb_split_database_session, but VOCAB is + an in-memory DuckDB database instead of a file, to exercise the + ephemeral-connection guard in _validate_cross_database_fk. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + stack = StackConfig.for_session( + connections={ + "primary": ConnectionConfig( + dialect="duckdb", database_name=str(Path(tmp_dir) / "primary.duckdb"), test_only=True + ), + "vocab": ConnectionConfig(dialect="duckdb", database_name=":memory:", test_only=True), + }, + databases={"cdm": CDMDatabaseConfig(connection="primary", vocab_connection="vocab")}, + ) + resolved = Resolver(stack).resolve_database("cdm") + primary_engine, _ = resolved.create_engines(schema_claims=[staging_schema_claim()]) + try: + with primary_engine.connect() as conn: + create_tables(conn, [VocabReferencingTable.__table__], resolved=resolved) + conn.commit() + with so.Session(bind=conn) as session: + yield resolved, session + finally: + primary_engine.dispose() + + def test_duckdb_backend_identity_and_capabilities(): backend = DuckDBBackend() @@ -647,6 +675,23 @@ def test_duckdb_backend_validate_deferred_fks_passes_a_real_cross_database_refer backend.validate_deferred_fks(_VocabReferencingTableCls, session) +@pytest.mark.duckdb +@pytest.mark.db_dialect +def test_duckdb_backend_validate_deferred_fks_refuses_an_ephemeral_remote_connection( + duckdb_split_database_session_with_ephemeral_vocab, +): + """A cross-database FK into an in-memory VOCAB must raise clearly + instead of rebuilding a disconnected, empty engine and reporting + every reference as missing. + """ + resolved, session = duckdb_split_database_session_with_ephemeral_vocab + backend = DuckDBBackend(resolved=resolved) + session.execute(sa.insert(VocabReferencingTable), [{"id": 1, "vocab_id": 1}]) + + with pytest.raises(ValueError, match="ephemeral"): + backend.validate_deferred_fks(_VocabReferencingTableCls, session) + + @pytest.mark.duckdb @pytest.mark.db_dialect def test_duckdb_backend_merge_insert_batch_validation_is_a_no_op_without_deferred_fks(duckdb_session): From ad023268b83c85cf2d01725bc7eadf0c7da37b01 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 28 Oct 2026 03:14:12 +0000 Subject: [PATCH 17/17] Pull shared pagination-index logic up into DatabaseBackend _pagination_column/_ensure_pagination_index were byte-for-byte identical in postgres.py and duckdb.py. Make them the base class's concrete default (ordering by staging's _rownum, indexed on demand), with SQLite's rowid-based override now explicit about why it skips indexing instead of relying on the old no-op default. Co-Authored-By: Claude Sonnet 5 --- src/orm_loader/backends/base.py | 20 +++++++++++++++----- src/orm_loader/backends/duckdb.py | 13 ------------- src/orm_loader/backends/postgres.py | 13 ------------- src/orm_loader/backends/sqlite.py | 6 ++++++ 4 files changed, 21 insertions(+), 31 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 1977286..74baa60 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -491,13 +491,18 @@ def merge_insert( See :meth:`merge_replace`. """ - @abstractmethod def _pagination_column( self, table_cls: Type["CSVTableProtocol"], staging: sa.Table, ) -> sa.ColumnElement[int]: - """Column this backend orders a staging-table scan by, for pagination.""" + """Column this backend orders a staging-table scan by, for pagination. + + Defaults to the explicit _rownum column every staging table carries. + Override for a backend with a cheaper, already-indexed alternative + (e.g. SQLite's implicit rowid). + """ + return staging.c._rownum def _ensure_pagination_index( self, @@ -505,10 +510,15 @@ def _ensure_pagination_index( staging: sa.Table, session: so.Session, ) -> None: - """Make sure _pagination_column() is indexed, if it needs to be. - No-op by default. Overwrite in a backend that needs an index to paginate efficiently (e.g. DuckDB). + """Index _pagination_column() so paginated batches can seek it efficiently. + + Overridden as a no-op by a backend whose _pagination_column() is + already indexed (e.g. SQLite's rowid). """ - return + staging_name = self.staging_name_for_table(table_cls.__tablename__) + idx = sa.Index(f"{staging_name}_rownum_idx", staging.c._rownum) + idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) + session.commit() @property @abstractmethod diff --git a/src/orm_loader/backends/duckdb.py b/src/orm_loader/backends/duckdb.py index 25928a4..2bfaaf0 100644 --- a/src/orm_loader/backends/duckdb.py +++ b/src/orm_loader/backends/duckdb.py @@ -296,19 +296,6 @@ def restore_fk_check( "DuckDB has no foreign-key-toggle mechanism (no PRAGMA/SET equivalent exists)." ) - def _pagination_column( - self, table_cls: type["CSVTableProtocol"], staging: sa.Table - ) -> sa.ColumnElement[int]: - return staging.c._rownum - - def _ensure_pagination_index( - self, table_cls: type["CSVTableProtocol"], staging: sa.Table, session: so.Session - ) -> None: - staging_name = self.staging_name_for_table(table_cls.__tablename__) - idx = sa.Index(f"{staging_name}_rownum_idx", staging.c._rownum) - idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) - session.commit() - def _paginate_batches( self, table_cls: type["CSVTableProtocol"], diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index e8c2ea9..43c3e23 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -202,19 +202,6 @@ def restore_fk_check( sa.text(f"SET session_replication_role = '{safe_state}'"), **execute_options ) - def _pagination_column( - self, table_cls: type["CSVTableProtocol"], staging: sa.Table - ) -> sa.ColumnElement[int]: - return staging.c._rownum - - def _ensure_pagination_index( - self, table_cls: type["CSVTableProtocol"], staging: sa.Table, session: so.Session - ) -> None: - staging_name = self.staging_name_for_table(table_cls.__tablename__) - idx = sa.Index(f"{staging_name}_rownum_idx", staging.c._rownum) - idx.create(bind=session.connection(bind_arguments={"mapper": table_cls}), checkfirst=True) - session.commit() - @property def _commits_between_pagination_batches(self) -> bool: """True: bounds the WAL growth, lock duration, and dead-tuple diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 4f8c7be..14be15f 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -182,6 +182,12 @@ def _pagination_column( """SQLite's implicit rowid is already gapless and indexed.""" return sa.literal_column("rowid") + def _ensure_pagination_index( + self, table_cls: type["CSVTableProtocol"], staging: sa.Table, session: so.Session + ) -> None: + """No-op: rowid needs no index, and the staging table has no _rownum column.""" + return + @property def _commits_between_pagination_batches(self) -> bool: """False: keeps the whole merge in the caller's transaction, so a