Move RelationshipClass/RelationshipMapping onto the CDM auxiliary schema

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,
             )
 
 

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:

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

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,
+        )

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:

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)