diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0e7d4c8..7e7a7f4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,18 +6,20 @@ on: jobs: label-gate: uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/label-gate.yml@main - build-test: - uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres.yml@main + build-test-sqlite: + uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test.yml@main + build-test-postgres: + uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres-v2.yml@main with: postgres-db: orm_loader_test setup-commands: | uv run omop-config configure orm_loader \ - --set test_orm_db.kind=cdm \ - --set test_orm_db.connection.dialect=postgresql+psycopg \ - --set test_orm_db.connection.host=localhost \ - --set test_orm_db.connection.port=5432 \ - --set test_orm_db.connection.user=test \ - --set test_orm_db.connection.password=test \ - --set test_orm_db.connection.database_name=orm_loader_test \ - --set test_orm_db.connection.test_only=true \ - --set test_orm_db.schema_name=public + --set test_orm_db_pg.kind=cdm \ + --set test_orm_db_pg.connection.dialect=postgresql+psycopg \ + --set test_orm_db_pg.connection.host=localhost \ + --set test_orm_db_pg.connection.port=5432 \ + --set test_orm_db_pg.connection.user=test \ + --set test_orm_db_pg.connection.password=test \ + --set test_orm_db_pg.connection.database_name=orm_loader_test \ + --set test_orm_db_pg.connection.test_only=true \ + --set test_orm_db_pg.cdm_schema=public diff --git a/docs/tables/mat_view.md b/docs/tables/mat_view.md index 30d0426..6bc463f 100644 --- a/docs/tables/mat_view.md +++ b/docs/tables/mat_view.md @@ -105,18 +105,22 @@ PatientSummaryMV.refresh_mv(engine, concurrently=True) This is fail-closed by design. An index created manually outside `__mv_indexes__` does not satisfy the mixin's declaration contract; declare it in the class even if another migration is responsible for creating it. Expressions, partial indexes, and other index forms are outside this simple contract and should not be represented as `MaterializedViewIndex` entries. -By default, `schema=None` leaves the view name unqualified. PostgreSQL resolves that name through the connection's `search_path`, matching the behavior of existing callers. Pass `schema="reporting"` only when the caller intentionally wants an explicit schema-qualified target. +The view's physical schema comes from the bound connection's own `schema_translate_map`, the same mechanism every other schema-aware table in this stack uses (`oa_configurator.schema_of`). `__mv_role__` declares which role a view resolves through (defaulting to `Role.PRIMARY`); pass `role=` to `create_mv()`/`refresh_mv()`/`drop_mv()` to override it for one call. ```python -# Existing/default behavior: search_path resolves the target. +class VocabSummaryMV(MaterializedViewMixin): + __mv_name__ = "vocab_summary" + __mv_select__ = ... + __mv_role__ = Role.VOCAB # resolves via the connection's vocab schema + +# Uses __mv_role__ (Role.PRIMARY by default) via the connection's own schema_translate_map. RecentObservationMV.create_mv(engine) -# Explicit schema: the target is quoted and schema-qualified. -RecentObservationMV.create_mv(engine, schema="reporting") -RecentObservationMV.refresh_mv(engine, schema="reporting") +# Override for one call. +RecentObservationMV.create_mv(engine, role=Role.VOCAB) ``` -Explicit schema targets are quoted component by component. This matters for embedded quotes, spaces, and mixed-case identifiers. It also means an unqualified mixed-case name and the same name passed with `schema=` can address different PostgreSQL relations. Keep schema selection at the call site and do not assume that this API provides `schema_translate_map`, role-token, or general multi-schema behavior. +Every generated identifier is quoted through `oa_configurator.qualified()`, which quotes each component only when the dialect actually requires it (reserved words, mixed case, embedded quotes or spaces) — the same behavior every other Core-built query in this stack has. ## Failure handling and backend support @@ -166,7 +170,3 @@ The built-in implementation is PostgreSQL-oriented. SQLite rejects materialized- ::: orm_loader.mappers.ConcurrentRefreshNotEligibleError options: heading_level: 3 - -::: orm_loader.mappers.UnsupportedMaterializationDialectError - options: - heading_level: 3 diff --git a/pyproject.toml b/pyproject.toml index 110bc99..1e68a9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,7 +82,7 @@ testpaths = ["tests"] python_files = ["test_*.py"] python_classes = ["Test*"] python_functions = ["test_*"] -addopts = "-ra" +addopts = "-ra -m 'not db_dialect'" [tool.pyright] reportMissingTypeStubs = false diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index 17534f8..e5dd07a 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -1,13 +1,17 @@ from .postgres import PostgresBackend from .resolve import resolve_backend from .sqlite import SQLiteBackend -from .base import BackendCapabilities, DatabaseBackend, STAGING_SCHEMA, Dialect +from .base import ( + BackendCapabilities, + DatabaseBackend, + Dialect, + STAGING_SCHEMA, +) from ..mappers.materialised_view_errors import ( ConcurrentRefreshNotEligibleError, MaterializationError, MaterializationFailure, MaterializationOperation, - UnsupportedMaterializationDialectError, ) __all__ = [ @@ -21,6 +25,5 @@ "PostgresBackend", "STAGING_SCHEMA", "SQLiteBackend", - "UnsupportedMaterializationDialectError", "resolve_backend", ] diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 71f2557..2a32c88 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -3,7 +3,6 @@ from abc import ABC, abstractmethod from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass -from enum import Enum from collections.abc import Generator from functools import wraps from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Type, TypeVar, cast @@ -13,6 +12,8 @@ from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer +from oa_configurator import Dialect, Role + if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext from ..mappers.materialised_view_contracts import MaterializedViewIndex @@ -34,13 +35,6 @@ class BackendCapabilities: supports_materialized_views: bool = False -class Dialect(str, Enum): - """Supported SQLAlchemy dialect names.""" - - SQLITE = "sqlite" - POSTGRESQL = "postgresql" - - STAGING_SCHEMA: str = "staging" P = ParamSpec("P") @@ -192,6 +186,41 @@ def _as_connection( self, bind: Engine | Connection, ) -> Generator[Connection, None, None]: + """ + Normalize a bind into an open connection, guarding its dialect. + + Every backend method that takes a ``bind`` should route it through + this context manager rather than opening a connection itself, so the + dialect guard below applies uniformly. + + Parameters + ---------- + bind : Engine or Connection + An Engine opens a new connection and transaction scoped to this + context manager, committing on a clean exit. A Connection is + forwarded as-is; its transaction is owned by the caller, and + passing the same Connection into several backend calls groups + them into one shared transaction. + + Yields + ------ + Connection + An open connection whose dialect matches ``self.dialect``. + + Raises + ------ + TypeError + If ``bind``'s dialect does not match ``self.dialect``. Guards + against a bind resolved through a different backend being + passed directly into a method on this one. + """ + if bind.dialect.name != self.dialect.value: + raise TypeError( + f"{self.name} backend received a {bind.dialect.name!r} connection; " + f"expected {self.dialect.value!r}. The bind passed to this method must " + "be the same one (or share the same dialect as) the bind resolve_backend() " + "was given." + ) if isinstance(bind, Engine): with bind.begin() as conn: yield conn @@ -268,7 +297,6 @@ def merge_replace( self, table_cls: Type["CSVTableProtocol"], session: so.Session, - target_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, @@ -280,7 +308,6 @@ def merge_upsert( self, table_cls: Type["CSVTableProtocol"], session: so.Session, - target_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, @@ -292,7 +319,6 @@ def merge_insert( self, table_cls: Type["CSVTableProtocol"], session: so.Session, - target_name: str, *, merge_batch_size: int | None = None, ) -> None: @@ -345,14 +371,16 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - schema: str | None = None, + role: Role = Role.PRIMARY, with_data: bool = True, if_not_exists: bool = True, ) -> None: """Create a materialized view for the supplied selectable. - ``schema`` defaults to ``None``, leaving the target unqualified for - the connection's ``search_path`` to resolve. + The view's schema is the bind's own ``schema_translate_map`` entry + for *role* (via ``oa_configurator.schema_of``), letting a view + built over vocab/results-role tables land in that role's own + schema instead of always primary. """ raise NotImplementedError( f"Backend '{self.name}' has not implemented create_materialized_view()" @@ -364,12 +392,17 @@ def refresh_materialized_view( bind: "Engine | Connection", name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, concurrently: bool = False, declared_indexes: tuple["MaterializedViewIndex", ...] = (), ) -> None: """Refresh a materialized view. + The view's schema is the bind's own ``schema_translate_map`` entry + for *role* (via ``oa_configurator.schema_of``), letting a view + built over vocab/results-role tables land in that role's own + schema instead of always primary. + ``declared_indexes`` lets supporting backends validate a concurrent refresh request without defining a second catalog-based eligibility rule. Other backends may ignore it. @@ -384,16 +417,18 @@ def drop_materialized_view( bind: "Engine | Connection", name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, if_exists: bool = True, cascade: bool = False, ) -> None: """Drop a materialized view. - This is deliberately non-abstract: the default implementation - requires the capability flag and then raises ``NotImplementedError``. - Older third-party backend subclasses need no override to receive a - clear error when they do not support this operation. + The view's schema is the bind's own ``schema_translate_map`` entry + for *role*, matching ``create_materialized_view``. This is + deliberately non-abstract: the default implementation requires the + capability flag and then raises ``NotImplementedError``. Older + third-party backend subclasses need no override to receive a clear + error when they do not support this operation. """ raise NotImplementedError( f"Backend '{self.name}' has not implemented drop_materialized_view()" @@ -406,13 +441,15 @@ def create_materialized_view_index( name: str, index: "MaterializedViewIndex", *, - schema: str | None = None, + role: Role = Role.PRIMARY, if_not_exists: bool = True, ) -> None: """Create an index on a materialized view. - This is deliberately non-abstract for the same compatibility reason - as :meth:`drop_materialized_view`. + The view's schema is the bind's own ``schema_translate_map`` entry + for *role*, matching ``create_materialized_view``. This is + deliberately non-abstract for the same compatibility reason as + :meth:`drop_materialized_view`. """ raise NotImplementedError( f"Backend '{self.name}' has not implemented " diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 19ca8b7..98e0ea2 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -6,17 +6,18 @@ import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so +from oa_configurator import autocommit_connection, qualified, schema_of, Dialect, Role +from ..helpers.sql import role_of_table from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError from sqlalchemy.sql.compiler import IdentifierPreparer -from .base import BackendCapabilities, DatabaseBackend, Dialect, requires_capability +from .base import BackendCapabilities, DatabaseBackend, requires_capability from ..mappers.materialised_view_errors import ( ConcurrentRefreshNotEligibleError, MaterializationError, MaterializationFailure, MaterializationOperation, - UnsupportedMaterializationDialectError, ) if TYPE_CHECKING: @@ -29,26 +30,6 @@ _VALID_PG_REPLICATION_ROLES = frozenset({"origin", "local", "replica"}) -def _require_postgres_dialect( - conn: "Connection", - *, - operation: MaterializationOperation, - schema: str | None, - name: str, -) -> None: - dialect = getattr(conn, "dialect", None) - if dialect is not None and dialect.name == "postgresql": - return - raise UnsupportedMaterializationDialectError( - MaterializationFailure( - operation=operation, - schema=schema, - name=name, - reason=f"received dialect {getattr(dialect, 'name', dialect)!r}", - ) - ) - - class PostgresBackend(DatabaseBackend): def __init__(self, *, staging_schema: str | None = None) -> None: super().__init__(staging_schema=staging_schema) @@ -86,7 +67,7 @@ def create_staging_table( table = table_cls.__table__ preparer = self.identifier_preparer staging_ref = self.qualified_staging_name(table_cls.__tablename__) - source_ref = preparer.quote_identifier(table.name) + source_ref = qualified(session, table.name, role=role_of_table(table)) session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};')) session.execute( sa.text( @@ -172,49 +153,46 @@ def restore_fk_check( safe_state = self._normalize_fk_check_state(previous_state) session.execute(sa.text(f"SET session_replication_role = '{safe_state}'")) + 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(), checkfirst=True) + session.commit() + def merge_replace( self, table_cls: type["CSVTableProtocol"], session: so.Session, - target_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: - preparer = self.identifier_preparer - staging_ref = self.qualified_staging_name(table_cls.__tablename__) - target_ref = preparer.quote_identifier(target_name) - pk_join = " AND ".join( - f't.{preparer.quote_identifier(c)} = s.{preparer.quote_identifier(c)}' for c in pk_cols - ) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) - non_paginated_replace = sa.text( - f'DELETE FROM {target_ref} t USING {staging_ref} s WHERE {pk_join}' - ) + 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.text(f'SELECT COUNT(*) FROM {staging_ref}')).scalar_one() + total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() if total <= merge_batch_size: session.execute(non_paginated_replace) return - staging_name = self.staging_name_for_table(table_cls.__tablename__) - idx_ref = preparer.quote_identifier(f"{staging_name}_rownum_idx") - session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS {idx_ref} ON {staging_ref} (_rownum)')) - session.commit() + self._staging_rownum_index(table_cls, staging, session) start = 0 while start < total: end = start + merge_batch_size session.execute( - sa.text( - f'DELETE FROM {target_ref} t USING {staging_ref} s' - f' WHERE {pk_join} AND s._rownum > :start AND s._rownum <= :end' - ), - {"start": start, "end": end}, + sa.delete(target).where( + pk_join, staging.c._rownum > start, staging.c._rownum <= end + ) ) session.commit() start = end @@ -223,50 +201,42 @@ def merge_upsert( self, table_cls: type["CSVTableProtocol"], session: so.Session, - target_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: - preparer = self.identifier_preparer - staging_ref = self.qualified_staging_name(table_cls.__tablename__) - target_ref = preparer.quote_identifier(target_name) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) insertable_cols = self._insertable_column_names(table_cls) - cols_str = ", ".join(preparer.quote_identifier(c) for c in insertable_cols) - conflict_cols = ", ".join(preparer.quote_identifier(c) for c in pk_cols) - non_paginated_upsert = sa.text( - f'INSERT INTO {target_ref} ({cols_str})' - f' SELECT {cols_str} FROM {staging_ref}' - f' ON CONFLICT ({conflict_cols}) DO NOTHING' - ) + def _upsert(select_: sa.sql.Select[Any]) -> sa.Insert: + # sa.insert() has no .on_conflict_do_nothing() + return ( + postgresql.insert(target) + .from_select(insertable_cols, select_) + .on_conflict_do_nothing(index_elements=pk_cols) + ) + + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) if merge_batch_size is None: - session.execute(non_paginated_upsert) + session.execute(_upsert(non_paginated_select)) return - total = session.execute(sa.text(f'SELECT COUNT(*) FROM {staging_ref}')).scalar_one() + total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() if total <= merge_batch_size: - session.execute(non_paginated_upsert) + session.execute(_upsert(non_paginated_select)) return - staging_name = self.staging_name_for_table(table_cls.__tablename__) - idx_ref = preparer.quote_identifier(f"{staging_name}_rownum_idx") - session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS {idx_ref} ON {staging_ref} (_rownum)')) - session.commit() + self._staging_rownum_index(table_cls, staging, session) start = 0 while start < total: end = start + merge_batch_size - session.execute( - sa.text( - f'INSERT INTO {target_ref} ({cols_str})' - f' SELECT {cols_str} FROM {staging_ref}' - f' WHERE _rownum > :start AND _rownum <= :end' - f' ON CONFLICT ({conflict_cols}) DO NOTHING' - ), - {"start": start, "end": end}, + batch_select = non_paginated_select.where( + staging.c._rownum > start, staging.c._rownum <= end ) + session.execute(_upsert(batch_select)) session.commit() start = end @@ -274,50 +244,39 @@ def merge_insert( self, table_cls: type["CSVTableProtocol"], session: so.Session, - target_name: str, *, merge_batch_size: int | None = None, ) -> None: - preparer = self.identifier_preparer - staging_ref = self.qualified_staging_name(table_cls.__tablename__) - target_ref = preparer.quote_identifier(target_name) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) insertable_cols = self._insertable_column_names(table_cls) - cols_str = ", ".join(preparer.quote_identifier(c) for c in insertable_cols) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) - non_paginated_insert = sa.text( - f'INSERT INTO {target_ref} ({cols_str})' - f' SELECT {cols_str} FROM {staging_ref}' - ) + def _insert(select_: sa.sql.Select[Any]) -> sa.Insert: + return sa.insert(target).from_select(insertable_cols, select_) if merge_batch_size is None: - session.execute(non_paginated_insert) + session.execute(_insert(non_paginated_select)) return - total = session.execute(sa.text(f'SELECT COUNT(*) FROM {staging_ref}')).scalar_one() + total = session.execute(sa.select(sa.func.count()).select_from(staging)).scalar_one() if total <= merge_batch_size: - session.execute(non_paginated_insert) + 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. - staging_name = self.staging_name_for_table(table_cls.__tablename__) - idx_ref = preparer.quote_identifier(f"{staging_name}_rownum_idx") - session.execute(sa.text(f'CREATE INDEX IF NOT EXISTS {idx_ref} ON {staging_ref} (_rownum)')) - session.commit() + self._staging_rownum_index(table_cls, staging, session) start = 0 while start < total: end = start + merge_batch_size - session.execute( - sa.text( - f'INSERT INTO {target_ref} ({cols_str})' - f' SELECT {cols_str} FROM {staging_ref}' - f' WHERE _rownum > :start AND _rownum <= :end' - ), - {"start": start, "end": end}, + batch_select = non_paginated_select.where( + staging.c._rownum > start, staging.c._rownum <= end ) + session.execute(_insert(batch_select)) session.commit() start = end @@ -335,23 +294,18 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - schema: str | None = None, + role: Role = Role.PRIMARY, with_data: bool = True, if_not_exists: bool = True, ) -> None: from ..mappers.materialised_view_mixin import CreateMaterializedView with self._as_connection(bind) as conn: - _require_postgres_dialect( - conn, - operation=MaterializationOperation.CREATE, - schema=schema, - name=name, - ) + schema = schema_of(conn, role=role) try: conn.execute( CreateMaterializedView( - self._mv_target(name, schema), + qualified(conn, name, schema=schema), selectable, with_data=with_data, if_not_exists=if_not_exists, @@ -374,17 +328,12 @@ def refresh_materialized_view( bind: Engine | Connection, name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, concurrently: bool = False, declared_indexes: tuple["MaterializedViewIndex", ...] = (), ) -> None: with self._as_connection(bind) as conn: - _require_postgres_dialect( - conn, - operation=MaterializationOperation.REFRESH, - schema=schema, - name=name, - ) + schema = schema_of(conn, role=role) if concurrently: if not any(index.unique for index in declared_indexes): raise ConcurrentRefreshNotEligibleError( @@ -396,7 +345,7 @@ def refresh_materialized_view( ) ) - safe_name = self._mv_target(name, schema) + safe_name = qualified(conn, name, schema=schema) concurrency = "CONCURRENTLY " if concurrently else "" try: conn.execute(sa.text(f"REFRESH MATERIALIZED VIEW {concurrency}{safe_name};")) @@ -426,23 +375,18 @@ def drop_materialized_view( bind: Engine | Connection, name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, if_exists: bool = True, cascade: bool = False, ) -> None: from ..mappers.materialised_view_contracts import DropMaterializedView with self._as_connection(bind) as conn: - _require_postgres_dialect( - conn, - operation=MaterializationOperation.DROP, - schema=schema, - name=name, - ) + schema = schema_of(conn, role=role) try: conn.execute( DropMaterializedView( - self._mv_target(name, schema), if_exists=if_exists, cascade=cascade + qualified(conn, name, schema=schema), if_exists=if_exists, cascade=cascade ) ) except Exception as error: @@ -463,22 +407,17 @@ def create_materialized_view_index( name: str, index: "MaterializedViewIndex", *, - schema: str | None = None, + role: Role = Role.PRIMARY, if_not_exists: bool = True, ) -> None: from ..mappers.materialised_view_contracts import CreateMaterializedViewIndex with self._as_connection(bind) as conn: - _require_postgres_dialect( - conn, - operation=MaterializationOperation.CREATE_INDEX, - schema=schema, - name=name, - ) + schema = schema_of(conn, role=role) try: conn.execute( CreateMaterializedViewIndex( - self._mv_target(name, schema), index, if_not_exists=if_not_exists + qualified(conn, name, schema=schema), index, if_not_exists=if_not_exists ) ) except Exception as error: @@ -493,11 +432,6 @@ def create_materialized_view_index( ) ) from error - def _mv_target(self, name: str, schema: str | None) -> str: - from ..helpers.sql import qualify_identifier - - return qualify_identifier(name, schema, self.identifier_preparer) - @contextmanager def engine_with_replica_role(self, engine: "Engine"): def _set_replica_role( @@ -514,9 +448,8 @@ def _set_replica_role( yield engine finally: sae.remove(engine, "connect", _set_replica_role) - with engine.connect() as conn: - conn = conn.execution_options(isolation_level="AUTOCOMMIT") - conn.execute(sa.text("SET session_replication_role = DEFAULT")) - role = conn.execute(sa.text("SHOW session_replication_role")).scalar() + with autocommit_connection(engine) as autocommit_conn: + autocommit_conn.execute(sa.text("SET session_replication_role = DEFAULT")) + role = autocommit_conn.execute(sa.text("SHOW session_replication_role")).scalar() if role != "origin": raise RuntimeError("Failed to restore session_replication_role") diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 445f8e3..42cf9ac 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -81,7 +81,7 @@ def _normalize_fk_check_state(previous_state: str | int) -> str: @property def name(self) -> str: - return "sqlite" + return Dialect.SQLITE @property def dialect(self) -> Dialect: @@ -151,93 +151,122 @@ def restore_fk_check( safe_state = self._normalize_fk_check_state(previous_state) session.execute(text(f"PRAGMA foreign_keys = {safe_state}")) + @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.""" + return sa.literal_column("rowid") + def merge_replace( self, table_cls: type["CSVTableProtocol"], session: so.Session, - target_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: - preparer = self.identifier_preparer - staging_name = self.staging_name_for_table(table_cls.__tablename__) - target_ref = preparer.quote_identifier(target_name) - staging_ref = preparer.quote_identifier(staging_name) - if len(pk_cols) == 1: - pk_ref = preparer.quote_identifier(pk_cols[0]) - session.execute( - sa.text( - f""" - DELETE FROM {target_ref} - WHERE {pk_ref} IN ( - SELECT {pk_ref} FROM {staging_ref} - ); - """ - ) - ) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + pk_match = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) + + # SQLite's DELETE has no USING/multi-table support (confirmed + # empirically: NotImplementedError on a plain multi-table WHERE), so + # this needs an EXISTS correlated subquery instead of Postgres's + # DELETE ... USING. + def _delete(extra: sa.ColumnElement[bool] | None = None) -> sa.Delete: + conditions = (pk_match,) if extra is None else (pk_match, extra) + return sa.delete(target).where(sa.exists().where(*conditions)) + + if merge_batch_size is None: + session.execute(_delete()) return - pk_match = " AND ".join( - f'{target_ref}.{preparer.quote_identifier(c)} = {staging_ref}.{preparer.quote_identifier(c)}' - for c in pk_cols - ) - session.execute( - sa.text( - f""" - DELETE FROM {target_ref} - WHERE EXISTS ( - SELECT 1 FROM {staging_ref} - WHERE {pk_match} - ); - """ - ) - ) + 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))) + session.commit() + start = end def merge_upsert( self, table_cls: type["CSVTableProtocol"], session: so.Session, - target_name: str, pk_cols: list[str], *, merge_batch_size: int | None = None, ) -> None: - preparer = self.identifier_preparer - staging_ref = preparer.quote_identifier(self.staging_name_for_table(table_cls.__tablename__)) - target_ref = preparer.quote_identifier(target_name) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) insertable_cols = self._insertable_column_names(table_cls) - cols_str = ", ".join(preparer.quote_identifier(c) for c in insertable_cols) - session.execute( - sa.text( - f""" - INSERT OR IGNORE INTO {target_ref} ({cols_str}) - SELECT {cols_str} FROM {staging_ref}; - """ + + def _upsert(select_: sa.Select[Any]) -> sa.Insert: + return ( + sqlite_dialect.insert(target) + .from_select(insertable_cols, select_) + .on_conflict_do_nothing(index_elements=pk_cols) ) - ) + + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + + if merge_batch_size is None: + # 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)) + session.commit() + start = end def merge_insert( self, table_cls: type["CSVTableProtocol"], session: so.Session, - target_name: str, *, merge_batch_size: int | None = None, ) -> None: - preparer = self.identifier_preparer - staging_ref = preparer.quote_identifier(self.staging_name_for_table(table_cls.__tablename__)) - target_ref = preparer.quote_identifier(target_name) + target = table_cls.__table__ + staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) insertable_cols = self._insertable_column_names(table_cls) - cols_str = ", ".join(preparer.quote_identifier(c) for c in insertable_cols) - session.execute( - sa.text( - f""" - INSERT INTO {target_ref} ({cols_str}) - SELECT {cols_str} FROM {staging_ref}; - """ - ) - ) + non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) + + def _insert(select_: sa.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)) + session.commit() + start = end def merge_context( self, @@ -272,7 +301,7 @@ def explain_fk_error( raise_error: bool = True, ) -> None: bind: Engine | Connection = session.get_bind() - if bind.dialect.name != "sqlite": + if bind.dialect.name != Dialect.SQLITE: raise exc with self._as_connection(bind) as conn: diff --git a/src/orm_loader/config.py b/src/orm_loader/config.py index 497c719..d712406 100644 --- a/src/orm_loader/config.py +++ b/src/orm_loader/config.py @@ -4,18 +4,35 @@ from typing import Annotated, ClassVar -from oa_configurator import CDMDatabaseConfig, PackageConfigBase, RefTo +from oa_configurator import CDMDatabaseConfig, PackageConfigBase, RefTo, register_reserved_schema +from pydantic import Field + +from .backends.base import STAGING_SCHEMA +# Guaranteed to be imported and registered if there is a config +register_reserved_schema(STAGING_SCHEMA, owner="orm-loader") class OrmLoaderConfig(PackageConfigBase): """oa-configurator config class for orm-loader. - orm-loader is connection-agnostic — it accepts SQLAlchemy sessions/engines + orm-loader is connection-agnostic: it accepts SQLAlchemy sessions/engines as parameters and owns no production database resource of its own. This class exists to register orm-loader in the oa-configurator ecosystem, provide a canonical ``configure_logging()`` entry point, and declare the test database used by the integration test suite. + Attributes + ---------- + test_orm_db_pg : str, optional + Name of the ``[databases.*]`` entry holding the test database. Must + resolve to a real PostgreSQL connection; used for real integration + testing of Postgres-only behavior. + test_orm_db_sqlite : str, optional + Same shape as ``test_orm_db_pg``, for tests that must always run + against SQLite specifically, regardless of what ``test_orm_db_pg`` + happens to be configured to. Left unconfigured by design in every + environment. + Notes ----- By design, this config is for internal use only and must not be @@ -25,4 +42,14 @@ class exists to register orm-loader in the oa-configurator ecosystem, tool_name: ClassVar[str] = "orm_loader" extra_logging_namespaces: ClassVar[tuple[str, ...]] = () - test_orm_db: Annotated[str | None, RefTo(CDMDatabaseConfig, is_test=True)] = None + test_orm_db_pg: Annotated[str | None, RefTo(CDMDatabaseConfig, is_test=True)] = Field( + default=None, + description="Real PostgreSQL test database, for Postgres-only integration testing.", + ) + test_orm_db_sqlite: Annotated[str | None, RefTo(CDMDatabaseConfig, is_test=True)] = Field( + default=None, + description=( + "Disposable SQLite test database; left unconfigured by design " + "(isolated_test_database(..., dialect='sqlite') provisions one automatically)." + ), + ) diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index 6ad5a70..e7757ae 100644 --- a/src/orm_loader/helpers/__init__.py +++ b/src/orm_loader/helpers/__init__.py @@ -9,7 +9,7 @@ from .metadata import Base from .discovery import get_model_by_tablename from .null_handlers import normalise_null -from .sql import qualify_identifier +from .sql import qualify_identifier, role_of_table __all__ = [ "IngestError", @@ -25,4 +25,5 @@ "get_model_by_tablename", "normalise_null", "qualify_identifier", + "role_of_table", ] diff --git a/src/orm_loader/helpers/sql.py b/src/orm_loader/helpers/sql.py index 6e30ba9..ea587d3 100644 --- a/src/orm_loader/helpers/sql.py +++ b/src/orm_loader/helpers/sql.py @@ -1,8 +1,27 @@ from __future__ import annotations +import sqlalchemy as sa +from oa_configurator import Role from sqlalchemy.sql.compiler import IdentifierPreparer +def role_of_table(table: sa.Table) -> Role: + """The ``Role`` a mapped table's own declared schema tag names. + + Every real CDM table is tagged ``schema=Role.X.value`` at class + definition time (Phase 2's schema-role parity work); reading it back off + the ``Table`` itself is the source of truth for which schema_translate_map + key a raw-SQL/reflection call site should resolve through, rather than + always defaulting to primary or reintroducing a manually-threaded + parameter that could disagree with what the table actually declares. + Falls back to ``Role.PRIMARY`` for a table with no schema tag at all + (schema=None), matching schema_of()'s own default. + """ + if table.schema is None: + return Role.PRIMARY + return Role(table.schema) + + def qualify_identifier(name: str, schema: str | None, preparer: IdentifierPreparer) -> str: """ Return a quoted, optionally schema-qualified SQL identifier. diff --git a/src/orm_loader/loaders/loading_helpers.py b/src/orm_loader/loaders/loading_helpers.py index 6bfdbf8..e30bb36 100644 --- a/src/orm_loader/loaders/loading_helpers.py +++ b/src/orm_loader/loaders/loading_helpers.py @@ -11,7 +11,7 @@ import pyarrow.csv as pv import io -from ..helpers.sql import qualify_identifier +from oa_configurator import qualified _SAFE_ENCODING = re.compile(r'^[A-Za-z][A-Za-z0-9_-]*$') @@ -274,7 +274,7 @@ def quick_load_pg( if not hasattr(raw_conn, "cursor"): raise RuntimeError("Expected DB-API connection for COPY") - table_ref = qualify_identifier(tablename, schema, session.get_bind().dialect.identifier_preparer) + table_ref = qualified(session, tablename, schema=schema) encoding = infer_encoding(path)['encoding'] or 'utf-8' if not _SAFE_ENCODING.match(encoding): diff --git a/src/orm_loader/mappers/__init__.py b/src/orm_loader/mappers/__init__.py index eb544c6..1fd7352 100644 --- a/src/orm_loader/mappers/__init__.py +++ b/src/orm_loader/mappers/__init__.py @@ -25,7 +25,6 @@ MaterializationError, MaterializationFailure, MaterializationOperation, - UnsupportedMaterializationDialectError, ) __all__ = [ @@ -38,7 +37,6 @@ "MaterializationOperation", "MaterializedViewIndex", "MaterializedViewMixin", - "UnsupportedMaterializationDialectError", "refresh_all_mvs", "resolve_mv_refresh_order", ] diff --git a/src/orm_loader/mappers/materialised_view_errors.py b/src/orm_loader/mappers/materialised_view_errors.py index ea84a5b..2c85650 100644 --- a/src/orm_loader/mappers/materialised_view_errors.py +++ b/src/orm_loader/mappers/materialised_view_errors.py @@ -40,13 +40,6 @@ def __init__(self, failure: MaterializationFailure) -> None: ) -class UnsupportedMaterializationDialectError(MaterializationError): - """Raised before executing Postgres-only DDL/catalog SQL against a - non-Postgres connection. Defense in depth: the normal ``resolve_backend`` - dispatch path already prevents this via ``_require_capability``; this - guards direct/manual ``PostgresBackend()`` use.""" - - class ConcurrentRefreshNotEligibleError(MaterializationError): """Raised when a concurrent materialized-view refresh is not eligible. diff --git a/src/orm_loader/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index ea7335c..ecdcbc2 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -4,6 +4,7 @@ import sqlalchemy as sa from sqlalchemy.ext import compiler from sqlalchemy.schema import DDLElement +from oa_configurator import Role from .materialised_view_contracts import MaterializedViewIndex @@ -21,7 +22,9 @@ class CreateMaterializedView(DDLElement): Parameters ---------- name - Name of the materialized view to be created. + Fully qualified, quoted name of the materialized view to be created + (see oa_configurator.qualified). The compiler has no live bindable + to qualify a bare name itself, so callers must qualify it first. selectable A SQLAlchemy Select construct defining the query backing the materialized view. @@ -46,8 +49,8 @@ def __init__( @compiler.compiles(CreateMaterializedView) def _create_view( - element: CreateMaterializedView, - compiler: sa.sql.compiler.SQLCompiler, + element: CreateMaterializedView, + compiler: sa.sql.compiler.SQLCompiler, **kwargs: Any ) -> str: @@ -82,6 +85,9 @@ class MaterializedViewMixin: - ``__mv_name__``: the name of the materialized view - ``__mv_select__``: a SQLAlchemy Select defining the view contents - optionally, ``__mv_dependencies__``: names of tables or materialized views this MV depends on + - optionally, ``__mv_role__``: the schema role the view itself lives under + (defaults to primary); set this on a vocab/results view so + :func:`refresh_all_mvs` resolves it to the right schema This mixin does not define ORM mappings; it is intended for schema-level helpers used during migrations, setup, or administrative workflows. @@ -175,6 +181,7 @@ class DailyObservationCountsMV(Base, MaterializedViewMixin): __mv_name__: str __mv_select__: sa.sql.Select[Any] __mv_dependencies__: set[str] = set() + __mv_role__: Role = Role.PRIMARY __mv_indexes__: tuple[MaterializedViewIndex, ...] = () @classmethod @@ -182,7 +189,7 @@ def create_mv( cls, bind: "sa.engine.Connection | sa.engine.Engine", *, - schema: str | None = None, + role: Role | None = None, with_data: bool = True, if_not_exists: bool = True, create_indexes: bool = True, @@ -194,9 +201,12 @@ def create_mv( ---------- bind A SQLAlchemy Engine or Connection used to execute the DDL. - schema - Explicit schema override. When omitted, the view name remains - unqualified for the connection's ``search_path`` to resolve. + role + Schema role the view's own physical schema resolves through. + Defaults to ``cls.__mv_role__`` when omitted. + Set ``__mv_role__`` on a vocab/results view to ensure it resolves + to the correct schema and can be refreshed by :func:`refresh_all_mvs` + without caller needing to know the role. create_indexes When True, create every index declared in ``__mv_indexes__``. @@ -236,13 +246,14 @@ def create_mv( from ..backends.resolve import resolve_backend backend = resolve_backend(bind) + role_ = role if role is not None else cls.__mv_role__ def create(connection: sa.engine.Connection | sa.engine.Engine) -> None: backend.create_materialized_view( connection, cls.__mv_name__, cls.__mv_select__, - schema=schema, + role=role_, with_data=with_data, if_not_exists=if_not_exists, ) @@ -252,7 +263,7 @@ def create(connection: sa.engine.Connection | sa.engine.Engine) -> None: connection, cls.__mv_name__, index, - schema=schema, + role=role_, if_not_exists=if_not_exists, ) @@ -269,7 +280,7 @@ def refresh_mv( cls, bind: "sa.engine.Connection | sa.engine.Engine", *, - schema: str | None = None, + role: Role | None = None, concurrently: bool = False, ) -> None: """ @@ -279,9 +290,9 @@ def refresh_mv( ---------- bind A SQLAlchemy Engine or Connection used to execute the refresh. - schema - Explicit schema override. When omitted, the view name remains - unqualified for the connection's ``search_path`` to resolve. + role + Schema role the view's own physical schema resolves through; + see :meth:`create_mv` for when to override the default. concurrently Request concurrent refresh, requiring a declared unique index. @@ -293,7 +304,7 @@ def refresh_mv( Examples -------- - ```python + ```python with engine.begin() as conn: RecentObservationMV.refresh_mv(conn) ``` @@ -301,10 +312,11 @@ def refresh_mv( from ..backends.resolve import resolve_backend backend = resolve_backend(bind) + role_ = role if role is not None else cls.__mv_role__ backend.refresh_materialized_view( bind, cls.__mv_name__, - schema=schema, + role=role_, concurrently=concurrently, declared_indexes=cls.__mv_indexes__, ) @@ -314,22 +326,28 @@ def drop_mv( cls, bind: "sa.engine.Connection | sa.engine.Engine", *, - schema: str | None = None, + role: Role | None = None, if_exists: bool = True, cascade: bool = False, ) -> None: """Drop the materialized view using the resolved backend. - When ``schema`` is omitted, the view name remains unqualified for the - connection's ``search_path`` to resolve. + Parameters + ---------- + bind + A SQLAlchemy Engine or Connection used to execute the drop. + role + Schema role the view's own physical schema resolves through; + see :meth:`create_mv` for when to override the default. """ from ..backends.resolve import resolve_backend backend = resolve_backend(bind) + role_ = role if role is not None else cls.__mv_role__ backend.drop_materialized_view( - bind, cls.__mv_name__, schema=schema, if_exists=if_exists, cascade=cascade + bind, cls.__mv_name__, role=role_, if_exists=if_exists, cascade=cascade ) - + def resolve_mv_refresh_order(mv_classes: list[type[MaterializedViewMixin]]) -> list[type]: """ diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 73f19ab..12dbcd3 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -2,6 +2,9 @@ import sqlalchemy as sa import sqlalchemy.orm as so import logging +from oa_configurator import schema_inspect + +from ..helpers.sql import role_of_table from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError from typing import Type, Any, Iterator @@ -120,9 +123,8 @@ def manage_indices( table_name = cls.__tablename__ indices = list(cls.__table__.indexes) if resolved_index_strategy == "drop_rebuild" else [] - inspector = sa.inspect(_require_bind(session)) - assert inspector is not None, "Failed to create inspector for index management" - + inspector = schema_inspect(session, role=role_of_table(cls.__table__)) + if indices: existing_in_db = {idx['name'] for idx in inspector.get_indexes(cls.__tablename__)} to_drop = [i for i in indices if i.name in existing_in_db] @@ -232,10 +234,22 @@ def get_staging_table( ------- sqlalchemy.Table The reflected staging table. + + Notes + ----- + Inspects and reflects via ``session.connection()``, not the bare + engine. Confirmed empirically: on SQLite's SingletonThreadPool, a + second connection opened straight from the engine is the same + underlying DBAPI connection, and closing that second wrapper resets + its perceived transaction state, silently discarding the session's + own uncommitted work. Using the session's own already-open + connection avoids ever opening a second one. Fetched fresh both + before and after the possible ``create_staging_table()`` call below, + since that call commits, which can invalidate an earlier reference. """ - engine = _require_bind(session) + _require_bind(session) backend = resolve_backend(session, staging_schema=staging_schema) - inspector = sa.inspect(engine) + inspector = sa.inspect(session.connection()) staging_name = backend.staging_name_for_table(cls.__tablename__) if not inspector.has_table(staging_name, schema=backend.staging_schema): @@ -245,7 +259,7 @@ def get_staging_table( return sa.Table( staging_name, sa.MetaData(), # throwaway — keeps staging table out of Base.metadata - autoload_with=engine, + autoload_with=session.connection(), schema=backend.staging_schema, ) @@ -405,10 +419,7 @@ def load_csv( f"Table `{cls.__tablename__}`: Checking whether target table is empty before staging load." ) check_started = perf_counter() - has_rows = cls._target_has_rows( - session=session, - target=cls.__tablename__, - ) + has_rows = cls._target_has_rows(session=session) logger.info( f"Table `{cls.__tablename__}`: Pre-load empty-table check completed in " f"{_format_elapsed(perf_counter() - check_started)}." @@ -459,20 +470,12 @@ def load_csv( def _target_has_rows( cls: Type[CSVTableProtocol], session: so.Session, - target: str, ) -> bool: """ Return whether the target table currently contains any rows. """ - table = cls.__table__ - if target not in {table.name, table.fullname}: - table = sa.Table( - target, - sa.MetaData(), - autoload_with=session.get_bind(), - ) row = session.execute( - sa.select(sa.literal(1)).select_from(table).limit(1) + sa.select(sa.literal(1)).select_from(cls.__table__).limit(1) ).first() return row is not None @@ -511,10 +514,7 @@ def merge_from_staging( f"Table `{target}`: Checking whether target table is empty for merge optimisation." ) check_started = perf_counter() - has_rows = cls._target_has_rows( - session=session, - target=target, - ) + has_rows = cls._target_has_rows(session=session) logger.info( f"Table `{target}`: Empty-table optimisation check completed in " f"{_format_elapsed(perf_counter() - check_started)}." @@ -530,14 +530,14 @@ 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, target, pk_cols, merge_batch_size=merge_batch_size) + backend.merge_replace(cls, session, pk_cols, merge_batch_size=merge_batch_size) 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, target, merge_batch_size=merge_batch_size) + backend.merge_insert(cls, session, merge_batch_size=merge_batch_size) logger.info( f"Table `{target}`: Merge insert phase completed in " f"{_format_elapsed(perf_counter() - insert_started)}." @@ -545,7 +545,7 @@ 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, target, pk_cols, merge_batch_size=merge_batch_size) + backend.merge_upsert(cls, session, pk_cols, merge_batch_size=merge_batch_size) logger.info( f"Table `{target}`: Merge upsert phase completed in " f"{_format_elapsed(perf_counter() - upsert_started)}." @@ -554,10 +554,7 @@ def merge_from_staging( if not target_empty_confirmed: logger.info(f"Table `{target}`: Checking whether target table is empty.") check_started = perf_counter() - has_rows = cls._target_has_rows( - session=session, - target=target, - ) + has_rows = cls._target_has_rows(session=session) logger.info( f"Table `{target}`: Empty-table check completed in " f"{_format_elapsed(perf_counter() - check_started)}." @@ -571,7 +568,7 @@ def merge_from_staging( logger.info(f"Table `{target}`: Merge insert-if-empty phase starting.") insert_started = perf_counter() - backend.merge_insert(cls, session, target, merge_batch_size=merge_batch_size) + backend.merge_insert(cls, session, merge_batch_size=merge_batch_size) 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 b08bda0..53dd0b6 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -109,7 +109,7 @@ def merge_from_staging( def drop_staging_table(cls, session: so.Session, *, staging_schema: str | None = None) -> None: ... @classmethod - def _target_has_rows(cls, session: so.Session, target: str) -> bool: ... + def _target_has_rows(cls, session: so.Session) -> bool: ... @classmethod def manage_indices( diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index 34154a2..d1583ce 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -12,6 +12,7 @@ import sqlalchemy.orm as so from sqlalchemy.engine import Connection, Engine +from oa_configurator import Role from orm_loader.backends import ( BackendCapabilities, DatabaseBackend, @@ -128,7 +129,7 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - schema: str | None = None, + role: Role = Role.PRIMARY, with_data: bool = True, if_not_exists: bool = True, ) -> None: @@ -139,7 +140,7 @@ def refresh_materialized_view( bind: Engine | Connection, name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, concurrently: bool = False, declared_indexes: tuple = (), ) -> None: diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index e9fa400..e79a5f9 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -7,32 +7,26 @@ import sqlalchemy as sa import sqlalchemy.orm as so from sqlalchemy.dialects import postgresql -from sqlalchemy.engine import Connection, Engine +from sqlalchemy.engine import Engine +from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role +from oa_configurator.testing import isolated_test_schema from orm_loader.backends import STAGING_SCHEMA, Dialect, PostgresBackend from orm_loader.helpers.sql import qualify_identifier +from tests.models import ComputedColumnTable -_TARGET_TABLE = "target_table" +_TARGET_TABLE = ComputedColumnTable.__tablename__ _STAGING_TABLE = f"_staging_{_TARGET_TABLE}" _PREPARER = postgresql.dialect().identifier_preparer _STAGING_TABLE_WITH_SCHEMA: str = qualify_identifier(_STAGING_TABLE, STAGING_SCHEMA, _PREPARER) +_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) + if TYPE_CHECKING: from orm_loader.tables.typing import CSVTableProtocol -class _ComputedTable: - __tablename__ = _TARGET_TABLE - __table__ = sa.Table( - _TARGET_TABLE, - sa.MetaData(), - sa.Column("id", sa.Integer, primary_key=True), - sa.Column("name", sa.String), - sa.Column("slug", sa.String, sa.Computed("lower(name)")), - ) - - class _FakeSession: dialect = postgresql.dialect() @@ -41,11 +35,20 @@ def __init__( scalar_result: str | int | bool = "origin", *, raise_on_execute: Exception | None = None, + schema_translate_map: dict[str, str | None] | None = None, ) -> None: self.statements: list[str] = [] self.scalar_result = scalar_result self.raise_on_execute = raise_on_execute self.commits = 0 + self._schema_translate_map = schema_translate_map + + def get_execution_options(self) -> dict: + """Minimal support for oa_configurator.schema_of(), which every + materialized-view backend method resolves its target schema through.""" + if self._schema_translate_map is None: + return {} + return {SCHEMA_TRANSLATE_MAP_KEY: self._schema_translate_map} def execute(self, statement, parameters=None): if hasattr(statement, "compile"): @@ -72,17 +75,10 @@ def commit(self) -> None: self.commits += 1 -_ComputedTableCls = cast("Type[CSVTableProtocol]", _ComputedTable) - - def _sess(s: _FakeSession) -> so.Session: return cast(so.Session, s) -def _as_engine(s: _FakeSession) -> Engine | Connection: - return cast(Engine, s) - - def test_postgres_backend_identity_and_capabilities(): backend = PostgresBackend() @@ -107,16 +103,14 @@ def test_postgres_backend_default_staging_schema_is_none(): assert backend.qualified_staging_name(_TARGET_TABLE) == _PREPARER.quote_identifier(_STAGING_TABLE) -def test_postgres_backend_create_staging_table_drops_computed_columns(): +def test_postgres_backend_create_staging_table_drops_computed_columns(pg_session): backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession() - backend.create_staging_table(_ComputedTableCls, _sess(session)) + backend.create_staging_table(_ComputedTableCls, pg_session) - assert any(f'DROP TABLE IF EXISTS {_STAGING_TABLE_WITH_SCHEMA}' in sql for sql in session.statements) - assert any(f'CREATE UNLOGGED TABLE {_STAGING_TABLE_WITH_SCHEMA}' in sql for sql in session.statements) - assert any(f'ALTER TABLE {_STAGING_TABLE_WITH_SCHEMA} DROP COLUMN "slug"' in sql for sql in session.statements) - assert session.commits == 1 + inspector = sa.inspect(pg_session.get_bind()) + cols = {c["name"] for c in inspector.get_columns(_STAGING_TABLE, schema=STAGING_SCHEMA)} + assert cols == {"id", "name", "_rownum"} # slug is computed, excluded def test_postgres_backend_drop_staging_table(): @@ -147,96 +141,88 @@ def test_postgres_backend_fk_methods_emit_expected_sql(): ] -def test_postgres_backend_merge_replace_uses_using_delete(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession(scalar_result=0) - - backend.merge_replace(_ComputedTableCls, _sess(session), _TARGET_TABLE, ["id", "name"]) - - sql = session.statements[0] - assert f'DELETE FROM "{_TARGET_TABLE}" t' in sql - assert f'USING {_STAGING_TABLE_WITH_SCHEMA} s' in sql - assert 't."id" = s."id" AND t."name" = s."name"' in sql - assert f'USING {qualify_identifier(_TARGET_TABLE, STAGING_SCHEMA, _PREPARER)}' not in sql - - -def test_postgres_backend_merge_insert_excludes_computed_columns(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession(scalar_result=0) - - backend.merge_insert(_ComputedTableCls, _sess(session), _TARGET_TABLE) +def test_postgres_backend_materialized_view_methods_work_end_to_end(pg_db): + """Real create + refresh + query, not just checking emitted SQL text. + The whole point is proving this DDL actually round-trips correctly.""" + backend = PostgresBackend() + conn = pg_db.connection + selectable = sa.select(sa.literal(1).label("n")) - sql = session.statements[0] - assert f'INSERT INTO "{_TARGET_TABLE}" ("id", "name")' in sql - assert f'SELECT "id", "name" FROM {_STAGING_TABLE_WITH_SCHEMA}' in sql + backend.create_materialized_view(conn, "mv_test", selectable) + backend.refresh_materialized_view(conn, "mv_test") + assert conn.execute(sa.text("SELECT n FROM mv_test")).scalar() == 1 -def test_postgres_backend_merge_upsert_excludes_computed_columns(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession(scalar_result=0) - backend.merge_upsert(_ComputedTableCls, _sess(session), _TARGET_TABLE, ["id"]) +def test_postgres_backend_materialized_view_respects_role(pg_db) -> None: + """create_materialized_view()/refresh_materialized_view() used to always + resolve schema=None -> schema_of(conn) with no role, which defaults to + Role.PRIMARY regardless of what role the view was actually built over. + A view over vocab-role tables must land in the vocab schema, not + wherever primary happens to be.""" + backend = PostgresBackend() + selectable = sa.select(sa.literal(1).label("n")) + engine = pg_db.connection.engine - sql = session.statements[0] - assert f'INSERT INTO "{_TARGET_TABLE}" ("id", "name")' in sql - assert 'ON CONFLICT ("id") DO NOTHING' in sql + with isolated_test_schema(engine, prefix="mv_primary") as primary_schema, \ + isolated_test_schema(engine, prefix="mv_vocab") as vocab_schema: + scoped = engine.execution_options( + schema_translate_map={Role.PRIMARY.value: primary_schema, "vocab": vocab_schema} + ) + with scoped.begin() as conn: + backend.create_materialized_view(conn, "mv_role_test", selectable, role=Role.VOCAB) + backend.refresh_materialized_view(conn, "mv_role_test", role=Role.VOCAB) + with engine.connect() as conn: + assert sa.inspect(conn).has_table("mv_role_test", schema=vocab_schema) + assert not sa.inspect(conn).has_table("mv_role_test", schema=primary_schema) + assert conn.execute( + sa.text(f'SELECT n FROM "{vocab_schema}".mv_role_test') + ).scalar() == 1 -def test_postgres_backend_merge_replace_paginated_path(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession(scalar_result=10) - backend.merge_replace( - _ComputedTableCls, _sess(session), _TARGET_TABLE, - ["id", "name"], merge_batch_size=3, - ) +def test_postgres_backend_materialized_view_methods_emit_expected_sql(): + backend = PostgresBackend() + session = _FakeSession() + selectable = sa.select(sa.literal(1).label("n")) - sqls = session.statements - assert any("CREATE INDEX IF NOT EXISTS" in s and "_rownum" in s for s in sqls) - assert any("_rownum >" in s and "DELETE" in s for s in sqls) - assert session.commits >= 4 # 1 for index + 4 batches (ceil(10/3)) + backend.create_materialized_view(_sess(session), "mv_test", selectable) + backend.refresh_materialized_view(_sess(session), "mv_test") + assert any('CREATE MATERIALIZED VIEW IF NOT EXISTS mv_test as SELECT' in sql for sql in session.statements) + assert any('REFRESH MATERIALIZED VIEW mv_test;' == sql for sql in session.statements) -def test_postgres_backend_merge_insert_paginated_path(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession(scalar_result=10) - backend.merge_insert( - _ComputedTableCls, _sess(session), _TARGET_TABLE, - merge_batch_size=3, - ) +def test_postgres_backend_rejects_mismatched_dialect_bind(): + from sqlalchemy.dialects import sqlite - sqls = session.statements - assert any("CREATE INDEX IF NOT EXISTS" in s and "_rownum" in s for s in sqls) - assert any("_rownum >" in s and "INSERT" in s for s in sqls) - assert session.commits >= 4 + backend = PostgresBackend() + session = _FakeSession() + session.dialect = sqlite.dialect() + selectable = sa.select(sa.literal(1).label("n")) + with pytest.raises(TypeError, match="received a 'sqlite' connection; expected 'postgresql'"): + backend.create_materialized_view(_sess(session), "mv_test", selectable) -def test_postgres_backend_merge_upsert_paginated_path(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession(scalar_result=10) + assert session.statements == [] - backend.merge_upsert( - _ComputedTableCls, _sess(session), _TARGET_TABLE, - ["id"], merge_batch_size=3, - ) - sqls = session.statements - assert any("CREATE INDEX IF NOT EXISTS" in s and "_rownum" in s for s in sqls) - assert any("_rownum >" in s and "INSERT" in s for s in sqls) - assert session.commits >= 4 +def test_postgres_backend_rejects_dialect_that_drifted_after_resolve_backend(): + """A bind resolved to PostgresBackend, then a differently-dialected bind + handed to one of its methods, must not run Postgres-only DDL against it.""" + from orm_loader.backends.resolve import resolve_backend + from sqlalchemy.dialects import sqlite + postgres_session = _FakeSession() + backend = resolve_backend(_sess(postgres_session)) + assert isinstance(backend, PostgresBackend) -def test_postgres_backend_materialized_view_methods_emit_expected_sql(): - backend = PostgresBackend() - session = _FakeSession() + sqlite_session = _FakeSession() + sqlite_session.dialect = sqlite.dialect() selectable = sa.select(sa.literal(1).label("n")) - backend.create_materialized_view(_as_engine(session), "mv_test", selectable) - backend.refresh_materialized_view(_as_engine(session), "mv_test") - - assert any('CREATE MATERIALIZED VIEW IF NOT EXISTS "mv_test" as SELECT' in sql for sql in session.statements) - assert any('REFRESH MATERIALIZED VIEW "mv_test";' == sql for sql in session.statements) + with pytest.raises(TypeError, match="received a 'sqlite' connection; expected 'postgresql'"): + backend.create_materialized_view(_sess(sqlite_session), "mv_test", selectable) def test_postgres_backend_quotes_unqualified_materialized_view_name(): @@ -255,39 +241,6 @@ def test_postgres_backend_quotes_unqualified_materialized_view_name(): assert any('ON "mv name" ("n")' in sql for sql in session.statements) -def test_postgres_backend_create_materialized_view_rejects_non_postgres_connection(): - from orm_loader.mappers.materialised_view_errors import ( - UnsupportedMaterializationDialectError, - ) - from sqlalchemy.dialects import sqlite - - backend = PostgresBackend() - session = _FakeSession() - session.dialect = sqlite.dialect() - selectable = sa.select(sa.literal(1).label("n")) - - with pytest.raises(UnsupportedMaterializationDialectError, match="received dialect 'sqlite'"): - backend.create_materialized_view(_sess(session), "mv_test", selectable) - - assert session.statements == [] - - -def test_postgres_backend_refresh_materialized_view_rejects_non_postgres_connection(): - from orm_loader.mappers.materialised_view_errors import ( - UnsupportedMaterializationDialectError, - ) - from sqlalchemy.dialects import sqlite - - backend = PostgresBackend() - session = _FakeSession() - session.dialect = sqlite.dialect() - - with pytest.raises(UnsupportedMaterializationDialectError, match="received dialect 'sqlite'"): - backend.refresh_materialized_view(_sess(session), "mv_test") - - assert session.statements == [] - - def test_postgres_backend_create_mv_quotes_name_for_legacy_search_path_resolution(): backend = PostgresBackend() session = _FakeSession() @@ -305,13 +258,13 @@ def test_postgres_backend_create_materialized_view_index_emits_expected_sql(): from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex backend = PostgresBackend() - session = _FakeSession() + session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) index = MaterializedViewIndex(name="mv_test_row_id_uq", columns=("row_id",), unique=True) - backend.create_materialized_view_index(_sess(session), "mv_test", index, schema="reporting") + backend.create_materialized_view_index(_sess(session), "mv_test", index) assert session.statements == [ - 'CREATE UNIQUE INDEX IF NOT EXISTS "mv_test_row_id_uq" ON "reporting"."mv_test" ("row_id")' + 'CREATE UNIQUE INDEX IF NOT EXISTS "mv_test_row_id_uq" ON reporting.mv_test ("row_id")' ] @@ -330,22 +283,20 @@ def test_postgres_backend_create_materialized_view_index_failure_mentions_index_ def test_postgres_backend_drop_materialized_view_default_args(): backend = PostgresBackend() - session = _FakeSession() + session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) - backend.drop_materialized_view(_sess(session), "mv_test", schema="reporting") + backend.drop_materialized_view(_sess(session), "mv_test") - assert session.statements == ['DROP MATERIALIZED VIEW IF EXISTS "reporting"."mv_test"'] + assert session.statements == ['DROP MATERIALIZED VIEW IF EXISTS reporting.mv_test'] def test_postgres_backend_drop_materialized_view_cascade_and_if_exists_false(): backend = PostgresBackend() - session = _FakeSession() + session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) - backend.drop_materialized_view( - _sess(session), "mv_test", schema="reporting", if_exists=False, cascade=True - ) + backend.drop_materialized_view(_sess(session), "mv_test", if_exists=False, cascade=True) - assert session.statements == ['DROP MATERIALIZED VIEW "reporting"."mv_test" CASCADE'] + assert session.statements == ['DROP MATERIALIZED VIEW reporting.mv_test CASCADE'] def test_postgres_backend_drop_materialized_view_failure_preserves_cause(): @@ -353,10 +304,12 @@ def test_postgres_backend_drop_materialized_view_failure_preserves_cause(): original = RuntimeError("boom") backend = PostgresBackend() - session = _FakeSession(raise_on_execute=original) + session = _FakeSession( + raise_on_execute=original, schema_translate_map={Role.PRIMARY.value: "reporting"} + ) with pytest.raises(MaterializationError) as exc_info: - backend.drop_materialized_view(_sess(session), "mv_test", schema="reporting") + backend.drop_materialized_view(_sess(session), "mv_test") assert exc_info.value.__cause__ is original assert exc_info.value.failure.cause is original @@ -384,12 +337,10 @@ def test_postgres_backend_refresh_concurrently_without_declared_unique_index_rai from orm_loader.mappers.materialised_view_errors import ConcurrentRefreshNotEligibleError backend = PostgresBackend() - session = _FakeSession() + session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) with pytest.raises(ConcurrentRefreshNotEligibleError, match="no simple unique index"): - backend.refresh_materialized_view( - _sess(session), "mv_test", schema="reporting", concurrently=True - ) + backend.refresh_materialized_view(_sess(session), "mv_test", concurrently=True) assert session.statements == [] @@ -406,21 +357,23 @@ def test_postgres_backend_refresh_concurrently_declared_but_database_rejects_it_ "columns of the materialized view." ) backend = PostgresBackend() - session = _FakeSession(raise_on_execute=sa.exc.OperationalError("REFRESH ...", {}, orig)) + session = _FakeSession( + raise_on_execute=sa.exc.OperationalError("REFRESH ...", {}, orig), + schema_translate_map={Role.PRIMARY.value: "reporting"}, + ) index = MaterializedViewIndex(name="mv_test_row_id_uq", columns=("row_id",), unique=True) with pytest.raises(ConcurrentRefreshNotEligibleError, match="cannot refresh") as exc_info: backend.refresh_materialized_view( _sess(session), "mv_test", - schema="reporting", concurrently=True, declared_indexes=(index,), ) assert exc_info.value.failure.cause.orig is orig assert session.statements == [ - 'REFRESH MATERIALIZED VIEW CONCURRENTLY "reporting"."mv_test";' + 'REFRESH MATERIALIZED VIEW CONCURRENTLY reporting.mv_test;' ] @@ -431,14 +384,16 @@ def test_postgres_backend_refresh_concurrently_unrelated_operational_error_propa orig = psycopg.errors.QueryCanceled("canceling statement due to statement timeout") backend = PostgresBackend() - session = _FakeSession(raise_on_execute=sa.exc.OperationalError("REFRESH ...", {}, orig)) + session = _FakeSession( + raise_on_execute=sa.exc.OperationalError("REFRESH ...", {}, orig), + schema_translate_map={Role.PRIMARY.value: "reporting"}, + ) index = MaterializedViewIndex(name="mv_test_row_id_uq", columns=("row_id",), unique=True) with pytest.raises(sa.exc.OperationalError) as exc_info: backend.refresh_materialized_view( _sess(session), "mv_test", - schema="reporting", concurrently=True, declared_indexes=(index,), ) @@ -450,25 +405,28 @@ def test_postgres_backend_refresh_concurrently_with_declared_index_emits_concurr from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex backend = PostgresBackend() - session = _FakeSession() + session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) index = MaterializedViewIndex(name="mv_test_row_id_uq", columns=("row_id",), unique=True) backend.refresh_materialized_view( _sess(session), "mv_test", - schema="reporting", concurrently=True, declared_indexes=(index,), ) assert session.statements[-1] == ( - 'REFRESH MATERIALIZED VIEW CONCURRENTLY "reporting"."mv_test";' + 'REFRESH MATERIALIZED VIEW CONCURRENTLY reporting.mv_test;' ) def test_postgres_backend_materialized_view_lifecycle_is_schema_isolated_with_adversarial_identifiers( - pg_session, pg_engine + pg_db, ): + """Two schemas, each addressed via its own scoped connection (role-based + resolution ties the schema to the connection, not to a per-call + override), must never bleed into each other even with adversarial, + quote-laden identifiers.""" from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex backend = PostgresBackend() @@ -476,76 +434,86 @@ def test_postgres_backend_materialized_view_lifecycle_is_schema_isolated_with_ad name = 'shared "view" name' index = MaterializedViewIndex(name="shared_name_row_id_uq", columns=("row_id",), unique=True) selectable = sa.select(sa.literal(1).label("row_id")) + engine = pg_db.connection.engine + preparer = postgresql.dialect().identifier_preparer - with pg_engine.begin() as conn: - for schema in (left_schema, right_schema): - preparer = postgresql.dialect().identifier_preparer - conn.execute(sa.text(f"CREATE SCHEMA {preparer.quote_identifier(schema)}")) - backend.create_materialized_view(conn, name, selectable, schema=schema) - backend.create_materialized_view_index(conn, name, index, schema=schema) - - backend.refresh_materialized_view( - conn, name, schema=left_schema, concurrently=True, declared_indexes=(index,) - ) - backend.drop_materialized_view(conn, name, schema=left_schema) - - assert conn.execute( - sa.text( - "SELECT EXISTS (SELECT 1 FROM pg_matviews " - "WHERE schemaname = :schema AND matviewname = :name)" - ), - {"schema": left_schema, "name": name}, - ).scalar() is False - assert conn.execute( - sa.text( - "SELECT EXISTS (SELECT 1 FROM pg_matviews " - "WHERE schemaname = :schema AND matviewname = :name)" - ), - {"schema": right_schema, "name": name}, - ).scalar() is True - - -def test_postgres_backend_refresh_concurrently_raises_when_declared_index_was_never_created( - pg_session, pg_engine -): + try: + with engine.begin() as setup_conn: + for schema in (left_schema, right_schema): + setup_conn.execute(sa.text(f"CREATE SCHEMA {preparer.quote_identifier(schema)}")) + + left = engine.execution_options(schema_translate_map={Role.PRIMARY.value: left_schema}) + right = engine.execution_options(schema_translate_map={Role.PRIMARY.value: right_schema}) + + for scoped in (left, right): + with scoped.begin() as conn: + backend.create_materialized_view(conn, name, selectable) + backend.create_materialized_view_index(conn, name, index) + + with left.begin() as conn: + backend.refresh_materialized_view(conn, name, concurrently=True, declared_indexes=(index,)) + backend.drop_materialized_view(conn, name) + + with engine.connect() as conn: + assert conn.execute( + sa.text( + "SELECT EXISTS (SELECT 1 FROM pg_matviews " + "WHERE schemaname = :schema AND matviewname = :name)" + ), + {"schema": left_schema, "name": name}, + ).scalar() is False + assert conn.execute( + sa.text( + "SELECT EXISTS (SELECT 1 FROM pg_matviews " + "WHERE schemaname = :schema AND matviewname = :name)" + ), + {"schema": right_schema, "name": name}, + ).scalar() is True + finally: + with engine.begin() as cleanup_conn: + for schema in (left_schema, right_schema): + cleanup_conn.execute( + sa.text(f"DROP SCHEMA IF EXISTS {preparer.quote_identifier(schema)} CASCADE") + ) + + +def test_postgres_backend_refresh_concurrently_raises_when_declared_index_was_never_created(pg_db): from orm_loader.mappers.materialised_view_errors import ConcurrentRefreshNotEligibleError from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex backend = PostgresBackend() index = MaterializedViewIndex(name="mv_missing_index_test_uq", columns=("row_id",), unique=True) + conn = pg_db.connection - with pg_engine.begin() as conn: - backend.create_materialized_view( - conn, "mv_missing_index_test", sa.select(sa.literal(1).label("row_id")) - ) + backend.create_materialized_view( + conn, "mv_missing_index_test", sa.select(sa.literal(1).label("row_id")) + ) - with pytest.raises(ConcurrentRefreshNotEligibleError) as exc_info: - backend.refresh_materialized_view( - conn, - "mv_missing_index_test", - concurrently=True, - declared_indexes=(index,), - ) + with pytest.raises(ConcurrentRefreshNotEligibleError) as exc_info: + backend.refresh_materialized_view( + conn, + "mv_missing_index_test", + concurrently=True, + declared_indexes=(index,), + ) - assert "concurrently" in str(exc_info.value).lower() - assert isinstance(exc_info.value.__cause__, sa.exc.OperationalError) + assert "concurrently" in str(exc_info.value).lower() + assert isinstance(exc_info.value.__cause__, sa.exc.OperationalError) -def test_postgres_backend_materialized_view_legacy_unqualified_path_still_round_trips( - pg_session, pg_engine -): +def test_postgres_backend_materialized_view_legacy_unqualified_path_still_round_trips(pg_db): backend = PostgresBackend() + conn = pg_db.connection - with pg_engine.begin() as conn: - backend.create_materialized_view( - conn, "mv_legacy_test", sa.select(sa.literal(1).label("n")) - ) - backend.refresh_materialized_view(conn, "mv_legacy_test") - assert conn.execute(sa.text("SELECT n FROM mv_legacy_test")).scalar() == 1 + backend.create_materialized_view( + conn, "mv_legacy_test", sa.select(sa.literal(1).label("n")) + ) + backend.refresh_materialized_view(conn, "mv_legacy_test") + assert conn.execute(sa.text("SELECT n FROM mv_legacy_test")).scalar() == 1 - backend.drop_materialized_view(conn, "mv_legacy_test") - with pytest.raises(sa.exc.ProgrammingError): - conn.execute(sa.text("SELECT n FROM mv_legacy_test")) + backend.drop_materialized_view(conn, "mv_legacy_test") + with pytest.raises(sa.exc.ProgrammingError): + conn.execute(sa.text("SELECT n FROM mv_legacy_test")) def test_postgres_backend_normalize_fk_check_state(): @@ -615,12 +583,25 @@ def __exit__(self, *_) -> None: def execution_options(self, **_): return self + def get_isolation_level(self): + return "READ COMMITTED" + + def rollback(self) -> None: + return None + + def close(self) -> None: + return None + def execute(self, statement): sql = str(statement.compile(dialect=postgresql.dialect())) statements.append(sql) return _Result() - class _Engine: + class _Engine(Engine): + def __init__(self) -> None: + # only exists for autocommit_connection() to route it into its real Engine branch + pass + def connect(self): events.append(("connect", self, "connect")) return _Conn() diff --git a/tests/backends/test_reserved_schema.py b/tests/backends/test_reserved_schema.py new file mode 100644 index 0000000..4ff133a --- /dev/null +++ b/tests/backends/test_reserved_schema.py @@ -0,0 +1,22 @@ +"""Confirms orm-loader's STAGING_SCHEMA registration (backends/base.py, +Phase 2.3) is actually picked up by oa-configurator's reserved-schema +check: resolving a CDM database configured with cdm_schema="staging" +must raise, proving the cross-package registration/enforcement wiring +works end to end, not just in isolation on either side. +""" + +from __future__ import annotations + +import pytest +from oa_configurator import CDMDatabaseConfig, ConnectionConfig, StackConfig +from pydantic import ValidationError + +from orm_loader.backends import STAGING_SCHEMA + + +def test_resolving_cdm_database_with_staging_schema_name_raises() -> None: + with pytest.raises(ValidationError, match=f"{STAGING_SCHEMA!r}.*orm-loader"): + StackConfig.for_session( + connections={"c": ConnectionConfig(dialect="sqlite", database_name=":memory:")}, + databases={"default": CDMDatabaseConfig(connection="c", cdm_schema=STAGING_SCHEMA)}, + ) diff --git a/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py new file mode 100644 index 0000000..381160a --- /dev/null +++ b/tests/backends/test_shared_backend.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Type, cast + +import pytest +import sqlalchemy as sa + +from oa_configurator.testing import DIALECT_PARAMS +from orm_loader.backends import STAGING_SCHEMA, DatabaseBackend, PostgresBackend, SQLiteBackend +from tests.models import ComputedColumnTable, CompositeTable + +if TYPE_CHECKING: + import sqlalchemy.orm as so + + from orm_loader.tables.typing import CSVTableProtocol + +_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) +_CompositeTableCls = cast("Type[CSVTableProtocol]", CompositeTable) + + +@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. + + 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). + """ + if request.param == "postgresql": + session = request.getfixturevalue("pg_session") + return PostgresBackend(staging_schema=STAGING_SCHEMA), session + session = request.getfixturevalue("session") + return SQLiteBackend(), session + + +def test_merge_replace_single_pk(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: + backend, session = merge_backend + backend.create_staging_table(_ComputedTableCls, session) + staging = _ComputedTableCls.get_staging_table(session, staging_schema=backend.staging_schema) + + session.execute( + sa.insert(ComputedColumnTable), + [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}], + ) + session.execute(sa.insert(staging), [{"id": 1, "name": "alpha-staged"}]) + + backend.merge_replace(_ComputedTableCls, session, ["id"]) + + remaining = session.execute(sa.select(ComputedColumnTable.id)).scalars().all() + assert remaining == [2] + + +def test_merge_replace_composite_pk(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: + backend, session = merge_backend + backend.create_staging_table(_CompositeTableCls, session) + staging = _CompositeTableCls.get_staging_table(session, staging_schema=backend.staging_schema) + + session.execute( + sa.insert(CompositeTable), + [{"a": 1, "b": 1, "value": "x"}, {"a": 2, "b": 2, "value": "y"}], + ) + session.execute(sa.insert(staging), [{"a": 1, "b": 1, "value": "staged"}]) + + backend.merge_replace(_CompositeTableCls, session, ["a", "b"]) + + remaining = session.execute(sa.select(CompositeTable.a, CompositeTable.b)).all() + assert remaining == [(2, 2)] + + +def test_merge_insert_excludes_computed_columns(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: + backend, session = merge_backend + backend.create_staging_table(_ComputedTableCls, session) + staging = _ComputedTableCls.get_staging_table(session, staging_schema=backend.staging_schema) + session.execute(sa.insert(staging), [{"id": 1, "name": "alpha"}]) + + backend.merge_insert(_ComputedTableCls, session) + + row = session.execute(sa.select(ComputedColumnTable)).scalars().one() + assert (row.id, row.name, row.slug) == (1, "alpha", "alpha") + + +def test_merge_upsert_excludes_computed_columns(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: + backend, session = merge_backend + backend.create_staging_table(_ComputedTableCls, session) + staging = _ComputedTableCls.get_staging_table(session, staging_schema=backend.staging_schema) + + session.execute(sa.insert(ComputedColumnTable), [{"id": 1, "name": "existing"}]) + session.execute( + sa.insert(staging), [{"id": 1, "name": "ignored"}, {"id": 2, "name": "new"}] + ) + + backend.merge_upsert(_ComputedTableCls, session, ["id"]) + + rows = {r.id: r.name for r in session.execute(sa.select(ComputedColumnTable)).scalars().all()} + assert rows == {1: "existing", 2: "new"} + + +def test_merge_replace_paginated_path(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: + backend, session = merge_backend + backend.create_staging_table(_ComputedTableCls, session) + staging = _ComputedTableCls.get_staging_table(session, staging_schema=backend.staging_schema) + + session.execute( + sa.insert(ComputedColumnTable), [{"id": i, "name": f"orig{i}"} for i in range(10)] + ) + session.execute(sa.insert(staging), [{"id": i, "name": f"staged{i}"} for i in range(10)]) + + backend.merge_replace(_ComputedTableCls, session, ["id"], merge_batch_size=3) + + remaining = session.execute( + sa.select(sa.func.count()).select_from(ComputedColumnTable.__table__) + ).scalar() + assert remaining == 0 + + +def test_merge_insert_paginated_path(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: + backend, session = merge_backend + backend.create_staging_table(_ComputedTableCls, session) + staging = _ComputedTableCls.get_staging_table(session, staging_schema=backend.staging_schema) + session.execute(sa.insert(staging), [{"id": i, "name": f"row{i}"} for i in range(10)]) + + backend.merge_insert(_ComputedTableCls, session, merge_batch_size=3) + + ids = sorted(session.execute(sa.select(ComputedColumnTable.id)).scalars().all()) + assert ids == list(range(10)) + + +def test_merge_upsert_paginated_path(merge_backend: tuple[DatabaseBackend, "so.Session"]) -> None: + backend, session = merge_backend + backend.create_staging_table(_ComputedTableCls, session) + staging = _ComputedTableCls.get_staging_table(session, staging_schema=backend.staging_schema) + + session.execute( + sa.insert(ComputedColumnTable), [{"id": i, "name": "kept"} for i in range(5)] + ) + session.execute( + sa.insert(staging), [{"id": i, "name": "should-not-overwrite"} for i in range(10)] + ) + + backend.merge_upsert(_ComputedTableCls, session, ["id"], merge_batch_size=3) + + rows = {r.id: r.name for r in session.execute(sa.select(ComputedColumnTable)).scalars().all()} + assert rows == {**{i: "kept" for i in range(5)}, **{i: "should-not-overwrite" for i in range(5, 10)}} diff --git a/tests/backends/test_sqlite_backend.py b/tests/backends/test_sqlite_backend.py index 9cc2fc0..09aabfe 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -10,25 +10,15 @@ from orm_loader.backends import Dialect, SQLiteBackend from orm_loader.helpers.sqlite import attach_sqlite_bulk_load_pragmas +from tests.models import ComputedColumnTable if TYPE_CHECKING: from orm_loader.tables.typing import CSVTableProtocol -_TARGET_TABLE = "target_table" +_TARGET_TABLE = ComputedColumnTable.__tablename__ _STAGING_TABLE = f"_staging_{_TARGET_TABLE}" -class _ComputedTable: - __tablename__ = _TARGET_TABLE - __table__ = sa.Table( - _TARGET_TABLE, - sa.MetaData(), - sa.Column("id", sa.Integer, primary_key=True), - sa.Column("name", sa.String), - sa.Column("slug", sa.String, sa.Computed("lower(name)")), - ) - - class _FakeSession: def __init__(self, scalar_result: int | str = 1) -> None: self.statements: list[str] = [] @@ -47,7 +37,7 @@ def scalar(self): return _Result(self.scalar_result) -_ComputedTableCls = cast("Type[CSVTableProtocol]", _ComputedTable) +_ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) def _sess(s: _FakeSession) -> so.Session: @@ -155,56 +145,6 @@ def test_sqlite_backend_normalize_fk_check_state(): raise AssertionError("Expected ValueError for unrecognised string") -def test_sqlite_backend_merge_replace_single_pk(): - backend = SQLiteBackend() - session = _FakeSession() - - backend.merge_replace( - _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id"] - ) - - sql = session.statements[0] - assert f'DELETE FROM "{_TARGET_TABLE}"' in sql - assert f'SELECT "id" FROM "{_STAGING_TABLE}"' in sql - - -def test_sqlite_backend_merge_replace_composite_pk(): - backend = SQLiteBackend() - session = _FakeSession() - - backend.merge_replace( - _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id", "name"] - ) - - sql = session.statements[0] - assert "WHERE EXISTS (" in sql - assert f'"{_TARGET_TABLE}"."id" = "{_STAGING_TABLE}"."id"' in sql - assert f'"{_TARGET_TABLE}"."name" = "{_STAGING_TABLE}"."name"' in sql - - -def test_sqlite_backend_merge_insert_excludes_computed_columns(): - backend = SQLiteBackend() - session = _FakeSession() - - backend.merge_insert(_ComputedTableCls, _sess(session), _TARGET_TABLE) - - sql = session.statements[0] - assert f'INSERT INTO "{_TARGET_TABLE}" ("id", "name")' in sql - assert f'SELECT "id", "name" FROM "{_STAGING_TABLE}"' in sql - - -def test_sqlite_backend_merge_upsert_excludes_computed_columns(): - backend = SQLiteBackend() - session = _FakeSession() - - backend.merge_upsert( - _ComputedTableCls, _sess(session), _TARGET_TABLE, ["id"] - ) - - sql = session.statements[0] - assert f'INSERT OR IGNORE INTO "{_TARGET_TABLE}" ("id", "name")' in sql - - def test_sqlite_backend_materialized_view_methods_raise(engine): backend = SQLiteBackend() selectable = sa.select(sa.literal(1).label("n")) diff --git a/tests/conftest.py b/tests/conftest.py index 7a6e01f..c8694d6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,3 @@ -import time from pathlib import Path import pytest @@ -6,7 +5,9 @@ import sqlalchemy.orm as so from dotenv import load_dotenv +from oa_configurator.testing import isolated_test_database from orm_loader.backends import STAGING_SCHEMA +from orm_loader.config import OrmLoaderConfig from tests.models import Base load_dotenv(Path(__file__).parent.parent / ".env") @@ -14,9 +15,12 @@ @pytest.fixture def engine(): - engine = sa.create_engine("sqlite:///:memory:", future=True) - Base.metadata.create_all(engine) - return engine + with isolated_test_database( + OrmLoaderConfig, "test_orm_db_sqlite", dialect="sqlite", future=True, + ) as db: + engine = db.connection.engine + Base.metadata.create_all(engine) + yield engine @pytest.fixture @@ -29,51 +33,33 @@ def session(engine): # Postgres fixtures # --------------------------------------------------------------------------- -@pytest.fixture(scope="session") -def pg_engine(): - from oa_configurator.pytest_plugin import ensure_test_db_exists, resolve_test_database - from orm_loader.config import OrmLoaderConfig - - url = resolve_test_database(OrmLoaderConfig, "test_orm_db") - - try: - ensure_test_db_exists(url) - except Exception as exc: - print(f"Could not ensure test DB exists, will try anyway: {exc}") - - last_err = None - for i in range(20): - engine: sa.Engine | None = None - try: - engine = sa.create_engine(url, future=True) - with engine.connect() as conn: - conn.execute(sa.text("SELECT 1")) - print("Postgres connection established") - yield engine - engine.dispose() - return - except Exception as exc: - if engine is not None: - engine.dispose() - last_err = exc - print(f"[{i}] Postgres not ready:", repr(exc)) - time.sleep(1) - - pytest.skip(f"PostgreSQL never became available: {last_err}") +@pytest.fixture +def pg_db(request): + """Isolated PostgreSQL test database. Everything done through + ``pg_db.connection``/``pg_db.session`` happens inside one transaction + that's rolled back on exit, so concurrent test runs can't collide and + nothing needs manual cleanup.""" + with isolated_test_database(OrmLoaderConfig, "test_orm_db_pg", request=request) as db: + yield db @pytest.fixture -def pg_session(pg_engine): - Session = so.sessionmaker(pg_engine, future=True) - with pg_engine.begin() as conn: - conn.execute(sa.text(f"DROP SCHEMA IF EXISTS {STAGING_SCHEMA} CASCADE")) - conn.execute(sa.text(f"CREATE SCHEMA {STAGING_SCHEMA}")) - Base.metadata.drop_all(conn) - Base.metadata.create_all(conn) - - session = Session() - try: - yield session - finally: - session.rollback() - session.close() +def pg_session(pg_db): + """The standard fixture for tests needing real tables ready to query: + creates the staging schema and Base.metadata inside pg_db's already-open, + rolled-back transaction, then returns pg_db.session.""" + conn = pg_db.connection + conn.execute(sa.text(f"CREATE SCHEMA IF NOT EXISTS {STAGING_SCHEMA}")) + Base.metadata.create_all(conn) + return pg_db.session + + +def schema_scoped_session( + conn: sa.Connection, table: sa.Table, schema_translate_map: dict +) -> so.Session: + """A Session scoped to *schema_translate_map*, with *table* already + created through it. + """ + scoped_conn = conn.execution_options(schema_translate_map=schema_translate_map) + table.create(scoped_conn, checkfirst=True) + return so.Session(bind=scoped_conn) diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 886e71b..6588821 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -15,7 +15,17 @@ from orm_loader.loaders.loader_interface import PandasLoader from orm_loader.tables.loadable_table import CSVLoadableTableInterface from orm_loader.tables.typing import CSVTableProtocol -from tests.models import Base, CompositeTable, EnumTable, Flag, ImpliedEnumTable, RequiredTable, Role, SimpleTable +from tests.models import ( + Base, + CompositeTable, + EnumTable, + Flag, + ImpliedEnumTable, + RequiredTable, + Role, + SimpleTable, + VocabRoleTable, +) # Typed aliases: Pylance cannot verify SQLAlchemy metaclass-generated attrs # satisfy CSVTableProtocol structurally, so we cast once per class here. @@ -24,6 +34,7 @@ _CompositeTable = cast(Type[CSVTableProtocol], CompositeTable) _EnumTable = cast(Type[CSVTableProtocol], EnumTable) _ImpliedEnumTable = cast(Type[CSVTableProtocol], ImpliedEnumTable) +_VocabRoleTable = cast(Type[CSVTableProtocol], VocabRoleTable) @pytest.fixture(autouse=True) @@ -65,6 +76,33 @@ def test_initial_csv_load(session, tmp_path): ] +def test_initial_csv_load_for_a_non_primary_role_table(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 + the table's declared role differs from primary. This is the SQLite + counterpart to test_schema_translate_map.py's Postgres-only, non-primary- + role coverage.""" + csv_path = tmp_path / "test_vocab_role_table.csv" + + pd.DataFrame( + [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}] + ).to_csv(csv_path, index=False, sep="\t") + + inserted = _VocabRoleTable.load_csv( + session, csv_path, dedupe=False, loader=PandasLoader() + ) + session.commit() + + assert inserted == 2 + + rows = session.execute( + sa.select(VocabRoleTable).order_by(VocabRoleTable.id) + ).scalars().all() + + assert [(r.id, r.name) for r in rows] == [(1, "alpha"), (2, "beta")] + + def test_replace_merge_strategy(session, tmp_path): csv_path = tmp_path / "test_table.csv" diff --git a/tests/loaders/test_pg_loader.py b/tests/loaders/test_pg_loader.py index 4e4c945..abafcc7 100644 --- a/tests/loaders/test_pg_loader.py +++ b/tests/loaders/test_pg_loader.py @@ -8,7 +8,6 @@ from tests.models import EnumTable, Role, SimpleTable -@pytest.mark.requires_database("test_orm_db") def test_copy_into_staging_with_extra_identity_column(pg_session, tmp_path): """COPY must succeed when the staging table has a _rownum identity column.""" csv = tmp_path / "test_table.csv" @@ -32,7 +31,6 @@ def test_copy_into_staging_with_extra_identity_column(pg_session, tmp_path): assert rownums == [1, 2], "_rownum must be auto-populated by IDENTITY sequence" -@pytest.mark.requires_database("test_orm_db") def test_copy_and_orm_path_equivalence(pg_session, tmp_path): csv = tmp_path / "test_table.csv" @@ -54,7 +52,6 @@ def test_copy_and_orm_path_equivalence(pg_session, tmp_path): -@pytest.mark.requires_database("test_orm_db") def test_postgres_copy_fast_path(pg_session, tmp_path): csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) @@ -64,7 +61,6 @@ def test_postgres_copy_fast_path(pg_session, tmp_path): assert inserted == 1 -@pytest.mark.requires_database("test_orm_db") def test_postgres_copy_fast_path_is_used(pg_session, tmp_path, monkeypatch): csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) @@ -84,7 +80,6 @@ def fake_quick_load_pg(*args, **kwargs): assert called["copy"] is True assert inserted == 1 -@pytest.mark.requires_database("test_orm_db") def test_copy_failure_falls_back_to_orm(pg_session, tmp_path, monkeypatch): csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) @@ -107,7 +102,6 @@ def broken_copy(*args, **kwargs): assert [(r.id, r.name) for r in rows] == [(1, "alpha")] -@pytest.mark.requires_database("test_orm_db") def test_postgres_upsert_does_not_update(pg_session, tmp_path): csv = tmp_path / "test_table.csv" @@ -124,7 +118,6 @@ def test_postgres_upsert_does_not_update(pg_session, tmp_path): assert [(r.id, r.name) for r in rows] == [(1, "alpha")] -@pytest.mark.requires_database("test_orm_db") def test_postgres_insert_if_empty(pg_session, tmp_path): csv = tmp_path / "test_table.csv" @@ -152,7 +145,6 @@ def test_postgres_insert_if_empty(pg_session, tmp_path): ] -@pytest.mark.requires_database("test_orm_db") def test_postgres_insert_if_empty_raises_on_non_empty_target(pg_session, tmp_path): csv = tmp_path / "test_table.csv" @@ -171,7 +163,6 @@ def test_postgres_insert_if_empty_raises_on_non_empty_target(pg_session, tmp_pat ) -@pytest.mark.requires_database("test_orm_db") def test_postgres_copy_large_batch(pg_session, tmp_path): csv = tmp_path / "test_table.csv" @@ -188,7 +179,6 @@ def test_postgres_copy_large_batch(pg_session, tmp_path): assert inserted == 9999 -@pytest.mark.requires_database("test_orm_db") def test_staging_schema_matches_target(pg_session, tmp_path): csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) @@ -259,7 +249,6 @@ def test_check_line_ending_unknown(caplog): assert "Unable to detect line ending" in caplog.text -@pytest.mark.requires_database("test_orm_db") def test_quick_load_pg_basic(pg_session, tmp_path): csv = tmp_path / "test_table.csv" csv.write_text("id,name\n1,alpha\n2,beta\n") @@ -275,7 +264,6 @@ def test_quick_load_pg_basic(pg_session, tmp_path): assert rows == [(1, "alpha"), (2, "beta")] -@pytest.mark.requires_database("test_orm_db") def test_quick_load_pg_lowercases_header(pg_session, tmp_path): csv = tmp_path / "test_table.csv" csv.write_text("ID,NAME\n1,alpha\n") @@ -287,7 +275,6 @@ def test_quick_load_pg_lowercases_header(pg_session, tmp_path): assert row == (1, "alpha") -@pytest.mark.requires_database("test_orm_db") def test_quick_load_pg_strips_literally_quoted_header(pg_session, tmp_path): """A header row with literal quote characters around each column name (a common CSV-export convention) used to round-trip into an invalid @@ -303,7 +290,6 @@ def test_quick_load_pg_strips_literally_quoted_header(pg_session, tmp_path): assert rows == [(1, "alpha"), (2, "beta")] -@pytest.mark.requires_database("test_orm_db") def test_quick_load_pg_tab_delimiter(pg_session, tmp_path): csv = tmp_path / "test_table.csv" csv.write_text("id\tname\n1\talpha\n2\tbeta\n") @@ -315,7 +301,6 @@ def test_quick_load_pg_tab_delimiter(pg_session, tmp_path): assert rows == [(1, "alpha"), (2, "beta")] -@pytest.mark.requires_database("test_orm_db") def test_quick_load_pg_rollback_on_error(pg_session, tmp_path): csv = tmp_path / "test_table.csv" csv.write_text("id,name\n1,alpha\n2,\n") # violates NOT NULL @@ -327,7 +312,6 @@ def test_quick_load_pg_rollback_on_error(pg_session, tmp_path): assert rows == 0 -@pytest.mark.requires_database("test_orm_db") def test_quick_load_pg_equivalence_with_orm(pg_session, tmp_path): csv = tmp_path / "test_table.csv" csv.write_text("id,name\n1,alpha\n2,beta\n") @@ -351,7 +335,6 @@ def test_quick_load_pg_equivalence_with_orm(pg_session, tmp_path): assert rows_pg == rows_orm -@pytest.mark.requires_database("test_orm_db") def test_quick_load_pg_trailing_blank_lines(pg_session, tmp_path): csv = tmp_path / "test_table.csv" @@ -370,7 +353,6 @@ def test_quick_load_pg_trailing_blank_lines(pg_session, tmp_path): assert total == 2 assert rows == [(1, "alpha"), (2, "beta")] -@pytest.mark.requires_database("test_orm_db") def test_copy_fails_with_raw_carriage_returns_but_succeeds_after_normalisation(pg_session, tmp_path): csv = tmp_path / "test_table.csv" @@ -424,7 +406,6 @@ def _clear_column_cast_rules(): _COLUMN_CAST_RULES.clear() -@pytest.mark.requires_database("test_orm_db") def test_enum_column_cast_rule_round_trips_on_real_postgres(pg_session, tmp_path): # The merge step that moves rows from staging to the target table is a # plain SQL copy with no Python-level type translation, so whatever text diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py new file mode 100644 index 0000000..e0eaffb --- /dev/null +++ b/tests/loaders/test_schema_translate_map.py @@ -0,0 +1,151 @@ +"""End-to-end proof that load_csv() respects schema_translate_map with no +caller-side workaround. This is the actual regression test for the bug this +whole plan exists to fix, distinct from the backend unit tests in +tests/backends/, which exercise the merge methods directly but never against +a genuinely non-default schema. + +Only Postgres is covered here. SQLite has no real schema concept (confirmed +in the plan's own audit), so there is no non-default-schema behavior to +regress there; SQLite's own dialect-specific correctness (the +postgresql.insert() vs sqlite.insert() upsert constructor split in +particular) is already covered by tests/backends/test_sqlite_backend.py and +the default-schema tests in test_loader_e2e.py. + +Not create_mock_engine: MockConnection.schema_for_object ignores +schema_translate_map entirely, which would make this test pass whether or +not translation actually works. Real Postgres, via pg_db. +""" + +from __future__ import annotations + +import uuid + +import pandas as pd +import sqlalchemy as sa +from oa_configurator import ensure_schema +from oa_configurator import Role as SchemaRole + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.loaders.loader_interface import PandasLoader + +from tests.conftest import schema_scoped_session +from tests.models import SimpleTable, VocabRoleTable + + +def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): + schema = f"test_schema_{uuid.uuid4().hex[:8]}" + conn = pg_db.connection + ensure_schema(conn, schema) + ensure_schema(conn, STAGING_SCHEMA) + + # Override the connection's default schema for this session only. This + # is the caller-side setup a real deployment does once at engine + # construction (ResolvedCDMDatabase.create_engine()), not a workaround + # threaded through load_csv() itself. + session = schema_scoped_session( + conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: schema} + ) + + csv_path = tmp_path / "test_table.csv" + pd.DataFrame( + [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}, {"id": 3, "name": "gamma"}] + ).to_csv(csv_path, index=False, sep="\t") + + inserted = SimpleTable.load_csv( + session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema=STAGING_SCHEMA + ) + session.commit() + + assert inserted == 3 + + # Read back through the schema-qualified name directly, not through + # schema_translate_map, to prove the rows are really there. + rows = conn.execute( + sa.text(f'SELECT id, name FROM "{schema}"."test_table" ORDER BY id') + ).fetchall() + assert rows == [(1, "alpha"), (2, "beta"), (3, "gamma")] + + # And that nothing leaked into the default/public schema. That was the + # exact failure mode the original bug caused: raw text() bypassing + # schema_translate_map, resolving through the connection's search_path + # instead. + leaked = conn.execute(sa.text("SELECT to_regclass('public.test_table')")).scalar() + assert leaked is None + + +def test_replace_merge_respects_non_default_schema_end_to_end(pg_db, tmp_path): + """A second load_csv() call with merge_strategy="replace" against the + same non-default schema. Proves the merge path itself, not just the + initial insert-if-empty fast path, qualifies correctly.""" + schema = f"test_schema_{uuid.uuid4().hex[:8]}" + conn = pg_db.connection + ensure_schema(conn, schema) + ensure_schema(conn, STAGING_SCHEMA) + + session = schema_scoped_session( + conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: schema} + ) + + def _write_and_load(rows: list[dict], path_name: str) -> int: + path = tmp_path / path_name + pd.DataFrame(rows).to_csv(path, index=False, sep="\t") + return SimpleTable.load_csv( + session, + path, + dedupe=False, + loader=PandasLoader(), + merge_strategy="replace", + staging_schema=STAGING_SCHEMA, + ) + + _write_and_load( + [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}], "test_table.csv" + ) + session.commit() + + _write_and_load([{"id": 1, "name": "alpha-updated"}], "test_table.csv") + session.commit() + + rows = conn.execute( + sa.text(f'SELECT id, name FROM "{schema}"."test_table" ORDER BY id') + ).fetchall() + assert rows == [(1, "alpha-updated"), (2, "beta")] + + +def test_load_csv_respects_non_primary_role_end_to_end(pg_db, tmp_path): + """Checks if the derivation of the role from the table's own + __table_role__ attribute works correctly for each role.""" + primary_schema = f"test_primary_{uuid.uuid4().hex[:8]}" + vocab_schema = f"test_vocab_{uuid.uuid4().hex[:8]}" + conn = pg_db.connection + ensure_schema(conn, primary_schema) + ensure_schema(conn, vocab_schema) + ensure_schema(conn, STAGING_SCHEMA) + + session = schema_scoped_session( + conn, + VocabRoleTable.__table__, + {SchemaRole.PRIMARY.value: primary_schema, SchemaRole.VOCAB.value: vocab_schema}, + ) + + csv_path = tmp_path / "test_vocab_role_table.csv" + pd.DataFrame([{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]).to_csv( + csv_path, index=False, sep="\t" + ) + + inserted = VocabRoleTable.load_csv( + session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema=STAGING_SCHEMA + ) + session.commit() + + assert inserted == 2 + + rows = conn.execute( + sa.text(f'SELECT id, name FROM "{vocab_schema}"."test_vocab_role_table" ORDER BY id') + ).fetchall() + assert rows == [(1, "alpha"), (2, "beta")] + + leaked = conn.execute( + sa.text(f"SELECT to_regclass('{primary_schema}.test_vocab_role_table')") + ).scalar() + assert leaked is None diff --git a/tests/loaders/test_split_connection.py b/tests/loaders/test_split_connection.py new file mode 100644 index 0000000..cc2a4ce --- /dev/null +++ b/tests/loaders/test_split_connection.py @@ -0,0 +1,62 @@ +"""Tests split CDM/vocab connection instances. +""" + +from __future__ import annotations + +import uuid + +import pandas as pd +import sqlalchemy as sa +from oa_configurator import ensure_schema +from oa_configurator import Role as SchemaRole + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.loaders.loader_interface import PandasLoader + +from tests.conftest import schema_scoped_session +from tests.models import SimpleTable, VocabRoleTable + + +def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_path): + """``pg_db`` (real Postgres, primary role) and ``session`` (real SQLite, + vocab role, via the module-level ``engine``/``session`` fixtures) are two + entirely different engines against two entirely different database + systems.""" + primary_schema = f"test_primary_{uuid.uuid4().hex[:8]}" + conn = pg_db.connection + ensure_schema(conn, primary_schema) + ensure_schema(conn, STAGING_SCHEMA) + primary_session = schema_scoped_session( + conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: primary_schema} + ) + + primary_csv = tmp_path / "test_table.csv" + pd.DataFrame([{"id": 1, "name": "primary-alpha"}]).to_csv(primary_csv, index=False, sep="\t") + + vocab_csv = tmp_path / "test_vocab_role_table.csv" + pd.DataFrame([{"id": 1, "name": "vocab-alpha"}]).to_csv(vocab_csv, index=False, sep="\t") + + # Interleaved on purpose: primary, then vocab, then primary again, so a + # module-level cache keyed wrong (or reused across calls) would surface + # as data landing in the wrong database. + SimpleTable.load_csv(primary_session, primary_csv, dedupe=False, loader=PandasLoader()) + primary_session.commit() + + VocabRoleTable.load_csv(session, vocab_csv, dedupe=False, loader=PandasLoader()) + session.commit() + + primary_rows = conn.execute( + sa.text(f'SELECT id, name FROM "{primary_schema}"."test_table"') + ).fetchall() + assert primary_rows == [(1, "primary-alpha")] + + vocab_rows = session.execute( + sa.select(VocabRoleTable).order_by(VocabRoleTable.id) + ).scalars().all() + assert [(r.id, r.name) for r in vocab_rows] == [(1, "vocab-alpha")] + + # Neither database saw the other's table/data at all. + leaked_vocab_table_in_pg = conn.execute( + sa.text(f"SELECT to_regclass('{primary_schema}.test_vocab_role_table')") + ).scalar() + assert leaked_vocab_table_in_pg is None diff --git a/tests/mappers/test_materialised_view_mixin.py b/tests/mappers/test_materialised_view_mixin.py index 20c5a51..fee8f10 100644 --- a/tests/mappers/test_materialised_view_mixin.py +++ b/tests/mappers/test_materialised_view_mixin.py @@ -6,6 +6,8 @@ import pytest import sqlalchemy as sa +from oa_configurator import Role +from oa_configurator.testing import isolated_test_schema from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex from orm_loader.mappers.materialised_view_mixin import ( MaterializedViewMixin, @@ -14,6 +16,38 @@ ) +class _PrimaryRoleMV(MaterializedViewMixin): + __mv_name__ = "mv_primary_role_test" + __mv_select__ = sa.select(sa.literal(1).label("n")) + + +class _VocabRoleMV(MaterializedViewMixin): + __mv_name__ = "mv_vocab_role_test" + __mv_select__ = sa.select(sa.literal(2).label("n")) + __mv_role__ = Role.VOCAB + + +def test_refresh_all_mvs_resolves_each_views_own_role(pg_db) -> None: + engine = pg_db.connection.engine + + with isolated_test_schema(engine, prefix="mv_primary") as primary_schema, \ + isolated_test_schema(engine, prefix="mv_vocab") as vocab_schema: + scoped = engine.execution_options( + schema_translate_map={Role.PRIMARY.value: primary_schema, Role.VOCAB.value: vocab_schema} + ) + with scoped.begin() as conn: + _PrimaryRoleMV.create_mv(conn) + _VocabRoleMV.create_mv(conn) + refresh_all_mvs(conn, [_PrimaryRoleMV, _VocabRoleMV]) + + with engine.connect() as conn: + inspector = sa.inspect(conn) + assert inspector.has_table("mv_primary_role_test", schema=primary_schema) + assert not inspector.has_table("mv_primary_role_test", schema=vocab_schema) + assert inspector.has_table("mv_vocab_role_test", schema=vocab_schema) + assert not inspector.has_table("mv_vocab_role_test", schema=primary_schema) + + class _FakeBackend: """Capture lifecycle calls without touching a real database.""" @@ -71,7 +105,7 @@ def test_create_mv_forwards_default_args_to_backend(fake_backend: _FakeBackend, ( "create_materialized_view", (bind, "mv_no_index", _SELECT), - {"schema": None, "with_data": True, "if_not_exists": True}, + {"role": Role.PRIMARY, "with_data": True, "if_not_exists": True}, ) ] @@ -91,7 +125,7 @@ def test_create_mv_creates_declared_indexes_after_the_view(fake_backend: _FakeBa "create_materialized_view_index", ] assert fake_backend.calls[1][1] == (bind, "mv_indexed", _INDEX) - assert fake_backend.calls[1][2] == {"schema": None, "if_not_exists": True} + assert fake_backend.calls[1][2] == {"role": Role.PRIMARY, "if_not_exists": True} def test_create_mv_create_indexes_false_skips_index_creation(fake_backend: _FakeBackend, bind): @@ -100,13 +134,13 @@ def test_create_mv_create_indexes_false_skips_index_creation(fake_backend: _Fake assert [call[0] for call in fake_backend.calls] == ["create_materialized_view"] -def test_create_mv_forwards_schema_with_data_and_if_not_exists_overrides( +def test_create_mv_forwards_role_with_data_and_if_not_exists_overrides( fake_backend: _FakeBackend, bind ): - _NoIndexMv.create_mv(bind, schema="reporting", with_data=False, if_not_exists=False) + _NoIndexMv.create_mv(bind, role=Role.VOCAB, with_data=False, if_not_exists=False) assert fake_backend.calls[0][2] == { - "schema": "reporting", + "role": Role.VOCAB, "with_data": False, "if_not_exists": False, } @@ -115,7 +149,7 @@ def test_create_mv_forwards_schema_with_data_and_if_not_exists_overrides( def test_create_mv_forwards_if_not_exists_to_declared_indexes(fake_backend: _FakeBackend, bind): _IndexedMv.create_mv(bind, if_not_exists=False) - assert fake_backend.calls[1][2] == {"schema": None, "if_not_exists": False} + assert fake_backend.calls[1][2] == {"role": Role.PRIMARY, "if_not_exists": False} def test_create_mv_engine_uses_one_transaction_for_view_and_indexes(monkeypatch: pytest.MonkeyPatch): @@ -177,16 +211,16 @@ def test_refresh_mv_forwards_default_args_and_declared_indexes(fake_backend: _Fa ( "refresh_materialized_view", (bind, "mv_indexed"), - {"schema": None, "concurrently": False, "declared_indexes": (_INDEX,)}, + {"role": Role.PRIMARY, "concurrently": False, "declared_indexes": (_INDEX,)}, ) ] -def test_refresh_mv_forwards_schema_and_concurrently(fake_backend: _FakeBackend, bind): - _IndexedMv.refresh_mv(bind, schema="reporting", concurrently=True) +def test_refresh_mv_forwards_role_and_concurrently(fake_backend: _FakeBackend, bind): + _IndexedMv.refresh_mv(bind, role=Role.VOCAB, concurrently=True) assert fake_backend.calls[0][2] == { - "schema": "reporting", + "role": Role.VOCAB, "concurrently": True, "declared_indexes": (_INDEX,), } @@ -199,16 +233,16 @@ def test_drop_mv_forwards_default_args(fake_backend: _FakeBackend, bind): ( "drop_materialized_view", (bind, "mv_no_index"), - {"schema": None, "if_exists": True, "cascade": False}, + {"role": Role.PRIMARY, "if_exists": True, "cascade": False}, ) ] -def test_drop_mv_forwards_schema_if_exists_and_cascade(fake_backend: _FakeBackend, bind): - _NoIndexMv.drop_mv(bind, schema="reporting", if_exists=False, cascade=True) +def test_drop_mv_forwards_role_if_exists_and_cascade(fake_backend: _FakeBackend, bind): + _NoIndexMv.drop_mv(bind, role=Role.VOCAB, if_exists=False, cascade=True) assert fake_backend.calls[0][2] == { - "schema": "reporting", + "role": Role.VOCAB, "if_exists": False, "cascade": True, } diff --git a/tests/models.py b/tests/models.py index 7c92fc9..898a3b0 100644 --- a/tests/models.py +++ b/tests/models.py @@ -2,6 +2,7 @@ from enum import Enum import sqlalchemy as sa +from oa_configurator import Role as SchemaRole from sqlalchemy.orm import declarative_base import sqlalchemy.orm as so from orm_loader.tables import CSVLoadableTableInterface @@ -20,6 +21,7 @@ class Flag(str, Enum): class PandasLoaderTable(CSVLoadableTableInterface, Base): __tablename__ = "test_pandas_loader" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} id = sa.Column(sa.Integer, primary_key=True) value = sa.Column(sa.String, nullable=False) @@ -28,6 +30,7 @@ class SimpleTable(Base, CSVLoadableTableInterface): __tablename__ = "test_table" __table_args__ = ( sa.Index("ix_test_table_name", "name"), + {"schema": SchemaRole.PRIMARY.value}, ) id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) @@ -36,6 +39,7 @@ class SimpleTable(Base, CSVLoadableTableInterface): class RequiredTable(Base, CSVLoadableTableInterface): __tablename__ = "required_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, nullable=False) @@ -43,6 +47,7 @@ class RequiredTable(Base, CSVLoadableTableInterface): class CompositeTable(Base, CSVLoadableTableInterface): __tablename__ = "composite_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} a: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) b: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) @@ -58,11 +63,26 @@ class EnumTable(Base, CSVLoadableTableInterface): """ __tablename__ = "enum_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) role: so.Mapped[Role | None] = so.mapped_column(sa.Enum(Role), nullable=True) +class ComputedColumnTable(Base, CSVLoadableTableInterface): + """A real, registered table with a computed column, for merge-method + tests that need get_staging_table() to work. Unlike a bare + __tablename__/__table__ pair, this actually implements + CSVLoadableTableInterface.""" + + __tablename__ = "computed_column_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) + slug: so.Mapped[str] = so.mapped_column(sa.String, sa.Computed("lower(name)")) + + class ImpliedEnumTable(Base, CSVLoadableTableInterface): """A plain String column with no type-level enum signal at all -- the OMOP CDM concept.standard_concept/invalid_reason shape register_column_cast_rule @@ -70,6 +90,23 @@ class ImpliedEnumTable(Base, CSVLoadableTableInterface): """ __tablename__ = "implied_enum_table" + __table_args__ = {"schema": SchemaRole.PRIMARY.value} id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) flag: so.Mapped[str | None] = so.mapped_column(sa.String(1), nullable=True) + + +class VocabRoleTable(Base, CSVLoadableTableInterface): + """A VOCAB-tagged table, so tests can prove the staging/index role + derivation (role_of_table(), threaded through create_staging_table()/ + manage_indices()) actually resolves a non-primary role correctly, + instead of only ever exercising the PRIMARY-tagged default.""" + + __tablename__ = "test_vocab_role_table" + __table_args__ = ( + sa.Index("ix_test_vocab_role_table_name", "name"), + {"schema": SchemaRole.VOCAB.value}, + ) + + id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) + name: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) diff --git a/tests/pg_db.py b/tests/pg_db.py deleted file mode 100644 index d0aacd5..0000000 --- a/tests/pg_db.py +++ /dev/null @@ -1,43 +0,0 @@ -import time -import pytest -import sqlalchemy as sa -from sqlalchemy.orm import sessionmaker - -from tests.models import Base - -POSTGRES_URL = "postgresql+psycopg://test:test@localhost:55432/test" - -@pytest.fixture(scope="session") -def pg_engine(): - # wait for container - for _ in range(20): - try: - engine = sa.create_engine(POSTGRES_URL, future=True) - with engine.connect() as conn: - conn.execute(sa.text("select 1")) - break - except Exception: - time.sleep(1) - else: - raise RuntimeError("Postgres never became available") - - yield engine - - engine.dispose() - - -@pytest.fixture -def pg_session(pg_engine): - Session = sessionmaker(bind=pg_engine, future=True) - with pg_engine.begin() as conn: - conn.execute(sa.text('DROP SCHEMA public CASCADE')) - conn.execute(sa.text('CREATE SCHEMA public')) - Base.metadata.drop_all(conn) - Base.metadata.create_all(conn) - - session = Session() - try: - yield session - finally: - session.rollback() - session.close()