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/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/pyproject.toml b/pyproject.toml index 4ca8cba..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", @@ -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..74baa60 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 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, SchemaClaim, open_connection, qualified - if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext from ..mappers.materialised_view_contracts import MaterializedViewIndex @@ -33,6 +32,23 @@ 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): + """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" @@ -99,7 +115,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 ---------- @@ -115,9 +138,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: @@ -269,6 +302,76 @@ 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: + """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_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_alias).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, @@ -324,6 +427,21 @@ 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], + 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. session is only used as read-only + to check whether target or a referencing table already has rows + """ + @abstractmethod def merge_replace( self, @@ -332,8 +450,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( @@ -343,8 +468,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( @@ -353,17 +483,135 @@ 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. + + validate_deferred_fks + See :meth:`merge_replace`. + """ + + 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. + + 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, + table_cls: Type["CSVTableProtocol"], + staging: sa.Table, + session: so.Session, + ) -> None: + """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). + """ + 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 + def _commits_between_pagination_batches(self) -> bool: + """Whether _paginate_batches()'s caller should commit after each batch.""" + + 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 _commit_batch_if_needed( + self, + table_cls: Type["CSVTableProtocol"], + session: so.Session, + batch_range: sa.ColumnElement[bool] | None, + *, + validate_deferred_fks: bool = True, ) -> None: - """Insert all staging rows into the target table.""" + """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 _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. 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_context( self, 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..2bfaaf0 --- /dev/null +++ b/src/orm_loader/backends/duckdb.py @@ -0,0 +1,590 @@ +from __future__ import annotations + +from collections.abc import Iterator +from typing import TYPE_CHECKING, Any, Type, cast + +import sqlalchemy as sa +import sqlalchemy.orm as so +from oa_configurator import ( + ResolvedCDMDatabase, + Role, + UnenforcedForeignKey, + is_ephemeral_url, + physical_schema_of, + qualified, + unenforced_foreign_keys, +) +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: + 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: + import duckdb_sqlalchemy + + 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, + supports_replace_with_kept_indexes=False, + ) + + 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. + + 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(...)." + ) + return unenforced_foreign_keys( + tables, resolved=self.resolved, physical_schemas=self.physical_schemas + ) + + 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``). + 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 = {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. + 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. + """ + 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 + 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 + 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: + dispose_engines(primary_engine, vocab_engine) + + 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. + + 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 = {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: + """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: + """Fully schema-qualified, quoted name of tablename's staging _rownum sequence. + + 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) + + def create_staging_table( + self, + table_cls: type["CSVTableProtocol"], + 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__) + 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) + + 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.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, + ) + session.commit() + + def drop_staging_table( + self, + table_cls: type["CSVTableProtocol"], + session: so.Session, + ) -> None: + bind_arguments = {"mapper": table_cls} + staging_name = self.staging_name_for_table(table_cls.__tablename__) + session.execute( + DropTable(sa.Table(staging_name, sa.MetaData(), schema=self.staging_schema), if_exists=True), + bind_arguments=bind_arguments, + ) + session.execute( + DropSequence( + sa.Sequence(f"{staging_name}_rownum_seq", schema=self.staging_schema), if_exists=True + ), + bind_arguments=bind_arguments, + ) + + @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)." + ) + + 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 + 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 + + 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. + + Call before any replace load touches table_cls, so a later table + failing this check doesn't leave earlier tables already replaced. + + Raises + ------ + ValueError + 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] + 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 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, + table_cls: type["CSVTableProtocol"], + session: so.Session, + 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 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. + + 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, batch by batch (see _paginate_batches()), draining + staging. The merge_insert() call that follows has nothing left to + do. + + Raises + ------ + ValueError + See check_replace_feasible(). + """ + self._validate_merge_batch_size(merge_batch_size) + 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] + 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 + + 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_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): + _process_batch(batch_range) + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks + ) + + def merge_upsert( + self, + table_cls: type["CSVTableProtocol"], + session: so.Session, + pk_cols: list[str], + *, + merge_batch_size: int | None = None, + validate_deferred_fks: bool = True, + ) -> 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) + 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)) + 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 + ) + + def merge_insert( + self, + table_cls: type["CSVTableProtocol"], + 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__ + 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_) + + 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/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 94eb39a..43c3e23 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 @@ -14,6 +15,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 ( @@ -25,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 @@ -35,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: @@ -61,6 +76,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( @@ -70,13 +86,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''' @@ -110,8 +130,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}, ) @@ -181,36 +202,18 @@ def restore_fk_check( sa.text(f"SET session_replication_role = '{safe_state}'"), **execute_options ) - def _staging_rownum_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() - - @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. + @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 session.execute( - sa.select(sa.func.coalesce(sa.func.max(staging.c._rownum), 0)) - ).scalar_one() + return True def merge_replace( self, @@ -219,43 +222,20 @@ 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. - - 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) 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 - ) + 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)) + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks ) - session.commit() - start = end def merge_upsert( self, @@ -264,6 +244,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__ @@ -280,27 +261,12 @@ 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 + 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_)) + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks ) - session.execute(_upsert(batch_select)) - session.commit() - start = end def merge_insert( self, @@ -308,7 +274,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. + + 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) @@ -318,31 +290,12 @@ 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 + 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_)) + self._commit_batch_if_needed( + table_cls, session, batch_range, validate_deferred_fks=validate_deferred_fks ) - session.execute(_insert(batch_select)) - session.commit() - start = end @requires_capability("supports_materialized_views", "materialized views") def create_materialized_view( @@ -405,12 +358,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/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/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 6cc2753..14be15f 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 @@ -10,11 +11,13 @@ 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 if TYPE_CHECKING: + from oa_configurator import ResolvedDatabase from sqlalchemy.engine import Connection, Engine from ..tables.typing import CSVTableProtocol @@ -36,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, @@ -45,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 @@ -97,6 +107,7 @@ def capabilities(self) -> BackendCapabilities: supports_unlogged_staging=False, supports_fk_toggle=True, supports_materialized_views=False, + supports_replace_with_kept_indexes=True, ) @property @@ -110,7 +121,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}, ) @@ -132,9 +143,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}, ) @@ -165,12 +176,26 @@ 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") + 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 + 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"], @@ -178,16 +203,15 @@ 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. - 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. + 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__ @@ -206,21 +230,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, @@ -229,6 +240,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__ @@ -244,24 +256,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, @@ -269,6 +267,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__ @@ -279,22 +278,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/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/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 77ccf14..b1b0566 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -12,10 +12,12 @@ declared_schema_tags, is_ephemeral_url, open_connection, + physical_schema_of, referred_schema_tag, without_cross_engine_foreign_keys, ) +from .engines import dispose_engines from .metadata import Base logger = logging.getLogger(__name__) @@ -91,9 +93,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 +131,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) @@ -199,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() diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 1f29b6d..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 @@ -373,6 +374,9 @@ def load_csv( index_strategy: str = "auto", 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: """ @@ -413,6 +417,23 @@ 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()). + 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 ------- @@ -444,6 +465,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, @@ -465,15 +494,29 @@ 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, - ) - - cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) + try: + 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, + 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, + resolved=resolved, + physical_schemas=physical_schemas, + ) + 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 @@ -502,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. @@ -516,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( @@ -544,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)}." @@ -559,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)}." @@ -582,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 b1a3b14..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): @@ -87,6 +90,9 @@ def load_csv( index_strategy: str = "auto", 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 @@ -103,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_base_backend.py b/tests/backends/test_base_backend.py index 34154a2..1a2db64 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -107,6 +107,15 @@ 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") + + @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 @@ -223,28 +232,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 +262,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 +289,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..7043220 --- /dev/null +++ b/tests/backends/test_duckdb_backend.py @@ -0,0 +1,776 @@ +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, 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, + SelfReferentialTable, + UnreferencedGrandparentTable, + VocabReferencingTable, + VocabSchemaTable, +) + +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) +_ReferencedParentTableCls = cast("Type[CSVTableProtocol]", ReferencedParentTable) +_SelfReferentialTableCls = cast("Type[CSVTableProtocol]", SelfReferentialTable) +_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() + + +@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() + + +@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() + + 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.capabilities.supports_replace_with_kept_indexes 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", "_rownum"] + assert all(c["nullable"] is True for c in cols if c["name"] != "_rownum") + + +@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 + + +@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. + """ + 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() + + 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, 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"] + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +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}]) + + 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_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"): + 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 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) + 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)} + + +@pytest.mark.duckdb +@pytest.mark.db_dialect +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=duckdb_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"): + 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_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 batches 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_replace_updates_and_inserts_self_referential_rows_together( + duckdb_session, duckdb_backend, +): + """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) + 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_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. Rolling back must + restore the original row. + """ + 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) + + 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_catches_a_genuine_orphan_in_its_own_batch( + duckdb_session_with_committed_schema, duckdb_backend, +): + """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) + 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) + + 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): + """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_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): + """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/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_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..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): @@ -62,9 +67,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 +77,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 +90,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 +113,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 +126,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 +164,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 +219,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 +237,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 +262,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 +274,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 +286,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/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() 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..5cf7bf2 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,69 @@ 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, 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 + 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") + + with pytest.raises(DeferredForeignKeyViolationError, match="cycle_a_table"): + _CycleATable.load_csv( + 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( + _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", + ) 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..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() @@ -54,6 +55,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 +104,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): @@ -96,6 +121,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.""" @@ -108,3 +167,37 @@ 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") + ) + + +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), + ) 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