From c9dea8c68d7eace6524ccf5bd50f1667cc678522 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 28 Aug 2026 01:21:41 +0000 Subject: [PATCH 01/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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/29] 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) From 808d11fbfb2d4694411f5af858456604d883038d Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 23:40:41 +0000 Subject: [PATCH 12/29] Remove unused import --- src/orm_loader/mappers/materialised_view_mixin.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/orm_loader/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index 6d56e25..47e3c8b 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -1,4 +1,3 @@ -from docutils.parsers.rst.languages.cs import roles from sqlalchemy.ext import compiler from sqlalchemy.schema import DDLElement import sqlalchemy as sa From d2ea59f816f4dec1496729ed5e4ca612a0755c21 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 16 Sep 2026 00:04:12 +0000 Subject: [PATCH 13/29] Have a scoped_session fixture --- tests/conftest.py | 11 ++++++++++ tests/loaders/test_schema_translate_map.py | 25 ++++++++++------------ tests/loaders/test_split_connection.py | 8 +++---- 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 0dacdb4..c8694d6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -52,3 +52,14 @@ def pg_session(pg_db): conn.execute(sa.text(f"CREATE SCHEMA IF NOT EXISTS {STAGING_SCHEMA}")) Base.metadata.create_all(conn) return pg_db.session + + +def schema_scoped_session( + conn: sa.Connection, table: sa.Table, schema_translate_map: dict +) -> so.Session: + """A Session scoped to *schema_translate_map*, with *table* already + created through it. + """ + scoped_conn = conn.execution_options(schema_translate_map=schema_translate_map) + table.create(scoped_conn, checkfirst=True) + return so.Session(bind=scoped_conn) diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index 16e275d..e0eaffb 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -22,13 +22,13 @@ 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.conftest import schema_scoped_session from tests.models import SimpleTable, VocabRoleTable @@ -42,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={SchemaRole.PRIMARY.value: schema}) - session = so.Session(bind=scoped_conn) - SimpleTable.__table__.create(scoped_conn, checkfirst=True) + session = schema_scoped_session( + conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: schema} + ) csv_path = tmp_path / "test_table.csv" pd.DataFrame( @@ -82,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={SchemaRole.PRIMARY.value: schema}) - session = so.Session(bind=scoped_conn) - SimpleTable.__table__.create(scoped_conn, checkfirst=True) + session = schema_scoped_session( + conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: schema} + ) def _write_and_load(rows: list[dict], path_name: str) -> int: path = tmp_path / path_name @@ -122,14 +122,11 @@ def test_load_csv_respects_non_primary_role_end_to_end(pg_db, tmp_path): 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 = schema_scoped_session( + conn, + VocabRoleTable.__table__, + {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( diff --git a/tests/loaders/test_split_connection.py b/tests/loaders/test_split_connection.py index 0e90978..cc2a4ce 100644 --- a/tests/loaders/test_split_connection.py +++ b/tests/loaders/test_split_connection.py @@ -7,13 +7,13 @@ 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.conftest import schema_scoped_session from tests.models import SimpleTable, VocabRoleTable @@ -26,9 +26,9 @@ def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_ 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_session = schema_scoped_session( + conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: primary_schema} + ) primary_csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "primary-alpha"}]).to_csv(primary_csv, index=False, sep="\t") From 7f14a91533e0e0da835cd6888a31763918b24ed8 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 16 Sep 2026 23:08:29 +0000 Subject: [PATCH 14/29] Remove redundant check for dialect and tie it into contextmanager _as_connection --- docs/tables/mat_view.md | 4 -- src/orm_loader/backends/__init__.py | 2 - src/orm_loader/backends/base.py | 35 ++++++++++++ src/orm_loader/backends/postgres.py | 33 ----------- src/orm_loader/mappers/__init__.py | 2 - .../mappers/materialised_view_errors.py | 7 --- tests/backends/test_postgres_backend.py | 55 +++++++++---------- 7 files changed, 62 insertions(+), 76 deletions(-) diff --git a/docs/tables/mat_view.md b/docs/tables/mat_view.md index 052b7c9..6bc463f 100644 --- a/docs/tables/mat_view.md +++ b/docs/tables/mat_view.md @@ -170,7 +170,3 @@ The built-in implementation is PostgreSQL-oriented. SQLite rejects materialized- ::: orm_loader.mappers.ConcurrentRefreshNotEligibleError options: heading_level: 3 - -::: orm_loader.mappers.UnsupportedMaterializationDialectError - options: - heading_level: 3 diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index 30defb8..e5dd07a 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -12,7 +12,6 @@ MaterializationError, MaterializationFailure, MaterializationOperation, - UnsupportedMaterializationDialectError, ) __all__ = [ @@ -26,6 +25,5 @@ "PostgresBackend", "STAGING_SCHEMA", "SQLiteBackend", - "UnsupportedMaterializationDialectError", "resolve_backend", ] diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index fa6fb8b..2a32c88 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -186,6 +186,41 @@ def _as_connection( self, bind: Engine | Connection, ) -> Generator[Connection, None, None]: + """ + Normalize a bind into an open connection, guarding its dialect. + + Every backend method that takes a ``bind`` should route it through + this context manager rather than opening a connection itself, so the + dialect guard below applies uniformly. + + Parameters + ---------- + bind : Engine or Connection + An Engine opens a new connection and transaction scoped to this + context manager, committing on a clean exit. A Connection is + forwarded as-is; its transaction is owned by the caller, and + passing the same Connection into several backend calls groups + them into one shared transaction. + + Yields + ------ + Connection + An open connection whose dialect matches ``self.dialect``. + + Raises + ------ + TypeError + If ``bind``'s dialect does not match ``self.dialect``. Guards + against a bind resolved through a different backend being + passed directly into a method on this one. + """ + if bind.dialect.name != self.dialect.value: + raise TypeError( + f"{self.name} backend received a {bind.dialect.name!r} connection; " + f"expected {self.dialect.value!r}. The bind passed to this method must " + "be the same one (or share the same dialect as) the bind resolve_backend() " + "was given." + ) if isinstance(bind, Engine): with bind.begin() as conn: yield conn diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index d184ea4..98e0ea2 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -18,7 +18,6 @@ MaterializationError, MaterializationFailure, MaterializationOperation, - UnsupportedMaterializationDialectError, ) if TYPE_CHECKING: @@ -31,26 +30,6 @@ _VALID_PG_REPLICATION_ROLES = frozenset({"origin", "local", "replica"}) -def _require_postgres_dialect( - conn: "Connection", - *, - operation: MaterializationOperation, - schema: str | None, - name: str, -) -> None: - dialect = getattr(conn, "dialect", None) - if dialect is not None and dialect.name == "postgresql": - return - raise UnsupportedMaterializationDialectError( - MaterializationFailure( - operation=operation, - schema=schema, - name=name, - reason=f"received dialect {getattr(dialect, 'name', dialect)!r}", - ) - ) - - class PostgresBackend(DatabaseBackend): def __init__(self, *, staging_schema: str | None = None) -> None: super().__init__(staging_schema=staging_schema) @@ -323,9 +302,6 @@ def create_materialized_view( with self._as_connection(bind) as conn: schema = schema_of(conn, role=role) - _require_postgres_dialect( - conn, operation=MaterializationOperation.CREATE, schema=schema, name=name - ) try: conn.execute( CreateMaterializedView( @@ -358,9 +334,6 @@ def refresh_materialized_view( ) -> None: with self._as_connection(bind) as conn: schema = schema_of(conn, role=role) - _require_postgres_dialect( - conn, operation=MaterializationOperation.REFRESH, schema=schema, name=name - ) if concurrently: if not any(index.unique for index in declared_indexes): raise ConcurrentRefreshNotEligibleError( @@ -410,9 +383,6 @@ def drop_materialized_view( with self._as_connection(bind) as conn: schema = schema_of(conn, role=role) - _require_postgres_dialect( - conn, operation=MaterializationOperation.DROP, schema=schema, name=name - ) try: conn.execute( DropMaterializedView( @@ -444,9 +414,6 @@ def create_materialized_view_index( with self._as_connection(bind) as conn: schema = schema_of(conn, role=role) - _require_postgres_dialect( - conn, operation=MaterializationOperation.CREATE_INDEX, schema=schema, name=name - ) try: conn.execute( CreateMaterializedViewIndex( diff --git a/src/orm_loader/mappers/__init__.py b/src/orm_loader/mappers/__init__.py index eb544c6..1fd7352 100644 --- a/src/orm_loader/mappers/__init__.py +++ b/src/orm_loader/mappers/__init__.py @@ -25,7 +25,6 @@ MaterializationError, MaterializationFailure, MaterializationOperation, - UnsupportedMaterializationDialectError, ) __all__ = [ @@ -38,7 +37,6 @@ "MaterializationOperation", "MaterializedViewIndex", "MaterializedViewMixin", - "UnsupportedMaterializationDialectError", "refresh_all_mvs", "resolve_mv_refresh_order", ] diff --git a/src/orm_loader/mappers/materialised_view_errors.py b/src/orm_loader/mappers/materialised_view_errors.py index ea84a5b..2c85650 100644 --- a/src/orm_loader/mappers/materialised_view_errors.py +++ b/src/orm_loader/mappers/materialised_view_errors.py @@ -40,13 +40,6 @@ def __init__(self, failure: MaterializationFailure) -> None: ) -class UnsupportedMaterializationDialectError(MaterializationError): - """Raised before executing Postgres-only DDL/catalog SQL against a - non-Postgres connection. Defense in depth: the normal ``resolve_backend`` - dispatch path already prevents this via ``_require_capability``; this - guards direct/manual ``PostgresBackend()`` use.""" - - class ConcurrentRefreshNotEligibleError(MaterializationError): """Raised when a concurrent materialized-view refresh is not eligible. diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 3abf88a..e79a5f9 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -193,53 +193,52 @@ def test_postgres_backend_materialized_view_methods_emit_expected_sql(): assert any('REFRESH MATERIALIZED VIEW mv_test;' == sql for sql in session.statements) -def test_postgres_backend_quotes_unqualified_materialized_view_name(): - from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex +def test_postgres_backend_rejects_mismatched_dialect_bind(): + from sqlalchemy.dialects import sqlite backend = PostgresBackend() session = _FakeSession() + session.dialect = sqlite.dialect() selectable = sa.select(sa.literal(1).label("n")) - backend.create_materialized_view(_sess(session), 'mv name', selectable) - backend.create_materialized_view_index( - _sess(session), 'mv name', MaterializedViewIndex(name="mv_name_idx", columns=("n",)) - ) + with pytest.raises(TypeError, match="received a 'sqlite' connection; expected 'postgresql'"): + backend.create_materialized_view(_sess(session), "mv_test", selectable) - assert any('CREATE MATERIALIZED VIEW IF NOT EXISTS "mv name" as SELECT' in sql for sql in session.statements) - assert any('ON "mv name" ("n")' in sql for sql in session.statements) + assert session.statements == [] -def test_postgres_backend_create_materialized_view_rejects_non_postgres_connection(): - from orm_loader.mappers.materialised_view_errors import ( - UnsupportedMaterializationDialectError, - ) +def test_postgres_backend_rejects_dialect_that_drifted_after_resolve_backend(): + """A bind resolved to PostgresBackend, then a differently-dialected bind + handed to one of its methods, must not run Postgres-only DDL against it.""" + from orm_loader.backends.resolve import resolve_backend from sqlalchemy.dialects import sqlite - backend = PostgresBackend() - session = _FakeSession() - session.dialect = sqlite.dialect() - selectable = sa.select(sa.literal(1).label("n")) + postgres_session = _FakeSession() + backend = resolve_backend(_sess(postgres_session)) + assert isinstance(backend, PostgresBackend) - with pytest.raises(UnsupportedMaterializationDialectError, match="received dialect 'sqlite'"): - backend.create_materialized_view(_sess(session), "mv_test", selectable) + sqlite_session = _FakeSession() + sqlite_session.dialect = sqlite.dialect() + selectable = sa.select(sa.literal(1).label("n")) - assert session.statements == [] + with pytest.raises(TypeError, match="received a 'sqlite' connection; expected 'postgresql'"): + backend.create_materialized_view(_sess(sqlite_session), "mv_test", selectable) -def test_postgres_backend_refresh_materialized_view_rejects_non_postgres_connection(): - from orm_loader.mappers.materialised_view_errors import ( - UnsupportedMaterializationDialectError, - ) - from sqlalchemy.dialects import sqlite +def test_postgres_backend_quotes_unqualified_materialized_view_name(): + from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex backend = PostgresBackend() session = _FakeSession() - session.dialect = sqlite.dialect() + selectable = sa.select(sa.literal(1).label("n")) - with pytest.raises(UnsupportedMaterializationDialectError, match="received dialect 'sqlite'"): - backend.refresh_materialized_view(_sess(session), "mv_test") + backend.create_materialized_view(_sess(session), 'mv name', selectable) + backend.create_materialized_view_index( + _sess(session), 'mv name', MaterializedViewIndex(name="mv_name_idx", columns=("n",)) + ) - assert session.statements == [] + assert any('CREATE MATERIALIZED VIEW IF NOT EXISTS "mv name" as SELECT' in sql for sql in session.statements) + assert any('ON "mv name" ("n")' in sql for sql in session.statements) def test_postgres_backend_create_mv_quotes_name_for_legacy_search_path_resolution(): From 9a333a215479bf90430f888223ec879953bfcd57 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 21 Sep 2026 00:58:18 +0000 Subject: [PATCH 15/29] Obtain role_of_table from oa-configurator, remove explicit primary fallback --- src/orm_loader/backends/postgres.py | 10 ++++++++-- src/orm_loader/helpers/__init__.py | 3 +-- src/orm_loader/helpers/sql.py | 19 ------------------- src/orm_loader/tables/loadable_table.py | 3 +-- tests/loaders/test_dedupe.py | 2 ++ tests/loaders/test_loader_e2e.py | 4 ++++ 6 files changed, 16 insertions(+), 25 deletions(-) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 98e0ea2..388e915 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -6,8 +6,14 @@ import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so -from oa_configurator import autocommit_connection, qualified, schema_of, Dialect, Role -from ..helpers.sql import role_of_table +from oa_configurator import ( + autocommit_connection, + qualified, + role_of_table, + schema_of, + Dialect, + Role +) from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError from sqlalchemy.sql.compiler import IdentifierPreparer diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index e7757ae..6ad5a70 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, role_of_table +from .sql import qualify_identifier __all__ = [ "IngestError", @@ -25,5 +25,4 @@ "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 ea587d3..6e30ba9 100644 --- a/src/orm_loader/helpers/sql.py +++ b/src/orm_loader/helpers/sql.py @@ -1,27 +1,8 @@ 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/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 12dbcd3..82178d1 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -2,9 +2,8 @@ import sqlalchemy as sa import sqlalchemy.orm as so import logging -from oa_configurator import schema_inspect +from oa_configurator import role_of_table, schema_inspect -from ..helpers.sql import role_of_table from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError from typing import Type, Any, Iterator diff --git a/tests/loaders/test_dedupe.py b/tests/loaders/test_dedupe.py index c84a6b3..34a6c31 100644 --- a/tests/loaders/test_dedupe.py +++ b/tests/loaders/test_dedupe.py @@ -1,5 +1,6 @@ import pyarrow as pa from typing import cast, Type +from oa_configurator import Role from orm_loader.loaders.loading_helpers import arrow_drop_duplicates import pandas as pd import sqlalchemy as sa @@ -16,6 +17,7 @@ class Base(DeclarativeBase): class DedupTable(Base, CSVLoadableTableInterface): __tablename__ = "dedup_table" + __table_args__ = {"schema": Role.PRIMARY.value} id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) value: so.Mapped[str] = so.mapped_column(sa.String, nullable=False) diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index 6588821..b9df959 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -7,6 +7,7 @@ import sqlalchemy as sa import sqlalchemy.event as sae import sqlalchemy.orm as so +from oa_configurator import Role as SchemaRole from orm_loader.backends import resolve_backend from orm_loader.helpers import IngestError @@ -539,6 +540,7 @@ def test_clean_nulls_passthrough(): def test_nullable_column_with_nan_does_not_crash(session, engine, tmp_path): class NullableTable(Base, CSVLoadableTableInterface): __tablename__ = "nullable_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, nullable=True) @@ -575,6 +577,7 @@ class NullableTable(Base, CSVLoadableTableInterface): def test_embedded_newline_in_field_is_preserved(session, engine, tmp_path): class TextTable(Base, CSVLoadableTableInterface): __tablename__ = "text_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) @@ -602,6 +605,7 @@ class TextTable(Base, CSVLoadableTableInterface): def test_embedded_tab_in_field(session, engine, tmp_path): class TextTable2(Base, CSVLoadableTableInterface): __tablename__ = "tab_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) From 7ad3ad8f4ea08b67eb97f3951b359294a7de7c95 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 21 Sep 2026 01:04:38 +0000 Subject: [PATCH 16/29] Fix missing test fixes that weren't committed --- tests/loaders/test_parquet_loader.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/loaders/test_parquet_loader.py b/tests/loaders/test_parquet_loader.py index 253ae26..96b00fe 100644 --- a/tests/loaders/test_parquet_loader.py +++ b/tests/loaders/test_parquet_loader.py @@ -3,6 +3,7 @@ import pandas as pd import sqlalchemy as sa import sqlalchemy.orm as so +from oa_configurator import Role from sqlalchemy.orm import DeclarativeBase from typing import cast, Type @@ -17,6 +18,7 @@ class Base(DeclarativeBase): class ParquetTable(Base, CSVLoadableTableInterface): __tablename__ = "parquet_table" + __table_args__ = {"schema": Role.PRIMARY.value} id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) value: so.Mapped[int] = so.mapped_column(sa.Integer, nullable=False) @@ -54,6 +56,7 @@ class NullSentinelTable(Base, CSVLoadableTableInterface): treats text null sentinels the same way the scalar path does.""" __tablename__ = "null_sentinel_table" + __table_args__ = {"schema": Role.PRIMARY.value} id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) name: so.Mapped[str | None] = so.mapped_column(sa.String, nullable=True) @@ -130,6 +133,7 @@ def test_arrow_path_nulls_float_nan(session, engine, tmp_path): class FloatTable(Base, CSVLoadableTableInterface): __tablename__ = "float_sentinel_table" + __table_args__ = {"schema": Role.PRIMARY.value} id: so.Mapped[int] = so.mapped_column(sa.Integer, primary_key=True) score: so.Mapped[float | None] = so.mapped_column(sa.Float, nullable=True) From a1e22b4fde917088d0fc91ccc48cc5908e94619a Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 21 Sep 2026 05:43:16 +0000 Subject: [PATCH 17/29] Addressed PR comments --- src/orm_loader/tables/loadable_table.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 82178d1..c51cb5b 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -122,9 +122,9 @@ def manage_indices( table_name = cls.__tablename__ indices = list(cls.__table__.indexes) if resolved_index_strategy == "drop_rebuild" else [] - inspector = schema_inspect(session, role=role_of_table(cls.__table__)) if indices: + inspector = schema_inspect(session, role=role_of_table(cls.__table__)) 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] @@ -180,7 +180,7 @@ def manage_indices( if indices: logger.info(f"Table `{table_name}`: Verifying/Rebuilding indices.") rebuild_started = perf_counter() - inspector.clear_cache() # Required to ensure we get the current state of the database after potential changes + inspector = schema_inspect(session, role=role_of_table(cls.__table__)) existing_idx_names = {idx['name'] for idx in inspector.get_indexes(table_name)} for idx in indices: From 53de2cfc0b2ecee8f092b2f6375a71989f45dd17 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 22 Sep 2026 23:58:16 +0000 Subject: [PATCH 18/29] Pull out role following oa-configurator changes --- docs/tables/mat_view.md | 24 +++++- src/orm_loader/backends/base.py | 42 +++++----- src/orm_loader/backends/postgres.py | 27 +++---- src/orm_loader/loaders/loading_helpers.py | 2 +- .../mappers/materialised_view_mixin.py | 70 +++++++++------- src/orm_loader/tables/loadable_table.py | 12 +-- tests/backends/test_base_backend.py | 5 +- tests/backends/test_postgres_backend.py | 63 +++++++-------- tests/loaders/test_loader_e2e.py | 16 ++-- tests/loaders/test_schema_translate_map.py | 12 +-- tests/loaders/test_split_connection.py | 14 ++-- tests/mappers/test_materialised_view_mixin.py | 80 +++++++++++++++---- tests/models.py | 11 +-- 13 files changed, 226 insertions(+), 152 deletions(-) diff --git a/docs/tables/mat_view.md b/docs/tables/mat_view.md index 6bc463f..30e9bbd 100644 --- a/docs/tables/mat_view.md +++ b/docs/tables/mat_view.md @@ -105,19 +105,35 @@ PatientSummaryMV.refresh_mv(engine, concurrently=True) This is fail-closed by design. An index created manually outside `__mv_indexes__` does not satisfy the mixin's declaration contract; declare it in the class even if another migration is responsible for creating it. Expressions, partial indexes, and other index forms are outside this simple contract and should not be represented as `MaterializedViewIndex` entries. -The view's physical schema comes from the bound connection's own `schema_translate_map`, the same mechanism every other schema-aware table in this stack uses (`oa_configurator.schema_of`). `__mv_role__` declares which role a view resolves through (defaulting to `Role.PRIMARY`); pass `role=` to `create_mv()`/`refresh_mv()`/`drop_mv()` to override it for one call. +The view's physical schema comes from the bound connection's own `schema_translate_map`. `create_mv()`/`refresh_mv()`/`drop_mv()` resolve it via `oa_configurator.schema_of()` before ever reaching the backend, which only ever sees an already-resolved physical schema string. + +The schema tag itself comes from one of three places, in order: + +1. an explicit `schema_tag=` argument, for a one-off override; +2. when the class is also declaratively mapped (combined with a `Base`, with a real `__table__`), that table's own `schema` -- set the schema the same way as any other mapped table, via `__table_args__ = {"schema": ...}`, and the materialized view follows it automatically; +3. `__mv_schema_tag__` (defaults to `"primary"`), only consulted for a Core-only declaration with no mapped table to derive anything from. ```python +# Core-only: no Base, no mapped table, so __mv_schema_tag__ is what resolves it. class VocabSummaryMV(MaterializedViewMixin): __mv_name__ = "vocab_summary" __mv_select__ = ... - __mv_role__ = Role.VOCAB # resolves via the connection's vocab schema + __mv_schema_tag__ = Role.VOCAB.value # resolves via the connection's vocab schema -# Uses __mv_role__ (Role.PRIMARY by default) via the connection's own schema_translate_map. +# Uses __mv_schema_tag__ ("primary" by default) via the connection's own schema_translate_map. RecentObservationMV.create_mv(engine) # Override for one call. -RecentObservationMV.create_mv(engine, role=Role.VOCAB) +RecentObservationMV.create_mv(engine, schema_tag=Role.VOCAB.value) + +# Declaratively mapped: the mapped table's own schema is authoritative, no +# __mv_schema_tag__ needed (or consulted) at all. +class VocabPatientSummaryMV(Base, MaterializedViewMixin): + __mv_name__ = "vocab_patient_summary" + __mv_select__ = ... + __tablename__ = "vocab_patient_summary" + __table_args__ = {"schema": Role.VOCAB.value} + patient_id = sa.Column(sa.Integer, primary_key=True) ``` Every generated identifier is quoted through `oa_configurator.qualified()`, which quotes each component only when the dialect actually requires it (reserved words, mixed case, embedded quotes or spaces) — the same behavior every other Core-built query in this stack has. diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 2a32c88..073635f 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -12,7 +12,7 @@ from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer -from oa_configurator import Dialect, Role +from oa_configurator import Dialect if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext @@ -371,16 +371,15 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - role: Role = Role.PRIMARY, + schema: str | None = None, with_data: bool = True, if_not_exists: bool = True, ) -> None: """Create a materialized view for the supplied selectable. - 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. + *schema* is the view's already-resolved physical schema (or None for + the connection's own default); callers resolve which + schema_translate_map key to use before calling. """ raise NotImplementedError( f"Backend '{self.name}' has not implemented create_materialized_view()" @@ -392,16 +391,14 @@ def refresh_materialized_view( bind: "Engine | Connection", name: str, *, - role: Role = Role.PRIMARY, + schema: str | None = None, concurrently: bool = False, declared_indexes: tuple["MaterializedViewIndex", ...] = (), ) -> None: """Refresh a materialized view. - The view's schema is the bind's own ``schema_translate_map`` entry - for *role* (via ``oa_configurator.schema_of``), letting a view - built over vocab/results-role tables land in that role's own - schema instead of always primary. + *schema* is the view's already-resolved physical schema, matching + ``create_materialized_view``. ``declared_indexes`` lets supporting backends validate a concurrent refresh request without defining a second catalog-based eligibility @@ -417,18 +414,18 @@ def drop_materialized_view( bind: "Engine | Connection", name: str, *, - role: Role = Role.PRIMARY, + schema: str | None = None, if_exists: bool = True, cascade: bool = False, ) -> None: """Drop a materialized view. - The view's schema is the bind's own ``schema_translate_map`` entry - for *role*, matching ``create_materialized_view``. This is - deliberately non-abstract: the default implementation requires the - capability flag and then raises ``NotImplementedError``. Older - third-party backend subclasses need no override to receive a clear - error when they do not support this operation. + *schema* is the view's already-resolved physical schema, matching + ``create_materialized_view``. This is deliberately non-abstract: the + default implementation requires the capability flag and then raises + ``NotImplementedError``. Older third-party backend subclasses need + no override to receive a clear error when they do not support this + operation. """ raise NotImplementedError( f"Backend '{self.name}' has not implemented drop_materialized_view()" @@ -441,15 +438,14 @@ def create_materialized_view_index( name: str, index: "MaterializedViewIndex", *, - role: Role = Role.PRIMARY, + schema: str | None = None, if_not_exists: bool = True, ) -> None: """Create an index on a materialized view. - The view's schema is the bind's own ``schema_translate_map`` entry - for *role*, matching ``create_materialized_view``. This is - deliberately non-abstract for the same compatibility reason as - :meth:`drop_materialized_view`. + *schema* is the view's already-resolved physical schema, matching + ``create_materialized_view``. This is deliberately non-abstract for + the same compatibility reason as :meth:`drop_materialized_view`. """ raise NotImplementedError( f"Backend '{self.name}' has not implemented " diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index 388e915..c19ecc5 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -9,10 +9,9 @@ from oa_configurator import ( autocommit_connection, qualified, - role_of_table, schema_of, + validate_schema_tag, Dialect, - Role ) from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError @@ -73,7 +72,9 @@ 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, role=role_of_table(table)) + source_ref = qualified( + session, table.name, physical_schema=schema_of(session, schema_tag=validate_schema_tag(table)) + ) session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};')) session.execute( sa.text( @@ -300,18 +301,17 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - role: Role = Role.PRIMARY, + schema: str | None = None, with_data: bool = True, if_not_exists: bool = True, ) -> None: from ..mappers.materialised_view_mixin import CreateMaterializedView with self._as_connection(bind) as conn: - schema = schema_of(conn, role=role) try: conn.execute( CreateMaterializedView( - qualified(conn, name, schema=schema), + qualified(conn, name, physical_schema=schema), selectable, with_data=with_data, if_not_exists=if_not_exists, @@ -334,12 +334,11 @@ def refresh_materialized_view( bind: Engine | Connection, name: str, *, - role: Role = Role.PRIMARY, + schema: str | None = None, concurrently: bool = False, declared_indexes: tuple["MaterializedViewIndex", ...] = (), ) -> None: with self._as_connection(bind) as conn: - schema = schema_of(conn, role=role) if concurrently: if not any(index.unique for index in declared_indexes): raise ConcurrentRefreshNotEligibleError( @@ -351,7 +350,7 @@ def refresh_materialized_view( ) ) - safe_name = qualified(conn, name, schema=schema) + safe_name = qualified(conn, name, physical_schema=schema) concurrency = "CONCURRENTLY " if concurrently else "" try: conn.execute(sa.text(f"REFRESH MATERIALIZED VIEW {concurrency}{safe_name};")) @@ -381,18 +380,17 @@ def drop_materialized_view( bind: Engine | Connection, name: str, *, - role: Role = Role.PRIMARY, + schema: str | None = None, if_exists: bool = True, cascade: bool = False, ) -> None: from ..mappers.materialised_view_contracts import DropMaterializedView with self._as_connection(bind) as conn: - schema = schema_of(conn, role=role) try: conn.execute( DropMaterializedView( - qualified(conn, name, schema=schema), if_exists=if_exists, cascade=cascade + qualified(conn, name, physical_schema=schema), if_exists=if_exists, cascade=cascade ) ) except Exception as error: @@ -413,17 +411,16 @@ def create_materialized_view_index( name: str, index: "MaterializedViewIndex", *, - role: Role = Role.PRIMARY, + schema: str | None = None, if_not_exists: bool = True, ) -> None: from ..mappers.materialised_view_contracts import CreateMaterializedViewIndex with self._as_connection(bind) as conn: - schema = schema_of(conn, role=role) try: conn.execute( CreateMaterializedViewIndex( - qualified(conn, name, schema=schema), index, if_not_exists=if_not_exists + qualified(conn, name, physical_schema=schema), index, if_not_exists=if_not_exists ) ) except Exception as error: diff --git a/src/orm_loader/loaders/loading_helpers.py b/src/orm_loader/loaders/loading_helpers.py index e30bb36..6a385cb 100644 --- a/src/orm_loader/loaders/loading_helpers.py +++ b/src/orm_loader/loaders/loading_helpers.py @@ -274,7 +274,7 @@ def quick_load_pg( if not hasattr(raw_conn, "cursor"): raise RuntimeError("Expected DB-API connection for COPY") - table_ref = qualified(session, tablename, schema=schema) + table_ref = qualified(session, tablename, physical_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 ecdcbc2..7e9d064 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -4,7 +4,7 @@ import sqlalchemy as sa from sqlalchemy.ext import compiler from sqlalchemy.schema import DDLElement -from oa_configurator import Role +from oa_configurator import Role, schema_of from .materialised_view_contracts import MaterializedViewIndex @@ -85,9 +85,11 @@ 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 + - optionally, ``__mv_schema_tag__``: the schema_translate_map key a Core-only + declaration (no ``Base``/mapped ``Table``) resolves through. + When the class is also declaratively mapped, its own mapped + table's ``schema`` is used instead and ``__mv_schema_tag__`` is ignored + (see :meth:`_resolve_schema_tag`.) This mixin does not define ORM mappings; it is intended for schema-level helpers used during migrations, setup, or administrative workflows. @@ -181,15 +183,28 @@ class DailyObservationCountsMV(Base, MaterializedViewMixin): __mv_name__: str __mv_select__: sa.sql.Select[Any] __mv_dependencies__: set[str] = set() - __mv_role__: Role = Role.PRIMARY + __mv_schema_tag__: str = Role.PRIMARY.value __mv_indexes__: tuple[MaterializedViewIndex, ...] = () + @classmethod + def _resolve_schema_tag(cls, schema_tag: str | None) -> str | None: + """Effective schema_tag for one lifecycle call. + A declratively mapped class's own mapped table's schema is authoritative, + overriding any ``__mv_schema_tag__``. + """ + if schema_tag is not None: + return schema_tag + table = getattr(cls, "__table__", None) + if table is not None: + return table.schema + return cls.__mv_schema_tag__ + @classmethod def create_mv( cls, bind: "sa.engine.Connection | sa.engine.Engine", *, - role: Role | None = None, + schema_tag: str | None = None, with_data: bool = True, if_not_exists: bool = True, create_indexes: bool = True, @@ -201,12 +216,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 ``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. + schema_tag + schema_translate_map key the view's own physical schema resolves + through. See :meth:`_resolve_schema_tag` for the default when + omitted: the mapped table's own schema if the class is + declaratively mapped, else ``cls.__mv_schema_tag__``. create_indexes When True, create every index declared in ``__mv_indexes__``. @@ -246,14 +260,15 @@ def create_mv( from ..backends.resolve import resolve_backend backend = resolve_backend(bind) - role_ = role if role is not None else cls.__mv_role__ + tag = cls._resolve_schema_tag(schema_tag) + schema = schema_of(bind, schema_tag=tag) def create(connection: sa.engine.Connection | sa.engine.Engine) -> None: backend.create_materialized_view( connection, cls.__mv_name__, cls.__mv_select__, - role=role_, + schema=schema, with_data=with_data, if_not_exists=if_not_exists, ) @@ -263,7 +278,7 @@ def create(connection: sa.engine.Connection | sa.engine.Engine) -> None: connection, cls.__mv_name__, index, - role=role_, + schema=schema, if_not_exists=if_not_exists, ) @@ -280,7 +295,7 @@ def refresh_mv( cls, bind: "sa.engine.Connection | sa.engine.Engine", *, - role: Role | None = None, + schema_tag: str | None = None, concurrently: bool = False, ) -> None: """ @@ -290,9 +305,9 @@ def refresh_mv( ---------- 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. + schema_tag + schema_translate_map key the view's own physical schema resolves + through; see :meth:`_resolve_schema_tag` for the default. concurrently Request concurrent refresh, requiring a declared unique index. @@ -312,11 +327,11 @@ def refresh_mv( from ..backends.resolve import resolve_backend backend = resolve_backend(bind) - role_ = role if role is not None else cls.__mv_role__ + tag = cls._resolve_schema_tag(schema_tag) backend.refresh_materialized_view( bind, cls.__mv_name__, - role=role_, + schema=schema_of(bind, schema_tag=tag), concurrently=concurrently, declared_indexes=cls.__mv_indexes__, ) @@ -326,7 +341,7 @@ def drop_mv( cls, bind: "sa.engine.Connection | sa.engine.Engine", *, - role: Role | None = None, + schema_tag: str | None = None, if_exists: bool = True, cascade: bool = False, ) -> None: @@ -336,16 +351,17 @@ def drop_mv( ---------- bind A SQLAlchemy Engine or Connection used to execute the drop. - role - Schema role the view's own physical schema resolves through; - see :meth:`create_mv` for when to override the default. + schema_tag + schema_translate_map key the view's own physical schema resolves + through; see :meth:`_resolve_schema_tag` for the default. """ from ..backends.resolve import resolve_backend backend = resolve_backend(bind) - role_ = role if role is not None else cls.__mv_role__ + tag = cls._resolve_schema_tag(schema_tag) backend.drop_materialized_view( - bind, cls.__mv_name__, role=role_, if_exists=if_exists, cascade=cascade + bind, cls.__mv_name__, schema=schema_of(bind, schema_tag=tag), + if_exists=if_exists, cascade=cascade, ) diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index c51cb5b..9ded33c 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -2,7 +2,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so import logging -from oa_configurator import role_of_table, schema_inspect +from oa_configurator import schema_of, validate_schema_tag from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError @@ -124,8 +124,9 @@ def manage_indices( indices = list(cls.__table__.indexes) if resolved_index_strategy == "drop_rebuild" else [] if indices: - inspector = schema_inspect(session, role=role_of_table(cls.__table__)) - existing_in_db = {idx['name'] for idx in inspector.get_indexes(cls.__tablename__)} + inspector = sa.inspect(session.connection()) + schema = schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) + existing_in_db = {idx['name'] for idx in inspector.get_indexes(cls.__tablename__, schema=schema)} to_drop = [i for i in indices if i.name in existing_in_db] if to_drop: @@ -180,8 +181,9 @@ def manage_indices( if indices: logger.info(f"Table `{table_name}`: Verifying/Rebuilding indices.") rebuild_started = perf_counter() - inspector = schema_inspect(session, role=role_of_table(cls.__table__)) - existing_idx_names = {idx['name'] for idx in inspector.get_indexes(table_name)} + inspector = sa.inspect(session.connection()) + schema = schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) + existing_idx_names = {idx['name'] for idx in inspector.get_indexes(table_name, schema=schema)} for idx in indices: if idx.name not in existing_idx_names: diff --git a/tests/backends/test_base_backend.py b/tests/backends/test_base_backend.py index d1583ce..34154a2 100644 --- a/tests/backends/test_base_backend.py +++ b/tests/backends/test_base_backend.py @@ -12,7 +12,6 @@ import sqlalchemy.orm as so from sqlalchemy.engine import Connection, Engine -from oa_configurator import Role from orm_loader.backends import ( BackendCapabilities, DatabaseBackend, @@ -129,7 +128,7 @@ def create_materialized_view( name: str, selectable: sa.sql.Select[Any], *, - role: Role = Role.PRIMARY, + schema: str | None = None, with_data: bool = True, if_not_exists: bool = True, ) -> None: @@ -140,7 +139,7 @@ def refresh_materialized_view( bind: Engine | Connection, name: str, *, - role: Role = Role.PRIMARY, + schema: str | None = None, concurrently: bool = False, declared_indexes: tuple = (), ) -> None: diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index e79a5f9..8dbfead 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -154,30 +154,27 @@ 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.""" +def test_postgres_backend_materialized_view_respects_schema(pg_db) -> None: + """Resolving a schema_tag to a physical schema is the caller's job + now (see MaterializedViewMixin.create_mv, which calls schema_of() before + ever reaching the backend). This proves the backend itself honors + whatever already-resolved schema it's given, placing the view there and + nowhere else.""" 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.begin() as conn: + backend.create_materialized_view(conn, "mv_schema_test", selectable, schema=vocab_schema) + backend.refresh_materialized_view(conn, "mv_schema_test", schema=vocab_schema) 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 sa.inspect(conn).has_table("mv_schema_test", schema=vocab_schema) + assert not sa.inspect(conn).has_table("mv_schema_test", schema=primary_schema) assert conn.execute( - sa.text(f'SELECT n FROM "{vocab_schema}".mv_role_test') + sa.text(f'SELECT n FROM "{vocab_schema}".mv_schema_test') ).scalar() == 1 @@ -261,7 +258,7 @@ def test_postgres_backend_create_materialized_view_index_emits_expected_sql(): session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) index = MaterializedViewIndex(name="mv_test_row_id_uq", columns=("row_id",), unique=True) - backend.create_materialized_view_index(_sess(session), "mv_test", index) + backend.create_materialized_view_index(_sess(session), "mv_test", index, schema="reporting") assert session.statements == [ 'CREATE UNIQUE INDEX IF NOT EXISTS "mv_test_row_id_uq" ON reporting.mv_test ("row_id")' @@ -285,7 +282,7 @@ def test_postgres_backend_drop_materialized_view_default_args(): backend = PostgresBackend() session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) - backend.drop_materialized_view(_sess(session), "mv_test") + backend.drop_materialized_view(_sess(session), "mv_test", schema="reporting") assert session.statements == ['DROP MATERIALIZED VIEW IF EXISTS reporting.mv_test'] @@ -294,7 +291,9 @@ def test_postgres_backend_drop_materialized_view_cascade_and_if_exists_false(): backend = PostgresBackend() session = _FakeSession(schema_translate_map={Role.PRIMARY.value: "reporting"}) - backend.drop_materialized_view(_sess(session), "mv_test", if_exists=False, cascade=True) + backend.drop_materialized_view( + _sess(session), "mv_test", schema="reporting", if_exists=False, cascade=True + ) assert session.statements == ['DROP MATERIALIZED VIEW reporting.mv_test CASCADE'] @@ -367,6 +366,7 @@ def test_postgres_backend_refresh_concurrently_declared_but_database_rejects_it_ backend.refresh_materialized_view( _sess(session), "mv_test", + schema="reporting", concurrently=True, declared_indexes=(index,), ) @@ -411,6 +411,7 @@ def test_postgres_backend_refresh_concurrently_with_declared_index_emits_concurr backend.refresh_materialized_view( _sess(session), "mv_test", + schema="reporting", concurrently=True, declared_indexes=(index,), ) @@ -423,10 +424,9 @@ def test_postgres_backend_refresh_concurrently_with_declared_index_emits_concurr def test_postgres_backend_materialized_view_lifecycle_is_schema_isolated_with_adversarial_identifiers( pg_db, ): - """Two schemas, each addressed via its own scoped connection (role-based - resolution ties the schema to the connection, not to a per-call - override), must never bleed into each other even with adversarial, - quote-laden identifiers.""" + """Two identically-named views, each placed in its own schema via an + explicit schema= argument, must never bleed into each other even with + adversarial, quote-laden identifiers.""" from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex backend = PostgresBackend() @@ -442,17 +442,16 @@ def test_postgres_backend_materialized_view_lifecycle_is_schema_isolated_with_ad for schema in (left_schema, right_schema): setup_conn.execute(sa.text(f"CREATE SCHEMA {preparer.quote_identifier(schema)}")) - left = engine.execution_options(schema_translate_map={Role.PRIMARY.value: left_schema}) - right = engine.execution_options(schema_translate_map={Role.PRIMARY.value: right_schema}) - - for scoped in (left, right): - with scoped.begin() as conn: - backend.create_materialized_view(conn, name, selectable) - backend.create_materialized_view_index(conn, name, index) + for schema in (left_schema, right_schema): + with engine.begin() as conn: + backend.create_materialized_view(conn, name, selectable, schema=schema) + backend.create_materialized_view_index(conn, name, index, schema=schema) - with left.begin() as conn: - backend.refresh_materialized_view(conn, name, concurrently=True, declared_indexes=(index,)) - backend.drop_materialized_view(conn, name) + with engine.begin() as conn: + backend.refresh_materialized_view( + conn, name, schema=left_schema, concurrently=True, declared_indexes=(index,) + ) + backend.drop_materialized_view(conn, name, schema=left_schema) with engine.connect() as conn: assert conn.execute( diff --git a/tests/loaders/test_loader_e2e.py b/tests/loaders/test_loader_e2e.py index b9df959..69ee5d4 100644 --- a/tests/loaders/test_loader_e2e.py +++ b/tests/loaders/test_loader_e2e.py @@ -25,7 +25,7 @@ RequiredTable, Role, SimpleTable, - VocabRoleTable, + VocabSchemaTable, ) # Typed aliases: Pylance cannot verify SQLAlchemy metaclass-generated attrs @@ -35,7 +35,7 @@ _CompositeTable = cast(Type[CSVTableProtocol], CompositeTable) _EnumTable = cast(Type[CSVTableProtocol], EnumTable) _ImpliedEnumTable = cast(Type[CSVTableProtocol], ImpliedEnumTable) -_VocabRoleTable = cast(Type[CSVTableProtocol], VocabRoleTable) +_VocabSchemaTable = cast(Type[CSVTableProtocol], VocabSchemaTable) @pytest.fixture(autouse=True) @@ -77,20 +77,20 @@ 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 +def test_initial_csv_load_for_a_non_primary_schema_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 + the table's declared schema tag differs from primary. This is the SQLite counterpart to test_schema_translate_map.py's Postgres-only, non-primary- - role coverage.""" + schema-tag 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( + inserted = _VocabSchemaTable.load_csv( session, csv_path, dedupe=False, loader=PandasLoader() ) session.commit() @@ -98,7 +98,7 @@ def test_initial_csv_load_for_a_non_primary_role_table(session, tmp_path): assert inserted == 2 rows = session.execute( - sa.select(VocabRoleTable).order_by(VocabRoleTable.id) + sa.select(VocabSchemaTable).order_by(VocabSchemaTable.id) ).scalars().all() assert [(r.id, r.name) for r in rows] == [(1, "alpha"), (2, "beta")] diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index e0eaffb..a843a6b 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -29,7 +29,7 @@ from orm_loader.loaders.loader_interface import PandasLoader from tests.conftest import schema_scoped_session -from tests.models import SimpleTable, VocabRoleTable +from tests.models import SimpleTable, VocabSchemaTable def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): @@ -112,9 +112,9 @@ def _write_and_load(rows: list[dict], path_name: str) -> int: 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.""" +def test_load_csv_respects_non_primary_schema_tag_end_to_end(pg_db, tmp_path): + """Checks that the schema tag is correctly derived from the table's own + schema (table.schema, via validate_schema_tag()) for a non-primary tag.""" primary_schema = f"test_primary_{uuid.uuid4().hex[:8]}" vocab_schema = f"test_vocab_{uuid.uuid4().hex[:8]}" conn = pg_db.connection @@ -124,7 +124,7 @@ def test_load_csv_respects_non_primary_role_end_to_end(pg_db, tmp_path): session = schema_scoped_session( conn, - VocabRoleTable.__table__, + VocabSchemaTable.__table__, {SchemaRole.PRIMARY.value: primary_schema, SchemaRole.VOCAB.value: vocab_schema}, ) @@ -133,7 +133,7 @@ def test_load_csv_respects_non_primary_role_end_to_end(pg_db, tmp_path): csv_path, index=False, sep="\t" ) - inserted = VocabRoleTable.load_csv( + inserted = VocabSchemaTable.load_csv( session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema=STAGING_SCHEMA ) session.commit() diff --git a/tests/loaders/test_split_connection.py b/tests/loaders/test_split_connection.py index cc2a4ce..6207552 100644 --- a/tests/loaders/test_split_connection.py +++ b/tests/loaders/test_split_connection.py @@ -14,14 +14,14 @@ from orm_loader.loaders.loader_interface import PandasLoader from tests.conftest import schema_scoped_session -from tests.models import SimpleTable, VocabRoleTable +from tests.models import SimpleTable, VocabSchemaTable 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.""" + """``pg_db`` (real Postgres, primary connection) and ``session`` (real + SQLite, vocab connection, 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) @@ -42,7 +42,7 @@ def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_ SimpleTable.load_csv(primary_session, primary_csv, dedupe=False, loader=PandasLoader()) primary_session.commit() - VocabRoleTable.load_csv(session, vocab_csv, dedupe=False, loader=PandasLoader()) + VocabSchemaTable.load_csv(session, vocab_csv, dedupe=False, loader=PandasLoader()) session.commit() primary_rows = conn.execute( @@ -51,7 +51,7 @@ def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_ assert primary_rows == [(1, "primary-alpha")] vocab_rows = session.execute( - sa.select(VocabRoleTable).order_by(VocabRoleTable.id) + sa.select(VocabSchemaTable).order_by(VocabSchemaTable.id) ).scalars().all() assert [(r.id, r.name) for r in vocab_rows] == [(1, "vocab-alpha")] diff --git a/tests/mappers/test_materialised_view_mixin.py b/tests/mappers/test_materialised_view_mixin.py index fee8f10..6bd3bbe 100644 --- a/tests/mappers/test_materialised_view_mixin.py +++ b/tests/mappers/test_materialised_view_mixin.py @@ -5,6 +5,7 @@ import pytest import sqlalchemy as sa +import sqlalchemy.orm as so from oa_configurator import Role from oa_configurator.testing import isolated_test_schema @@ -24,10 +25,10 @@ class _PrimaryRoleMV(MaterializedViewMixin): class _VocabRoleMV(MaterializedViewMixin): __mv_name__ = "mv_vocab_role_test" __mv_select__ = sa.select(sa.literal(2).label("n")) - __mv_role__ = Role.VOCAB + __mv_schema_tag__ = Role.VOCAB.value -def test_refresh_all_mvs_resolves_each_views_own_role(pg_db) -> None: +def test_refresh_all_mvs_resolves_each_views_own_schema_tag(pg_db) -> None: engine = pg_db.connection.engine with isolated_test_schema(engine, prefix="mv_primary") as primary_schema, \ @@ -105,7 +106,7 @@ def test_create_mv_forwards_default_args_to_backend(fake_backend: _FakeBackend, ( "create_materialized_view", (bind, "mv_no_index", _SELECT), - {"role": Role.PRIMARY, "with_data": True, "if_not_exists": True}, + {"schema": "primary", "with_data": True, "if_not_exists": True}, ) ] @@ -125,7 +126,7 @@ def test_create_mv_creates_declared_indexes_after_the_view(fake_backend: _FakeBa "create_materialized_view_index", ] assert fake_backend.calls[1][1] == (bind, "mv_indexed", _INDEX) - assert fake_backend.calls[1][2] == {"role": Role.PRIMARY, "if_not_exists": True} + assert fake_backend.calls[1][2] == {"schema": "primary", "if_not_exists": True} def test_create_mv_create_indexes_false_skips_index_creation(fake_backend: _FakeBackend, bind): @@ -134,22 +135,69 @@ def test_create_mv_create_indexes_false_skips_index_creation(fake_backend: _Fake assert [call[0] for call in fake_backend.calls] == ["create_materialized_view"] -def test_create_mv_forwards_role_with_data_and_if_not_exists_overrides( +def test_create_mv_forwards_schema_tag_with_data_and_if_not_exists_overrides( fake_backend: _FakeBackend, bind ): - _NoIndexMv.create_mv(bind, role=Role.VOCAB, with_data=False, if_not_exists=False) + _NoIndexMv.create_mv(bind, schema_tag=Role.VOCAB, with_data=False, if_not_exists=False) assert fake_backend.calls[0][2] == { - "role": Role.VOCAB, + "schema": "vocab", "with_data": False, "if_not_exists": False, } +_MappedMvBase = so.declarative_base() + + +class _MappedVocabMv(_MappedMvBase, MaterializedViewMixin): + """Declaratively mapped, schema set via __table_args__ (not __mv_schema_tag__).""" + + __mv_name__ = "mv_mapped_vocab" + __mv_select__ = _SELECT + __tablename__ = "mv_mapped_vocab" + __table_args__ = {"schema": Role.VOCAB.value} + + row_id = sa.Column(sa.Integer, primary_key=True) + + +class _MappedNoSchemaMv(_MappedMvBase, MaterializedViewMixin): + """Declaratively mapped, no schema set at all -- should resolve to None, + not fall back to __mv_schema_tag__'s "primary" default.""" + + __mv_name__ = "mv_mapped_no_schema" + __mv_select__ = _SELECT + __tablename__ = "mv_mapped_no_schema" + + row_id = sa.Column(sa.Integer, primary_key=True) + + +def test_create_mv_defers_to_the_mapped_tables_own_schema(fake_backend: _FakeBackend, bind): + _MappedVocabMv.create_mv(bind) + + assert fake_backend.calls[0][2]["schema"] == "vocab" + + +def test_create_mv_mapped_table_with_no_schema_resolves_to_none_not_primary( + fake_backend: _FakeBackend, bind +): + _MappedNoSchemaMv.create_mv(bind) + + assert fake_backend.calls[0][2]["schema"] is None + + +def test_create_mv_explicit_schema_tag_overrides_the_mapped_tables_own_schema( + fake_backend: _FakeBackend, bind +): + _MappedVocabMv.create_mv(bind, schema_tag=Role.PRIMARY) + + assert fake_backend.calls[0][2]["schema"] == "primary" + + def test_create_mv_forwards_if_not_exists_to_declared_indexes(fake_backend: _FakeBackend, bind): _IndexedMv.create_mv(bind, if_not_exists=False) - assert fake_backend.calls[1][2] == {"role": Role.PRIMARY, "if_not_exists": False} + assert fake_backend.calls[1][2] == {"schema": "primary", "if_not_exists": False} def test_create_mv_engine_uses_one_transaction_for_view_and_indexes(monkeypatch: pytest.MonkeyPatch): @@ -211,16 +259,16 @@ def test_refresh_mv_forwards_default_args_and_declared_indexes(fake_backend: _Fa ( "refresh_materialized_view", (bind, "mv_indexed"), - {"role": Role.PRIMARY, "concurrently": False, "declared_indexes": (_INDEX,)}, + {"schema": "primary", "concurrently": False, "declared_indexes": (_INDEX,)}, ) ] -def test_refresh_mv_forwards_role_and_concurrently(fake_backend: _FakeBackend, bind): - _IndexedMv.refresh_mv(bind, role=Role.VOCAB, concurrently=True) +def test_refresh_mv_forwards_schema_tag_and_concurrently(fake_backend: _FakeBackend, bind): + _IndexedMv.refresh_mv(bind, schema_tag=Role.VOCAB, concurrently=True) assert fake_backend.calls[0][2] == { - "role": Role.VOCAB, + "schema": "vocab", "concurrently": True, "declared_indexes": (_INDEX,), } @@ -233,16 +281,16 @@ def test_drop_mv_forwards_default_args(fake_backend: _FakeBackend, bind): ( "drop_materialized_view", (bind, "mv_no_index"), - {"role": Role.PRIMARY, "if_exists": True, "cascade": False}, + {"schema": "primary", "if_exists": True, "cascade": False}, ) ] -def test_drop_mv_forwards_role_if_exists_and_cascade(fake_backend: _FakeBackend, bind): - _NoIndexMv.drop_mv(bind, role=Role.VOCAB, if_exists=False, cascade=True) +def test_drop_mv_forwards_schema_tag_if_exists_and_cascade(fake_backend: _FakeBackend, bind): + _NoIndexMv.drop_mv(bind, schema_tag=Role.VOCAB, if_exists=False, cascade=True) assert fake_backend.calls[0][2] == { - "role": Role.VOCAB, + "schema": "vocab", "if_exists": False, "cascade": True, } diff --git a/tests/models.py b/tests/models.py index 898a3b0..8b47569 100644 --- a/tests/models.py +++ b/tests/models.py @@ -96,11 +96,12 @@ class ImpliedEnumTable(Base, CSVLoadableTableInterface): 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.""" +class VocabSchemaTable(Base, CSVLoadableTableInterface): + """A VOCAB-tagged table, so tests can prove the staging/index schema + derivation (table.schema resolved via validate_schema_tag()/schema_of(), + threaded through create_staging_table()/manage_indices()) actually + resolves a non-primary schema tag correctly, instead of only ever + exercising the PRIMARY-tagged default.""" __tablename__ = "test_vocab_role_table" __table_args__ = ( From 7bd1f5842cd8ed0796939ffcd403684a6627ae72 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 23 Sep 2026 01:15:13 +0000 Subject: [PATCH 19/29] Quick edit to docstring --- tests/mappers/test_materialised_view_mixin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mappers/test_materialised_view_mixin.py b/tests/mappers/test_materialised_view_mixin.py index 6bd3bbe..4d8a752 100644 --- a/tests/mappers/test_materialised_view_mixin.py +++ b/tests/mappers/test_materialised_view_mixin.py @@ -162,7 +162,7 @@ class _MappedVocabMv(_MappedMvBase, MaterializedViewMixin): class _MappedNoSchemaMv(_MappedMvBase, MaterializedViewMixin): - """Declaratively mapped, no schema set at all -- should resolve to None, + """Declaratively mapped, no schema set at all. Should resolve to None, not fall back to __mv_schema_tag__'s "primary" default.""" __mv_name__ = "mv_mapped_no_schema" From f52bb2aa4bb7aee391dc7338686eba6baa0b7bae Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 23 Sep 2026 04:30:44 +0000 Subject: [PATCH 20/29] Follow-up from internal review --- src/orm_loader/backends/base.py | 15 +++----- src/orm_loader/backends/sqlite.py | 4 +- src/orm_loader/helpers/__init__.py | 2 - src/orm_loader/helpers/bootstrap.py | 37 ++++++++++++++++--- src/orm_loader/helpers/sql.py | 30 --------------- .../mappers/materialised_view_contracts.py | 6 +-- .../mappers/materialised_view_mixin.py | 18 ++++----- tests/backends/test_postgres_backend.py | 10 ++--- 8 files changed, 55 insertions(+), 67 deletions(-) delete mode 100644 src/orm_loader/helpers/sql.py diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 073635f..c8809c9 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -12,7 +12,7 @@ from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer -from oa_configurator import Dialect +from oa_configurator import Dialect, open_connection, qualified if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext @@ -120,10 +120,8 @@ def qualified_staging_name(self, tablename: str) -> str: str e.g. '"staging"."_staging_concept"' or '"_staging_concept"'. """ - from ..helpers.sql import qualify_identifier - - return qualify_identifier( - self.staging_name_for_table(tablename), self.staging_schema, self.identifier_preparer + return qualified( + self.identifier_preparer, self.staging_name_for_table(tablename), physical_schema=self.staging_schema ) @property @@ -221,11 +219,8 @@ def _as_connection( "be the same one (or share the same dialect as) the bind resolve_backend() " "was given." ) - if isinstance(bind, Engine): - with bind.begin() as conn: - yield conn - else: - yield bind + with open_connection(bind) as conn: + yield conn def _insertable_column_names( self, diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index 42cf9ac..b35ce19 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 schema_if_supported + from .base import BackendCapabilities, DatabaseBackend, Dialect if TYPE_CHECKING: @@ -40,7 +42,7 @@ def __init__( journal_mode: str = "WAL", defer_foreign_keys: bool = True, ) -> None: - if staging_schema is not None: + if staging_schema is not None and schema_if_supported(staging_schema, Dialect.SQLITE) is None: logger.warning( "SQLite does not support schema-qualified staging tables; " f"got staging_schema={staging_schema!r}. Setting staging_schema=None." diff --git a/src/orm_loader/helpers/__init__.py b/src/orm_loader/helpers/__init__.py index 6ad5a70..651ab2f 100644 --- a/src/orm_loader/helpers/__init__.py +++ b/src/orm_loader/helpers/__init__.py @@ -9,7 +9,6 @@ from .metadata import Base from .discovery import get_model_by_tablename from .null_handlers import normalise_null -from .sql import qualify_identifier __all__ = [ "IngestError", @@ -24,5 +23,4 @@ "Base", "get_model_by_tablename", "normalise_null", - "qualify_identifier", ] diff --git a/src/orm_loader/helpers/bootstrap.py b/src/orm_loader/helpers/bootstrap.py index 08f7760..b27c42f 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -1,13 +1,40 @@ -from .metadata import Base +from contextlib import ExitStack import logging + import sqlalchemy as sa +from oa_configurator import ResolvedDatabase, guard_schema_provenance_for, open_connection, validate_schema_tag + +from .metadata import Base + logger = logging.getLogger(__name__) -def create_db(engine: sa.engine.Engine) -> None: + +def create_db( + bindable: sa.engine.Engine | sa.engine.Connection, *, resolved: ResolvedDatabase | None = None +) -> None: logger.debug("Creating database schema") - Base.metadata.create_all(engine) + tables_by_schema_tag: dict[str, list[sa.Table]] = {} + for table in Base.metadata.tables.values(): + schema_tag = validate_schema_tag(table) + if schema_tag is not None: + tables_by_schema_tag.setdefault(schema_tag, []).append(table) + + with open_connection(bindable) as connection, ExitStack() as guard_stack: + # One provenance guard per schema_tag (count only known at runtime); ExitStack defers every write until the block below succeeds. + for schema_tag, tables in tables_by_schema_tag.items(): + # A same-named table could already exist under a drifted schema, attached to an unrelated table. + guard_stack.enter_context( + guard_schema_provenance_for(connection, resolved, schema_tag=schema_tag, tables=tables) + ) + Base.metadata.create_all(connection) + -def bootstrap(engine: sa.engine.Engine, *, create: bool = True) -> None: +def bootstrap( + bindable: sa.engine.Engine | sa.engine.Connection, + *, + create: bool = True, + resolved: ResolvedDatabase | None = None, +) -> None: logger.info("Bootstrapping schema (create=%s)", create) if create: - create_db(engine) + create_db(bindable, resolved=resolved) diff --git a/src/orm_loader/helpers/sql.py b/src/orm_loader/helpers/sql.py deleted file mode 100644 index 6e30ba9..0000000 --- a/src/orm_loader/helpers/sql.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -from sqlalchemy.sql.compiler import IdentifierPreparer - - -def qualify_identifier(name: str, schema: str | None, preparer: IdentifierPreparer) -> str: - """ - Return a quoted, optionally schema-qualified SQL identifier. - - Parameters - ---------- - name - The SQL identifier to qualify (e.g. a table name). - schema - Schema name to prefix. If None, returns only the quoted identifier. - Useful for backends that do not support schema-qualified identifiers (e.g. SQLite). - preparer - The dialect-specific identifier preparer used to quote and escape each - component. Delegating to SQLAlchemy here (rather than hand-rolled - f-string quoting) ensures embedded quote characters are escaped - correctly for the target dialect. - - Returns - ------- - str - e.g. '"staging"."_staging_foo"' or '"_staging_foo"'. - """ - if schema: - return f"{preparer.quote_identifier(schema)}.{preparer.quote_identifier(name)}" - return preparer.quote_identifier(name) diff --git a/src/orm_loader/mappers/materialised_view_contracts.py b/src/orm_loader/mappers/materialised_view_contracts.py index e91e8a2..cbba3a2 100644 --- a/src/orm_loader/mappers/materialised_view_contracts.py +++ b/src/orm_loader/mappers/materialised_view_contracts.py @@ -2,7 +2,7 @@ Schema qualification for these DDL elements follows the same convention as ``CreateMaterializedView``: callers must build a fully qualified, quoted -target string themselves (see ``orm_loader.helpers.sql.qualify_identifier``) +target string themselves (see ``oa_configurator.qualified``) before constructing one of these elements. The compiler has no live bindable to qualify a bare name itself. """ @@ -53,7 +53,7 @@ class DropMaterializedView(DDLElement): ---------- name Fully qualified, quoted name of the materialized view to drop (see - ``orm_loader.helpers.sql.qualify_identifier``). + ``oa_configurator.qualified``). if_exists Emit ``IF EXISTS`` so dropping an already-absent view is a no-op rather than an error. @@ -85,7 +85,7 @@ class CreateMaterializedViewIndex(DDLElement): ---------- target Fully qualified, quoted name of the materialized view to index (see - ``orm_loader.helpers.sql.qualify_identifier``). + ``oa_configurator.qualified``). index The index to create. if_not_exists diff --git a/src/orm_loader/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index 7e9d064..ba128e8 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -4,7 +4,12 @@ import sqlalchemy as sa from sqlalchemy.ext import compiler from sqlalchemy.schema import DDLElement -from oa_configurator import Role, schema_of +from oa_configurator import ( + Role, + open_connection, + schema_of, + validate_schema_tag +) from .materialised_view_contracts import MaterializedViewIndex @@ -196,7 +201,7 @@ def _resolve_schema_tag(cls, schema_tag: str | None) -> str | None: return schema_tag table = getattr(cls, "__table__", None) if table is not None: - return table.schema + return validate_schema_tag(table) return cls.__mv_schema_tag__ @classmethod @@ -282,13 +287,8 @@ def create(connection: sa.engine.Connection | sa.engine.Engine) -> None: if_not_exists=if_not_exists, ) - if isinstance(bind, sa.engine.Engine): - with bind.begin() as connection: - create(connection) - elif isinstance(bind, sa.engine.Connection): - create(bind) - else: # pragma: no cover - guarded above; keeps the runtime contract explicit. - raise TypeError("bind must be a SQLAlchemy Engine or Connection") + with open_connection(bind) as connection: + create(connection) @classmethod def refresh_mv( diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 8dbfead..5b450f0 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -9,16 +9,15 @@ from sqlalchemy.dialects import postgresql from sqlalchemy.engine import Engine -from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role +from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role, qualified from oa_configurator.testing import isolated_test_schema from orm_loader.backends import STAGING_SCHEMA, Dialect, PostgresBackend -from orm_loader.helpers.sql import qualify_identifier from tests.models import ComputedColumnTable _TARGET_TABLE = 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) +_STAGING_TABLE_WITH_SCHEMA: str = qualified(_PREPARER, _STAGING_TABLE, physical_schema=STAGING_SCHEMA) _ComputedTableCls = cast("Type[CSVTableProtocol]", ComputedColumnTable) @@ -91,16 +90,13 @@ def test_postgres_backend_identity_and_capabilities(): assert backend.capabilities.supports_materialized_views is True -def test_qualify_identifier_escapes_embedded_quotes(): - assert qualify_identifier("table", 'schema"name', _PREPARER) == '"schema""name"."table"' - assert qualify_identifier('ta"ble', None, _PREPARER) == '"ta""ble"' def test_postgres_backend_default_staging_schema_is_none(): backend = PostgresBackend() assert backend.staging_schema is None - assert backend.qualified_staging_name(_TARGET_TABLE) == _PREPARER.quote_identifier(_STAGING_TABLE) + assert backend.qualified_staging_name(_TARGET_TABLE) == qualified(_PREPARER, _STAGING_TABLE, physical_schema=None) def test_postgres_backend_create_staging_table_drops_computed_columns(pg_session): From a24e45367470f37cb7f7d315f297e1ef26580016 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 23 Sep 2026 04:48:59 +0000 Subject: [PATCH 21/29] Disambiguate physical schema from schema tag --- src/orm_loader/backends/postgres.py | 4 ++-- src/orm_loader/mappers/materialised_view_mixin.py | 8 ++++---- src/orm_loader/tables/loadable_table.py | 6 +++--- tests/backends/test_postgres_backend.py | 4 ++-- tests/models.py | 2 +- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index c19ecc5..bc20c15 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -9,7 +9,7 @@ from oa_configurator import ( autocommit_connection, qualified, - schema_of, + physical_schema_of, validate_schema_tag, Dialect, ) @@ -73,7 +73,7 @@ def create_staging_table( preparer = self.identifier_preparer staging_ref = self.qualified_staging_name(table_cls.__tablename__) source_ref = qualified( - session, table.name, physical_schema=schema_of(session, schema_tag=validate_schema_tag(table)) + session, table.name, physical_schema=physical_schema_of(session, schema_tag=validate_schema_tag(table)) ) session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};')) session.execute( diff --git a/src/orm_loader/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index ba128e8..fb77fb3 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -7,7 +7,7 @@ from oa_configurator import ( Role, open_connection, - schema_of, + physical_schema_of, validate_schema_tag ) @@ -266,7 +266,7 @@ def create_mv( backend = resolve_backend(bind) tag = cls._resolve_schema_tag(schema_tag) - schema = schema_of(bind, schema_tag=tag) + schema = physical_schema_of(bind, schema_tag=tag) def create(connection: sa.engine.Connection | sa.engine.Engine) -> None: backend.create_materialized_view( @@ -331,7 +331,7 @@ def refresh_mv( backend.refresh_materialized_view( bind, cls.__mv_name__, - schema=schema_of(bind, schema_tag=tag), + schema=physical_schema_of(bind, schema_tag=tag), concurrently=concurrently, declared_indexes=cls.__mv_indexes__, ) @@ -360,7 +360,7 @@ def drop_mv( backend = resolve_backend(bind) tag = cls._resolve_schema_tag(schema_tag) backend.drop_materialized_view( - bind, cls.__mv_name__, schema=schema_of(bind, schema_tag=tag), + bind, cls.__mv_name__, schema=physical_schema_of(bind, schema_tag=tag), if_exists=if_exists, cascade=cascade, ) diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 9ded33c..6a354cb 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -2,7 +2,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so import logging -from oa_configurator import schema_of, validate_schema_tag +from oa_configurator import physical_schema_of, validate_schema_tag from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError @@ -125,7 +125,7 @@ def manage_indices( if indices: inspector = sa.inspect(session.connection()) - schema = schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) + schema = physical_schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) existing_in_db = {idx['name'] for idx in inspector.get_indexes(cls.__tablename__, schema=schema)} to_drop = [i for i in indices if i.name in existing_in_db] @@ -182,7 +182,7 @@ def manage_indices( logger.info(f"Table `{table_name}`: Verifying/Rebuilding indices.") rebuild_started = perf_counter() inspector = sa.inspect(session.connection()) - schema = schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) + schema = physical_schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) existing_idx_names = {idx['name'] for idx in inspector.get_indexes(table_name, schema=schema)} for idx in indices: diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index 5b450f0..bf77a45 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -43,7 +43,7 @@ def __init__( self._schema_translate_map = schema_translate_map def get_execution_options(self) -> dict: - """Minimal support for oa_configurator.schema_of(), which every + """Minimal support for oa_configurator.physical_schema_of(), which every materialized-view backend method resolves its target schema through.""" if self._schema_translate_map is None: return {} @@ -152,7 +152,7 @@ def test_postgres_backend_materialized_view_methods_work_end_to_end(pg_db): def test_postgres_backend_materialized_view_respects_schema(pg_db) -> None: """Resolving a schema_tag to a physical schema is the caller's job - now (see MaterializedViewMixin.create_mv, which calls schema_of() before + now (see MaterializedViewMixin.create_mv, which calls physical_schema_of() before ever reaching the backend). This proves the backend itself honors whatever already-resolved schema it's given, placing the view there and nowhere else.""" diff --git a/tests/models.py b/tests/models.py index 8b47569..03a1645 100644 --- a/tests/models.py +++ b/tests/models.py @@ -98,7 +98,7 @@ class ImpliedEnumTable(Base, CSVLoadableTableInterface): class VocabSchemaTable(Base, CSVLoadableTableInterface): """A VOCAB-tagged table, so tests can prove the staging/index schema - derivation (table.schema resolved via validate_schema_tag()/schema_of(), + derivation (table.schema resolved via validate_schema_tag()/physical_schema_of(), threaded through create_staging_table()/manage_indices()) actually resolves a non-primary schema tag correctly, instead of only ever exercising the PRIMARY-tagged default.""" From f5083aaadda7111f63fa79382365af02e01a2784 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 24 Sep 2026 00:17:25 +0000 Subject: [PATCH 22/29] Docstring adaptations --- src/orm_loader/loaders/data_classes.py | 3 +++ src/orm_loader/tables/loadable_table.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/orm_loader/loaders/data_classes.py b/src/orm_loader/loaders/data_classes.py index b092be7..f272c4f 100644 --- a/src/orm_loader/loaders/data_classes.py +++ b/src/orm_loader/loaders/data_classes.py @@ -61,6 +61,9 @@ class LoaderContext: Whether to apply type casting / normalisation. dedupe Whether to perform deduplication (pre-insertion for source issues) + quote_mode + Quoting mode, resolved to the same concrete mode on both the + PostgreSQL COPY fast-path and the pandas ORM fallback. staging_schema Schema the staging table lives in, passed to resolve_backend() so every backend resolution within this load shares the same schema. diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 6a354cb..2265dee 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -392,7 +392,9 @@ def load_csv( Merge strategy to apply (e.g. ``replace``, ``upsert``, or ``insert_if_empty``). quote_mode - Quoting mode used by the PostgreSQL fast-path loader. + Quoting mode. Governs parsing on both the PostgreSQL COPY + fast-path and the pandas ORM fallback, so a file that falls back + to the ORM path is still parsed identically. index_strategy Index handling strategy during merge. Use ``"auto"`` to let the backend choose a sensible default. From fcb6a334f54e42c721fd6ed528671df55ac7dc2e Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 24 Sep 2026 02:48:49 +0000 Subject: [PATCH 23/29] Updated Docs --- docs/loaders/context.md | 2 +- docs/registry/validation.md | 12 ++++++++++++ docs/tables/mat_view.md | 4 ++-- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/docs/loaders/context.md b/docs/loaders/context.md index 29418fd..59a4abe 100644 --- a/docs/loaders/context.md +++ b/docs/loaders/context.md @@ -25,7 +25,7 @@ on globals or implicit configuration. | `chunksize` | Optional chunk size | | `normalise` | Whether to cast values to ORM types | | `dedupe` | Whether to deduplicate incoming data | -| `quote_mode` | CSV quoting mode for PostgreSQL fast-path loading | +| `quote_mode` | CSV quoting mode, resolved to the same concrete mode on both the PostgreSQL COPY fast-path and the pandas ORM fallback | ::: orm_loader.loaders.data_classes.LoaderContext diff --git a/docs/registry/validation.md b/docs/registry/validation.md index de875ae..3cf0e13 100644 --- a/docs/registry/validation.md +++ b/docs/registry/validation.md @@ -45,3 +45,15 @@ Validates primary key presence, nullability, and alignment with specs. Validates structural correctness of foreign key definitions. ::: orm_loader.registry.validation.ForeignKeyShapeValidator + +--- + +## `always_on_validators()` + +Returns the standard bundle of all four validators above +(`ColumnPresenceValidator`, `ColumnNullabilityValidator`, +`PrimaryKeyValidator`, `ForeignKeyShapeValidator`) as a ready-made list, for +callers that want the default validation set without constructing each one +individually. + +::: orm_loader.registry.validation_presets.always_on_validators diff --git a/docs/tables/mat_view.md b/docs/tables/mat_view.md index 30e9bbd..f03a7be 100644 --- a/docs/tables/mat_view.md +++ b/docs/tables/mat_view.md @@ -105,7 +105,7 @@ PatientSummaryMV.refresh_mv(engine, concurrently=True) This is fail-closed by design. An index created manually outside `__mv_indexes__` does not satisfy the mixin's declaration contract; declare it in the class even if another migration is responsible for creating it. Expressions, partial indexes, and other index forms are outside this simple contract and should not be represented as `MaterializedViewIndex` entries. -The view's physical schema comes from the bound connection's own `schema_translate_map`. `create_mv()`/`refresh_mv()`/`drop_mv()` resolve it via `oa_configurator.schema_of()` before ever reaching the backend, which only ever sees an already-resolved physical schema string. +The view's physical schema comes from the bound connection's own `schema_translate_map`. `create_mv()`/`refresh_mv()`/`drop_mv()` resolve it via `oa_configurator.physical_schema_of()` before ever reaching the backend, which only ever sees an already-resolved physical schema string. The schema tag itself comes from one of three places, in order: @@ -142,7 +142,7 @@ Every generated identifier is quoted through `oa_configurator.qualified()`, whic The built-in implementation is PostgreSQL-oriented. SQLite rejects materialized-view operations with `NotImplementedError`; this is intentional, not an emulation using ordinary views. -`drop_mv()` and declared-index creation wrap execution failures in `MaterializationError`. +`create_mv()`, `drop_mv()`, and declared-index creation all wrap execution failures in `MaterializationError`. ## API reference From 64fe41aab8579ac572c350b16601e548dc427df2 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 1 Oct 2026 00:31:40 +0000 Subject: [PATCH 24/29] Update to most recent oa-config changes, catch unregistered staging schema, expose registering of staging schema to caller --- src/orm_loader/backends/__init__.py | 2 + src/orm_loader/backends/base.py | 28 ++++- src/orm_loader/backends/postgres.py | 20 +++- src/orm_loader/config.py | 6 +- src/orm_loader/helpers/bootstrap.py | 80 ++++++++++---- .../mappers/materialised_view_mixin.py | 9 +- src/orm_loader/tables/loadable_table.py | 6 +- tests/backends/test_reserved_schema.py | 75 ++++++++++--- tests/conftest.py | 11 +- tests/mappers/test_materialised_view_mixin.py | 100 +++++++++++++++--- 10 files changed, 271 insertions(+), 66 deletions(-) diff --git a/src/orm_loader/backends/__init__.py b/src/orm_loader/backends/__init__.py index e5dd07a..2c9c799 100644 --- a/src/orm_loader/backends/__init__.py +++ b/src/orm_loader/backends/__init__.py @@ -6,6 +6,7 @@ DatabaseBackend, Dialect, STAGING_SCHEMA, + staging_schema_claim, ) from ..mappers.materialised_view_errors import ( ConcurrentRefreshNotEligibleError, @@ -26,4 +27,5 @@ "STAGING_SCHEMA", "SQLiteBackend", "resolve_backend", + "staging_schema_claim", ] diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index c8809c9..2a8eac4 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -12,7 +12,7 @@ from sqlalchemy.engine import Connection, Engine from sqlalchemy.sql.compiler import IdentifierPreparer -from oa_configurator import Dialect, open_connection, qualified +from oa_configurator import Dialect, SchemaClaim, open_connection, qualified if TYPE_CHECKING: from ..loaders.data_classes import LoaderContext @@ -37,6 +37,32 @@ class BackendCapabilities: STAGING_SCHEMA: str = "staging" + +def staging_schema_claim(physical_schema: str | None = None, *, reserved: bool = True) -> SchemaClaim: + """SchemaClaim reserving orm-loader's staging schema. + Since orm-loader doesn't have an engine creation step of its own, + this claim needs to be registered by the caller ingesting tables into + the database through orm-loader. + + Add to your own ``create_engine(schema_claims=[...])`` call before + handing a session into this package's staging path. + + Parameters + ---------- + physical_schema : str, optional + Physical schema to reserve. Defaults to ``STAGING_SCHEMA`` itself. + reserved : bool, optional + Whether no other owner may claim this physical schema on the same + connection. Defaults to True. + """ + return SchemaClaim( + schema_tag=STAGING_SCHEMA, + physical_schema=physical_schema or STAGING_SCHEMA, + reserved=reserved, + owner="orm_loader", + ) + + P = ParamSpec("P") R = TypeVar("R") diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index bc20c15..a089a3a 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -10,14 +10,19 @@ autocommit_connection, qualified, physical_schema_of, - validate_schema_tag, Dialect, + UnregisteredSchemaTagError, ) from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError from sqlalchemy.sql.compiler import IdentifierPreparer -from .base import BackendCapabilities, DatabaseBackend, requires_capability +from .base import ( + STAGING_SCHEMA, + BackendCapabilities, + DatabaseBackend, + requires_capability +) from ..mappers.materialised_view_errors import ( ConcurrentRefreshNotEligibleError, MaterializationError, @@ -69,11 +74,20 @@ def create_staging_table( table_cls: type["CSVTableProtocol"], session: so.Session, ) -> None: + if self.staging_schema == STAGING_SCHEMA: + try: + physical_schema_of(session, schema_tag=STAGING_SCHEMA) + except UnregisteredSchemaTagError as exc: + raise UnregisteredSchemaTagError( + "orm-loader needs its staging schema reserved on this engine. Add " + "orm_loader.staging_schema_claim() to your own " + "create_engine(schema_claims=[...]) call." + ) from exc table = table_cls.__table__ preparer = self.identifier_preparer staging_ref = self.qualified_staging_name(table_cls.__tablename__) source_ref = qualified( - session, table.name, physical_schema=physical_schema_of(session, schema_tag=validate_schema_tag(table)) + session, table.name, physical_schema=physical_schema_of(session, schema_tag=table.schema) ) session.execute(sa.text(f'DROP TABLE IF EXISTS {staging_ref};')) session.execute( diff --git a/src/orm_loader/config.py b/src/orm_loader/config.py index d712406..dc1e524 100644 --- a/src/orm_loader/config.py +++ b/src/orm_loader/config.py @@ -4,13 +4,9 @@ from typing import Annotated, ClassVar -from oa_configurator import CDMDatabaseConfig, PackageConfigBase, RefTo, register_reserved_schema +from oa_configurator import CDMDatabaseConfig, PackageConfigBase, RefTo 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/src/orm_loader/helpers/bootstrap.py b/src/orm_loader/helpers/bootstrap.py index b27c42f..c4d091a 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -2,39 +2,81 @@ import logging import sqlalchemy as sa -from oa_configurator import ResolvedDatabase, guard_schema_provenance_for, open_connection, validate_schema_tag +from oa_configurator import ( + ResolvedCDMDatabase, + ResolvedDatabase, + guard_schema_provenance_for, + is_ephemeral_url, + open_connection, +) from .metadata import Base logger = logging.getLogger(__name__) +Bindable = sa.engine.Engine | sa.engine.Connection -def create_db( - bindable: sa.engine.Engine | sa.engine.Connection, *, resolved: ResolvedDatabase | None = None -) -> None: + +def _resolve_binds(resolved: ResolvedDatabase, bindable: Bindable | None) -> tuple[Bindable, Bindable]: + """Return (primary_bind, vocab_bind) to run DDL against. + + Built from resolved.create_engine()/create_engines() when bindable is + omitted, keeping schema-claim registration and vocab/primary routing in + one place. This is the default, enforced path for a real database! + + bindable is required instead whenever resolved can't be safely + reconnected to: an ephemeral SQLite connection (':memory:', including + shared-cache mode) gets a separate, empty database from a second, + independently-built engine, since nothing else shares its data. bindable + is also accepted for a non-ephemeral resolved when the caller already + holds the one true engine and resolved can't be trusted to describe it + faithfully. + """ + if is_ephemeral_url(resolved.connection.safe_url) and bindable is None: + raise ValueError( + f"{resolved.name!r} resolves to an ephemeral connection (SQLite ':memory:' " + "or shared-cache mode): a freshly built engine would be a separate, empty " + "database. Pass the caller's own already-open engine/connection via `bindable`." + ) + if bindable is not None: + return bindable, bindable + if isinstance(resolved, ResolvedCDMDatabase): + return resolved.create_engines() + engine = resolved.create_engine() + return engine, engine + + +def create_db(resolved: ResolvedDatabase, *, bindable: Bindable | None = None) -> None: logger.debug("Creating database schema") - tables_by_schema_tag: dict[str, list[sa.Table]] = {} + primary_bind, vocab_bind = _resolve_binds(resolved, bindable) + + tables_by_bind: dict[int, tuple[Bindable, dict[str, list[sa.Table]]]] = {} for table in Base.metadata.tables.values(): - schema_tag = validate_schema_tag(table) - if schema_tag is not None: - tables_by_schema_tag.setdefault(schema_tag, []).append(table) + schema_tag = table.schema + if schema_tag is None: + continue + bind = resolved.route_for_schema_tag(schema_tag, vocab=vocab_bind, primary=primary_bind) + _, tables_by_tag = tables_by_bind.setdefault(id(bind), (bind, {})) + tables_by_tag.setdefault(schema_tag, []).append(table) - with open_connection(bindable) as connection, ExitStack() as guard_stack: - # One provenance guard per schema_tag (count only known at runtime); ExitStack defers every write until the block below succeeds. - for schema_tag, tables in tables_by_schema_tag.items(): - # A same-named table could already exist under a drifted schema, attached to an unrelated table. - guard_stack.enter_context( - guard_schema_provenance_for(connection, resolved, schema_tag=schema_tag, tables=tables) - ) - Base.metadata.create_all(connection) + with ExitStack() as guard_stack: + for bind, tables_by_tag in tables_by_bind.values(): + connection = guard_stack.enter_context(open_connection(bind)) + # One provenance guard per schema_tag on this connection; ExitStack defers every write until the block below succeeds. + for schema_tag in tables_by_tag: + guard_stack.enter_context( + guard_schema_provenance_for(connection, resolved, schema_tag=schema_tag) + ) + tables = [table for tables in tables_by_tag.values() for table in tables] + Base.metadata.create_all(connection, tables=tables, checkfirst=True) def bootstrap( - bindable: sa.engine.Engine | sa.engine.Connection, + resolved: ResolvedDatabase, *, create: bool = True, - resolved: ResolvedDatabase | None = None, + bindable: Bindable | None = None, ) -> None: logger.info("Bootstrapping schema (create=%s)", create) if create: - create_db(bindable, resolved=resolved) + create_db(resolved, bindable=bindable) diff --git a/src/orm_loader/mappers/materialised_view_mixin.py b/src/orm_loader/mappers/materialised_view_mixin.py index fb77fb3..39a2447 100644 --- a/src/orm_loader/mappers/materialised_view_mixin.py +++ b/src/orm_loader/mappers/materialised_view_mixin.py @@ -5,10 +5,9 @@ from sqlalchemy.ext import compiler from sqlalchemy.schema import DDLElement from oa_configurator import ( - Role, - open_connection, - physical_schema_of, - validate_schema_tag + Role, + open_connection, + physical_schema_of, ) from .materialised_view_contracts import MaterializedViewIndex @@ -201,7 +200,7 @@ def _resolve_schema_tag(cls, schema_tag: str | None) -> str | None: return schema_tag table = getattr(cls, "__table__", None) if table is not None: - return validate_schema_tag(table) + return table.schema return cls.__mv_schema_tag__ @classmethod diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 2265dee..1278c60 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -2,7 +2,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so import logging -from oa_configurator import physical_schema_of, validate_schema_tag +from oa_configurator import physical_schema_of from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError @@ -125,7 +125,7 @@ def manage_indices( if indices: inspector = sa.inspect(session.connection()) - schema = physical_schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) + schema = physical_schema_of(session, schema_tag=cls.__table__.schema) existing_in_db = {idx['name'] for idx in inspector.get_indexes(cls.__tablename__, schema=schema)} to_drop = [i for i in indices if i.name in existing_in_db] @@ -182,7 +182,7 @@ def manage_indices( logger.info(f"Table `{table_name}`: Verifying/Rebuilding indices.") rebuild_started = perf_counter() inspector = sa.inspect(session.connection()) - schema = physical_schema_of(session, schema_tag=validate_schema_tag(cls.__table__)) + schema = physical_schema_of(session, schema_tag=cls.__table__.schema) existing_idx_names = {idx['name'] for idx in inspector.get_indexes(table_name, schema=schema)} for idx in indices: diff --git a/tests/backends/test_reserved_schema.py b/tests/backends/test_reserved_schema.py index 4ff133a..91db0c5 100644 --- a/tests/backends/test_reserved_schema.py +++ b/tests/backends/test_reserved_schema.py @@ -1,22 +1,69 @@ -"""Confirms orm-loader's STAGING_SCHEMA registration (backends/base.py, -Phase 2.3) is actually picked up by oa-configurator's reserved-schema -check: resolving a CDM database configured with cdm_schema="staging" -must raise, proving the cross-package registration/enforcement wiring -works end to end, not just in isolation on either side. +"""Confirms orm-loader's STAGING_SCHEMA reservation (registered via +orm_loader.staging_schema_claim(), passed into a consumer's own +create_engine(schema_claims=[...]) call) is actually picked up by +oa-configurator's own conflict check: a CDM database resolving +cdm_schema=STAGING_SCHEMA on the same connection must raise, proving the +cross-package registration/enforcement wiring works end to end, not just +in isolation on either side. + +Registration happens once, at create_engine() time, via the sanctioned +schema_claims= path -- create_staging_table() only reads the already- +registered schema (see backends/postgres.py), it never registers it +itself. Uses a fresh, genuinely committing engine (never the +rollback-protected pg_db.connection), since the registration write and +the later conflict check are two separate connections and must actually +see each other. """ from __future__ import annotations +from oa_configurator import ( + CDMDatabaseConfig, + ConnectionConfig, + GenericDatabaseConfig, + Resolver, + SchemaOwnershipError, + StackConfig, +) +from oa_configurator.domains.resources.schema_registry import SchemaRegistry +from sqlalchemy.engine import make_url +from sqlalchemy import Table +from typing import cast import pytest -from oa_configurator import CDMDatabaseConfig, ConnectionConfig, StackConfig -from pydantic import ValidationError -from orm_loader.backends import STAGING_SCHEMA +from orm_loader.backends import STAGING_SCHEMA, staging_schema_claim + +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + + +def _connection_config(pg_db) -> ConnectionConfig: + connection = pg_db.resolved.connection + url = make_url(connection.url) + return ConnectionConfig( + dialect=url.drivername, host=url.host, port=url.port, + user=url.username, password=url.password, database_name=url.database, + test_only=connection.test_only, + ) -def test_resolving_cdm_database_with_staging_schema_name_raises() -> None: - with pytest.raises(ValidationError, match=f"{STAGING_SCHEMA!r}.*orm-loader"): - StackConfig.for_session( - connections={"c": ConnectionConfig(dialect="sqlite", database_name=":memory:")}, - databases={"default": CDMDatabaseConfig(connection="c", cdm_schema=STAGING_SCHEMA)}, - ) +def test_resolving_cdm_database_with_staging_schema_name_raises(pg_db) -> None: + connection_config = _connection_config(pg_db) + stack = StackConfig.for_session( + connections={"c": connection_config}, + databases={ + "loader": GenericDatabaseConfig(connection="c"), + "default": CDMDatabaseConfig(connection="c", cdm_schema=STAGING_SCHEMA), + }, + ) + resolver = Resolver(stack) + engine = resolver.resolve_database("loader").create_engine( + schema_claims=[staging_schema_claim()] + ) + try: + with pytest.raises(SchemaOwnershipError, match=f"{STAGING_SCHEMA!r}.*orm_loader"): + resolver.resolve_database("default").create_engine() + finally: + engine.dispose() + with pg_db.committing_engine.begin() as connection: + table = cast(Table, SchemaRegistry.__table__) + connection.execute(table.delete().where(table.c.physical_schema == STAGING_SCHEMA)) diff --git a/tests/conftest.py b/tests/conftest.py index c8694d6..a358fe9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,8 +5,9 @@ import sqlalchemy.orm as so from dotenv import load_dotenv +from oa_configurator import ensure_schema from oa_configurator.testing import isolated_test_database -from orm_loader.backends import STAGING_SCHEMA +from orm_loader.backends import STAGING_SCHEMA, staging_schema_claim from orm_loader.config import OrmLoaderConfig from tests.models import Base @@ -16,7 +17,7 @@ @pytest.fixture def engine(): with isolated_test_database( - OrmLoaderConfig, "test_orm_db_sqlite", dialect="sqlite", future=True, + OrmLoaderConfig, "test_orm_db_sqlite", dialect="sqlite", ) as db: engine = db.connection.engine Base.metadata.create_all(engine) @@ -39,7 +40,9 @@ def pg_db(request): ``pg_db.connection``/``pg_db.session`` happens inside one transaction that's rolled back on exit, so concurrent test runs can't collide and nothing needs manual cleanup.""" - with isolated_test_database(OrmLoaderConfig, "test_orm_db_pg", request=request) as db: + with isolated_test_database( + OrmLoaderConfig, "test_orm_db_pg", request=request, schema_claims=[staging_schema_claim()], + ) as db: yield db @@ -49,7 +52,7 @@ def pg_session(pg_db): 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}")) + ensure_schema(conn, STAGING_SCHEMA) Base.metadata.create_all(conn) return pg_db.session diff --git a/tests/mappers/test_materialised_view_mixin.py b/tests/mappers/test_materialised_view_mixin.py index 4d8a752..a43408d 100644 --- a/tests/mappers/test_materialised_view_mixin.py +++ b/tests/mappers/test_materialised_view_mixin.py @@ -49,6 +49,64 @@ def test_refresh_all_mvs_resolves_each_views_own_schema_tag(pg_db) -> None: assert not inspector.has_table("mv_vocab_role_test", schema=primary_schema) +class _IndexedPgMv(MaterializedViewMixin): + __mv_name__ = "mv_indexed_pg_test" + __mv_select__ = sa.select(sa.literal(1).label("row_id")) + __mv_indexes__ = (MaterializedViewIndex(name="mv_indexed_pg_test_row_id_uq", columns=("row_id",), unique=True),) + + +def test_create_mv_with_declared_index_against_postgres(pg_db) -> None: + """create_mv() through the mixin creates both the view and its declared + index against a real database.test_create_mv_creates_declared_indexes_after_the_view + only proves the right backend calls are made, via a fake backend.""" + engine = pg_db.connection.engine + + with isolated_test_schema(engine, prefix="mv_indexed") as schema: + scoped = engine.execution_options(schema_translate_map={Role.PRIMARY.value: schema}) + with scoped.begin() as conn: + _IndexedPgMv.create_mv(conn) + + with engine.connect() as conn: + inspector = sa.inspect(conn) + assert inspector.has_table("mv_indexed_pg_test", schema=schema) + index_names = {idx["name"] for idx in inspector.get_indexes("mv_indexed_pg_test", schema=schema)} + assert "mv_indexed_pg_test_row_id_uq" in index_names + + +def test_drop_mv_cascade_against_postgres(pg_db) -> None: + """drop_mv() through the mixin actually drops the view against a real + database. test_drop_mv_forwards_default_args only proves the right + backend call is made, via a fake backend.""" + engine = pg_db.connection.engine + + with isolated_test_schema(engine, prefix="mv_drop") as schema: + scoped = engine.execution_options(schema_translate_map={Role.PRIMARY.value: schema}) + with scoped.begin() as conn: + _PrimaryRoleMV.create_mv(conn) + with engine.connect() as conn: + assert sa.inspect(conn).has_table("mv_primary_role_test", schema=schema) + + with scoped.begin() as conn: + _PrimaryRoleMV.drop_mv(conn, cascade=True) + with engine.connect() as conn: + assert not sa.inspect(conn).has_table("mv_primary_role_test", schema=schema) + + +def test_refresh_mv_concurrently_against_postgres(pg_db) -> None: + """refresh_mv(concurrently=True) through the mixin succeeds against a + real database when a unique index is declared. + The eligibility check and error translation are already proven at the + backend level directly in test_postgres_backend.py. + This proves the mixin wires into that correctly end to end.""" + engine = pg_db.connection.engine + + with isolated_test_schema(engine, prefix="mv_refresh") as schema: + scoped = engine.execution_options(schema_translate_map={Role.PRIMARY.value: schema}) + with scoped.begin() as conn: + _IndexedPgMv.create_mv(conn) + _IndexedPgMv.refresh_mv(conn, concurrently=True) + + class _FakeBackend: """Capture lifecycle calls without touching a real database.""" @@ -106,7 +164,8 @@ def test_create_mv_forwards_default_args_to_backend(fake_backend: _FakeBackend, ( "create_materialized_view", (bind, "mv_no_index", _SELECT), - {"schema": "primary", "with_data": True, "if_not_exists": True}, + # sqlite has no schema concept, so physical_schema_of() always folds to None here. + {"schema": None, "with_data": True, "if_not_exists": True}, ) ] @@ -126,7 +185,8 @@ def test_create_mv_creates_declared_indexes_after_the_view(fake_backend: _FakeBa "create_materialized_view_index", ] assert fake_backend.calls[1][1] == (bind, "mv_indexed", _INDEX) - assert fake_backend.calls[1][2] == {"schema": "primary", "if_not_exists": True} + # sqlite has no schema concept, so physical_schema_of() always folds to None here. + assert fake_backend.calls[1][2] == {"schema": None, "if_not_exists": True} def test_create_mv_create_indexes_false_skips_index_creation(fake_backend: _FakeBackend, bind): @@ -140,8 +200,10 @@ def test_create_mv_forwards_schema_tag_with_data_and_if_not_exists_overrides( ): _NoIndexMv.create_mv(bind, schema_tag=Role.VOCAB, with_data=False, if_not_exists=False) + # sqlite has no schema concept, so physical_schema_of() always folds to None here, + # regardless of which schema_tag was requested assert fake_backend.calls[0][2] == { - "schema": "vocab", + "schema": None, "with_data": False, "if_not_exists": False, } @@ -172,10 +234,19 @@ class _MappedNoSchemaMv(_MappedMvBase, MaterializedViewMixin): row_id = sa.Column(sa.Integer, primary_key=True) -def test_create_mv_defers_to_the_mapped_tables_own_schema(fake_backend: _FakeBackend, bind): +def test_create_mv_defers_to_the_mapped_tables_own_schema( + fake_backend: _FakeBackend, bind, monkeypatch: pytest.MonkeyPatch +): + # sqlite folds every schema to None regardless of tag, which would make "mapped + # table's own schema" and "mixin default" indistinguishable here. Patch + # physical_schema_of() to pass the resolved tag through unchanged, isolating this + # test to the mixin's own schema_tag routing. + mixin_module = importlib.import_module("orm_loader.mappers.materialised_view_mixin") + monkeypatch.setattr(mixin_module, "physical_schema_of", lambda bind, schema_tag: schema_tag) + _MappedVocabMv.create_mv(bind) - assert fake_backend.calls[0][2]["schema"] == "vocab" + assert fake_backend.calls[0][2]["schema"] == Role.VOCAB.value def test_create_mv_mapped_table_with_no_schema_resolves_to_none_not_primary( @@ -187,17 +258,21 @@ def test_create_mv_mapped_table_with_no_schema_resolves_to_none_not_primary( def test_create_mv_explicit_schema_tag_overrides_the_mapped_tables_own_schema( - fake_backend: _FakeBackend, bind + fake_backend: _FakeBackend, bind, monkeypatch: pytest.MonkeyPatch ): + # See test_create_mv_defers_to_the_mapped_tables_own_schema for why this is patched. + mixin_module = importlib.import_module("orm_loader.mappers.materialised_view_mixin") + monkeypatch.setattr(mixin_module, "physical_schema_of", lambda bind, schema_tag: schema_tag) + _MappedVocabMv.create_mv(bind, schema_tag=Role.PRIMARY) - assert fake_backend.calls[0][2]["schema"] == "primary" + assert fake_backend.calls[0][2]["schema"] == Role.VOCAB.value def test_create_mv_forwards_if_not_exists_to_declared_indexes(fake_backend: _FakeBackend, bind): _IndexedMv.create_mv(bind, if_not_exists=False) - assert fake_backend.calls[1][2] == {"schema": "primary", "if_not_exists": False} + assert fake_backend.calls[1][2] == {"schema": None, "if_not_exists": False} def test_create_mv_engine_uses_one_transaction_for_view_and_indexes(monkeypatch: pytest.MonkeyPatch): @@ -259,7 +334,7 @@ def test_refresh_mv_forwards_default_args_and_declared_indexes(fake_backend: _Fa ( "refresh_materialized_view", (bind, "mv_indexed"), - {"schema": "primary", "concurrently": False, "declared_indexes": (_INDEX,)}, + {"schema": None, "concurrently": False, "declared_indexes": (_INDEX,)}, ) ] @@ -268,7 +343,7 @@ def test_refresh_mv_forwards_schema_tag_and_concurrently(fake_backend: _FakeBack _IndexedMv.refresh_mv(bind, schema_tag=Role.VOCAB, concurrently=True) assert fake_backend.calls[0][2] == { - "schema": "vocab", + "schema": None, "concurrently": True, "declared_indexes": (_INDEX,), } @@ -281,7 +356,8 @@ def test_drop_mv_forwards_default_args(fake_backend: _FakeBackend, bind): ( "drop_materialized_view", (bind, "mv_no_index"), - {"schema": "primary", "if_exists": True, "cascade": False}, + # sqlite has no schema concept, so physical_schema_of() always folds to None here. + {"schema": None, "if_exists": True, "cascade": False}, ) ] @@ -290,7 +366,7 @@ def test_drop_mv_forwards_schema_tag_if_exists_and_cascade(fake_backend: _FakeBa _NoIndexMv.drop_mv(bind, schema_tag=Role.VOCAB, if_exists=False, cascade=True) assert fake_backend.calls[0][2] == { - "schema": "vocab", + "schema": None, "if_exists": False, "cascade": True, } From b533dac3ffc0d6a5ea97c645f41d72b4c5a87a80 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 2 Oct 2026 02:51:41 +0000 Subject: [PATCH 25/29] Updated tests --- tests/backends/test_reserved_schema.py | 33 +--- tests/conftest.py | 27 ++- tests/loaders/test_schema_translate_map.py | 186 ++++++++---------- tests/loaders/test_split_connection.py | 74 +++---- tests/mappers/test_materialised_view_mixin.py | 49 ++--- 5 files changed, 152 insertions(+), 217 deletions(-) diff --git a/tests/backends/test_reserved_schema.py b/tests/backends/test_reserved_schema.py index 91db0c5..79a9caa 100644 --- a/tests/backends/test_reserved_schema.py +++ b/tests/backends/test_reserved_schema.py @@ -19,16 +19,11 @@ from oa_configurator import ( CDMDatabaseConfig, - ConnectionConfig, GenericDatabaseConfig, Resolver, SchemaOwnershipError, - StackConfig, ) -from oa_configurator.domains.resources.schema_registry import SchemaRegistry -from sqlalchemy.engine import make_url -from sqlalchemy import Table -from typing import cast +from oa_configurator.testing import reset_schema_registry_rows import pytest from orm_loader.backends import STAGING_SCHEMA, staging_schema_claim @@ -36,26 +31,15 @@ pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] -def _connection_config(pg_db) -> ConnectionConfig: - connection = pg_db.resolved.connection - url = make_url(connection.url) - return ConnectionConfig( - dialect=url.drivername, host=url.host, port=url.port, - user=url.username, password=url.password, database_name=url.database, - test_only=connection.test_only, - ) - - -def test_resolving_cdm_database_with_staging_schema_name_raises(pg_db) -> None: - connection_config = _connection_config(pg_db) - stack = StackConfig.for_session( - connections={"c": connection_config}, +def test_resolving_cdm_database_with_staging_schema_name_raises(pg_db, cleanup_after_test) -> None: + reset_schema_registry_rows(cleanup_after_test, pg_db.committing_engine, [STAGING_SCHEMA]) + connection = pg_db.resolved.connection.name + resolver = Resolver.from_active_config().with_overrides( databases={ - "loader": GenericDatabaseConfig(connection="c"), - "default": CDMDatabaseConfig(connection="c", cdm_schema=STAGING_SCHEMA), + "loader": GenericDatabaseConfig(connection=connection), + "default": CDMDatabaseConfig(connection=connection, cdm_schema=STAGING_SCHEMA), }, ) - resolver = Resolver(stack) engine = resolver.resolve_database("loader").create_engine( schema_claims=[staging_schema_claim()] ) @@ -64,6 +48,3 @@ def test_resolving_cdm_database_with_staging_schema_name_raises(pg_db) -> None: resolver.resolve_database("default").create_engine() finally: engine.dispose() - with pg_db.committing_engine.begin() as connection: - table = cast(Table, SchemaRegistry.__table__) - connection.execute(table.delete().where(table.c.physical_schema == STAGING_SCHEMA)) diff --git a/tests/conftest.py b/tests/conftest.py index a358fe9..0711e20 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,5 @@ +from collections.abc import Iterable, Iterator +from contextlib import contextmanager from pathlib import Path import pytest @@ -5,8 +7,8 @@ import sqlalchemy.orm as so from dotenv import load_dotenv -from oa_configurator import ensure_schema -from oa_configurator.testing import isolated_test_database +from oa_configurator import Role, ensure_schema +from oa_configurator.testing import ScopedTestSchema, isolated_test_database, scoped_test_schema from orm_loader.backends import STAGING_SCHEMA, staging_schema_claim from orm_loader.config import OrmLoaderConfig from tests.models import Base @@ -57,12 +59,17 @@ def pg_session(pg_db): return pg_db.session +@contextmanager def schema_scoped_session( - conn: sa.Connection, table: sa.Table, schema_translate_map: dict -) -> so.Session: - """A Session scoped to *schema_translate_map*, with *table* already - created through it. - """ - scoped_conn = conn.execution_options(schema_translate_map=schema_translate_map) - table.create(scoped_conn, checkfirst=True) - return so.Session(bind=scoped_conn) + pg_db, table: sa.Table, *, prefix: str, split_roles: Iterable[Role] = () +) -> Iterator[tuple[ScopedTestSchema, so.Session]]: + """A Session on committed scoped test schemas, with the staging schema + claimed and *table* already created.""" + with scoped_test_schema( + pg_db.resolved, prefix=prefix, split_roles=split_roles, schema_claims=[staging_schema_claim()] + ) as scoped: + with scoped.engine.begin() as conn: + ensure_schema(conn, STAGING_SCHEMA) + table.create(conn, checkfirst=True) + with so.Session(scoped.engine) as session: + yield scoped, session diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index a843a6b..735610c 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -18,11 +18,8 @@ from __future__ import annotations -import uuid - import pandas as pd import sqlalchemy as sa -from oa_configurator import ensure_schema from oa_configurator import Role as SchemaRole from orm_loader.backends import STAGING_SCHEMA @@ -33,119 +30,90 @@ def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): - schema = f"test_schema_{uuid.uuid4().hex[:8]}" - conn = pg_db.connection - ensure_schema(conn, schema) - ensure_schema(conn, STAGING_SCHEMA) - - # Override the connection's default schema for this session only. This - # is the caller-side setup a real deployment does once at engine - # construction (ResolvedCDMDatabase.create_engine()), not a workaround - # threaded through load_csv() itself. - session = schema_scoped_session( - conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: schema} - ) - - csv_path = tmp_path / "test_table.csv" - pd.DataFrame( - [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}, {"id": 3, "name": "gamma"}] - ).to_csv(csv_path, index=False, sep="\t") - - inserted = SimpleTable.load_csv( - session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema=STAGING_SCHEMA - ) - session.commit() - - assert inserted == 3 - - # Read back through the schema-qualified name directly, not through - # schema_translate_map, to prove the rows are really there. - rows = conn.execute( - sa.text(f'SELECT id, name FROM "{schema}"."test_table" ORDER BY id') - ).fetchall() - assert rows == [(1, "alpha"), (2, "beta"), (3, "gamma")] - - # And that nothing leaked into the default/public schema. That was the - # exact failure mode the original bug caused: raw text() bypassing - # schema_translate_map, resolving through the connection's search_path - # instead. - leaked = conn.execute(sa.text("SELECT to_regclass('public.test_table')")).scalar() - assert leaked is None + with schema_scoped_session(pg_db, SimpleTable.__table__, prefix="loader_primary") as (scoped, session): + schema = scoped.schemas[SchemaRole.PRIMARY] + 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 + + with scoped.engine.connect() as conn: + # 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() + # Nothing leaked into public: raw text() bypassing + # schema_translate_map would resolve through search_path instead. + leaked = conn.execute(sa.text("SELECT to_regclass('public.test_table')")).scalar() + assert rows == [(1, "alpha"), (2, "beta"), (3, "gamma")] + assert leaked is None def test_replace_merge_respects_non_default_schema_end_to_end(pg_db, tmp_path): """A second load_csv() call with merge_strategy="replace" against the same non-default schema. Proves the merge path itself, not just the initial insert-if-empty fast path, qualifies correctly.""" - schema = f"test_schema_{uuid.uuid4().hex[:8]}" - conn = pg_db.connection - ensure_schema(conn, schema) - ensure_schema(conn, STAGING_SCHEMA) - - session = schema_scoped_session( - conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: schema} - ) - - def _write_and_load(rows: list[dict], path_name: str) -> int: - path = tmp_path / path_name - pd.DataFrame(rows).to_csv(path, index=False, sep="\t") - return SimpleTable.load_csv( - session, - path, - dedupe=False, - loader=PandasLoader(), - merge_strategy="replace", - staging_schema=STAGING_SCHEMA, - ) - - _write_and_load( - [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}], "test_table.csv" - ) - session.commit() - - _write_and_load([{"id": 1, "name": "alpha-updated"}], "test_table.csv") - session.commit() - - rows = conn.execute( - sa.text(f'SELECT id, name FROM "{schema}"."test_table" ORDER BY id') - ).fetchall() - assert rows == [(1, "alpha-updated"), (2, "beta")] + with schema_scoped_session(pg_db, SimpleTable.__table__, prefix="loader_replace") as (scoped, session): + schema = scoped.schemas[SchemaRole.PRIMARY] + + 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() + + with scoped.engine.connect() as conn: + rows = conn.execute( + sa.text(f'SELECT id, name FROM "{schema}"."test_table" ORDER BY id') + ).fetchall() + assert rows == [(1, "alpha-updated"), (2, "beta")] def test_load_csv_respects_non_primary_schema_tag_end_to_end(pg_db, tmp_path): """Checks that the schema tag is correctly derived from the table's own - schema (table.schema, via validate_schema_tag()) for a non-primary tag.""" - primary_schema = f"test_primary_{uuid.uuid4().hex[:8]}" - vocab_schema = f"test_vocab_{uuid.uuid4().hex[:8]}" - conn = pg_db.connection - ensure_schema(conn, primary_schema) - ensure_schema(conn, vocab_schema) - ensure_schema(conn, STAGING_SCHEMA) - - session = schema_scoped_session( - conn, - VocabSchemaTable.__table__, - {SchemaRole.PRIMARY.value: primary_schema, SchemaRole.VOCAB.value: vocab_schema}, - ) - - csv_path = tmp_path / "test_vocab_role_table.csv" - pd.DataFrame([{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]).to_csv( - csv_path, index=False, sep="\t" - ) - - inserted = VocabSchemaTable.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 + schema (table.schema) for a non-primary tag.""" + with schema_scoped_session( + pg_db, VocabSchemaTable.__table__, prefix="loader_vocab", split_roles=[SchemaRole.VOCAB] + ) as (scoped, session): + primary_schema = scoped.schemas[SchemaRole.PRIMARY] + vocab_schema = scoped.schemas[SchemaRole.VOCAB] + 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 = VocabSchemaTable.load_csv( + session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema=STAGING_SCHEMA + ) + session.commit() + + assert inserted == 2 + + with scoped.engine.connect() as conn: + rows = conn.execute( + sa.text(f'SELECT id, name FROM "{vocab_schema}"."test_vocab_role_table" ORDER BY id') + ).fetchall() + leaked = conn.execute( + sa.text(f"SELECT to_regclass('{primary_schema}.test_vocab_role_table')") + ).scalar() + assert rows == [(1, "alpha"), (2, "beta")] + assert leaked is None diff --git a/tests/loaders/test_split_connection.py b/tests/loaders/test_split_connection.py index 6207552..931b51a 100644 --- a/tests/loaders/test_split_connection.py +++ b/tests/loaders/test_split_connection.py @@ -3,14 +3,10 @@ from __future__ import annotations -import uuid - import pandas as pd import sqlalchemy as sa -from oa_configurator import ensure_schema from oa_configurator import Role as SchemaRole -from orm_loader.backends import STAGING_SCHEMA from orm_loader.loaders.loader_interface import PandasLoader from tests.conftest import schema_scoped_session @@ -22,41 +18,35 @@ def test_load_csv_across_two_genuinely_separate_connections(pg_db, session, tmp_ SQLite, vocab connection, via the module-level ``engine``/``session`` fixtures) are two entirely different engines against two entirely different database systems.""" - primary_schema = f"test_primary_{uuid.uuid4().hex[:8]}" - conn = pg_db.connection - ensure_schema(conn, primary_schema) - ensure_schema(conn, STAGING_SCHEMA) - primary_session = schema_scoped_session( - conn, SimpleTable.__table__, {SchemaRole.PRIMARY.value: primary_schema} - ) - - primary_csv = tmp_path / "test_table.csv" - pd.DataFrame([{"id": 1, "name": "primary-alpha"}]).to_csv(primary_csv, index=False, sep="\t") - - vocab_csv = tmp_path / "test_vocab_role_table.csv" - pd.DataFrame([{"id": 1, "name": "vocab-alpha"}]).to_csv(vocab_csv, index=False, sep="\t") - - # Interleaved on purpose: primary, then vocab, then primary again, so a - # module-level cache keyed wrong (or reused across calls) would surface - # as data landing in the wrong database. - SimpleTable.load_csv(primary_session, primary_csv, dedupe=False, loader=PandasLoader()) - primary_session.commit() - - VocabSchemaTable.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(VocabSchemaTable).order_by(VocabSchemaTable.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 + with schema_scoped_session(pg_db, SimpleTable.__table__, prefix="loader_split") as (scoped, primary_session): + primary_schema = scoped.schemas[SchemaRole.PRIMARY] + 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() + + VocabSchemaTable.load_csv(session, vocab_csv, dedupe=False, loader=PandasLoader()) + session.commit() + + with scoped.engine.connect() as conn: + primary_rows = conn.execute( + sa.text(f'SELECT id, name FROM "{primary_schema}"."test_table"') + ).fetchall() + # 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 primary_rows == [(1, "primary-alpha")] + assert leaked_vocab_table_in_pg is None + + vocab_rows = session.execute( + sa.select(VocabSchemaTable).order_by(VocabSchemaTable.id) + ).scalars().all() + assert [(r.id, r.name) for r in vocab_rows] == [(1, "vocab-alpha")] diff --git a/tests/mappers/test_materialised_view_mixin.py b/tests/mappers/test_materialised_view_mixin.py index a43408d..e4cf523 100644 --- a/tests/mappers/test_materialised_view_mixin.py +++ b/tests/mappers/test_materialised_view_mixin.py @@ -8,7 +8,7 @@ import sqlalchemy.orm as so from oa_configurator import Role -from oa_configurator.testing import isolated_test_schema +from oa_configurator.testing import scoped_test_schema from orm_loader.mappers.materialised_view_contracts import MaterializedViewIndex from orm_loader.mappers.materialised_view_mixin import ( MaterializedViewMixin, @@ -29,19 +29,15 @@ class _VocabRoleMV(MaterializedViewMixin): def test_refresh_all_mvs_resolves_each_views_own_schema_tag(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: + with scoped_test_schema(pg_db.resolved, prefix="mv", split_roles=[Role.VOCAB]) as scoped: + primary_schema = scoped.schemas[Role.PRIMARY] + vocab_schema = scoped.schemas[Role.VOCAB] + with scoped.engine.begin() as conn: _PrimaryRoleMV.create_mv(conn) _VocabRoleMV.create_mv(conn) refresh_all_mvs(conn, [_PrimaryRoleMV, _VocabRoleMV]) - with engine.connect() as conn: + with scoped.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) @@ -59,14 +55,12 @@ def test_create_mv_with_declared_index_against_postgres(pg_db) -> None: """create_mv() through the mixin creates both the view and its declared index against a real database.test_create_mv_creates_declared_indexes_after_the_view only proves the right backend calls are made, via a fake backend.""" - engine = pg_db.connection.engine - - with isolated_test_schema(engine, prefix="mv_indexed") as schema: - scoped = engine.execution_options(schema_translate_map={Role.PRIMARY.value: schema}) - with scoped.begin() as conn: + with scoped_test_schema(pg_db.resolved, prefix="mv_indexed") as scoped: + schema = scoped.schemas[Role.PRIMARY] + with scoped.engine.begin() as conn: _IndexedPgMv.create_mv(conn) - with engine.connect() as conn: + with scoped.engine.connect() as conn: inspector = sa.inspect(conn) assert inspector.has_table("mv_indexed_pg_test", schema=schema) index_names = {idx["name"] for idx in inspector.get_indexes("mv_indexed_pg_test", schema=schema)} @@ -77,18 +71,16 @@ def test_drop_mv_cascade_against_postgres(pg_db) -> None: """drop_mv() through the mixin actually drops the view against a real database. test_drop_mv_forwards_default_args only proves the right backend call is made, via a fake backend.""" - engine = pg_db.connection.engine - - with isolated_test_schema(engine, prefix="mv_drop") as schema: - scoped = engine.execution_options(schema_translate_map={Role.PRIMARY.value: schema}) - with scoped.begin() as conn: + with scoped_test_schema(pg_db.resolved, prefix="mv_drop") as scoped: + schema = scoped.schemas[Role.PRIMARY] + with scoped.engine.begin() as conn: _PrimaryRoleMV.create_mv(conn) - with engine.connect() as conn: + with scoped.engine.connect() as conn: assert sa.inspect(conn).has_table("mv_primary_role_test", schema=schema) - with scoped.begin() as conn: + with scoped.engine.begin() as conn: _PrimaryRoleMV.drop_mv(conn, cascade=True) - with engine.connect() as conn: + with scoped.engine.connect() as conn: assert not sa.inspect(conn).has_table("mv_primary_role_test", schema=schema) @@ -98,11 +90,8 @@ def test_refresh_mv_concurrently_against_postgres(pg_db) -> None: The eligibility check and error translation are already proven at the backend level directly in test_postgres_backend.py. This proves the mixin wires into that correctly end to end.""" - engine = pg_db.connection.engine - - with isolated_test_schema(engine, prefix="mv_refresh") as schema: - scoped = engine.execution_options(schema_translate_map={Role.PRIMARY.value: schema}) - with scoped.begin() as conn: + with scoped_test_schema(pg_db.resolved, prefix="mv_refresh") as scoped: + with scoped.engine.begin() as conn: _IndexedPgMv.create_mv(conn) _IndexedPgMv.refresh_mv(conn, concurrently=True) @@ -266,7 +255,7 @@ def test_create_mv_explicit_schema_tag_overrides_the_mapped_tables_own_schema( _MappedVocabMv.create_mv(bind, schema_tag=Role.PRIMARY) - assert fake_backend.calls[0][2]["schema"] == Role.VOCAB.value + assert fake_backend.calls[0][2]["schema"] == Role.PRIMARY.value def test_create_mv_forwards_if_not_exists_to_declared_indexes(fake_backend: _FakeBackend, bind): From 3a6ec134fa50cfccbcc03ac230a6695e57befa72 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 7 Oct 2026 00:35:33 +0000 Subject: [PATCH 26/29] Schema_tag routing as a result of private review, smaller helpers --- src/orm_loader/backends/base.py | 22 ++++--- src/orm_loader/backends/postgres.py | 25 +++----- src/orm_loader/backends/resolve.py | 41 +++++++++--- src/orm_loader/backends/sqlite.py | 20 +++--- src/orm_loader/helpers/bootstrap.py | 69 +++++++++++++++----- src/orm_loader/loaders/data_classes.py | 9 +-- src/orm_loader/tables/loadable_table.py | 73 +++++++++++----------- src/orm_loader/tables/typing.py | 12 ++-- tests/backends/test_postgres_backend.py | 3 +- tests/backends/test_reserved_schema.py | 10 ++- tests/backends/test_shared_backend.py | 21 ++++--- tests/conftest.py | 6 +- tests/loaders/test_pg_loader.py | 24 +++---- tests/loaders/test_schema_translate_map.py | 6 +- 14 files changed, 201 insertions(+), 140 deletions(-) diff --git a/src/orm_loader/backends/base.py b/src/orm_loader/backends/base.py index 2a8eac4..16b01df 100644 --- a/src/orm_loader/backends/base.py +++ b/src/orm_loader/backends/base.py @@ -99,20 +99,24 @@ class DatabaseBackend(ABC): without changing existing loader orchestration yet. """ - def __init__(self, staging_schema: str | None = None) -> None: + def __init__(self, *, staging_schema_tag: str | None = None, staging_schema: str | None = None) -> None: """ Parameters ---------- + staging_schema_tag + schema_translate_map tag the staging schema was resolved from. + ``None`` means no schema qualification was requested at all. + Stored so a caller re-resolving the backend later (e.g. after + the engine's claims may have changed) re-derives the same tag + rather than reusing an already-resolved physical name. staging_schema - Schema in which staging tables are created. ``None`` means no - schema qualification — staging tables land in whatever schema the - connection's search_path resolves to. Callers that want - schema-isolated staging can pass an explicit schema; backends must - not enable it implicitly because the schema may not have been - provisioned. ``STAGING_SCHEMA`` provides a shared convention for - callers that opt in. SQLite has no schema concept and always uses - ``None``. + Already-resolved physical schema staging tables are created in, + via ``resolve_backend()``. ``None`` means no schema + qualification — staging tables land in whatever schema the + connection's search_path resolves to. SQLite has no schema + concept and always resolves to ``None``. """ + self.staging_schema_tag = staging_schema_tag self.staging_schema = staging_schema @staticmethod diff --git a/src/orm_loader/backends/postgres.py b/src/orm_loader/backends/postgres.py index a089a3a..eddefd9 100644 --- a/src/orm_loader/backends/postgres.py +++ b/src/orm_loader/backends/postgres.py @@ -11,16 +11,14 @@ qualified, physical_schema_of, Dialect, - UnregisteredSchemaTagError, ) from sqlalchemy.dialects import postgresql from sqlalchemy.exc import OperationalError from sqlalchemy.sql.compiler import IdentifierPreparer from .base import ( - STAGING_SCHEMA, - BackendCapabilities, - DatabaseBackend, + BackendCapabilities, + DatabaseBackend, requires_capability ) from ..mappers.materialised_view_errors import ( @@ -41,8 +39,8 @@ class PostgresBackend(DatabaseBackend): - def __init__(self, *, staging_schema: str | None = None) -> None: - super().__init__(staging_schema=staging_schema) + def __init__(self, *, staging_schema_tag: str | None = None, staging_schema: str | None = None) -> None: + super().__init__(staging_schema_tag=staging_schema_tag, staging_schema=staging_schema) @staticmethod def staging_name_for_table(tablename: str) -> str: @@ -74,15 +72,6 @@ def create_staging_table( table_cls: type["CSVTableProtocol"], session: so.Session, ) -> None: - if self.staging_schema == STAGING_SCHEMA: - try: - physical_schema_of(session, schema_tag=STAGING_SCHEMA) - except UnregisteredSchemaTagError as exc: - raise UnregisteredSchemaTagError( - "orm-loader needs its staging schema reserved on this engine. Add " - "orm_loader.staging_schema_claim() to your own " - "create_engine(schema_claims=[...]) call." - ) from exc table = table_cls.__table__ preparer = self.identifier_preparer staging_ref = self.qualified_staging_name(table_cls.__tablename__) @@ -191,7 +180,7 @@ def merge_replace( merge_batch_size: int | None = None, ) -> None: target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) pk_join = sa.and_(*(target.c[c] == staging.c[c] for c in pk_cols)) non_paginated_replace = sa.delete(target).where(pk_join) @@ -227,7 +216,7 @@ def merge_upsert( merge_batch_size: int | None = None, ) -> None: target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls) def _upsert(select_: sa.sql.Select[Any]) -> sa.Insert: @@ -269,7 +258,7 @@ def merge_insert( merge_batch_size: int | None = None, ) -> None: target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls) non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) diff --git a/src/orm_loader/backends/resolve.py b/src/orm_loader/backends/resolve.py index 94ee0d2..c604e89 100644 --- a/src/orm_loader/backends/resolve.py +++ b/src/orm_loader/backends/resolve.py @@ -1,15 +1,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING import sqlalchemy.orm as so +from oa_configurator import Bindable, UnregisteredSchemaTagError, physical_schema_of + from .base import DatabaseBackend, Dialect from .postgres import PostgresBackend from .sqlite import SQLiteBackend -if TYPE_CHECKING: - from sqlalchemy.engine import Connection, Engine - _BACKEND_TYPES: dict[Dialect, type[DatabaseBackend]] = { Dialect.POSTGRESQL: PostgresBackend, @@ -17,7 +15,7 @@ } -def _dialect(bindable: "so.Session | Engine | Connection") -> Dialect: +def _dialect(bindable: Bindable) -> Dialect: if isinstance(bindable, so.Session): bind = bindable.get_bind() dialect_name = bind.dialect.name @@ -34,10 +32,37 @@ def _dialect(bindable: "so.Session | Engine | Connection") -> Dialect: ) from exc -def resolve_backend(bindable: "so.Session | Engine | Connection", **kwargs) -> DatabaseBackend: - """Resolve a concrete backend from a SQLAlchemy session, engine, or connection.""" +def resolve_backend( + bindable: Bindable, + *, + staging_schema_tag: str | None = None, + **kwargs, +) -> DatabaseBackend: + """Resolve a concrete backend from a SQLAlchemy session, engine, or connection. + + staging_schema_tag is resolved to its physical schema here, as it is the common + entry point for all backends. + + Raises + ------ + oa_configurator.UnregisteredSchemaTagError + If staging_schema_tag was never reserved on bindable via + ``create_engine(schema_claims=[staging_schema_claim()])``. + """ dialect = _dialect(bindable) + staging_schema: str | None = None + if staging_schema_tag is not None: + try: + staging_schema = physical_schema_of(bindable, schema_tag=staging_schema_tag) + except UnregisteredSchemaTagError as exc: + raise UnregisteredSchemaTagError( + "orm-loader needs its staging schema reserved on this engine. Add " + "orm_loader.staging_schema_claim() to your own " + "create_engine(schema_claims=[...]) call." + ) from exc try: - return _BACKEND_TYPES[dialect](**kwargs) + return _BACKEND_TYPES[dialect]( + staging_schema_tag=staging_schema_tag, staging_schema=staging_schema, **kwargs + ) except KeyError: raise NotImplementedError(f"No backend registered for dialect '{dialect.value}'") diff --git a/src/orm_loader/backends/sqlite.py b/src/orm_loader/backends/sqlite.py index b35ce19..b3db94e 100644 --- a/src/orm_loader/backends/sqlite.py +++ b/src/orm_loader/backends/sqlite.py @@ -13,8 +13,6 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.sql.compiler import IdentifierPreparer -from oa_configurator import schema_if_supported - from .base import BackendCapabilities, DatabaseBackend, Dialect if TYPE_CHECKING: @@ -37,18 +35,18 @@ def staging_name_for_table(tablename: str) -> str: def __init__( self, *, + staging_schema_tag: str | None = None, staging_schema: str | None = None, busy_timeout_ms: int = 60000, journal_mode: str = "WAL", defer_foreign_keys: bool = True, ) -> None: - if staging_schema is not None and schema_if_supported(staging_schema, Dialect.SQLITE) is None: - logger.warning( - "SQLite does not support schema-qualified staging tables; " - f"got staging_schema={staging_schema!r}. Setting staging_schema=None." + if staging_schema is not None: + raise ValueError( + "SQLite has no schema concept, so staging_schema must be None. " + "physical_schema_of() already resolves it to None for this dialect." ) - staging_schema = None - super().__init__(staging_schema=staging_schema) + super().__init__(staging_schema_tag=staging_schema_tag, staging_schema=staging_schema) self.busy_timeout_ms = busy_timeout_ms self.journal_mode = self._validate_journal_mode(journal_mode) self.defer_foreign_keys = defer_foreign_keys @@ -168,7 +166,7 @@ def merge_replace( merge_batch_size: int | None = None, ) -> None: target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) 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 @@ -205,7 +203,7 @@ def merge_upsert( merge_batch_size: int | None = None, ) -> None: target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls) def _upsert(select_: sa.Select[Any]) -> sa.Insert: @@ -245,7 +243,7 @@ def merge_insert( merge_batch_size: int | None = None, ) -> None: target = table_cls.__table__ - staging = table_cls.get_staging_table(session, staging_schema=self.staging_schema) + staging = table_cls.get_staging_table(session, staging_schema_tag=self.staging_schema_tag) insertable_cols = self._insertable_column_names(table_cls) non_paginated_select = sa.select(*(staging.c[c] for c in insertable_cols)) diff --git a/src/orm_loader/helpers/bootstrap.py b/src/orm_loader/helpers/bootstrap.py index c4d091a..e8bc5b6 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -1,11 +1,13 @@ -from contextlib import ExitStack import logging +from typing import cast import sqlalchemy as sa from oa_configurator import ( ResolvedCDMDatabase, ResolvedDatabase, - guard_schema_provenance_for, + UnregisteredSchemaTagError, + claimed_schema_tags, + declared_schema_tags, is_ephemeral_url, open_connection, ) @@ -14,10 +16,15 @@ logger = logging.getLogger(__name__) -Bindable = sa.engine.Engine | sa.engine.Connection +# Narrower than oa_configurator.Bindable: +# open_connection(), the one thing every bindable here is routed through, +# only accepts Engine | Connection. +_EngineOrConnection = sa.engine.Engine | sa.engine.Connection -def _resolve_binds(resolved: ResolvedDatabase, bindable: Bindable | None) -> tuple[Bindable, Bindable]: +def _resolve_binds( + resolved: ResolvedDatabase, bindable: _EngineOrConnection | None +) -> tuple[_EngineOrConnection, _EngineOrConnection]: """Return (primary_bind, vocab_bind) to run DDL against. Built from resolved.create_engine()/create_engines() when bindable is @@ -46,36 +53,66 @@ def _resolve_binds(resolved: ResolvedDatabase, bindable: Bindable | None) -> tup return engine, engine -def create_db(resolved: ResolvedDatabase, *, bindable: Bindable | None = None) -> None: +def create_db(resolved: ResolvedDatabase, *, bindable: _EngineOrConnection | None = None) -> None: + """Create every schema-tagged table in Base.metadata. + + Schema creation (CREATE SCHEMA) happens inside create_engine() when + this function builds its own engine(s). This function's only job is + metadata.create_all() for the tables themselves, grouped by schema + tag. Because the engine(s) used here are always freshly built for this + one call (or explicitly handed in), there's no window for schema drift + between the engine's creation and this call's create_all(), so no + separate provenance guard is needed here, unlike a long-lived engine + reused across many later calls. + """ logger.debug("Creating database schema") + owned = bindable is None primary_bind, vocab_bind = _resolve_binds(resolved, bindable) - tables_by_bind: dict[int, tuple[Bindable, dict[str, list[sa.Table]]]] = {} + tables_by_bind: dict[int, tuple[_EngineOrConnection, dict[str, list[sa.Table]]]] = {} + untagged: list[str] = [] for table in Base.metadata.tables.values(): schema_tag = table.schema if schema_tag is None: + untagged.append(table.name) continue bind = resolved.route_for_schema_tag(schema_tag, vocab=vocab_bind, primary=primary_bind) _, tables_by_tag = tables_by_bind.setdefault(id(bind), (bind, {})) tables_by_tag.setdefault(schema_tag, []).append(table) - with ExitStack() as guard_stack: + if untagged: + raise ValueError( + f"create_db() found table(s) with no schema tag (Table.schema is None): " + f"{sorted(untagged)}. Every table must declare its schema_tag; an untagged " + "table would otherwise be silently excluded from create_db()." + ) + + try: for bind, tables_by_tag in tables_by_bind.values(): - connection = guard_stack.enter_context(open_connection(bind)) - # One provenance guard per schema_tag on this connection; ExitStack defers every write until the block below succeeds. - for schema_tag in tables_by_tag: - guard_stack.enter_context( - guard_schema_provenance_for(connection, resolved, schema_tag=schema_tag) - ) - tables = [table for tables in tables_by_tag.values() for table in tables] - Base.metadata.create_all(connection, tables=tables, checkfirst=True) + with open_connection(bind) as connection: + tables = [table for tables in tables_by_tag.values() for table in tables] + missing = declared_schema_tags(tables) - claimed_schema_tags(connection) + if missing: + raise UnregisteredSchemaTagError( + f"create_db(): table(s) declare schema tag(s) {sorted(missing)} that " + "aren't claimed on this connection. Add them to your own " + "create_engine(schema_claims=[...]) call." + ) + Base.metadata.create_all(connection, tables=tables, checkfirst=True) + finally: + if owned: + # _resolve_binds() only returns Engine instances on this path + # (bindable is None): create_engine()/create_engines(). + cast(sa.engine.Engine, primary_bind).dispose() + if vocab_bind is not primary_bind: + cast(sa.engine.Engine, vocab_bind).dispose() def bootstrap( resolved: ResolvedDatabase, *, create: bool = True, - bindable: Bindable | None = None, + bindable: _EngineOrConnection | None = None, ) -> None: logger.info("Bootstrapping schema (create=%s)", create) if create: diff --git a/src/orm_loader/loaders/data_classes.py b/src/orm_loader/loaders/data_classes.py index f272c4f..1236c97 100644 --- a/src/orm_loader/loaders/data_classes.py +++ b/src/orm_loader/loaders/data_classes.py @@ -64,9 +64,10 @@ class LoaderContext: quote_mode Quoting mode, resolved to the same concrete mode on both the PostgreSQL COPY fast-path and the pandas ORM fallback. - staging_schema - Schema the staging table lives in, passed to resolve_backend() so - every backend resolution within this load shares the same schema. + staging_schema_tag + schema_translate_map tag the staging table lives in, passed to + resolve_backend() so every backend resolution within this load + resolves the same schema. """ tableclass: Type["CSVTableProtocol"] session: so.Session @@ -77,7 +78,7 @@ class LoaderContext: normalise: bool = True dedupe: bool = True quote_mode: str = "auto" - staging_schema: str | None = None + staging_schema_tag: str | None = None class LoaderInterface: diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 1278c60..64a8e08 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -73,7 +73,7 @@ def create_staging_table( cls: Type[CSVTableProtocol], session: so.Session, *, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ): """ Create a fresh staging table for ingestion. @@ -85,9 +85,10 @@ def create_staging_table( ---------- session An active SQLAlchemy session bound to an engine. - staging_schema - Schema the staging table should be created in. ``None`` means no - schema qualification (backend-default behavior). + staging_schema_tag + schema_translate_map tag the staging table should be created + under. ``None`` means no schema qualification (backend-default + behavior). Raises ------ @@ -97,7 +98,7 @@ def create_staging_table( If the database dialect is unsupported. """ _require_bind(session) - backend = resolve_backend(session, staging_schema=staging_schema) + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag) backend.create_staging_table(cls, session) @classmethod @@ -107,7 +108,7 @@ def manage_indices( session: so.Session, index_strategy: str = "auto", *, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ) -> Iterator[None]: """ Manage non-primary-key indexes around a staged merge. @@ -117,7 +118,7 @@ def manage_indices( moment SQLite keeps indexes by default, while PostgreSQL drops and rebuilds them. """ - backend = resolve_backend(session, staging_schema=staging_schema) + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag) resolved_index_strategy = backend.resolve_index_strategy(index_strategy) table_name = cls.__tablename__ @@ -218,7 +219,7 @@ def get_staging_table( cls: Type[CSVTableProtocol], session: so.Session, *, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ) -> sa.Table: """ Return the reflected staging table, creating it if necessary. @@ -227,9 +228,9 @@ def get_staging_table( ---------- session An active SQLAlchemy session bound to an engine. - staging_schema - Schema the staging table lives in. ``None`` means no schema - qualification (backend-default behavior). + staging_schema_tag + schema_translate_map tag the staging table lives in. ``None`` + means no schema qualification (backend-default behavior). Returns ------- @@ -249,13 +250,13 @@ def get_staging_table( since that call commits, which can invalidate an earlier reference. """ _require_bind(session) - backend = resolve_backend(session, staging_schema=staging_schema) + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag) 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): logger.debug(f"Staging table {staging_name} does not exist; recreating",) - cls.create_staging_table(session, staging_schema=staging_schema) + cls.create_staging_table(session, staging_schema_tag=staging_schema_tag) return sa.Table( staging_name, @@ -290,10 +291,10 @@ def load_staging( """ _require_bind(loader_context.session) - backend = resolve_backend(loader_context.session, staging_schema=loader_context.staging_schema) + backend = resolve_backend(loader_context.session, staging_schema_tag=loader_context.staging_schema_tag) total = 0 - cls.create_staging_table(loader_context.session, staging_schema=loader_context.staging_schema) + cls.create_staging_table(loader_context.session, staging_schema_tag=loader_context.staging_schema_tag) try: total = backend.load_staging_fast(loader_context=loader_context) @@ -362,7 +363,7 @@ def load_csv( quote_mode: str = "auto", index_strategy: str = "auto", merge_batch_size: int | None = None, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ) -> int: """ @@ -398,11 +399,11 @@ def load_csv( index_strategy Index handling strategy during merge. Use ``"auto"`` to let the backend choose a sensible default. - staging_schema - Schema the staging table lives in. ``None`` means no schema - qualification (backend-default behavior). Threaded through - every internal step of the load lifecycle so they all resolve - the same backend/schema. + staging_schema_tag + schema_translate_map tag the staging table lives in. ``None`` + means no schema qualification (backend-default behavior). + Threaded through every internal step of the load lifecycle so + they all resolve the same backend/schema. Returns ------- @@ -438,12 +439,12 @@ def load_csv( tableclass=cls, session=session, path=path, - staging_table=cls.get_staging_table(session, staging_schema=staging_schema), + staging_table=cls.get_staging_table(session, staging_schema_tag=staging_schema_tag), chunksize=chunksize, normalise=normalise, dedupe=dedupe, quote_mode=quote_mode, - staging_schema=staging_schema, + staging_schema_tag=staging_schema_tag, ) if loader is None: @@ -455,15 +456,15 @@ def load_csv( # Merge staging to target (Wrapped in our index dropper!) logger.info(f"Table `{cls.__tablename__}`: Merging staging data into target table") - with cls.manage_indices(session, index_strategy=index_strategy, staging_schema=staging_schema): + with cls.manage_indices(session, index_strategy=index_strategy, staging_schema_tag=staging_schema_tag): cls.merge_from_staging( session, merge_strategy=merge_strategy, merge_batch_size=merge_batch_size, - staging_schema=staging_schema, + staging_schema_tag=staging_schema_tag, ) - cls.drop_staging_table(session, staging_schema=staging_schema) + cls.drop_staging_table(session, staging_schema_tag=staging_schema_tag) logger.info(f"Table `{cls.__tablename__}`: Successfully finished ingestion. Total rows: {total}") return total @@ -490,7 +491,7 @@ def merge_from_staging( merge_strategy: str = "replace", *, merge_batch_size: int | None = None, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ): """ Merge data from the staging table into the target table. @@ -502,15 +503,15 @@ def merge_from_staging( merge_strategy Merge strategy to apply (for example ``replace``, ``upsert``, or ``insert_if_empty``). - staging_schema - Schema the staging table lives in. ``None`` means no schema - qualification (backend-default behavior). + staging_schema_tag + schema_translate_map tag the staging table lives in. ``None`` + means no schema qualification (backend-default behavior). """ target = cls.__tablename__ pk_cols = cls.pk_names() _require_bind(session) - backend = resolve_backend(session, staging_schema=staging_schema) + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag) target_empty_confirmed = False if merge_strategy in {"replace", "upsert"}: logger.info( @@ -584,7 +585,7 @@ def drop_staging_table( cls: Type[CSVTableProtocol], session: so.Session, *, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ): """ Drop the staging table if it exists. @@ -593,11 +594,11 @@ def drop_staging_table( ---------- session An active SQLAlchemy session bound to an engine. - staging_schema - Schema the staging table lives in. ``None`` means no schema - qualification (backend-default behavior). + staging_schema_tag + schema_translate_map tag the staging table lives in. ``None`` + means no schema qualification (backend-default behavior). """ - backend = resolve_backend(session, staging_schema=staging_schema) + backend = resolve_backend(session, staging_schema_tag=staging_schema_tag) backend.drop_staging_table(cls, session) @classmethod diff --git a/src/orm_loader/tables/typing.py b/src/orm_loader/tables/typing.py index 53dd0b6..b1a3b14 100644 --- a/src/orm_loader/tables/typing.py +++ b/src/orm_loader/tables/typing.py @@ -67,7 +67,7 @@ class CSVTableProtocol(ORMTableProtocol, Protocol): def _select_loader(cls, path: Path) -> "LoaderInterface": ... @classmethod - def create_staging_table(cls, session: so.Session, *, staging_schema: str | None = None) -> None: ... + def create_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> None: ... @classmethod def load_staging(cls: Type["CSVTableProtocol"], loader: "LoaderInterface", loader_context: "LoaderContext") -> int: ... @@ -86,14 +86,14 @@ def load_csv( quote_mode: str = "csv", index_strategy: str = "auto", merge_batch_size: int | None = None, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ) -> int: ... @classmethod def orm_staging_load(cls, loader: "LoaderInterface", loader_context: "LoaderContext") -> int: ... @classmethod - def get_staging_table(cls, session: so.Session, *, staging_schema: str | None = None) -> sa.Table: ... + def get_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> sa.Table: ... @classmethod def merge_from_staging( @@ -102,18 +102,18 @@ def merge_from_staging( merge_strategy: str = "replace", *, merge_batch_size: int | None = None, - staging_schema: str | None = None, + staging_schema_tag: str | None = None, ) -> None: ... @classmethod - def drop_staging_table(cls, session: so.Session, *, staging_schema: str | None = None) -> None: ... + def drop_staging_table(cls, session: so.Session, *, staging_schema_tag: str | None = None) -> None: ... @classmethod def _target_has_rows(cls, session: so.Session) -> bool: ... @classmethod def manage_indices( - cls, session: so.Session, index_strategy: str = "auto", *, staging_schema: str | None = None + cls, session: so.Session, index_strategy: str = "auto", *, staging_schema_tag: str | None = None ) -> AbstractContextManager[None]: ... diff --git a/tests/backends/test_postgres_backend.py b/tests/backends/test_postgres_backend.py index bf77a45..8bc5fe6 100644 --- a/tests/backends/test_postgres_backend.py +++ b/tests/backends/test_postgres_backend.py @@ -12,6 +12,7 @@ from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role, qualified from oa_configurator.testing import isolated_test_schema from orm_loader.backends import STAGING_SCHEMA, Dialect, PostgresBackend +from orm_loader.backends.resolve import resolve_backend from tests.models import ComputedColumnTable _TARGET_TABLE = ComputedColumnTable.__tablename__ @@ -100,7 +101,7 @@ def test_postgres_backend_default_staging_schema_is_none(): def test_postgres_backend_create_staging_table_drops_computed_columns(pg_session): - backend = PostgresBackend(staging_schema=STAGING_SCHEMA) + backend = resolve_backend(pg_session, staging_schema_tag=STAGING_SCHEMA) backend.create_staging_table(_ComputedTableCls, pg_session) diff --git a/tests/backends/test_reserved_schema.py b/tests/backends/test_reserved_schema.py index 79a9caa..4f76560 100644 --- a/tests/backends/test_reserved_schema.py +++ b/tests/backends/test_reserved_schema.py @@ -32,12 +32,18 @@ def test_resolving_cdm_database_with_staging_schema_name_raises(pg_db, cleanup_after_test) -> None: + """ + Notes + ----- + Overrides pg_db's own CDM entry in place, since two + CDM entries sharing one physical connection is itself unconfigurable.""" reset_schema_registry_rows(cleanup_after_test, pg_db.committing_engine, [STAGING_SCHEMA]) connection = pg_db.resolved.connection.name + cdm_name = pg_db.resolved.name resolver = Resolver.from_active_config().with_overrides( databases={ "loader": GenericDatabaseConfig(connection=connection), - "default": CDMDatabaseConfig(connection=connection, cdm_schema=STAGING_SCHEMA), + cdm_name: CDMDatabaseConfig(connection=connection, cdm_schema=STAGING_SCHEMA), }, ) engine = resolver.resolve_database("loader").create_engine( @@ -45,6 +51,6 @@ def test_resolving_cdm_database_with_staging_schema_name_raises(pg_db, cleanup_a ) try: with pytest.raises(SchemaOwnershipError, match=f"{STAGING_SCHEMA!r}.*orm_loader"): - resolver.resolve_database("default").create_engine() + resolver.resolve_database(cdm_name).create_engine() finally: engine.dispose() diff --git a/tests/backends/test_shared_backend.py b/tests/backends/test_shared_backend.py index 381160a..a100339 100644 --- a/tests/backends/test_shared_backend.py +++ b/tests/backends/test_shared_backend.py @@ -6,7 +6,8 @@ import sqlalchemy as sa from oa_configurator.testing import DIALECT_PARAMS -from orm_loader.backends import STAGING_SCHEMA, DatabaseBackend, PostgresBackend, SQLiteBackend +from orm_loader.backends import STAGING_SCHEMA, DatabaseBackend +from orm_loader.backends.resolve import resolve_backend from tests.models import ComputedColumnTable, CompositeTable if TYPE_CHECKING: @@ -32,15 +33,15 @@ def merge_backend(request: pytest.FixtureRequest) -> tuple[DatabaseBackend, "so. """ if request.param == "postgresql": session = request.getfixturevalue("pg_session") - return PostgresBackend(staging_schema=STAGING_SCHEMA), session + return resolve_backend(session, staging_schema_tag=STAGING_SCHEMA), session session = request.getfixturevalue("session") - return SQLiteBackend(), session + return resolve_backend(session), 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) + staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) session.execute( sa.insert(ComputedColumnTable), @@ -57,7 +58,7 @@ def test_merge_replace_single_pk(merge_backend: tuple[DatabaseBackend, "so.Sessi 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) + staging = _CompositeTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) session.execute( sa.insert(CompositeTable), @@ -74,7 +75,7 @@ def test_merge_replace_composite_pk(merge_backend: tuple[DatabaseBackend, "so.Se 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) + staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) session.execute(sa.insert(staging), [{"id": 1, "name": "alpha"}]) backend.merge_insert(_ComputedTableCls, session) @@ -86,7 +87,7 @@ def test_merge_insert_excludes_computed_columns(merge_backend: tuple[DatabaseBac 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) + staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) session.execute(sa.insert(ComputedColumnTable), [{"id": 1, "name": "existing"}]) session.execute( @@ -102,7 +103,7 @@ def test_merge_upsert_excludes_computed_columns(merge_backend: tuple[DatabaseBac 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) + staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) session.execute( sa.insert(ComputedColumnTable), [{"id": i, "name": f"orig{i}"} for i in range(10)] @@ -120,7 +121,7 @@ def test_merge_replace_paginated_path(merge_backend: tuple[DatabaseBackend, "so. 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) + staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) 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) @@ -132,7 +133,7 @@ def test_merge_insert_paginated_path(merge_backend: tuple[DatabaseBackend, "so.S 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) + staging = _ComputedTableCls.get_staging_table(session, staging_schema_tag=backend.staging_schema_tag) session.execute( sa.insert(ComputedColumnTable), [{"id": i, "name": "kept"} for i in range(5)] diff --git a/tests/conftest.py b/tests/conftest.py index 0711e20..af3c066 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,9 +7,9 @@ import sqlalchemy.orm as so from dotenv import load_dotenv -from oa_configurator import Role, ensure_schema +from oa_configurator import Role from oa_configurator.testing import ScopedTestSchema, isolated_test_database, scoped_test_schema -from orm_loader.backends import STAGING_SCHEMA, staging_schema_claim +from orm_loader.backends import staging_schema_claim from orm_loader.config import OrmLoaderConfig from tests.models import Base @@ -54,7 +54,6 @@ def pg_session(pg_db): 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 - ensure_schema(conn, STAGING_SCHEMA) Base.metadata.create_all(conn) return pg_db.session @@ -69,7 +68,6 @@ def schema_scoped_session( pg_db.resolved, prefix=prefix, split_roles=split_roles, schema_claims=[staging_schema_claim()] ) as scoped: with scoped.engine.begin() as conn: - ensure_schema(conn, STAGING_SCHEMA) table.create(conn, checkfirst=True) with so.Session(scoped.engine) as session: yield scoped, session diff --git a/tests/loaders/test_pg_loader.py b/tests/loaders/test_pg_loader.py index abafcc7..20f37cf 100644 --- a/tests/loaders/test_pg_loader.py +++ b/tests/loaders/test_pg_loader.py @@ -13,7 +13,7 @@ def test_copy_into_staging_with_extra_identity_column(pg_session, tmp_path): csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]).to_csv(csv, index=False) - SimpleTable.create_staging_table(pg_session, staging_schema=STAGING_SCHEMA) + SimpleTable.create_staging_table(pg_session, staging_schema_tag=STAGING_SCHEMA) staging_name = resolve_backend(pg_session).staging_name_for_table(SimpleTable.__tablename__) cols = pg_session.execute(sa.text( @@ -41,7 +41,7 @@ def test_copy_and_orm_path_equivalence(pg_session, tmp_path): ] ).to_csv(csv, index=False, sep="\t") - SimpleTable.load_csv(pg_session, csv, staging_schema=STAGING_SCHEMA) + SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) pg_session.commit() rows = pg_session.execute(sa.select(SimpleTable).order_by(SimpleTable.id)).scalars().all() @@ -56,7 +56,7 @@ 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) - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema=STAGING_SCHEMA) + inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) pg_session.commit() assert inserted == 1 @@ -74,7 +74,7 @@ def fake_quick_load_pg(*args, **kwargs): import orm_loader.loaders.loading_helpers as loading_helpers monkeypatch.setattr(loading_helpers, "quick_load_pg", fake_quick_load_pg) - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema=STAGING_SCHEMA) + inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) pg_session.commit() assert called["copy"] is True @@ -94,7 +94,7 @@ def broken_copy(*args, **kwargs): monkeypatch.setattr(loading_helpers, "quick_load_pg", broken_copy) - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema=STAGING_SCHEMA) + inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) pg_session.commit() rows = pg_session.execute(sa.select(SimpleTable)).scalars().all() @@ -106,12 +106,12 @@ def test_postgres_upsert_does_not_update(pg_session, tmp_path): csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - SimpleTable.load_csv(pg_session, csv, staging_schema=STAGING_SCHEMA) + SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) pg_session.commit() pd.DataFrame([{"id": 1, "name": "alpha_updated"}]).to_csv(csv, index=False) - SimpleTable.load_csv(pg_session, csv, merge_strategy="upsert", staging_schema=STAGING_SCHEMA) + SimpleTable.load_csv(pg_session, csv, merge_strategy="upsert", staging_schema_tag=STAGING_SCHEMA) pg_session.commit() rows = pg_session.execute(sa.select(SimpleTable)).scalars().all() @@ -132,7 +132,7 @@ def test_postgres_insert_if_empty(pg_session, tmp_path): pg_session, csv, merge_strategy="insert_if_empty", - staging_schema=STAGING_SCHEMA, + staging_schema_tag=STAGING_SCHEMA, ) pg_session.commit() @@ -149,7 +149,7 @@ def test_postgres_insert_if_empty_raises_on_non_empty_target(pg_session, tmp_pat csv = tmp_path / "test_table.csv" pd.DataFrame([{"id": 1, "name": "alpha"}]).to_csv(csv, index=False) - SimpleTable.load_csv(pg_session, csv, staging_schema=STAGING_SCHEMA) + SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) pg_session.commit() pd.DataFrame([{"id": 2, "name": "beta"}]).to_csv(csv, index=False) @@ -159,7 +159,7 @@ def test_postgres_insert_if_empty_raises_on_non_empty_target(pg_session, tmp_pat pg_session, csv, merge_strategy="insert_if_empty", - staging_schema=STAGING_SCHEMA, + staging_schema_tag=STAGING_SCHEMA, ) @@ -171,7 +171,7 @@ def test_postgres_copy_large_batch(pg_session, tmp_path): ) df.to_csv(csv, index=False) - inserted = SimpleTable.load_csv(pg_session, csv, staging_schema=STAGING_SCHEMA) + inserted = SimpleTable.load_csv(pg_session, csv, staging_schema_tag=STAGING_SCHEMA) pg_session.commit() count = pg_session.execute(sa.text('SELECT COUNT(*) FROM test_table')).scalar() @@ -183,7 +183,7 @@ 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) - SimpleTable.create_staging_table(pg_session, staging_schema=STAGING_SCHEMA) + SimpleTable.create_staging_table(pg_session, staging_schema_tag=STAGING_SCHEMA) staging_name = resolve_backend(pg_session).staging_name_for_table(SimpleTable.__tablename__) cols = pg_session.execute(sa.text(""" diff --git a/tests/loaders/test_schema_translate_map.py b/tests/loaders/test_schema_translate_map.py index 735610c..9d6a15b 100644 --- a/tests/loaders/test_schema_translate_map.py +++ b/tests/loaders/test_schema_translate_map.py @@ -38,7 +38,7 @@ def test_load_csv_respects_non_default_schema_end_to_end(pg_db, tmp_path): ).to_csv(csv_path, index=False, sep="\t") inserted = SimpleTable.load_csv( - session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema=STAGING_SCHEMA + session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema_tag=STAGING_SCHEMA ) session.commit() @@ -73,7 +73,7 @@ def _write_and_load(rows: list[dict], path_name: str) -> int: dedupe=False, loader=PandasLoader(), merge_strategy="replace", - staging_schema=STAGING_SCHEMA, + staging_schema_tag=STAGING_SCHEMA, ) _write_and_load([{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}], "test_table.csv") @@ -102,7 +102,7 @@ def test_load_csv_respects_non_primary_schema_tag_end_to_end(pg_db, tmp_path): ) inserted = VocabSchemaTable.load_csv( - session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema=STAGING_SCHEMA + session, csv_path, dedupe=False, loader=PandasLoader(), staging_schema_tag=STAGING_SCHEMA ) session.commit() From 0b2fbb82532264dca7f492e0197ac3ef1b25b9ea Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 7 Oct 2026 00:40:32 +0000 Subject: [PATCH 27/29] Remove unncessary bind --- src/orm_loader/helpers/bootstrap.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/src/orm_loader/helpers/bootstrap.py b/src/orm_loader/helpers/bootstrap.py index e8bc5b6..5764939 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -3,6 +3,7 @@ import sqlalchemy as sa from oa_configurator import ( + Bindable, ResolvedCDMDatabase, ResolvedDatabase, UnregisteredSchemaTagError, @@ -16,15 +17,8 @@ logger = logging.getLogger(__name__) -# Narrower than oa_configurator.Bindable: -# open_connection(), the one thing every bindable here is routed through, -# only accepts Engine | Connection. -_EngineOrConnection = sa.engine.Engine | sa.engine.Connection - -def _resolve_binds( - resolved: ResolvedDatabase, bindable: _EngineOrConnection | None -) -> tuple[_EngineOrConnection, _EngineOrConnection]: +def _resolve_binds(resolved: ResolvedDatabase, bindable: Bindable | None) -> tuple[Bindable, Bindable]: """Return (primary_bind, vocab_bind) to run DDL against. Built from resolved.create_engine()/create_engines() when bindable is @@ -53,7 +47,7 @@ def _resolve_binds( return engine, engine -def create_db(resolved: ResolvedDatabase, *, bindable: _EngineOrConnection | None = None) -> None: +def create_db(resolved: ResolvedDatabase, *, bindable: Bindable | None = None) -> None: """Create every schema-tagged table in Base.metadata. Schema creation (CREATE SCHEMA) happens inside create_engine() when @@ -69,7 +63,7 @@ def create_db(resolved: ResolvedDatabase, *, bindable: _EngineOrConnection | Non owned = bindable is None primary_bind, vocab_bind = _resolve_binds(resolved, bindable) - tables_by_bind: dict[int, tuple[_EngineOrConnection, dict[str, list[sa.Table]]]] = {} + tables_by_bind: dict[int, tuple[Bindable, dict[str, list[sa.Table]]]] = {} untagged: list[str] = [] for table in Base.metadata.tables.values(): schema_tag = table.schema @@ -112,7 +106,7 @@ def bootstrap( resolved: ResolvedDatabase, *, create: bool = True, - bindable: _EngineOrConnection | None = None, + bindable: Bindable | None = None, ) -> None: logger.info("Bootstrapping schema (create=%s)", create) if create: From 5f5d5af6ff416a2a7f39030e75e5c90dbf93f253 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 7 Oct 2026 05:21:17 +0000 Subject: [PATCH 28/29] Polishing of docstrings --- src/orm_loader/helpers/bootstrap.py | 15 ++++++++++++--- tests/models.py | 5 +---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/src/orm_loader/helpers/bootstrap.py b/src/orm_loader/helpers/bootstrap.py index 5764939..290d72d 100644 --- a/src/orm_loader/helpers/bootstrap.py +++ b/src/orm_loader/helpers/bootstrap.py @@ -53,13 +53,22 @@ def create_db(resolved: ResolvedDatabase, *, bindable: Bindable | None = None) - Schema creation (CREATE SCHEMA) happens inside create_engine() when this function builds its own engine(s). This function's only job is metadata.create_all() for the tables themselves, grouped by schema - tag. Because the engine(s) used here are always freshly built for this + tag. + + Notes + ----- + Because the engine(s) used here are always freshly built for this one call (or explicitly handed in), there's no window for schema drift between the engine's creation and this call's create_all(), so no - separate provenance guard is needed here, unlike a long-lived engine - reused across many later calls. + separate provenance guard is needed here. """ logger.debug("Creating database schema") + if not Base.metadata.tables: + logger.warning( + "create_db(): Base.metadata declares no tables. If this is unexpected, " + "the module(s) defining your ORM models likely weren't imported yet. " + "SQLAlchemy only registers a model on Base when its module runs." + ) owned = bindable is None primary_bind, vocab_bind = _resolve_binds(resolved, bindable) diff --git a/tests/models.py b/tests/models.py index 03a1645..efc3b17 100644 --- a/tests/models.py +++ b/tests/models.py @@ -98,10 +98,7 @@ class ImpliedEnumTable(Base, CSVLoadableTableInterface): class VocabSchemaTable(Base, CSVLoadableTableInterface): """A VOCAB-tagged table, so tests can prove the staging/index schema - derivation (table.schema resolved via validate_schema_tag()/physical_schema_of(), - threaded through create_staging_table()/manage_indices()) actually - resolves a non-primary schema tag correctly, instead of only ever - exercising the PRIMARY-tagged default.""" + derivation actually resolves a non-primary schema tag correctly.""" __tablename__ = "test_vocab_role_table" __table_args__ = ( From 48d5d2ec69d7b660ae3fe725fa1b88dc4f51aa03 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 7 Oct 2026 21:47:54 +0000 Subject: [PATCH 29/29] Fix ty --- src/orm_loader/helpers/bulk.py | 6 +++--- src/orm_loader/tables/loadable_table.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/orm_loader/helpers/bulk.py b/src/orm_loader/helpers/bulk.py index a74824b..459ccf5 100644 --- a/src/orm_loader/helpers/bulk.py +++ b/src/orm_loader/helpers/bulk.py @@ -2,7 +2,7 @@ from contextlib import contextmanager from sqlalchemy import Engine from sqlalchemy.orm import Session -from typing import Iterator +from typing import Generator logger = logging.getLogger(__name__) @@ -38,7 +38,7 @@ def bulk_load_context( *, disable_fk: bool = True, no_autoflush: bool = True, -) -> Iterator[None]: +) -> Generator[None, None, None]: """ Wrap a trusted bulk operation in backend-aware session settings. @@ -56,7 +56,7 @@ def bulk_load_context( @contextmanager -def engine_with_replica_role(engine: Engine) -> Iterator[Engine]: +def engine_with_replica_role(engine: Engine) -> Generator[Engine, None, None]: """ Force ``session_replication_role=replica`` on PostgreSQL engine connections. diff --git a/src/orm_loader/tables/loadable_table.py b/src/orm_loader/tables/loadable_table.py index 64a8e08..bcbcb04 100644 --- a/src/orm_loader/tables/loadable_table.py +++ b/src/orm_loader/tables/loadable_table.py @@ -6,7 +6,7 @@ from sqlalchemy.exc import InvalidRequestError, UnboundExecutionError -from typing import Type, Any, Iterator +from typing import Type, Any, Generator from pathlib import Path from contextlib import contextmanager from time import perf_counter @@ -109,7 +109,7 @@ def manage_indices( index_strategy: str = "auto", *, staging_schema_tag: str | None = None, - ) -> Iterator[None]: + ) -> Generator[None, None, None]: """ Manage non-primary-key indexes around a staged merge.