From c9dea8c68d7eace6524ccf5bd50f1667cc678522 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 28 Aug 2026 01:21:41 +0000 Subject: [PATCH 01/11] Fix the schema not being properly bound to the engines by using direct SQLAlchemy core objects, update test suite and unify shared backend tests, update schema registration --- src/orm_loader/backends/__init__.py | 9 +- src/orm_loader/backends/base.py | 22 ++- src/orm_loader/backends/postgres.py | 133 +++++++--------- src/orm_loader/backends/sqlite.py | 147 +++++++++++------- src/orm_loader/loaders/loading_helpers.py | 4 +- .../mappers/materialised_view_mixin.py | 8 +- src/orm_loader/tables/loadable_table.py | 33 ++-- tests/backends/test_postgres_backend.py | 130 ++-------------- tests/backends/test_reserved_schema.py | 22 +++ tests/backends/test_shared_backend.py | 139 +++++++++++++++++ tests/backends/test_sqlite_backend.py | 66 +------- tests/conftest.py | 61 ++------ tests/loaders/test_pg_loader.py | 19 --- tests/loaders/test_schema_translate_map.py | 110 +++++++++++++ tests/models.py | 13 ++ tests/pg_db.py | 43 ----- 16 files changed, 523 insertions(+), 436 deletions(-) create mode 100644 tests/backends/test_reserved_schema.py create mode 100644 tests/backends/test_shared_backend.py create mode 100644 tests/loaders/test_schema_translate_map.py delete mode 100644 tests/pg_db.py diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index 35a99cf..785358a 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -1,14 +1,19 @@ 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, +) __all__ = [ "BackendCapabilities", "DatabaseBackend", - "STAGING_SCHEMA", "Dialect", "PostgresBackend", + "STAGING_SCHEMA", "SQLiteBackend", "resolve_backend", ] diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index e845d95..00b2c91 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -9,6 +9,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so +from oa_configurator import register_reserved_schema from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer @@ -41,6 +42,8 @@ class Dialect(str, Enum): STAGING_SCHEMA: str = "staging" +register_reserved_schema(STAGING_SCHEMA, owner="orm-loader") + class DatabaseBackend(ABC): """ @@ -239,7 +242,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, @@ -251,7 +253,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, @@ -263,7 +264,6 @@ def merge_insert( self, table_cls: Type["CSVTableProtocol"], session: so.Session, - target_name: str, *, merge_batch_size: int | None = None, ) -> None: @@ -315,13 +315,25 @@ def create_materialized_view( bind: "Engine | Connection", name: str, selectable: sa.sql.Select[Any], + *, + schema: str | None = None, ) -> None: - """Create a materialized view for the supplied selectable.""" + """Create a materialized view for the supplied selectable. + + *schema* defaults to the bind's own ``schema_translate_map`` (via + ``oa_configurator.schema_of``) when not given explicitly. + """ @abstractmethod def refresh_materialized_view( self, bind: "Engine | Connection", name: str, + *, + schema: str | None = None, ) -> None: - """Refresh a materialized view.""" + """Refresh a materialized view. + + *schema* defaults to the bind's own ``schema_translate_map`` (via + ``oa_configurator.schema_of``) when not given explicitly. + """ diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 8ea17a0..3b57972 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -6,6 +6,7 @@ import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so +from oa_configurator import qualified, schema_of from sqlalchemy.dialects import postgresql from sqlalchemy.sql.compiler import IdentifierPreparer @@ -57,7 +58,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) session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};')) session.execute( sa.text( @@ -143,49 +144,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 @@ -194,50 +192,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 @@ -245,50 +235,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 @@ -304,22 +283,26 @@ def create_materialized_view( bind: Engine | Connection, name: str, selectable: sa.sql.Select[Any], + *, + schema: str | None = None, ) -> None: from ..mappers.materialised_view_mixin import CreateMaterializedView with self._as_connection(bind) as conn: - conn.execute(CreateMaterializedView(name, selectable)) + effective_schema = schema if schema is not None else schema_of(conn) + qualified_name = qualified(conn, name, schema=effective_schema) + conn.execute(CreateMaterializedView(qualified_name, selectable)) def refresh_materialized_view( self, bind: Engine | Connection, name: str, + *, + schema: str | None = None, ) -> None: with self._as_connection(bind) as conn: - safe_name = name - dialect = getattr(conn, "dialect", None) - if dialect is not None: - safe_name = dialect.identifier_preparer.quote(name) + effective_schema = schema if schema is not None else schema_of(conn) + safe_name = qualified(conn, name, schema=effective_schema) conn.execute(sa.text(f"REFRESH MATERIALIZED VIEW {safe_name};")) @contextmanager diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index eb51f8d..ffd5d2c 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -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, @@ -251,6 +280,8 @@ def create_materialized_view( bind: "Engine | Connection", name: str, selectable: sa.sql.Select[Any], + *, + schema: str | None = None, ) -> None: self._require_capability("supports_materialized_views", "materialized views") @@ -258,6 +289,8 @@ def refresh_materialized_view( self, bind: "Engine | Connection", name: str, + *, + schema: str | None = None, ) -> None: self._require_capability("supports_materialized_views", "materialized views") 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/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index aa6e67b..1ff46e5 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -19,7 +19,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. @@ -31,8 +33,8 @@ def __init__(self, name: str, selectable: sa.sql.Select[Any]): @compiler.compiles(CreateMaterializedView) def _create_view( - element: CreateMaterializedView, - compiler: sa.sql.compiler.SQLCompiler, + element: CreateMaterializedView, + compiler: sa.sql.compiler.SQLCompiler, **kwargs: Any ) -> str: diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 73f19ab..eb335e6 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -2,6 +2,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so import logging +from oa_configurator import schema_inspect, schema_of from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError from typing import Type, Any, Iterator @@ -120,9 +121,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) + 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 +232,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 +257,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, ) @@ -470,6 +482,7 @@ def _target_has_rows( target, sa.MetaData(), autoload_with=session.get_bind(), + schema=schema_of(session), ) row = session.execute( sa.select(sa.literal(1)).select_from(table).limit(1) @@ -530,14 +543,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 +558,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)}." @@ -571,7 +584,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/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index ea0ef83..b94136c 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -6,32 +6,24 @@ 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 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: def __init__(self, scalar_result: str | int = "origin") -> None: self.statements: list[str] = [] @@ -61,17 +53,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() @@ -96,16 +81,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(): @@ -136,96 +119,17 @@ 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) - - 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 - - -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"]) - - sql = session.statements[0] - assert f'INSERT INTO "{_TARGET_TABLE}" ("id", "name")' in sql - assert 'ON CONFLICT ("id") DO NOTHING' in sql - - -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, - ) - - 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)) - - -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, - ) - - 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_merge_upsert_paginated_path(): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) - session = _FakeSession(scalar_result=10) - - 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_materialized_view_methods_emit_expected_sql(): +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() - session = _FakeSession() + conn = pg_db.connection 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") + backend.create_materialized_view(conn, "mv_test", selectable) + backend.refresh_materialized_view(conn, "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) + assert conn.execute(sa.text("SELECT n FROM mv_test")).scalar() == 1 def test_postgres_backend_normalize_fk_check_state(): diff --git a/tests/backends/test_reserved_schema.py b/tests/backends/test_reserved_schema.py new file mode 100644 index 0000000..09b7c53 --- /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 schema_name="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, Resolver, StackConfig + +from orm_loader.backends import STAGING_SCHEMA + + +def test_resolving_cdm_database_with_staging_schema_name_raises() -> None: + cfg = StackConfig.for_session( + connections={"c": ConnectionConfig(dialect="sqlite", database_name=":memory:")}, + databases={"default": CDMDatabaseConfig(connection="c", schema_name=STAGING_SCHEMA)}, + ) + with pytest.raises(RuntimeError, match=f"{STAGING_SCHEMA!r}.*orm-loader"): + Resolver(cfg).resolve_database("default") diff --git a/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py new file mode 100644 index 0000000..e1cf0e3 --- /dev/null +++ b/tests/backends/test_shared_backend.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Type, cast + +import pytest +import sqlalchemy as sa + +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=["postgres", "sqlite"]) +def merge_backend(request: pytest.FixtureRequest) -> tuple[DatabaseBackend, "so.Session"]: + """Same merge-method contract exercised against both real backends. + Only the postgres param ever requests pg_session, so the sqlite param + never needs a database.""" + if request.param == "postgres": + 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 d93e5e6..535e887 100644 --- a/tests/backends/test_sqlite_backend.py +++ b/tests/backends/test_sqlite_backend.py @@ -9,25 +9,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] = [] @@ -46,7 +36,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: @@ -154,56 +144,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..8fdec5e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,3 @@ -import time from pathlib import Path import pytest @@ -29,51 +28,25 @@ 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 +@pytest.fixture +def pg_db(): + """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.""" + from oa_configurator.testing import isolated_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}") + with isolated_test_database(OrmLoaderConfig, "test_orm_db") 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 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..e71a2b6 --- /dev/null +++ b/tests/loaders/test_schema_translate_map.py @@ -0,0 +1,110 @@ +"""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 +import sqlalchemy.orm as so +from oa_configurator import ensure_schema + +from orm_loader.backends import STAGING_SCHEMA +from orm_loader.loaders.loader_interface import PandasLoader +from tests.models import Base, SimpleTable + + +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. + scoped_conn = conn.execution_options(schema_translate_map={None: schema}) + session = so.Session(bind=scoped_conn) + Base.metadata.create_all(scoped_conn) + + 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) + + scoped_conn = conn.execution_options(schema_translate_map={None: schema}) + session = so.Session(bind=scoped_conn) + Base.metadata.create_all(scoped_conn) + + 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")] diff --git a/tests/models.py b/tests/models.py index 7c92fc9..51ae58b 100644 --- a/tests/models.py +++ b/tests/models.py @@ -63,6 +63,19 @@ class EnumTable(Base, CSVLoadableTableInterface): 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" + + 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 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() From 1c3477d05b65f443a6de513a1e54ac7ee096be66 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 31 Aug 2026 05:11:46 +0000 Subject: [PATCH 02/11] Adapt to newr test mechanism --- pyproject.toml | 2 +- src/orm_loader/config.py | 21 +++++++++++++++++++-- tests/backends/test_shared_backend.py | 18 ++++++++++++++++-- tests/conftest.py | 16 +++++++++------- 4 files changed, 45 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 110bc99..b29e68d 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 requires_process_isolation\"" [tool.pyright] reportMissingTypeStubs = false diff --git a/src/orm_loader/config.py b/src/orm_loader/config.py index 497c719..09324ce 100644 --- a/src/orm_loader/config.py +++ b/src/orm_loader/config.py @@ -5,17 +5,30 @@ from typing import Annotated, ClassVar from oa_configurator import CDMDatabaseConfig, PackageConfigBase, RefTo +from pydantic import Field 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 +38,8 @@ 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)] = None diff --git a/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py index e1cf0e3..0c9823c 100644 --- a/tests/backends/test_shared_backend.py +++ b/tests/backends/test_shared_backend.py @@ -17,11 +17,25 @@ _CompositeTableCls = cast("Type[CSVTableProtocol]", CompositeTable) -@pytest.fixture(params=["postgres", "sqlite"]) +@pytest.fixture( + params=[pytest.param("postgres", marks=pytest.mark.requires_process_isolation), "sqlite"] +) def merge_backend(request: pytest.FixtureRequest) -> tuple[DatabaseBackend, "so.Session"]: """Same merge-method contract exercised against both real backends. Only the postgres param ever requests pg_session, so the sqlite param - never needs a database.""" + never needs a database. + + The postgres param carries its own requires_process_isolation mark + directly (rather than relying on the usual pg_db-in-fixturenames + auto-detection): request.getfixturevalue("pg_session") is a dynamic, + runtime lookup, invisible to pytest's collection-time fixturenames + computation, so the auto-detection mechanism can't see it and would + silently leave these Postgres-touching runs in the default suite, + alongside SQLite tests in the same process. Confirmed via + `pytest -m requires_process_isolation --collect-only`: without this + explicit mark, this file's postgres-param tests were being deselected + from that run entirely. + """ if request.param == "postgres": session = request.getfixturevalue("pg_session") return PostgresBackend(staging_schema=STAGING_SCHEMA), session diff --git a/tests/conftest.py b/tests/conftest.py index 8fdec5e..7f0990a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,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") @@ -13,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 @@ -34,10 +39,7 @@ def pg_db(): ``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.""" - from oa_configurator.testing import isolated_test_database - from orm_loader.config import OrmLoaderConfig - - with isolated_test_database(OrmLoaderConfig, "test_orm_db") as db: + with isolated_test_database(OrmLoaderConfig, "test_orm_db_pg") as db: yield db From 33088b1fbfb7c05e66756c6e7fc5bc9c25e54a58 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 1 Sep 2026 04:41:52 +0000 Subject: [PATCH 03/11] Adhere to process isolation for postgres DBs --- pyproject.toml | 2 +- src/orm_loader/config.py | 8 +++++++- tests/backends/test_shared_backend.py | 28 +++++++++++---------------- tests/conftest.py | 4 ++-- 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b29e68d..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 -m \"not requires_process_isolation\"" +addopts = "-ra -m 'not db_dialect'" [tool.pyright] reportMissingTypeStubs = false diff --git a/src/orm_loader/config.py b/src/orm_loader/config.py index 09324ce..dc1e524 100644 --- a/src/orm_loader/config.py +++ b/src/orm_loader/config.py @@ -42,4 +42,10 @@ class exists to register orm-loader in the oa-configurator ecosystem, default=None, description="Real PostgreSQL test database, for Postgres-only integration testing.", ) - test_orm_db_sqlite: Annotated[str | None, RefTo(CDMDatabaseConfig, is_test=True)] = None + 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/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py index 0c9823c..381160a 100644 --- a/tests/backends/test_shared_backend.py +++ b/tests/backends/test_shared_backend.py @@ -5,6 +5,7 @@ import pytest import sqlalchemy as sa +from oa_configurator.testing import DIALECT_PARAMS from orm_loader.backends import STAGING_SCHEMA, DatabaseBackend, PostgresBackend, SQLiteBackend from tests.models import ComputedColumnTable, CompositeTable @@ -17,26 +18,19 @@ _CompositeTableCls = cast("Type[CSVTableProtocol]", CompositeTable) -@pytest.fixture( - params=[pytest.param("postgres", marks=pytest.mark.requires_process_isolation), "sqlite"] -) +@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 postgres param ever requests pg_session, so the sqlite param - never needs a database. - - The postgres param carries its own requires_process_isolation mark - directly (rather than relying on the usual pg_db-in-fixturenames - auto-detection): request.getfixturevalue("pg_session") is a dynamic, - runtime lookup, invisible to pytest's collection-time fixturenames - computation, so the auto-detection mechanism can't see it and would - silently leave these Postgres-touching runs in the default suite, - alongside SQLite tests in the same process. Confirmed via - `pytest -m requires_process_isolation --collect-only`: without this - explicit mark, this file's postgres-param tests were being deselected - from that run entirely. + 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 == "postgres": + if request.param == "postgresql": session = request.getfixturevalue("pg_session") return PostgresBackend(staging_schema=STAGING_SCHEMA), session session = request.getfixturevalue("session") diff --git a/tests/conftest.py b/tests/conftest.py index 7f0990a..0dacdb4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -34,12 +34,12 @@ def session(engine): # --------------------------------------------------------------------------- @pytest.fixture -def pg_db(): +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") as db: + with isolated_test_database(OrmLoaderConfig, "test_orm_db_pg", request=request) as db: yield db From 60f898da66db175f4105485d5aa49f356bd911bf Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 1 Sep 2026 05:32:24 +0000 Subject: [PATCH 04/11] Use outstanding autocommit --- src/orm_loader/backends/postgres.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 3b57972..fc679f4 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -6,7 +6,7 @@ import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so -from oa_configurator import qualified, schema_of +from oa_configurator import autocommit_connection, qualified, schema_of from sqlalchemy.dialects import postgresql from sqlalchemy.sql.compiler import IdentifierPreparer @@ -322,7 +322,7 @@ def _set_replica_role( finally: sae.remove(engine, "connect", _set_replica_role) with engine.connect() as conn: - conn = conn.execution_options(isolation_level="AUTOCOMMIT") + conn = autocommit_connection(conn) conn.execute(sa.text("SET session_replication_role = DEFAULT")) role = conn.execute(sa.text("SHOW session_replication_role")).scalar() if role != "origin": From 310d5c2af2766cfe7f448c029e3aa1b71162c6ea Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 1 Sep 2026 06:11:49 +0000 Subject: [PATCH 05/11] Update CI --- .github/workflows/ci.yml | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0e7d4c8..62deef1 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-sqlite: + uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test.yml@main build-test: uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres.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.schema_name=public From 5f27e38c1dc6970ac6d048e34061668b7333e8fd Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 3 Sep 2026 05:55:42 +0000 Subject: [PATCH 06/11] Updated CI --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 62deef1..2669506 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -8,8 +8,8 @@ jobs: uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/label-gate.yml@main build-test-sqlite: uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test.yml@main - build-test: - uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres.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: | From 260a9232442edf06bb3e13501a77513d57deaf1d Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 4 Sep 2026 06:11:40 +0000 Subject: [PATCH 07/11] Small fixes to autocommit and where the staging schema is registered --- src/orm_loader/backends/base.py | 3 --- src/orm_loader/backends/postgres.py | 7 +++---- src/orm_loader/config.py | 6 +++++- tests/backends/test_postgres_backend.py | 15 ++++++++++++++- tests/backends/test_reserved_schema.py | 14 +++++++------- 5 files changed, 29 insertions(+), 16 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 00b2c91..3d40581 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -9,7 +9,6 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import register_reserved_schema from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer @@ -42,8 +41,6 @@ class Dialect(str, Enum): STAGING_SCHEMA: str = "staging" -register_reserved_schema(STAGING_SCHEMA, owner="orm-loader") - class DatabaseBackend(ABC): """ diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index fc679f4..f59fd43 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -321,9 +321,8 @@ def _set_replica_role( yield engine finally: sae.remove(engine, "connect", _set_replica_role) - with engine.connect() as conn: - conn = autocommit_connection(conn) - 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/config.py b/src/orm_loader/config.py index dc1e524..d712406 100644 --- a/src/orm_loader/config.py +++ b/src/orm_loader/config.py @@ -4,9 +4,13 @@ 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. diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index b94136c..b76d945 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -199,12 +199,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 index 09b7c53..a62f78f 100644 --- a/tests/backends/test_reserved_schema.py +++ b/tests/backends/test_reserved_schema.py @@ -8,15 +8,15 @@ from __future__ import annotations import pytest -from oa_configurator import CDMDatabaseConfig, ConnectionConfig, Resolver, StackConfig +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: - cfg = StackConfig.for_session( - connections={"c": ConnectionConfig(dialect="sqlite", database_name=":memory:")}, - databases={"default": CDMDatabaseConfig(connection="c", schema_name=STAGING_SCHEMA)}, - ) - with pytest.raises(RuntimeError, match=f"{STAGING_SCHEMA!r}.*orm-loader"): - Resolver(cfg).resolve_database("default") + 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", schema_name=STAGING_SCHEMA)}, + ) From 2228f6b87655dcd42dfe064b22e2f8a9848e78e1 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 7 Sep 2026 05:18:33 +0000 Subject: [PATCH 08/11] oa-configurator dialect changes --- src/orm_loader/backends/base.py | 10 ++-------- src/orm_loader/backends/postgres.py | 4 ++-- src/orm_loader/backends/sqlite.py | 4 ++-- 3 files changed, 6 insertions(+), 12 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 3d40581..80d43b9 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 typing import TYPE_CHECKING, Type, Any @@ -12,6 +11,8 @@ from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer +from oa_configurator import Dialect + if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext from ..tables.typing import CSVTableProtocol @@ -32,13 +33,6 @@ class BackendCapabilities: supports_materialized_views: bool = False -class Dialect(str, Enum): - """Supported SQLAlchemy dialect names.""" - - SQLITE = "sqlite" - POSTGRESQL = "postgresql" - - STAGING_SCHEMA: str = "staging" diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index f59fd43..8609267 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -6,11 +6,11 @@ import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so -from oa_configurator import autocommit_connection, qualified, schema_of +from oa_configurator import autocommit_connection, qualified, schema_of, Dialect from sqlalchemy.dialects import postgresql from sqlalchemy.sql.compiler import IdentifierPreparer -from .base import BackendCapabilities, DatabaseBackend, Dialect +from .base import BackendCapabilities, DatabaseBackend if TYPE_CHECKING: from sqlalchemy.engine import Connection, Engine diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index ffd5d2c..68dbb99 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: @@ -320,7 +320,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: From ab813c8fbabd22fd3c194922a7342bf0f5fab68c Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 00:06:45 +0000 Subject: [PATCH 09/11] cdm_schema renames --- .github/workflows/ci.yml | 2 +- tests/backends/test_reserved_schema.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2669506..7e7a7f4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,4 +22,4 @@ jobs: --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.schema_name=public + --set test_orm_db_pg.cdm_schema=public diff --git a/tests/backends/test_reserved_schema.py b/tests/backends/test_reserved_schema.py index a62f78f..4ff133a 100644 --- a/tests/backends/test_reserved_schema.py +++ b/tests/backends/test_reserved_schema.py @@ -1,6 +1,6 @@ """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 schema_name="staging" +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. """ @@ -18,5 +18,5 @@ 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", schema_name=STAGING_SCHEMA)}, + databases={"default": CDMDatabaseConfig(connection="c", cdm_schema=STAGING_SCHEMA)}, ) From ba3139faf59729eca07b99078c889fea53ab0617 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 00:07:34 +0000 Subject: [PATCH 10/11] Funnel role into each method for correct schema resolution --- src/orm_loader/backends/base.py | 18 +++++++---- src/orm_loader/backends/postgres.py | 15 ++++----- src/orm_loader/backends/sqlite.py | 6 ++-- src/orm_loader/helpers/__init__.py | 3 +- src/orm_loader/helpers/sql.py | 19 +++++++++++ .../mappers/materialised_view_mixin.py | 22 ++++++++++--- src/orm_loader/tables/loadable_table.py | 32 +++++-------------- src/orm_loader/tables/typing.py | 2 +- tests/backends/test_postgres_backend.py | 29 +++++++++++++++++ tests/loaders/test_schema_translate_map.py | 12 ++++--- tests/models.py | 8 +++++ 11 files changed, 114 insertions(+), 52 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 80d43b9..cf5fa56 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -11,7 +11,7 @@ from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer -from oa_configurator import Dialect +from oa_configurator import Dialect, Role if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext @@ -307,12 +307,14 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - schema: str | None = None, + role: Role = Role.PRIMARY, ) -> None: """Create a materialized view for the supplied selectable. - *schema* defaults to the bind's own ``schema_translate_map`` (via - ``oa_configurator.schema_of``) when not given explicitly. + 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. """ @abstractmethod @@ -321,10 +323,12 @@ def refresh_materialized_view( bind: "Engine | Connection", name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, ) -> None: """Refresh a materialized view. - *schema* defaults to the bind's own ``schema_translate_map`` (via - ``oa_configurator.schema_of``) when not given explicitly. + 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. """ diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 8609267..354cd35 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -6,7 +6,8 @@ import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so -from oa_configurator import autocommit_connection, qualified, schema_of, Dialect +from oa_configurator import autocommit_connection, qualified, Dialect, Role +from ..helpers.sql import role_of_table from sqlalchemy.dialects import postgresql from sqlalchemy.sql.compiler import IdentifierPreparer @@ -58,7 +59,7 @@ def create_staging_table( table = table_cls.__table__ preparer = self.identifier_preparer staging_ref = self.qualified_staging_name(table_cls.__tablename__) - source_ref = qualified(session, 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( @@ -284,13 +285,12 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - schema: str | None = None, + role: Role = Role.PRIMARY, ) -> None: from ..mappers.materialised_view_mixin import CreateMaterializedView with self._as_connection(bind) as conn: - effective_schema = schema if schema is not None else schema_of(conn) - qualified_name = qualified(conn, name, schema=effective_schema) + qualified_name = qualified(conn, name, role=role) conn.execute(CreateMaterializedView(qualified_name, selectable)) def refresh_materialized_view( @@ -298,11 +298,10 @@ def refresh_materialized_view( bind: Engine | Connection, name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, ) -> None: with self._as_connection(bind) as conn: - effective_schema = schema if schema is not None else schema_of(conn) - safe_name = qualified(conn, name, schema=effective_schema) + safe_name = qualified(conn, name, role=role) conn.execute(sa.text(f"REFRESH MATERIALIZED VIEW {safe_name};")) @contextmanager diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 68dbb99..d5315cd 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -13,6 +13,8 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.sql.compiler import IdentifierPreparer +from oa_configurator import Role + from .base import BackendCapabilities, DatabaseBackend, Dialect if TYPE_CHECKING: @@ -281,7 +283,7 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - schema: str | None = None, + role: Role = Role.PRIMARY, ) -> None: self._require_capability("supports_materialized_views", "materialized views") @@ -290,7 +292,7 @@ def refresh_materialized_view( bind: "Engine | Connection", name: str, *, - schema: str | None = None, + role: Role = Role.PRIMARY, ) -> None: self._require_capability("supports_materialized_views", "materialized views") 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/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index 1ff46e5..d82608f 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -3,6 +3,7 @@ import sqlalchemy as sa from typing import Any from collections import defaultdict, deque +from oa_configurator import Role from ..backends.resolve import resolve_backend class CreateMaterializedView(DDLElement): @@ -162,7 +163,9 @@ class DailyObservationCountsMV(Base, MaterializedViewMixin): __mv_dependencies__: set[str] = set() @classmethod - def create_mv(cls, bind: "sa.engine.Connection | sa.engine.Engine") -> None: + def create_mv( + cls, bind: "sa.engine.Connection | sa.engine.Engine", *, role: Role = Role.PRIMARY + ) -> None: """ Create the materialized view if it does not already exist. @@ -170,6 +173,12 @@ def create_mv(cls, bind: "sa.engine.Connection | sa.engine.Engine") -> None: ---------- bind A SQLAlchemy Engine or Connection used to execute the DDL. + role + Schema role the view's own physical schema resolves through + (defaults to primary). Set this to the role of the tables + ``__mv_select__`` reads from when it's a vocab/results view, + not primary -- otherwise the view always lands in the primary + schema regardless of what it was actually built over. Notes ----- @@ -202,10 +211,12 @@ def create_mv(cls, bind: "sa.engine.Connection | sa.engine.Engine") -> None: ``` """ backend = resolve_backend(bind) - backend.create_materialized_view(bind, cls.__mv_name__, cls.__mv_select__) + backend.create_materialized_view(bind, cls.__mv_name__, cls.__mv_select__, role=role) @classmethod - def refresh_mv(cls, bind: "sa.engine.Connection | sa.engine.Engine") -> None: + def refresh_mv( + cls, bind: "sa.engine.Connection | sa.engine.Engine", *, role: Role = Role.PRIMARY + ) -> None: """ Refresh the contents of the materialized view. @@ -213,6 +224,9 @@ def refresh_mv(cls, bind: "sa.engine.Connection | sa.engine.Engine") -> None: ---------- bind A SQLAlchemy Engine or Connection used to execute the refresh. + role + Schema role the view's own physical schema resolves through; + see :meth:`create_mv` for when to override the default. Notes ----- @@ -228,7 +242,7 @@ def refresh_mv(cls, bind: "sa.engine.Connection | sa.engine.Engine") -> None: ``` """ backend = resolve_backend(bind) - backend.refresh_materialized_view(bind, cls.__mv_name__) + backend.refresh_materialized_view(bind, cls.__mv_name__, role=role) 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 eb335e6..12dbcd3 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -2,7 +2,9 @@ import sqlalchemy as sa import sqlalchemy.orm as so import logging -from oa_configurator import schema_inspect, schema_of +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 @@ -121,7 +123,7 @@ def manage_indices( table_name = cls.__tablename__ indices = list(cls.__table__.indexes) if resolved_index_strategy == "drop_rebuild" else [] - inspector = schema_inspect(session) + 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__)} @@ -417,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)}." @@ -471,21 +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(), - schema=schema_of(session), - ) 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 @@ -524,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)}." @@ -567,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)}." 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_postgres_backend.py b/tests/backends/test_postgres_backend.py index b76d945..471783b 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -8,6 +8,8 @@ from sqlalchemy.dialects import postgresql from sqlalchemy.engine import Engine +from oa_configurator import 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 @@ -132,6 +134,33 @@ def test_postgres_backend_materialized_view_methods_work_end_to_end(pg_db): assert conn.execute(sa.text("SELECT n FROM mv_test")).scalar() == 1 +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 + + 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_normalize_fk_check_state(): normalize = PostgresBackend._normalize_fk_check_state diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index e71a2b6..7488be6 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -24,10 +24,12 @@ import sqlalchemy as sa import sqlalchemy.orm as so 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.models import Base, SimpleTable + +from tests.models import SimpleTable def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): @@ -40,9 +42,9 @@ def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): # is the caller-side setup a real deployment does once at engine # construction (ResolvedCDMDatabase.create_engine()), not a workaround # threaded through load_csv() itself. - scoped_conn = conn.execution_options(schema_translate_map={None: schema}) + scoped_conn = conn.execution_options(schema_translate_map={SchemaRole.PRIMARY.value: schema}) session = so.Session(bind=scoped_conn) - Base.metadata.create_all(scoped_conn) + SimpleTable.__table__.create(scoped_conn, checkfirst=True) csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -80,9 +82,9 @@ def test_replace_merge_respects_non_default_schema_end_to_end(pg_db, tmp_path): ensure_schema(conn, schema) ensure_schema(conn, STAGING_SCHEMA) - scoped_conn = conn.execution_options(schema_translate_map={None: schema}) + scoped_conn = conn.execution_options(schema_translate_map={SchemaRole.PRIMARY.value: schema}) session = so.Session(bind=scoped_conn) - Base.metadata.create_all(scoped_conn) + SimpleTable.__table__.create(scoped_conn, checkfirst=True) def _write_and_load(rows: list[dict], path_name: str) -> int: path = tmp_path / path_name diff --git a/tests/models.py b/tests/models.py index 51ae58b..7910f36 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,6 +63,7 @@ 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) @@ -70,6 +76,7 @@ class ComputedColumnTable(Base, CSVLoadableTableInterface): 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) @@ -83,6 +90,7 @@ 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) From ede7911b51349197e8390c574682b350b22a7947 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 04:49:23 +0000 Subject: [PATCH 11/11] Extend to non-primary tagged tables, split connection test --- .../mappers/materialised_view_mixin.py | 32 +++++++--- tests/backends/test_base_backend.py | 12 +++- tests/loaders/test_loader_e2e.py | 40 +++++++++++- tests/loaders/test_schema_translate_map.py | 44 ++++++++++++- tests/loaders/test_split_connection.py | 62 +++++++++++++++++++ tests/mappers/test_materialised_view_mixin.py | 39 ++++++++++++ tests/models.py | 16 +++++ 7 files changed, 232 insertions(+), 13 deletions(-) create mode 100644 tests/loaders/test_split_connection.py create mode 100644 tests/mappers/test_materialised_view_mixin.py diff --git a/src/orm_loader/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index d82608f..6d56e25 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -1,3 +1,4 @@ +from docutils.parsers.rst.languages.cs import roles from sqlalchemy.ext import compiler from sqlalchemy.schema import DDLElement import sqlalchemy as sa @@ -68,6 +69,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. @@ -161,10 +165,14 @@ class DailyObservationCountsMV(Base, MaterializedViewMixin): __mv_name__: str __mv_select__: sa.sql.Select[Any] __mv_dependencies__: set[str] = set() + __mv_role__: Role = Role.PRIMARY @classmethod def create_mv( - cls, bind: "sa.engine.Connection | sa.engine.Engine", *, role: Role = Role.PRIMARY + cls, + bind: "sa.engine.Connection | sa.engine.Engine", + *, + role: Role | None = None ) -> None: """ Create the materialized view if it does not already exist. @@ -174,11 +182,11 @@ def create_mv( bind A SQLAlchemy Engine or Connection used to execute the DDL. role - Schema role the view's own physical schema resolves through - (defaults to primary). Set this to the role of the tables - ``__mv_select__`` reads from when it's a vocab/results view, - not primary -- otherwise the view always lands in the primary - schema regardless of what it was actually built over. + 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. Notes ----- @@ -211,11 +219,14 @@ def create_mv( ``` """ backend = resolve_backend(bind) - backend.create_materialized_view(bind, cls.__mv_name__, cls.__mv_select__, role=role) + role_ = role if role is not None else cls.__mv_role__ + backend.create_materialized_view( + bind, cls.__mv_name__, cls.__mv_select__, role=role_ + ) @classmethod def refresh_mv( - cls, bind: "sa.engine.Connection | sa.engine.Engine", *, role: Role = Role.PRIMARY + cls, bind: "sa.engine.Connection | sa.engine.Engine", *, role: Role | None = None ) -> None: """ Refresh the contents of the materialized view. @@ -242,7 +253,10 @@ def refresh_mv( ``` """ backend = resolve_backend(bind) - backend.refresh_materialized_view(bind, cls.__mv_name__, role=role) + role_ = role if role is not None else cls.__mv_role__ + backend.refresh_materialized_view( + bind, cls.__mv_name__, role=role_ + ) def resolve_mv_refresh_order(mv_classes: list[type[MaterializedViewMixin]]) -> list[type]: diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index ad50ded..ba1ed27 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, @@ -123,11 +124,18 @@ def restore_fk_check(self, session: so.Session, previous_state: str | int) -> No self.calls.append(("restore_fk_check", previous_state)) def create_materialized_view( - self, bind: Engine | Connection, name: str, selectable: sa.sql.Select[Any] + self, + bind: Engine | Connection, + name: str, + selectable: sa.sql.Select[Any], + *, + role: Role = Role.PRIMARY, ) -> None: return None - def refresh_materialized_view(self, bind: Engine | Connection, name: str) -> None: + def refresh_materialized_view( + self, bind: Engine | Connection, name: str, *, role: Role = Role.PRIMARY + ) -> None: return None 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_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index 7488be6..16e275d 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -29,7 +29,7 @@ from orm_loader.backends import STAGING_SCHEMA from orm_loader.loaders.loader_interface import PandasLoader -from tests.models import SimpleTable +from tests.models import SimpleTable, VocabRoleTable def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): @@ -110,3 +110,45 @@ def _write_and_load(rows: list[dict], path_name: str) -> int: 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) + + scoped_conn = conn.execution_options( + schema_translate_map={ + SchemaRole.PRIMARY.value: primary_schema, + SchemaRole.VOCAB.value: vocab_schema, + } + ) + session = so.Session(bind=scoped_conn) + VocabRoleTable.__table__.create(scoped_conn, checkfirst=True) + + 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..0e90978 --- /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 +import sqlalchemy.orm as so +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.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) + scoped_conn = conn.execution_options(schema_translate_map={SchemaRole.PRIMARY.value: primary_schema}) + primary_session = so.Session(bind=scoped_conn) + SimpleTable.__table__.create(scoped_conn, checkfirst=True) + + 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 new file mode 100644 index 0000000..8ff05f1 --- /dev/null +++ b/tests/mappers/test_materialised_view_mixin.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import sqlalchemy as sa + +from oa_configurator import Role +from oa_configurator.testing import isolated_test_schema +from orm_loader.mappers.materialised_view_mixin import MaterializedViewMixin, refresh_all_mvs + + +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) diff --git a/tests/models.py b/tests/models.py index 7910f36..898a3b0 100644 --- a/tests/models.py +++ b/tests/models.py @@ -94,3 +94,19 @@ class ImpliedEnumTable(Base, CSVLoadableTableInterface): 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)