diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index dff286e..11c0a88 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -9,7 +9,12 @@ import typer import sqlalchemy.orm as so -from oa_configurator import ResolvedCDMDatabase, open_connection +from oa_configurator import ( + ResolvedCDMDatabase, + claimed_schema_tags, + open_connection, + physical_schema_of, +) from orm_loader.backends import STAGING_SCHEMA, resolve_backend from orm_loader.helpers.metadata import Base @@ -20,7 +25,7 @@ from omop_graph.extensions.omop_alchemy import ( RelationshipClass, RelationshipMapping, - create_extension_tables, + create_auxiliary_tables, ) from omop_graph.cli_utils import populate_test_data @@ -134,7 +139,7 @@ def _load_relationship_classification( engine: sa.Engine, relationship_ids: set[str], ) -> None: - """Body of :func:`relationship_classification`, on *engine*, which hosts the extension tables.""" + """Body of :func:`relationship_classification`, on *engine*, which hosts the auxiliary tables.""" pred_class_dir_pl = ( Path(pred_class_dir) if pred_class_dir else packaged_predicate_csv_dir() ) @@ -206,7 +211,19 @@ def _load_relationship_classification( subset=["relationship_id", "predicate_kind", "predicate_subkind"] ) - loader_backend = resolve_backend(engine, staging_schema_tag=STAGING_SCHEMA) + # The auxiliary schema is a caller-claimed tag, not one of resolved's own + # configured Role tags, so a dialect that cannot cross schemas (DuckDB) + # needs its physical name explicitly to classify the FK to vocab as + # feasible or deferred. Same pattern as + # MaintenanceContext._foreign_keys_creatable()/create_tables(): derive it + # from every tag claimed on the engine, not just the one in play here. + physical_schemas = { + tag: physical_schema_of(engine, schema_tag=tag) for tag in claimed_schema_tags(engine) + } + + loader_backend = resolve_backend( + engine, staging_schema_tag=STAGING_SCHEMA, physical_schemas=physical_schemas + ) drop_staging_sql = ( sa.text( @@ -233,7 +250,7 @@ def _load_relationship_classification( ] with open_connection(engine) as connection: Base.metadata.drop_all(bind=connection, tables=tables_to_drop, checkfirst=True) # type: ignore - create_extension_tables(connection, resolved=resolved) + create_auxiliary_tables(connection, resolved=resolved) with tempfile.TemporaryDirectory() as tmp_dir, so.Session(engine) as session: for model, df in zip( @@ -253,12 +270,16 @@ def _load_relationship_classification( merge_strategy="replace", staging_schema_tag=STAGING_SCHEMA, resolved=resolved, + physical_schemas=physical_schemas, ) session.commit() for model in (RelationshipClass, RelationshipMapping): model.validate_foreign_keys( # ty: ignore[invalid-argument-type] - session, staging_schema_tag=STAGING_SCHEMA, resolved=resolved + session, + staging_schema_tag=STAGING_SCHEMA, + resolved=resolved, + physical_schemas=physical_schemas, ) diff --git a/src/omop_graph/db/session.py b/src/omop_graph/db/session.py index 826f088..6ed980e 100644 --- a/src/omop_graph/db/session.py +++ b/src/omop_graph/db/session.py @@ -8,8 +8,9 @@ import sqlalchemy.orm as so from sqlalchemy.engine import Engine from oa_configurator import ResolvedCDMDatabase, Resolver -from omop_alchemy.config import create_cdm_engines +from omop_alchemy.config import CDM_AUXILIARY_SCHEMA_CLAIM from omop_alchemy.cross_database import CDMSession, cdm_sessionmaker +from orm_loader.backends import staging_schema_claim from sqlalchemy.orm import sessionmaker from omop_graph.config import OmopGraphConfig @@ -48,9 +49,9 @@ def cdm_session_factory( ) -> sessionmaker[CDMSession]: """Read-only routed session factory on a CDM database from the active config. - Builds the engine pair with ``create_cdm_engines()``, checking its schema - claims without registering them. The engines live as long as the process, - so this suits scripts that run once. + Builds the engine pair directly, checking its schema claims without + registering them. The engines live as long as the process, so this + suits scripts that run once. Parameters ---------- @@ -62,34 +63,41 @@ def cdm_session_factory( def _create_cdm_sessions( - resolved: ResolvedCDMDatabase, *, register_claims: bool + resolved: ResolvedCDMDatabase, + *, + register_claims: bool, ) -> tuple[sessionmaker[CDMSession], Engine, Engine]: - """Build routed sessions and their owned engine pair through omop-alchemy.""" - primary, vocab = create_cdm_engines( - resolved, register_claims=register_claims + """Build routed sessions and their owned engine pair. + + Notes + ----- + Claims the following additional schemas: + - staging schema: for CSV bulk-loading of RelationshipClass/RelationshipMapping + - CDM_AUXILIARY_SCHEMA: for the tables those two live in + """ + + primary, vocab = resolved.create_engines( + schema_claims=[staging_schema_claim(), CDM_AUXILIARY_SCHEMA_CLAIM], + register_claims=register_claims, ) return cdm_sessionmaker(resolved, primary=primary, vocab=vocab), primary, vocab @contextmanager def open_cdm_sessions( - resolved: ResolvedCDMDatabase, *, register_claims: bool = True + resolved: ResolvedCDMDatabase, ) -> Generator[so.sessionmaker[CDMSession], None, None]: - """Yield a routed session factory on ``create_cdm_engines(resolved)``, disposing both engines on exit. + """Yield a routed session factory, disposing both engines on exit. Parameters ---------- resolved : ResolvedCDMDatabase - register_claims : bool, optional - As for :func:`create_cdm_engines`. Yields ------ sqlalchemy.orm.sessionmaker[CDMSession] """ - cdm_sessions, primary, vocab = _create_cdm_sessions( - resolved, register_claims=register_claims - ) + cdm_sessions, primary, vocab = _create_cdm_sessions(resolved, register_claims=True) try: yield cdm_sessions finally: diff --git a/src/omop_graph/extensions/omop_alchemy.py b/src/omop_graph/extensions/omop_alchemy.py index 984a768..4d4e934 100644 --- a/src/omop_graph/extensions/omop_alchemy.py +++ b/src/omop_graph/extensions/omop_alchemy.py @@ -1,14 +1,14 @@ # Extension to omop-alchemy package import sqlalchemy as sa import sqlalchemy.orm as so -from orm_loader.helpers import Base, create_tables +from orm_loader.helpers import create_tables from omop_alchemy.cdm.base import ( ReferenceTable, - cdm_table, - CDMTableBase, + CDMAuxiliaryTableBase, merge_table_args, role_fk, ) +from omop_alchemy.config import CDM_AUXILIARY_SCHEMA from oa_configurator import ResolvedCDMDatabase, Role from enum import Enum @@ -28,12 +28,12 @@ def _predicate_kind_values(enum_cls: type[PredicateKind]) -> list[str]: return [member.value for member in enum_cls] -def _predicate_kind_enum() -> sa.Enum: +def _predicate_kind_enum(*, schema: str) -> sa.Enum: """Use a varchar-backed enum on DuckDB and the native enum elsewhere.""" return sa.Enum( PredicateKind, values_callable=_predicate_kind_values, - schema=Role.PRIMARY.value, + schema=schema, ).with_variant( sa.Enum( PredicateKind, @@ -45,18 +45,17 @@ def _predicate_kind_enum() -> sa.Enum: ) -@cdm_table -class RelationshipClass(ReferenceTable, CDMTableBase, Base): +class RelationshipClass(ReferenceTable, CDMAuxiliaryTableBase): """ - Extensions table: Defines the semantic categories (Parent) + Auxiliary table: Defines the semantic categories (Parent) and entity types (Child). """ __tablename__ = "relationship_class" - # Must match RelationshipMapping's schema, or SQLAlchemy silently builds a second, unlinked Table object. - __table_args__ = merge_table_args({"schema": Role.PRIMARY.value}) predicate_kind: so.Mapped[PredicateKind] = so.mapped_column( - _predicate_kind_enum(), primary_key=True + # Must match CDMAuxiliaryTableBase's own schema to ensure that + # schema_translate_map routes the enum correctly. + _predicate_kind_enum(schema=CDM_AUXILIARY_SCHEMA), primary_key=True ) predicate_subkind: so.Mapped[str] = so.mapped_column( sa.String(20), primary_key=True @@ -66,10 +65,9 @@ class RelationshipClass(ReferenceTable, CDMTableBase, Base): inference: so.Mapped[str] = so.mapped_column(sa.String(40), nullable=False) -@cdm_table -class RelationshipMapping(ReferenceTable, CDMTableBase, Base): +class RelationshipMapping(ReferenceTable, CDMAuxiliaryTableBase): """ - Extensions table: Maps standard OMOP relationship_ids to + Auxiliary table: Maps standard OMOP relationship_ids to their parent (predicate_kind - one of PredicateKind) and more fine-grained subclasses . """ @@ -80,32 +78,35 @@ class RelationshipMapping(ReferenceTable, CDMTableBase, Base): primary_key=True, ) predicate_kind: so.Mapped[PredicateKind] = so.mapped_column( - _predicate_kind_enum(), primary_key=True + # Must match CDMAuxiliaryTableBase's own schema to ensure that + # schema_translate_map routes the enum correctly. + _predicate_kind_enum(schema=CDM_AUXILIARY_SCHEMA), primary_key=True ) predicate_subkind: so.Mapped[str] = so.mapped_column( sa.String(20), primary_key=True ) - # Must match RelationshipClass's schema, or SQLAlchemy silently builds a second, unlinked Table object. + # Qualified string to RelationshipClass as it lives in the same + # fixed auxiliary schema (CDMAuxiliaryTableBase's default). __table_args__ = merge_table_args( sa.ForeignKeyConstraint( ["predicate_kind", "predicate_subkind"], [ - role_fk(Role.PRIMARY, "relationship_class.predicate_kind"), - role_fk(Role.PRIMARY, "relationship_class.predicate_subkind"), + f"{CDM_AUXILIARY_SCHEMA}.relationship_class.predicate_kind", + f"{CDM_AUXILIARY_SCHEMA}.relationship_class.predicate_subkind", ], name="fk_rel_mapping_to_rel_class", ), - {"schema": Role.PRIMARY.value}, + {"schema": CDM_AUXILIARY_SCHEMA}, ) -def create_extension_tables( +def create_auxiliary_tables( connection: sa.Connection, *, resolved: ResolvedCDMDatabase, ) -> None: - """Create the extension tables as they can physically exist on *resolved*. + """Create the auxiliary tables as they can physically exist on *resolved*. Goes through ``orm_loader``'s ``create_tables``: on a split deployment the foreign key to the vocabulary is left out, since no dialect can express a diff --git a/tests/fixtures/helpers.py b/tests/fixtures/helpers.py index 7de6513..b9a5175 100644 --- a/tests/fixtures/helpers.py +++ b/tests/fixtures/helpers.py @@ -2,15 +2,19 @@ from __future__ import annotations -from typing import cast +from contextlib import contextmanager +from typing import Generator, cast import sqlalchemy as sa +from oa_configurator import SchemaClaim +from oa_configurator.testing import isolated_test_schema from omop_alchemy.cdm.model.vocabulary import ( Concept, Concept_Class, Domain, Vocabulary ) +from omop_alchemy.config import CDM_AUXILIARY_SCHEMA, CDM_AUXILIARY_SCHEMA_OWNER VOCAB_TABLES = cast( "tuple[sa.Table, ...]", @@ -21,3 +25,25 @@ needing more tables (e.g. Relationship, Concept_Relationship) extends this tuple rather than redefining its own copy of the shared core. """ + + +@contextmanager +def cdm_auxiliary_schema_claim(engine: sa.Engine, *, prefix: str) -> Generator[SchemaClaim, None, None]: + """A SchemaClaim for a fresh, test-isolated physical auxiliary schema. + + Notes + ----- + Needed before any Base.metadata.create_all() call reachable from a test, + even one scoped via tables=: a dialect-level custom type (e.g. + RelationshipClass's predicate_kind) is created metadata-wide, not + scoped to tables=, so it's attempted regardless of which tables the + call actually asked for. scoped_test_schema() doesn't scope the + auxiliary tag on its own (only Role tags), so this must be claimed + and physically created separately. + """ + with isolated_test_schema(engine, prefix=f"{prefix}_ext") as physical_schema: + yield SchemaClaim( + schema_tag=CDM_AUXILIARY_SCHEMA, + physical_schema=physical_schema, + owner=CDM_AUXILIARY_SCHEMA_OWNER, + ) diff --git a/tests/fixtures/mock_cdm.py b/tests/fixtures/mock_cdm.py index 5752771..2e8e83b 100644 --- a/tests/fixtures/mock_cdm.py +++ b/tests/fixtures/mock_cdm.py @@ -7,6 +7,7 @@ import sqlalchemy as sa from sqlalchemy.orm import Session, sessionmaker +from oa_configurator import Dialect from oa_configurator.testing import isolated_test_database, scoped_test_schema from orm_loader.helpers import Base from orm_loader.backends import resolve_backend @@ -18,6 +19,7 @@ from omop_alchemy.cdm.model.vocabulary.domain import Domain from omop_alchemy.cdm.model.vocabulary.relationship import Relationship from omop_alchemy.cdm.model.vocabulary.vocabulary import Vocabulary +from omop_alchemy.config import CDM_AUXILIARY_SCHEMA_CLAIM from omop_graph.config import OmopGraphConfig from omop_graph.extensions.omop_alchemy import ( @@ -27,6 +29,8 @@ ) from omop_graph.graph.kg import KnowledgeGraph +from fixtures.helpers import cdm_auxiliary_schema_claim + PARENT_CANCER_ID = 443392 CONCEPT_META_ID = 0 LANGUAGE_CONCEPT_ID = 1 @@ -35,21 +39,22 @@ @pytest.fixture(scope="module") def mock_cdm_engine() -> Iterator[sa.Engine]: with isolated_test_database(OmopGraphConfig, "test_cdm_db_pg") as db: - with scoped_test_schema(db.resolved, prefix="mock_cdm") as scoped: - _create_mock_cdm_tables(scoped.engine) - yield scoped.engine + with cdm_auxiliary_schema_claim(db.committing_engine, prefix="mock_cdm") as aux_claim: + with scoped_test_schema( + db.resolved, prefix="mock_cdm", schema_claims=[aux_claim], register_claims=False + ) as scoped: + _create_mock_cdm_tables(scoped.engine) + yield scoped.engine @pytest.fixture(scope="module") def mock_cdm_engine_sqlite() -> Iterator[sa.Engine]: - engine = sa.create_engine("sqlite+pysqlite:///:memory:").execution_options( - schema_translate_map={"primary": None, "vocab": None} - ) - _create_mock_cdm_tables(engine) - try: - yield engine - finally: - engine.dispose() + with isolated_test_database( + OmopGraphConfig, "test_cdm_db_sqlite", dialect=Dialect.SQLITE, + schema_claims=[CDM_AUXILIARY_SCHEMA_CLAIM], + ) as db: + _create_mock_cdm_tables(db.committing_engine) + yield db.committing_engine def _create_mock_cdm_tables(engine: sa.Engine) -> None: diff --git a/tests/test_concept_queries.py b/tests/test_concept_queries.py index cef9bc3..ee9383a 100644 --- a/tests/test_concept_queries.py +++ b/tests/test_concept_queries.py @@ -24,14 +24,19 @@ q_entities, ) -from fixtures.helpers import VOCAB_TABLES +from fixtures.helpers import VOCAB_TABLES, cdm_auxiliary_schema_claim _META_CONCEPT_ID = 0 @pytest.fixture() def concept_engine(pg_db) -> Iterator[sa.Engine]: - with scoped_test_schema(pg_db.resolved, prefix="concept_queries") as scoped: + with ( + cdm_auxiliary_schema_claim(pg_db.committing_engine, prefix="concept_queries") as aux_claim, + scoped_test_schema( + pg_db.resolved, prefix="concept_queries", schema_claims=[aux_claim], register_claims=False + ) as scoped, + ): engine = scoped.engine Base.metadata.create_all(engine, tables=VOCAB_TABLES, checkfirst=True) diff --git a/tests/test_edges_cross_domain_regression.py b/tests/test_edges_cross_domain_regression.py index 97c1435..551dcb9 100644 --- a/tests/test_edges_cross_domain_regression.py +++ b/tests/test_edges_cross_domain_regression.py @@ -22,6 +22,7 @@ from orm_loader.helpers import Base from omop_alchemy.cdm.model.vocabulary.concept import Concept from omop_alchemy.cdm.model.vocabulary.concept_relationship import Concept_Relationship +from omop_alchemy.config import CDM_AUXILIARY_SCHEMA_CLAIM from omop_graph.extensions.omop_alchemy import ( PredicateKind, @@ -54,7 +55,11 @@ @pytest.fixture(scope="module") def procedure_kg() -> Iterator[KnowledgeGraph]: engine = sa.create_engine("sqlite+pysqlite:///:memory:").execution_options( - schema_translate_map={"primary": None, "vocab": None} + schema_translate_map={ + "primary": None, + "vocab": None, + CDM_AUXILIARY_SCHEMA_CLAIM.schema_tag: None, + } ) tables = cast( list[sa.Table], diff --git a/tests/test_fulltext_vocab_schema_postgres.py b/tests/test_fulltext_vocab_schema_postgres.py index b9f0065..59fe4fc 100644 --- a/tests/test_fulltext_vocab_schema_postgres.py +++ b/tests/test_fulltext_vocab_schema_postgres.py @@ -17,13 +17,22 @@ from omop_graph.graph.queries import q_concept_name_fulltext -from fixtures.helpers import VOCAB_TABLES +from fixtures.helpers import VOCAB_TABLES, cdm_auxiliary_schema_claim pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] def test_fulltext_query_finds_the_tsvector_column_in_the_vocab_schema(pg_db): - with scoped_test_schema(pg_db.resolved, prefix="fulltext", split_roles=[Role.VOCAB]) as scoped: + with ( + cdm_auxiliary_schema_claim(pg_db.committing_engine, prefix="fulltext_a") as aux_claim, + scoped_test_schema( + pg_db.resolved, + prefix="fulltext", + split_roles=[Role.VOCAB], + schema_claims=[aux_claim], + register_claims=False, + ) as scoped, + ): engine = scoped.engine Base.metadata.create_all(bind=engine, tables=VOCAB_TABLES, checkfirst=True) @@ -50,7 +59,16 @@ def test_fulltext_query_finds_the_tsvector_column_in_the_vocab_schema(pg_db): def test_fulltext_query_still_raises_when_the_column_is_genuinely_absent(pg_db): - with scoped_test_schema(pg_db.resolved, prefix="fulltext", split_roles=[Role.VOCAB]) as scoped: + with ( + cdm_auxiliary_schema_claim(pg_db.committing_engine, prefix="fulltext_b") as aux_claim, + scoped_test_schema( + pg_db.resolved, + prefix="fulltext", + split_roles=[Role.VOCAB], + schema_claims=[aux_claim], + register_claims=False, + ) as scoped, + ): Base.metadata.create_all(bind=scoped.engine, tables=VOCAB_TABLES, checkfirst=True) with pytest.raises(FullTextError): diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index bca03d2..3cc0a1c 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -90,6 +90,7 @@ def _seed_one_concept(engine: sa.Engine, *, concept_id: int, name: str) -> None: def _populate(scoped, *, name: str) -> None: + """Seed one concept plus the relationship vocabulary/classification.""" Base.metadata.create_all(bind=scoped.engine, checkfirst=True) _seed_one_concept(scoped.engine, concept_id=_CONCEPT_ID, name=name) seed_relationship_vocabulary(scoped.engine) @@ -100,7 +101,10 @@ def _populate(scoped, *, name: str) -> None: def test_get_adapter_resolves_the_slug_as_the_database_name( pg_db, monkeypatch, descriptor, expected_name ) -> None: - with scoped_test_schema(pg_db.resolved, prefix="oaklib_get_adapter", schema_claims=[staging_schema_claim()]) as scoped: + with scoped_test_schema( + pg_db.resolved, prefix="oaklib_get_adapter", schema_claims=[staging_schema_claim()], + register_claims=False, + ) as scoped: _populate(scoped, name="Adapter concept") def fake_resolve_cdm_database(name): @@ -116,7 +120,7 @@ def fake_resolve_cdm_database(name): def _one_engine_for_both_roles(kg: KnowledgeGraph) -> bool: - """Do the graph's sessions send vocabulary and extension tables to one engine?""" + """Do the graph's sessions send vocabulary and auxiliary tables to one engine?""" with kg.session_factory() as session: return session.get_bind(Concept) is session.get_bind(RelationshipMapping) @@ -126,7 +130,10 @@ def test_resolve_cdm_database_resolves_a_named_entry(pg_db) -> None: def test_resolved_path_reuses_one_engine_without_a_vocab_split(pg_db) -> None: - with scoped_test_schema(pg_db.resolved, prefix="oaklib_resolved", schema_claims=[staging_schema_claim()]) as scoped: + with scoped_test_schema( + pg_db.resolved, prefix="oaklib_resolved", schema_claims=[staging_schema_claim()], + register_claims=False, + ) as scoped: _populate(scoped, name="Resolved-path concept") adapter = OMOPAlchemyImplementation(resolved=scoped.resolved) @@ -148,7 +155,13 @@ def test_two_connection_entries_for_one_database_collapse_to_one_engine(pg_db) - connections={"oaklib_split_vocab": resolver.config.connections[pg_db.resolved.connection.name]}, databases={name: resolver.config.databases[name].model_copy(update={"vocab_connection": "oaklib_split_vocab"})}, ) - with scoped_test_schema(resolver.resolve_database(name), prefix="oaklib_split", resolver=resolver, schema_claims=[staging_schema_claim()]) as scoped: + with scoped_test_schema( + resolver.resolve_database(name), + prefix="oaklib_split", + resolver=resolver, + schema_claims=[staging_schema_claim()], + register_claims=False, + ) as scoped: assert scoped.resolved.connection.name != scoped.resolved.vocab_connection.name assert scoped.resolved.foreign_key_can_span(Role.PRIMARY, Role.VOCAB) _populate(scoped, name="Split-wiring concept") @@ -160,7 +173,10 @@ def test_two_connection_entries_for_one_database_collapse_to_one_engine(pg_db) - def test_kg_injection_path_uses_the_injected_knowledge_graph(pg_db) -> None: - with scoped_test_schema(pg_db.resolved, prefix="oaklib_kg", schema_claims=[staging_schema_claim()]) as scoped: + with scoped_test_schema( + pg_db.resolved, prefix="oaklib_kg", schema_claims=[staging_schema_claim()], + register_claims=False, + ) as scoped: _populate(scoped, name="Injected concept") kg = KnowledgeGraph( cdm_sessionmaker(scoped.resolved, primary=scoped.engine, vocab=scoped.engine) diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index af1bb27..1882552 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -12,24 +12,43 @@ import logging import sqlalchemy as sa -from oa_configurator import Role +import sqlalchemy.orm as so from oa_configurator.testing import scoped_test_schema from orm_loader.backends import staging_schema_claim from orm_loader.helpers import Base -from omop_graph.cli import relationship_classification +from omop_alchemy.cdm.model.vocabulary.relationship import Relationship +from omop_graph.cli import _load_relationship_classification, relationship_classification from omop_graph.extensions.omop_alchemy import RelationshipClass, RelationshipMapping from fixtures.mock_cdm import seed_relationship_vocabulary +from fixtures.helpers import cdm_auxiliary_schema_claim + def test_relationship_classification_respects_the_configured_schema(pg_db): - with scoped_test_schema(pg_db.resolved, prefix="relclass_routing", schema_claims=[staging_schema_claim()]) as scoped: - schema = scoped.schemas[Role.PRIMARY] + """The loading body creates RelationshipClass/RelationshipMapping (and the + predicatekind enum) wherever its engine resolves CDM_AUXILIARY_SCHEMA to, + not a hardcoded literal schema name. + """ + with ( + cdm_auxiliary_schema_claim(pg_db.committing_engine, prefix="relclass_routing") as aux_claim, + scoped_test_schema( + pg_db.resolved, + prefix="relclass_routing", + schema_claims=[staging_schema_claim(), aux_claim], + register_claims=False, + ) as scoped, + ): + auxiliary_schema = aux_claim.physical_schema Base.metadata.create_all(bind=scoped.engine, checkfirst=True) seed_relationship_vocabulary(scoped.engine) - relationship_classification(resolved=scoped.resolved) + with so.Session(scoped.engine) as session: + relationship_ids = set(session.scalars(sa.select(Relationship.relationship_id)).all()) + _load_relationship_classification( + None, resolved=scoped.resolved, engine=scoped.engine, relationship_ids=relationship_ids + ) with scoped.engine.connect() as conn: n_class = conn.execute( @@ -43,18 +62,23 @@ def test_relationship_classification_respects_the_configured_schema(pg_db): "SELECT n.nspname FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace " "WHERE t.typname = 'predicatekind' AND n.nspname = :schema" ), - {"schema": schema}, + {"schema": auxiliary_schema}, ).scalar() assert n_class and n_class > 0 assert n_mapping and n_mapping > 0 - assert sa.inspect(scoped.engine).has_table("relationship_class", schema=schema) - assert enum_schema == schema + assert sa.inspect(scoped.engine).has_table("relationship_class", schema=auxiliary_schema) + assert enum_schema == auxiliary_schema def test_relationship_classification_is_idempotent(pg_db): """Re-running against the same schema, the real-world redeploy case the DROP TABLE/enum-drop cleanup exists for, must not fail.""" - with scoped_test_schema(pg_db.resolved, prefix="relclass_idempotent", schema_claims=[staging_schema_claim()]) as scoped: + with scoped_test_schema( + pg_db.resolved, + prefix="relclass_idempotent", + schema_claims=[staging_schema_claim()], + register_claims=False, + ) as scoped: Base.metadata.create_all(bind=scoped.engine, checkfirst=True) seed_relationship_vocabulary(scoped.engine) @@ -85,6 +109,7 @@ def test_relationship_classification_warns_and_drops_unknown_ids_colocated( pg_db.resolved, prefix="relclass_unknown", schema_claims=[staging_schema_claim()], + register_claims=False, ) as scoped: Base.metadata.create_all(bind=scoped.engine, checkfirst=True) seed_relationship_vocabulary(scoped.engine) diff --git a/tests/test_schema_provenance_guard.py b/tests/test_schema_provenance_guard.py index 98acaa3..b759c00 100644 --- a/tests/test_schema_provenance_guard.py +++ b/tests/test_schema_provenance_guard.py @@ -16,7 +16,7 @@ from __future__ import annotations import pytest -from oa_configurator import Role, SchemaDriftError +from oa_configurator import Role, SchemaClaim, SchemaDriftError, ensure_schema from oa_configurator.testing import ( guarded_resolver, reset_schema_registry_rows, @@ -26,6 +26,7 @@ from orm_loader.helpers import Base +from omop_alchemy.config import CDM_AUXILIARY_SCHEMA, CDM_AUXILIARY_SCHEMA_OWNER from omop_graph import cli as omop_graph_cli from fixtures.mock_cdm import seed_relationship_vocabulary @@ -39,7 +40,22 @@ def test_relationship_classification_guard_fires_on_reconfigured_schema( resolver = guarded_resolver(pg_db.resolved) resolved = resolver.resolve_database(pg_db.resolved.name) - with scoped_test_schema(resolved, prefix="graph_guard_a", resolver=resolver) as scoped_a: + # Literal CDM_AUXILIARY_SCHEMA, not a test-isolated name: relationship_classification() + # below runs its own resolved-only path, which resolves its engine via + # create_cdm_engines() and claims the same literal name and owner + # internally (matching production, like MAINTENANCE_SCHEMA) -- a + # test-isolated name or owner here would silently diverge from what + # that path actually uses. + auxiliary_claim = SchemaClaim( + schema_tag=CDM_AUXILIARY_SCHEMA, + physical_schema=CDM_AUXILIARY_SCHEMA, + reserved=True, + owner=CDM_AUXILIARY_SCHEMA_OWNER, + ) + with scoped_test_schema( + resolved, prefix="graph_guard_a", resolver=resolver, schema_claims=[auxiliary_claim] + ) as scoped_a: + ensure_schema(scoped_a.engine, CDM_AUXILIARY_SCHEMA) Base.metadata.create_all(bind=scoped_a.engine, checkfirst=True) seed_relationship_vocabulary(scoped_a.engine) monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: scoped_a.resolved) diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index e121da2..4c7d1a1 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -19,8 +19,13 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import ResolvedCDMDatabase, Resolver, Role -from oa_configurator.testing import isolated_test_database, resolve_with_role_schemas, scoped_test_schema +from oa_configurator import ResolvedCDMDatabase, Resolver, Role, SchemaClaim +from oa_configurator.testing import ( + isolated_test_database, + resolve_with_role_schemas, + scoped_test_schema, +) +from orm_loader.backends import staging_schema_claim from orm_loader.config import OrmLoaderConfig from orm_loader.helpers import Base from orm_loader.backends import resolve_backend @@ -33,9 +38,10 @@ Relationship, Vocabulary, ) +from omop_alchemy.config import CDM_AUXILIARY_SCHEMA, CDM_AUXILIARY_SCHEMA_OWNER from omop_alchemy.cross_database import cdm_sessionmaker -from omop_graph.cli import relationship_classification +from omop_graph.cli import _load_relationship_classification from omop_graph.config import OmopGraphConfig from omop_graph.extensions.omop_alchemy import ( PredicateKind, @@ -44,7 +50,7 @@ ) from omop_graph.graph.kg import KnowledgeGraph -from fixtures.helpers import VOCAB_TABLES +from fixtures.helpers import VOCAB_TABLES, cdm_auxiliary_schema_claim pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] @@ -71,7 +77,7 @@ sa.Enum( PredicateKind, values_callable=lambda obj: [e.value for e in obj], - schema=Role.PRIMARY.value, + schema=CDM_AUXILIARY_SCHEMA, ), primary_key=True, ), @@ -80,7 +86,7 @@ sa.Column("semantics", sa.String(40), nullable=False), sa.Column("inference", sa.String(40), nullable=False), # Must match RelationshipClass's schema, or SQLAlchemy silently builds a second, unlinked Table object. - schema=Role.PRIMARY.value, + schema=CDM_AUXILIARY_SCHEMA, ) _shadow_relationship_mapping = sa.Table( "relationship_mapping", @@ -91,12 +97,12 @@ sa.Enum( PredicateKind, values_callable=lambda obj: [e.value for e in obj], - schema=Role.PRIMARY.value, + schema=CDM_AUXILIARY_SCHEMA, ), primary_key=True, ), sa.Column("predicate_subkind", sa.String(20), primary_key=True), - schema=Role.PRIMARY.value, + schema=CDM_AUXILIARY_SCHEMA, ) @@ -114,6 +120,9 @@ def split_engines() -> Iterator[_Engines]: with ( isolated_test_database(OmopGraphConfig, "test_cdm_db_pg") as primary_db, isolated_test_database(OrmLoaderConfig, "test_orm_db_pg") as vocab_db, + # RelationshipClass/RelationshipMapping live in the auxiliary schema, + # on the same physical connection as primary. + cdm_auxiliary_schema_claim(primary_db.committing_engine, prefix="vocab_split") as aux_claim, ): name = primary_db.resolved.name resolver = Resolver.from_active_config() @@ -125,21 +134,38 @@ def split_engines() -> Iterator[_Engines]: } ) split = resolver.resolve_database(name) - with scoped_test_schema(split, prefix="vocab_split", resolver=resolver) as scoped: - primary_engine, vocab_engine = scoped.resolved.create_engines() - # Vocab tables are created through the vocab server's own resolved: - # create_all on the split vocab engine also emits the PRIMARY-tagged - # predicatekind type, whose schema only exists on the primary server. + with scoped_test_schema( + split, prefix="vocab_split", resolver=resolver, schema_claims=[aux_claim], register_claims=False + ) as scoped: + primary_engine, vocab_engine = scoped.resolved.create_engines( + schema_claims=[staging_schema_claim(), aux_claim], register_claims=False + ) + # create_all() on the vocab engine also emits the auxiliary-tagged + # predicatekind type; co-locate it with vocab_schema since its own + # schema only exists on the primary server. vocab_schema = scoped.schemas[Role.VOCAB] vocab_only = resolve_with_role_schemas( vocab_db.resolved, {role: vocab_schema for role in vocab_db.resolved.schema_tags()} ) - vocab_setup_engine, _ = vocab_only.create_engines() + vocab_setup_engine, _ = vocab_only.create_engines( + schema_claims=[ + SchemaClaim( + schema_tag=CDM_AUXILIARY_SCHEMA, + physical_schema=vocab_schema, + owner=CDM_AUXILIARY_SCHEMA_OWNER, + ) + ], + register_claims=False, + ) try: _shadow_metadata.create_all(primary_engine) Base.metadata.create_all(vocab_setup_engine, tables=_VOCAB_TABLES, checkfirst=True) _seed(primary_engine, vocab_setup_engine) - yield _Engines(primary=primary_engine, vocab=vocab_engine, resolved=scoped.resolved) + yield _Engines( + primary=primary_engine, + vocab=vocab_engine, + resolved=scoped.resolved, + ) finally: vocab_setup_engine.dispose() primary_engine.dispose() @@ -381,9 +407,9 @@ def test_relationships_invert_swaps_subjects_and_objects(split_engines: _Engines def test_relationship_classification_succeeds_on_a_genuinely_split_vocab_connection( split_engines: _Engines, tmp_path ) -> None: - """relationship_classification() creates RelationshipMapping without its - FK to the vocabulary. A relationship ID the vocabulary lacks is warned - about and dropped.""" + """The loading body creates RelationshipMapping without its FK to the + vocabulary when they're on separate connections. A relationship ID the + vocabulary lacks is warned about and dropped.""" (tmp_path / "predicate_classification.csv").write_text( "class,subclass,description,semantics,inference\n" "Identity,mapping,Identity mapping,identity,none\n" @@ -394,7 +420,12 @@ def test_relationship_classification_succeeds_on_a_genuinely_split_vocab_connect "Identity,mapping,unknown_rel,Unknown,,,\n" ) - relationship_classification(str(tmp_path), resolved=split_engines.resolved) + with so.Session(split_engines.vocab) as session: + relationship_ids = set(session.scalars(sa.select(Relationship.relationship_id)).all()) + _load_relationship_classification( + str(tmp_path), resolved=split_engines.resolved, engine=split_engines.primary, + relationship_ids=relationship_ids, + ) with so.Session(split_engines.primary) as session: relationship_ids = set(