40 fix schema translate map rel class

.github/workflows/ci.yml

@@ -6,5 +6,29 @@ on:
 jobs:
   label-gate:
     uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/label-gate.yml@main
-  build-test:
+  build-test-sqlite:
     uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test.yml@main
+  build-test-postgres:
+    uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres-v2.yml@main
+    with:
+      postgres-db: omop_graph_test
+      setup-commands: |
+        uv run omop-config configure omop_graph \
+          --set cdm_db.kind=cdm \
+          --set cdm_db.connection.dialect=postgresql+psycopg \
+          --set cdm_db.connection.host=localhost \
+          --set cdm_db.connection.port=5432 \
+          --set cdm_db.connection.user=test \
+          --set cdm_db.connection.password=test \
+          --set cdm_db.connection.database_name=omop_graph_ci_placeholder \
+          --set cdm_db.connection.test_only=false \
+          --set cdm_db.schema_name=public \
+          --set test_cdm_db_pg.kind=cdm \
+          --set test_cdm_db_pg.connection.dialect=postgresql+psycopg \
+          --set test_cdm_db_pg.connection.host=localhost \
+          --set test_cdm_db_pg.connection.port=5432 \
+          --set test_cdm_db_pg.connection.user=test \
+          --set test_cdm_db_pg.connection.password=test \
+          --set test_cdm_db_pg.connection.database_name=omop_graph_test \
+          --set test_cdm_db_pg.connection.test_only=true \
+          --set test_cdm_db_pg.schema_name=public

docs/oaklib/interface.md

@@ -31,8 +31,8 @@ The primary adapter class that inherits from multiple OAK interfaces:
 ### Resource Management
 To initialize a connection, `omop-graph` uses a specialized resource factory:
 
-* **`OMOPOntologyResource`**: A dataclass that wraps the SQLAlchemy connection URL, treating the database as a live ontology source.
-* **`omop_resource()`**: A factory function that resolves database credentials from an explicit URL, or from the active oa-configurator stack config (`OmopGraphConfig.cdm_db`) when no URL is given.
+* **`OMOPOntologyResource`**: A dataclass that wraps the SQLAlchemy connection URL, treating the database as a live ontology source. When the resolved CDM database has a genuinely separate `vocab_connection` configured, it also carries a second URL (`vocab_url`) for the vocabulary server.
+* **`omop_resource()`**: A factory function that resolves database credentials from an explicit URL, or from the active oa-configurator stack config (`OmopGraphConfig.cdm_db`) when no URL is given. Populates `vocab_url` automatically from the resolved config's `vocab_connection`, when configured; an explicit `url=` has no vocabulary split to carry.
 
 ---
 

pytest.toml

@@ -1,6 +1,6 @@
 [pytest]
 testpaths = ["tests"]
-addopts = ["-rf", "-rx", "--disable-pytest-warnings"]
+addopts = ["-rf", "-rx", "--disable-pytest-warnings", "-m", "not db_dialect"]
 log_cli = true
 log_cli_level = "DEBUG"
 log_cli_format = "%(asctime)s | %(name)s | %(levelname)s | %(message)s"

src/omop_graph/cli.py

@@ -1,5 +1,7 @@
 import logging
 import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager
 from importlib import resources
 from pathlib import Path
 from typing import Annotated, Optional, cast
@@ -9,13 +11,15 @@
 import typer
 from sqlalchemy.orm import sessionmaker
 
-from orm_loader.backends import resolve_backend
+from oa_configurator import ResolvedCDMDatabase, Role, ensure_schema, guard_schema_provenance, schema_of
+
+from orm_loader.backends import STAGING_SCHEMA, resolve_backend
 from orm_loader.helpers import bulk_load_context
 from orm_loader.helpers.metadata import Base
 from orm_loader.loaders.loader_interface import PandasLoader
 
 from omop_graph.config import OmopGraphConfig
-from omop_graph.db.session import make_engine
+from omop_graph.db.session import make_engine, resolve_cdm_database
 from omop_graph.extensions.omop_alchemy import RelationshipClass, RelationshipMapping
 from omop_graph.cli_utils import populate_test_data
 
@@ -57,20 +61,42 @@ def packaged_predicate_csv_dir() -> Path:
     return Path(str(resources.files("omop_graph") / "data"))
 
 
-@app.command()
+@contextmanager
+def _open_connection(bindable: sa.Engine | sa.Connection) -> Iterator[sa.Connection]:
+    """Yield a Connection: opens its own transaction for an Engine, or uses
+    an already-open Connection directly, participating in the caller's own
+    transaction (needed by the rollback-based pg_db test fixture).
+    """
+    if isinstance(bindable, sa.Engine):
+        with bindable.begin() as connection:
+            yield connection
+    else:
+        yield bindable
+
+
 def relationship_classification(
-    pred_class_dir: Annotated[
-        Optional[str],
-        typer.Option(
-            help=(
-                "Path to the directory containing `predicate_classification.csv` "
-                "and `predicate_mapping.csv`. Defaults to the copies shipped with "
-                "omop-graph; pass a directory to override them."
-            )
-        ),
-    ] = None,
-):
-    """Load pre-classified predicates into the database."""
+    pred_class_dir: Optional[str] = None,
+    *,
+    engine: sa.Engine | sa.Connection | None = None,
+    resolved: ResolvedCDMDatabase | None = None,
+) -> None:
+    """Load pre-classified predicates into the database.
+
+    Parameters
+    ----------
+    pred_class_dir : str, optional
+        Path to the directory containing `predicate_classification.csv` and
+        `predicate_mapping.csv`. Defaults to the copies shipped with
+        omop-graph.
+    engine : sqlalchemy.Engine or sqlalchemy.Connection, optional
+        Bindable to run against. Defaults to the active oa-configurator
+        config's resolved CDM engine, in which case resolved is also
+        resolved internally and any value passed here is ignored.
+    resolved : ResolvedCDMDatabase, optional
+        Enables the schema-provenance guard. Only meaningful together with
+        an explicitly injected engine/connection, since the engine=None
+        path always resolves its own regardless of what's passed here.
+    """
     pred_class_dir_pl = (
         Path(pred_class_dir) if pred_class_dir else packaged_predicate_csv_dir()
     )
@@ -139,32 +165,46 @@ def relationship_classification(
         subset=["relationship_id", "predicate_kind", "predicate_subkind"]
     )
 
-    engine = make_engine()
+    if engine is None:
+        resolved = resolve_cdm_database()
+        engine = resolved.create_engine()
+    db_schema = schema_of(engine)
+    ensure_schema(engine, db_schema)
+    ensure_schema(engine, STAGING_SCHEMA)
+
     Session = sessionmaker(bind=engine, future=True)
     session = Session()
-    loader_backend = resolve_backend(engine)
-
-    with engine.begin() as conn:
-        conn.execute(
-            sa.text(
-                "DROP TABLE IF EXISTS "
-                f"{loader_backend.qualified_staging_name(RelationshipMapping.__tablename__)} CASCADE"
-            )
-        )
-        conn.execute(
-            sa.text(
-                "DROP TABLE IF EXISTS "
-                f"{loader_backend.qualified_staging_name(RelationshipClass.__tablename__)} CASCADE"
-            )
-        )
-        conn.execute(sa.text("DROP TYPE IF EXISTS predicatekindenum CASCADE;"))
+    loader_backend = resolve_backend(engine, staging_schema=STAGING_SCHEMA)
 
+    drop_staging_sql = (
+        sa.text(
+            "DROP TABLE IF EXISTS "
+            f"{loader_backend.qualified_staging_name(RelationshipMapping.__tablename__)} CASCADE"
+        ),
+        sa.text(
+            "DROP TABLE IF EXISTS "
+            f"{loader_backend.qualified_staging_name(RelationshipClass.__tablename__)} CASCADE"
+        ),
+    )
+    with _open_connection(engine) as connection:
+        for stmt in drop_staging_sql:
+            connection.execute(stmt)
+
+    # DROP TYPE IF EXISTS predicatekindenum was dead code: the Enum column
+    # never set an explicit name=, so SQLAlchemy's generated type name is
+    # actually "predicatekind". drop_all(tables=[...]) already drops the
+    # shared Enum type exactly once, deduped, since both tables using it
+    # are always in the same tables= list.
     tables_to_drop = [
         RelationshipMapping.__table__,
         RelationshipClass.__table__,
     ]
-    Base.metadata.drop_all(bind=engine, tables=tables_to_drop, checkfirst=True)  # type: ignore
-    Base.metadata.create_all(bind=engine, tables=tables_to_drop)  # type: ignore
+    # Both tables live in the primary schema (only RelationshipMapping's FK
+    # target is vocab-tagged, via role_fk), so the guard checks Role.PRIMARY.
+    with _open_connection(engine) as connection:
+        with guard_schema_provenance(connection, resolved, role=Role.PRIMARY):
+            Base.metadata.drop_all(bind=connection, tables=tables_to_drop, checkfirst=True)  # type: ignore
+            Base.metadata.create_all(bind=connection, tables=tables_to_drop)  # type: ignore
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         for model, df in zip(
@@ -184,9 +224,27 @@ def relationship_classification(
                     dedupe=True,
                     merge_strategy="replace",
                     loader=PandasLoader(),
+                    staging_schema=STAGING_SCHEMA,
                 )
                 session.commit()
 
 
+@app.command(name="relationship-classification")
+def relationship_classification_cmd(
+    pred_class_dir: Annotated[
+        Optional[str],
+        typer.Option(
+            help=(
+                "Path to the directory containing `predicate_classification.csv` "
+                "and `predicate_mapping.csv`. Defaults to the copies shipped with "
+                "omop-graph; pass a directory to override them."
+            )
+        ),
+    ] = None,
+):
+    """Load pre-classified predicates into the database."""
+    relationship_classification(pred_class_dir)
+
+
 if __name__ == "__main__":
     app()

src/omop_graph/config.py

@@ -37,6 +37,21 @@ class OmopGraphConfig(PackageConfigBase):
     )
 
     cdm_db: Annotated[str, RefTo(CDMDatabaseConfig)] = "cdm_db"
+    test_cdm_db_pg: Annotated[
+        str | None, RefTo(CDMDatabaseConfig, is_test=True)
+    ] = Field(
+        default=None,
+        description="Real PostgreSQL test CDM database, for Postgres-only integration testing.",
+    )
+    test_cdm_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)."
+        ),
+    )
     embedding_model_name: Annotated[str | None, RefTo(ModelConfig)] = Field(
         default=None,
         description=(

src/omop_graph/db/session.py

@@ -7,10 +7,26 @@
 from sqlalchemy import create_engine, URL, Engine
 from sqlalchemy.orm import sessionmaker, Session
 
-from oa_configurator import Resolver
+from oa_configurator import ResolvedCDMDatabase, Resolver
 from omop_graph.config import OmopGraphConfig
 
 
+def resolve_cdm_database() -> ResolvedCDMDatabase:
+    """Resolve the active oa-configurator config's CDM database.
+
+    Split out from make_engine() for callers that need the resolved object
+    itself (e.g. schema-provenance guarding, Phase 9), not just an engine.
+    """
+    resolver = Resolver.from_active_config()
+    db_name = resolver.resolve_package_config(OmopGraphConfig).cdm_db
+    resolved = resolver.resolve_database(db_name)
+    if not isinstance(resolved, ResolvedCDMDatabase):
+        raise TypeError(
+            f"OmopGraphConfig.cdm_db must resolve to a CDM database, got {type(resolved).__name__}"
+        )
+    return resolved
+
+
 def make_engine(
     url: Optional[Union[URL, str]] = None,
     *,
@@ -31,9 +47,11 @@ def make_engine(
         Keyword arguments forwarded to ``sqlalchemy.create_engine`` in both paths.
         Common keys: ``echo``, ``connect_args``, ``pool_size``.
     execution_options : dict, optional
-        Options forwarded to ``engine.execution_options()``. In the resolver path these
-        are merged with the auto-generated ``schema_translate_map`` (resolver wins on
-        that key via ``setdefault``).
+        Options forwarded to ``engine.execution_options()``. In the resolver
+        path, a ``schema_translate_map`` here may add keys the resolver
+        doesn't define, but may not include ``None``: that key is always set
+        from the resolved config, and ``create_engine()`` raises
+        ``ValueError`` if it is overridden here.
 
     Returns
     -------
@@ -42,9 +60,7 @@ def make_engine(
     """
     engine_kwargs = engine_kwargs or {}
     if url is None:
-        resolver = Resolver.from_active_config()
-        db_name = resolver.resolve_package_config(OmopGraphConfig).cdm_db
-        database = resolver.resolve_database(db_name)
+        database = resolve_cdm_database()
         return database.create_engine(execution_options=execution_options, **engine_kwargs)
 
     from sqlalchemy import make_url as _make_url

src/omop_graph/extensions/omop_alchemy.py

@@ -2,7 +2,8 @@
 import sqlalchemy as sa
 import sqlalchemy.orm as so
 from orm_loader.helpers import Base
-from omop_alchemy.cdm.base import ReferenceTable, cdm_table, CDMTableBase
+from omop_alchemy.cdm.base import ReferenceTable, cdm_table, CDMTableBase, role_fk
+from oa_configurator import Role
 
 from enum import Enum
 from dataclasses import dataclass
@@ -51,7 +52,8 @@ class RelationshipMapping(ReferenceTable, CDMTableBase, Base):
     __tablename__ = "relationship_mapping"
 
     relationship_id: so.Mapped[str] = so.mapped_column(
-        sa.ForeignKey("relationship.relationship_id"), primary_key=True
+        sa.ForeignKey(role_fk(Role.VOCAB, "relationship.relationship_id")),
+        primary_key=True,
     )
     predicate_kind: so.Mapped[PredicateKind] = so.mapped_column(
         sa.Enum(