From 01d4e74e7409cce4cde6b295aed7d052db05d5d7 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 1 Sep 2026 05:00:32 +0000 Subject: [PATCH 1/5] Utilise oa-configurator test interface, new schema_translate_map carrying methods --- Dockerfile | 3 - pytest.toml | 2 +- src/omop_graph/cli.py | 100 ++++-- src/omop_graph/config.py | 15 + src/omop_graph/extensions/omop_alchemy.py | 6 +- src/omop_graph/graph/kg.py | 130 +++++-- src/omop_graph/graph/queries.py | 157 ++++++--- .../oaklib_interface/omop_factory.py | 10 +- .../oaklib_interface/omop_implementation.py | 6 +- .../oaklib_interface/omop_resource.py | 6 + tests/conftest.py | 22 ++ tests/fixtures/mock_cdm.py | 22 +- tests/test_concept_queries.py | 85 +++-- .../test_edges_same_connection_regression.py | 34 ++ tests/test_oaklib_schema_awareness.py | 194 +++++++++++ tests/test_pg_db_fixture.py | 23 ++ tests/test_relationship_classification.py | 71 ++++ tests/test_vocab_split_connection.py | 318 ++++++++++++++++++ 18 files changed, 1056 insertions(+), 148 deletions(-) delete mode 100644 Dockerfile create mode 100644 tests/test_edges_same_connection_regression.py create mode 100644 tests/test_oaklib_schema_awareness.py create mode 100644 tests/test_pg_db_fixture.py create mode 100644 tests/test_relationship_classification.py create mode 100644 tests/test_vocab_split_connection.py diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index 208b49a..0000000 --- a/Dockerfile +++ /dev/null @@ -1,3 +0,0 @@ -FROM python:3.12-slim -RUN pip install --no-cache-dir ".[postgres,emb,pgvector,faiss-cpu]" -WORKDIR /workspace diff --git a/pytest.toml b/pytest.toml index b2320fb..ba0e085 100644 --- a/pytest.toml +++ b/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" diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index c7a2038..06b32de 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -9,7 +9,9 @@ import typer from sqlalchemy.orm import sessionmaker -from orm_loader.backends import resolve_backend +from oa_configurator import ensure_schema, 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 @@ -57,20 +59,23 @@ def packaged_predicate_csv_dir() -> Path: return Path(str(resources.files("omop_graph") / "data")) -@app.command() 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, +) -> 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. + """ pred_class_dir_pl = ( Path(pred_class_dir) if pred_class_dir else packaged_predicate_csv_dir() ) @@ -139,26 +144,41 @@ def relationship_classification( subset=["relationship_id", "predicate_kind", "predicate_subkind"] ) - engine = make_engine() + if engine is None: + engine = make_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" + ), + ) + if isinstance(engine, sa.Engine): + with engine.begin() as conn: + for stmt in drop_staging_sql: + conn.execute(stmt) + else: + for stmt in drop_staging_sql: + engine.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", meaning this line never matched anything, + # with IF EXISTS silently no-op'ing every run. drop_all(tables=[...]) already + # drops a shared Enum type exactly once, correctly deduped, once every + # table using it is in the same tables= list (true here, both tables + # always move together), so no manual DROP TYPE is needed at all. tables_to_drop = [ RelationshipMapping.__table__, RelationshipClass.__table__, @@ -184,9 +204,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() diff --git a/src/omop_graph/config.py b/src/omop_graph/config.py index 54f7381..42b157a 100644 --- a/src/omop_graph/config.py +++ b/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=( diff --git a/src/omop_graph/extensions/omop_alchemy.py b/src/omop_graph/extensions/omop_alchemy.py index c456cb5..f2a6cba 100644 --- a/src/omop_graph/extensions/omop_alchemy.py +++ b/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( diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index cf1401d..174e685 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -22,7 +22,7 @@ from typing import Dict, Optional, Tuple, Literal, Generator, TYPE_CHECKING from dataclasses import dataclass -from sqlalchemy import Engine +from sqlalchemy import Engine, Row from sqlalchemy.orm import Session, sessionmaker from omop_alchemy.backends import FullTextError from omop_alchemy.cdm.query import ConceptFilter @@ -73,6 +73,8 @@ q_children, q_predicate_name, q_predicate_row_with_ancestry, + q_relationship_mapping_all, + q_relationship_mapping_row, q_roots, q_singletons, q_entities, @@ -145,6 +147,39 @@ def provider_type(self) -> str: return self.resolved_model.provider.provider +def _relationship_mapping_lookup(session: Session) -> dict[str, Row]: + """RelationshipMapping rows keyed by relationship_id. + + RelationshipMapping is an omop-graph extension table, not vocab-role, so + it never lives on a split ``vocab_engine``. This always runs against + the primary connection. + """ + return { + row.relationship_id: row + for row in session.execute(q_relationship_mapping_all()).all() + } + + +def _predicate_from_rows(ancestry_row: Row, mapping_row: Row) -> Predicate: + """Build a Predicate from a Relationship-ancestry row and a RelationshipMapping row. + + The two rows come from the same query in a same-connection deployment + (pass the row twice), or from two separately-fetched engines in a + split-connection one. This is the one place that shape difference + collapses back into a single code path. + """ + return Predicate( + relationship_id=ancestry_row.relationship_id, + name=ancestry_row.relationship_name, + reverse_id=ancestry_row.reverse_relationship_id, + is_hierarchical=bool(ancestry_row.is_hierarchical), + anc_up=bool(ancestry_row.anc_up), + anc_down=bool(ancestry_row.anc_down), + predicate_kind=PredicateKind(mapping_row.predicate_kind), + predicate_subkind=mapping_row.predicate_subkind, + ) + + class KnowledgeGraph(GraphBackend): """ The main entry point for interacting with the OMOP Graph. @@ -156,16 +191,32 @@ class KnowledgeGraph(GraphBackend): ---------- cdm_engine : Engine The SQLAlchemy engine for the OMOP CDM database. + vocab_engine : Engine, optional + A separate engine for the vocabulary connection, for a deployment + where ``vocab_connection`` names a physically different server than + ``connection``. Omit (the common case) when vocabulary tables sit on + the same connection as everything else: same-connection queries + stay a single eager join. When given and different from + ``cdm_engine``, the three queries that join a vocab-role table + (Concept/Concept_Relationship/Relationship) against + RelationshipMapping (not vocab-role, since it's an omop-graph + extension table) fetch each side from its own engine and merge in + Python, since a SQL join cannot span two physical connections. """ def __init__( self, cdm_engine: Engine, + vocab_engine: Optional[Engine] = None, emb_config: Optional[KnowledgeGraphEmbeddingConfiguration] = None, ): self.cdm_engine = cdm_engine self.session_factory = sessionmaker(bind=self.cdm_engine, future=True) + self.vocab_engine = vocab_engine if vocab_engine is not None else cdm_engine + self._vocab_split = self.vocab_engine is not self.cdm_engine + self.vocab_session_factory = sessionmaker(bind=self.vocab_engine, future=True) + try: with self.session_factory() as session: self._relationship_mapping: dict[str, RelationshipMappingElement] = ( @@ -416,18 +467,22 @@ def predicate(self, relationship_id: str) -> Predicate: Predicate The predicate definition. """ + if self._vocab_split: + with self.vocab_session_factory() as vsession: + ancestry_row = vsession.execute( + q_predicate_row_with_ancestry( + relationship_id, include_classification=False + ) + ).one() + with self.session_factory() as session: + mapping_row = session.execute( + q_relationship_mapping_row(relationship_id) + ).one() + return _predicate_from_rows(ancestry_row, mapping_row) + with self.session_factory() as session: row = session.execute(q_predicate_row_with_ancestry(relationship_id)).one() - return Predicate( - relationship_id=row.relationship_id, - name=row.relationship_name, - reverse_id=row.reverse_relationship_id, - is_hierarchical=bool(row.is_hierarchical), - anc_up=bool(row.anc_up), - anc_down=bool(row.anc_down), - predicate_kind=PredicateKind(row.predicate_kind), - predicate_subkind=row.predicate_subkind, - ) + return _predicate_from_rows(row, row) def predicate_name(self, relationship_id: str) -> str: """ @@ -573,6 +628,32 @@ def iter_edges( within_domain: bool = True, ) -> Generator[EdgeView, None, None]: + if self._vocab_split: + with self.vocab_session_factory() as vsession: + vocab_rows = vsession.execute( + q_edges( + concept_ids=concept_ids, + predicate_ids=predicate_ids, + direction=direction, + active_only=active_only, + on=on, + within_domain=within_domain, + include_classification=False, + ) + ).all() + mapping_by_id = _relationship_mapping_lookup(session) + for vrow in vocab_rows: + mapping = mapping_by_id.get(vrow.predicate_id) + if mapping is None: + continue + if predicate_kinds and PredicateKind(mapping.predicate_kind) not in predicate_kinds: + continue + data = dict(vrow._mapping) + data["predicate_kind"] = PredicateKind(mapping.predicate_kind) + data["predicate_subkind"] = mapping.predicate_subkind + yield EdgeView(**data) + return + stmt = q_edges( concept_ids=concept_ids, predicate_ids=predicate_ids, @@ -682,21 +763,22 @@ def predicates(self) -> tuple[Predicate, ...]: """ Return all predicates known to the knowledge graph. """ + if self._vocab_split: + with self.vocab_session_factory() as vsession: + ancestry_rows = vsession.execute( + q_all_predicates_with_ancestry(include_classification=False) + ).all() + with self.session_factory() as session: + mapping_by_id = _relationship_mapping_lookup(session) + return tuple( + _predicate_from_rows(row, mapping_by_id[row.relationship_id]) + for row in ancestry_rows + if row.relationship_id in mapping_by_id + ) + with self.session_factory() as session: rows = session.execute(q_all_predicates_with_ancestry()).all() - return tuple( - Predicate( - relationship_id=row.relationship_id, - name=row.relationship_name, - reverse_id=row.reverse_relationship_id, - is_hierarchical=bool(row.is_hierarchical), - anc_up=bool(row.anc_up), - anc_down=bool(row.anc_down), - predicate_kind=PredicateKind(row.predicate_kind), - predicate_subkind=row.predicate_subkind, - ) - for row in rows - ) + return tuple(_predicate_from_rows(row, row) for row in rows) @functools.cached_property def _valid_domains(self) -> frozenset[str]: diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 8bb8763..4118918 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -26,12 +26,13 @@ or_, select, Engine, - inspect, column, ) from sqlalchemy.orm import aliased from sqlalchemy.sql import Select +from oa_configurator import schema_inspect + from omop_alchemy.backends import ( CONCEPT_NAME_TSVECTOR_COLUMN, CONCEPT_SYNONYM_NAME_TSVECTOR_COLUMN, @@ -352,7 +353,7 @@ def q_concept_name_fulltext( Concept_Synonym.concept_synonym_name if synonym else Concept.concept_name ) - inspector = inspect(engine) + inspector = schema_inspect(engine) target_table = Concept_Synonym if synonym else Concept target_col = ( CONCEPT_SYNONYM_NAME_TSVECTOR_COLUMN @@ -423,61 +424,108 @@ def q_predicate_row(relationship_id: str) -> Select: ).where(Relationship.relationship_id == relationship_id) -def q_predicate_row_with_ancestry(relationship_id: str) -> Select: +def q_predicate_row_with_ancestry( + relationship_id: str, *, include_classification: bool = True +) -> Select: """ Query a predicate and its reverse to determine directionality. This joins the Relationship table with itself to determine if the relationship points 'up' (towards ancestors) or 'down' (towards descendants). + Parameters + ---------- + include_classification : bool, optional + Join in RelationshipMapping's predicate_kind/predicate_subkind. Set to + False for a split-connection deployment (Relationship is vocab-role, + RelationshipMapping is not, so they can live on different physical + connections). The caller fetches RelationshipMapping separately via + :func:`q_relationship_mapping_row` and merges in Python. + Returns ------- Select Columns: relationship_id, relationship_name, reverse_relationship_id, - is_hierarchical, anc_down, anc_up. + is_hierarchical, anc_down, anc_up, plus predicate_kind/predicate_subkind + when include_classification is True. """ Rel = Relationship Rev = aliased(Relationship) - Rm = aliased(RelationshipMapping) - return ( - select( - Rel.relationship_id, - Rel.relationship_name, - Rel.reverse_relationship_id, - Rel.is_hierarchical_relationship_expr().label("is_hierarchical"), - Rel.is_ancestry_defining_expr().label("anc_down"), - Rev.is_ancestry_defining_expr().label("anc_up"), - Rm.predicate_kind, - Rm.predicate_subkind, - ) - .join( - Rev, - Rel.reverse_relationship_id == Rev.relationship_id, - ) - .join(Rm, Rel.relationship_id == Rm.relationship_id) # Match string IDs - .where(Rel.relationship_id == relationship_id) + stmt = select( + Rel.relationship_id, + Rel.relationship_name, + Rel.reverse_relationship_id, + Rel.is_hierarchical_relationship_expr().label("is_hierarchical"), + Rel.is_ancestry_defining_expr().label("anc_down"), + Rev.is_ancestry_defining_expr().label("anc_up"), + ).join( + Rev, + Rel.reverse_relationship_id == Rev.relationship_id, ) + if include_classification: + Rm = aliased(RelationshipMapping) + stmt = stmt.add_columns(Rm.predicate_kind, Rm.predicate_subkind).join( + Rm, Rel.relationship_id == Rm.relationship_id + ) -def q_all_predicates_with_ancestry() -> Select: - """Query all predicates with derived ancestry direction flags and classification.""" + return stmt.where(Rel.relationship_id == relationship_id) + + +def q_all_predicates_with_ancestry(*, include_classification: bool = True) -> Select: + """Query all predicates with derived ancestry direction flags and classification. + + Parameters + ---------- + include_classification : bool, optional + See :func:`q_predicate_row_with_ancestry`. + """ Rel = Relationship Rev = aliased(Relationship) - Rm = aliased(RelationshipMapping) - return ( - select( - Rel.relationship_id, - Rel.relationship_name, - Rel.reverse_relationship_id, - Rel.is_hierarchical_relationship_expr().label("is_hierarchical"), - Rel.is_ancestry_defining_expr().label("anc_down"), - Rev.is_ancestry_defining_expr().label("anc_up"), - Rm.predicate_kind, - Rm.predicate_subkind, + + stmt = select( + Rel.relationship_id, + Rel.relationship_name, + Rel.reverse_relationship_id, + Rel.is_hierarchical_relationship_expr().label("is_hierarchical"), + Rel.is_ancestry_defining_expr().label("anc_down"), + Rev.is_ancestry_defining_expr().label("anc_up"), + ).join(Rev, Rel.reverse_relationship_id == Rev.relationship_id) + + if include_classification: + Rm = aliased(RelationshipMapping) + stmt = stmt.add_columns(Rm.predicate_kind, Rm.predicate_subkind).join( + Rm, Rel.relationship_id == Rm.relationship_id ) - .join(Rev, Rel.reverse_relationship_id == Rev.relationship_id) - .join(Rm, Rel.relationship_id == Rm.relationship_id) + + return stmt + + +def q_relationship_mapping_row(relationship_id: str) -> Select: + """Query one RelationshipMapping row by relationship_id. + + The primary-role half of a split-connection predicate lookup, pairing + with :func:`q_predicate_row_with_ancestry`'s ``include_classification=False``. + """ + return select( + RelationshipMapping.relationship_id, + RelationshipMapping.predicate_kind, + RelationshipMapping.predicate_subkind, + ).where(RelationshipMapping.relationship_id == relationship_id) + + +def q_relationship_mapping_all() -> Select: + """Query every RelationshipMapping row, keyed by relationship_id. + + The primary-role half of a split-connection edges/predicates lookup. + RelationshipMapping is a small reference table, so callers merge it as a + plain dict rather than joining across connections. + """ + return select( + RelationshipMapping.relationship_id, + RelationshipMapping.predicate_kind, + RelationshipMapping.predicate_subkind, ) @@ -489,11 +537,30 @@ def q_edges( active_only: bool = False, on: Optional[date] = None, within_domain: bool = False, + include_classification: bool = True, ) -> Select: - """Query outgoing edges for a batch of concept IDs.""" + """Query outgoing edges for a batch of concept IDs. + + Parameters + ---------- + include_classification : bool, optional + Join in RelationshipMapping's predicate_kind/predicate_subkind. + Concept_Relationship is vocab-role, RelationshipMapping is not, so + for a split-connection deployment set this to False and merge + RelationshipMapping (via :func:`q_relationship_mapping_all`) + in Python instead. ``predicate_kinds`` cannot be applied in SQL + when this is False (the column isn't joined); the caller must + filter after merging. + """ if isinstance(concept_ids, int): concept_ids = (concept_ids,) + if not include_classification and predicate_kinds: + raise ValueError( + "predicate_kinds requires include_classification=True; filter " + "after merging RelationshipMapping in Python instead." + ) + Subj = aliased(Concept) Obj = aliased(Concept) @@ -504,13 +571,17 @@ def q_edges( Concept_Relationship.valid_start_date, Concept_Relationship.valid_end_date, Concept_Relationship.invalid_reason, - RelationshipMapping.predicate_kind, - RelationshipMapping.predicate_subkind, - ).join( - RelationshipMapping, - Concept_Relationship.relationship_id == RelationshipMapping.relationship_id, ) + if include_classification: + stmt = stmt.add_columns( + RelationshipMapping.predicate_kind, RelationshipMapping.predicate_subkind + ).join( + RelationshipMapping, + Concept_Relationship.relationship_id + == RelationshipMapping.relationship_id, + ) + if active_only: stmt = stmt.where(Concept_Relationship.is_valid_expr()) if on is not None: diff --git a/src/omop_graph/oaklib_interface/omop_factory.py b/src/omop_graph/oaklib_interface/omop_factory.py index e4bef90..e1cdbe4 100644 --- a/src/omop_graph/oaklib_interface/omop_factory.py +++ b/src/omop_graph/oaklib_interface/omop_factory.py @@ -7,7 +7,7 @@ from sqlalchemy.engine import URL from .omop_resource import OMOPOntologyResource -from oa_configurator import Resolver +from oa_configurator import ResolvedCDMDatabase, Resolver from omop_graph.config import OmopGraphConfig @@ -33,13 +33,21 @@ def omop_resource( ------- OMOPOntologyResource """ + execution_options = None if url is None: resolver = Resolver.from_active_config() db_name = resolver.resolve_package_config(OmopGraphConfig).cdm_db database = resolver.resolve_database(db_name) + if not isinstance(database, ResolvedCDMDatabase): + raise TypeError( + f"OmopGraphConfig.cdm_db must resolve to a CDM database, got " + f"{type(database).__name__}" + ) url = database.connection.url + execution_options = {"schema_translate_map": database.schema_translate_map()} return OMOPOntologyResource( slug=slug, url=url, + execution_options=execution_options, ) diff --git a/src/omop_graph/oaklib_interface/omop_implementation.py b/src/omop_graph/oaklib_interface/omop_implementation.py index d11b232..30714dc 100644 --- a/src/omop_graph/oaklib_interface/omop_implementation.py +++ b/src/omop_graph/oaklib_interface/omop_implementation.py @@ -896,7 +896,11 @@ def __init__( "No database URL provided for OMOPAlchemyImplementation" ) - engine = make_engine(self.engine_string, engine_kwargs={"echo": False, "future": True}) + engine = make_engine( + self.engine_string, + engine_kwargs={"echo": False, "future": True}, + execution_options=self.resource.execution_options, + ) self._connection = None diff --git a/src/omop_graph/oaklib_interface/omop_resource.py b/src/omop_graph/oaklib_interface/omop_resource.py index c5c73b1..3fce3e2 100644 --- a/src/omop_graph/oaklib_interface/omop_resource.py +++ b/src/omop_graph/oaklib_interface/omop_resource.py @@ -27,6 +27,11 @@ class OMOPOntologyResource(OntologyResource): Whether the resource is in-memory. Defaults to False. readonly : bool, optional Whether the resource is read-only. Defaults to True. + execution_options : dict, optional + Forwarded to the engine built from ``url`` (e.g. a + ``schema_translate_map``). Not carried by ``url`` itself, so a + caller resolving through oa-configurator needs this to keep the + configured schema past this resource object. """ url: Optional[Union[str, URL]] = None # type: ignore[assignment] @@ -35,6 +40,7 @@ class OMOPOntologyResource(OntologyResource): local: bool = False # type: ignore[assignment] in_memory: bool = False # type: ignore[assignment] readonly: bool = True # type: ignore[assignment] + execution_options: Optional[dict] = None # type: ignore[assignment] def _parsed_url(self) -> Optional[URL]: """ diff --git a/tests/conftest.py b/tests/conftest.py index 4d5115d..887403f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,28 @@ pytest_plugins = ("fixtures.mock_cdm",) +@pytest.fixture +def pg_db(request): + """Canonical isolated PostgreSQL test database (Phase 0 of the + schema_translate_map fix). + + Resolves via OA_Configurator resource 'test_cdm_db_pg' in ~/.config/omop/config.toml. + Run: omop-config configure omop_graph (answer Y when asked to configure test database). + + Everything a test does through ``pg_db.connection``/``pg_db.session`` + happens inside one transaction that's rolled back on exit. Nothing + here is ever committed to the shared server, so concurrent test runs + can't collide and no manual cleanup is needed. omop-graph has no + existing Postgres test fixture; built from scratch here, matching the + convention already used by OMOP_Alchemy/orm-loader/omop-emb. + """ + from oa_configurator.testing import isolated_test_database + from omop_graph.config import OmopGraphConfig + + with isolated_test_database(OmopGraphConfig, "test_cdm_db_pg", request=request) as db: + yield db + + class WhitelistFilter(logging.Filter): def __init__(self, whitelist): self.whitelist = whitelist diff --git a/tests/fixtures/mock_cdm.py b/tests/fixtures/mock_cdm.py index 2527d5e..7783e2e 100644 --- a/tests/fixtures/mock_cdm.py +++ b/tests/fixtures/mock_cdm.py @@ -1,12 +1,13 @@ from __future__ import annotations from datetime import date -from typing import cast +from typing import Iterator, cast import pytest import sqlalchemy as sa from sqlalchemy.orm import Session, sessionmaker +from oa_configurator.testing import isolated_test_database from orm_loader.helpers import Base from omop_alchemy.cdm.model.vocabulary.concept import Concept from omop_alchemy.cdm.model.vocabulary.concept_ancestor import Concept_Ancestor @@ -17,6 +18,7 @@ from omop_alchemy.cdm.model.vocabulary.relationship import Relationship from omop_alchemy.cdm.model.vocabulary.vocabulary import Vocabulary +from omop_graph.config import OmopGraphConfig from omop_graph.extensions.omop_alchemy import ( PredicateKind, RelationshipClass, @@ -30,8 +32,20 @@ @pytest.fixture(scope="module") -def mock_cdm_engine() -> sa.Engine: - engine = sa.create_engine("sqlite+pysqlite:///:memory:", future=True) +def mock_cdm_engine() -> Iterator[sa.Engine]: + with isolated_test_database( + OmopGraphConfig, + "test_cdm_db_sqlite", + dialect="sqlite", + future=True, + execution_options={"schema_translate_map": {None: None, "vocab": None, "results": None}}, + ) as db: + engine = db.connection.engine + _create_mock_cdm_tables(engine) + yield engine + + +def _create_mock_cdm_tables(engine: sa.Engine) -> None: tables = cast( list[sa.Table], [ @@ -54,8 +68,6 @@ def mock_cdm_engine() -> sa.Engine: with session_local() as session: seed_mock_cdm(session) - return engine - @pytest.fixture() def mock_cdm_kg( diff --git a/tests/test_concept_queries.py b/tests/test_concept_queries.py index 57e28b8..2a55904 100644 --- a/tests/test_concept_queries.py +++ b/tests/test_concept_queries.py @@ -3,14 +3,18 @@ from __future__ import annotations from datetime import date +from typing import Iterator import pytest import sqlalchemy as sa from sqlalchemy.orm import Session +from oa_configurator.testing import isolated_test_database + from omop_alchemy.cdm.model.vocabulary import Concept from omop_alchemy.cdm.query import ConceptFilter +from omop_graph.config import OmopGraphConfig from omop_graph.graph.nodes import ConceptView from omop_graph.graph.queries import ( q_concept_filtered, @@ -21,45 +25,52 @@ @pytest.fixture() -def concept_engine() -> sa.Engine: - engine = sa.create_engine("sqlite+pysqlite:///:memory:", future=True) - Concept.__table__.create(engine) - - valid_from = date(2000, 1, 1) - valid_until = date(2099, 12, 31) - - def concept( - concept_id: int, - *, - standard_concept: str | None, - invalid_reason: str | None, - ) -> Concept: - return Concept( - concept_id=concept_id, - concept_name="Shared label", - domain_id="Condition", - vocabulary_id="SNOMED", - concept_class_id="Clinical Finding", - standard_concept=standard_concept, - concept_code=f"TEST-{concept_id}", - valid_start_date=valid_from, - valid_end_date=valid_until, - invalid_reason=invalid_reason, - ) +def concept_engine() -> Iterator[sa.Engine]: + with isolated_test_database( + OmopGraphConfig, + "test_cdm_db_sqlite", + dialect="sqlite", + future=True, + execution_options={"schema_translate_map": {None: None, "vocab": None, "results": None}}, + ) as db: + engine = db.connection.engine + Concept.__table__.create(engine) + + valid_from = date(2000, 1, 1) + valid_until = date(2099, 12, 31) + + def concept( + concept_id: int, + *, + standard_concept: str | None, + invalid_reason: str | None, + ) -> Concept: + return Concept( + concept_id=concept_id, + concept_name="Shared label", + domain_id="Condition", + vocabulary_id="SNOMED", + concept_class_id="Clinical Finding", + standard_concept=standard_concept, + concept_code=f"TEST-{concept_id}", + valid_start_date=valid_from, + valid_end_date=valid_until, + invalid_reason=invalid_reason, + ) - with Session(engine) as session: - session.add_all( - [ - concept(1, standard_concept="S", invalid_reason=None), - concept(2, standard_concept="C", invalid_reason=" "), - concept(3, standard_concept=None, invalid_reason=None), - concept(4, standard_concept="S", invalid_reason="U"), - concept(5, standard_concept=" ", invalid_reason="X"), - ] - ) - session.commit() + with Session(engine) as session: + session.add_all( + [ + concept(1, standard_concept="S", invalid_reason=None), + concept(2, standard_concept="C", invalid_reason=" "), + concept(3, standard_concept=None, invalid_reason=None), + concept(4, standard_concept="S", invalid_reason="U"), + concept(5, standard_concept=" ", invalid_reason="X"), + ] + ) + session.commit() - return engine + yield engine def test_concept_filter_applies_canonical_graph_constraints( diff --git a/tests/test_edges_same_connection_regression.py b/tests/test_edges_same_connection_regression.py new file mode 100644 index 0000000..3fc64f0 --- /dev/null +++ b/tests/test_edges_same_connection_regression.py @@ -0,0 +1,34 @@ +"""Same-connection regression coverage for kg.py's split-vocab merge (Phase 3.2). + +``KnowledgeGraph.iter_edges``/``predicate``/``predicates`` gained a +split-connection branch (see ``test_vocab_split_connection.py``). This pins +the default, unsplit path -- the one every existing deployment actually +uses -- stays on the original single eager join, protecting against a +future edit accidentally forcing the split-path branch unconditionally. +""" + +from __future__ import annotations + +from omop_graph.extensions.omop_alchemy import PredicateKind +from omop_graph.graph.kg import KnowledgeGraph + + +def test_edges_use_single_eager_join_when_no_split_is_configured( + mock_cdm_kg: KnowledgeGraph, +) -> None: + assert mock_cdm_kg._vocab_split is False + + edges = mock_cdm_kg.edges( + concept_ids=900001, + direction="out", + active_only=False, + within_domain=False, + ) + + assert len(edges) == 1 + edge = edges[0] + assert edge.subject_id == 900001 + assert edge.object_id == 196653 + assert edge.predicate_id == "maps to" + assert edge.predicate_kind == PredicateKind.IDENTITY + assert edge.predicate_subkind == "mapping" diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py new file mode 100644 index 0000000..fad54b5 --- /dev/null +++ b/tests/test_oaklib_schema_awareness.py @@ -0,0 +1,194 @@ +"""OAK-lib adapter schema-awareness gap (Phase 4). + +`omop_resource()` used to resolve a full `ResolvedCDMDatabase` internally +but discard it after extracting `database.connection.url`, so +`OMOPAlchemyImplementation`'s internally-built engine never carried +`schema_translate_map`. There are two genuinely different construction +paths here, tested separately: + +1. `kg=`-injected construction (what any caller in this stack that can + reach a `Resolver` should use): the internal `make_engine()` call still + runs but its result is discarded, so this path was never actually + broken by the bug. A caller building its own schema-aware engine and + passing `kg=` already worked. Tested here anyway, since it's the + pattern this stack's own production code should use and had no + coverage at all. +2. Bare `engine_string=`/`resource=`-only construction (OAK-lib's own + generic `materialize()` invocation, which only ever gets a URL string, + never a live connection): this is the one path the bug actually broke, + and the only one that needed the `execution_options` fix. +""" + +from __future__ import annotations + +from datetime import date + +import sqlalchemy as sa + +from oa_configurator import qualified +from oa_configurator.testing import isolated_test_schema +from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary +from orm_loader.helpers import Base + +from omop_graph.cli import relationship_classification +from omop_graph.db.session import make_engine +from omop_graph.graph.kg import KnowledgeGraph +from omop_graph.oaklib_interface.omop_factory import omop_resource +from omop_graph.oaklib_interface.omop_implementation import OMOPAlchemyImplementation +from omop_graph.oaklib_interface.omop_resource import OMOPOntologyResource + +_META_CONCEPT_ID = 0 +_CONCEPT_ID = 1001 +_TODAY = date(2020, 1, 1) +_FAR_FUTURE = date(2099, 12, 31) +_VOCAB_TABLES = (Domain.__table__, Vocabulary.__table__, Concept_Class.__table__, Concept.__table__) + + +def _seed_one_concept(bindable: sa.Engine | sa.Connection, *, concept_id: int, name: str) -> None: + """Minimal, real vocab bootstrap: Domain/Vocabulary/Concept_Class/Concept + form a genuine insert cycle in Postgres (each references-row's own + *_concept_id FK requires a Concept row to exist, and that Concept row's + domain_id/vocabulary_id/concept_class_id FKs require the reference rows + to exist), the same cycle production bulk-loads handle by disabling FK + triggers for the load, then re-enabling them. Accepts either an Engine + or an already-open Connection: an Engine has no .execute() of its own, + so this opens one short-lived connection for the trigger toggles. + """ + opened_here = isinstance(bindable, sa.Engine) + conn = bindable.connect() if opened_here else bindable + try: + for table in _VOCAB_TABLES: + conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name)} DISABLE TRIGGER ALL")) + # Only commit a connection opened here: pg_db's own Connection is + # already inside an explicit, rollback-based outer transaction, and + # calling .commit() on it directly would end that transaction for + # real, defeating the isolation the fixture exists to provide. A + # freshly-opened connection has no such transaction to protect, and + # DDL needs to actually persist for the Session below (a genuinely + # separate connection from the pool) to see it. + if opened_here: + conn.commit() + finally: + if opened_here: + conn.close() + + with sa.orm.Session(bindable) as session: + session.add_all( + [ + Concept( + concept_id=_META_CONCEPT_ID, + concept_name="Meta concept", + domain_id="Metadata", + vocabulary_id="OMOP", + concept_class_id="Metadata", + standard_concept="S", + concept_code="META", + valid_start_date=_TODAY, + valid_end_date=_FAR_FUTURE, + ), + Concept( + concept_id=concept_id, + concept_name=name, + domain_id="Metadata", + vocabulary_id="OMOP", + concept_class_id="Metadata", + standard_concept="S", + concept_code=str(concept_id), + valid_start_date=_TODAY, + valid_end_date=_FAR_FUTURE, + ), + Domain(domain_id="Metadata", domain_name="Metadata", domain_concept_id=_META_CONCEPT_ID), + Vocabulary( + vocabulary_id="OMOP", + vocabulary_name="OMOP", + vocabulary_reference="local", + vocabulary_version="test", + vocabulary_concept_id=_META_CONCEPT_ID, + ), + Concept_Class( + concept_class_id="Metadata", + concept_class_name="Metadata", + concept_class_concept_id=_META_CONCEPT_ID, + ), + ] + ) + session.commit() + + conn = bindable.connect() if opened_here else bindable + try: + for table in _VOCAB_TABLES: + conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name)} ENABLE TRIGGER ALL")) + if opened_here: + conn.commit() + finally: + if opened_here: + conn.close() + + +def test_omop_resource_execution_options_carry_the_configured_schema() -> None: + """Object inspection only: no query, no data, no database connection + at all. omop_resource() resolves the active config's schema_translate_map + purely from typed config data, and make_engine() with an explicit url= + never opens a connection either (Engine construction is lazy).""" + resource = omop_resource() + + engine = make_engine(resource.url, execution_options=resource.execution_options) + + assert resource.execution_options is not None + assert ( + engine.get_execution_options()["schema_translate_map"] + == resource.execution_options["schema_translate_map"] + ) + + +def test_kg_injection_path_resolves_against_the_configured_schema(pg_db) -> None: + """The kg= injection path this stack's own production code should + prefer: build a schema-aware engine externally, wrap it, pass kg=. + The internal make_engine(engine_string, ...) call still runs but its + result is discarded. engine_string must still be a resolvable dialect, + just never actually connected to, so a bare "sqlite:///:memory:" + placeholder is fine here.""" + schema = "phase4_oaklib_kg_injection" + conn = pg_db.connection + conn.execute(sa.text(f"CREATE SCHEMA {schema}")) + scoped = conn.execution_options( + schema_translate_map={None: schema, "vocab": schema, "results": schema} + ) + Base.metadata.create_all(bind=scoped, checkfirst=True) + _seed_one_concept(scoped, concept_id=_CONCEPT_ID, name="Test concept") + relationship_classification(engine=scoped) + + kg = KnowledgeGraph(cdm_engine=scoped) + adapter = OMOPAlchemyImplementation(engine_string="sqlite:///:memory:", kg=kg) + + assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Test concept" + + +def test_bare_engine_string_path_resolves_against_the_configured_schema(pg_db) -> None: + """The one path that can't be dependency-injected: OAK-lib's own + generic materialize() mechanism only ever hands a URL string to + OMOPAlchemyImplementation, never a live connection. This is the only + remaining legitimate use of isolated_test_schema() in this whole plan, + since it's the only caller that genuinely can't accept pg_db's + rolled-back Connection. Construction goes through omop_resource(), + which needs a real, committed, independently-connectable schema.""" + with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_bare") as schema: + engine = pg_db.connection.engine.execution_options( + schema_translate_map={None: schema, "vocab": schema, "results": schema} + ) + Base.metadata.create_all(bind=engine, checkfirst=True) + _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Bare-string concept") + relationship_classification(engine=engine) + + resource = OMOPOntologyResource( + # str(url) masks the password by default (renders "***"), and + # this is the one place that string actually needs to be usable + # to open a real connection, not just for display. + url=pg_db.connection.engine.url.render_as_string(hide_password=False), + execution_options={ + "schema_translate_map": {None: schema, "vocab": schema, "results": schema} + }, + ) + adapter = OMOPAlchemyImplementation(resource=resource) + + assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Bare-string concept" diff --git a/tests/test_pg_db_fixture.py b/tests/test_pg_db_fixture.py new file mode 100644 index 0000000..0754ddc --- /dev/null +++ b/tests/test_pg_db_fixture.py @@ -0,0 +1,23 @@ +"""Smoke test for the pg_db fixture (Phase 0 of the schema_translate_map fix). + +omop-graph had no Postgres test fixture at all before this -- this proves +the newly-added one actually works, not just that it's wired up. +""" + +import sqlalchemy as sa + + +def test_pg_db_yields_a_working_connection_and_session(pg_db): + assert pg_db.connection.execute(sa.text("SELECT 1")).scalar() == 1 + assert pg_db.session.connection() is pg_db.connection + + +def test_pg_db_rolls_back_between_tests(pg_db): + """A second, independent test using the same fixture must not see + anything from a prior test -- proving isolation, not just connectivity.""" + exists = pg_db.connection.execute( + sa.text("SELECT to_regclass('pg_db_fixture_smoke_test')") + ).scalar() + assert exists is None + pg_db.connection.execute(sa.text("CREATE TABLE pg_db_fixture_smoke_test (id INT)")) + # Never committed -- rolled back automatically when this test ends. diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py new file mode 100644 index 0000000..6751b0a --- /dev/null +++ b/tests/test_relationship_classification.py @@ -0,0 +1,71 @@ +"""Regression test for the originally reported bug: relationship-classification +silently ignored the configured CDM schema. + +Runs entirely on Phase 0's rollback-based ``pg_db`` fixture: real Postgres, +a non-default schema created inside the test's own already-open transaction, +nothing ever committed. Also covers the DROP TYPE naming-mismatch fix +(Phase 4): the enum column never set an explicit ``name=``, so the real +generated type is ``predicatekind``, not the ``predicatekindenum`` the old +raw SQL referenced. Confirmed here rather than assumed. +""" + +from __future__ import annotations + +import sqlalchemy as sa + +from orm_loader.helpers import Base + +from omop_graph.cli import relationship_classification +from omop_graph.extensions.omop_alchemy import RelationshipClass, RelationshipMapping + + +def _scoped_connection(pg_db, schema: str) -> sa.Connection: + conn = pg_db.connection + conn.execute(sa.text(f"CREATE SCHEMA {schema}")) + return conn.execution_options( + schema_translate_map={None: schema, "vocab": schema, "results": schema} + ) + + +def test_relationship_classification_respects_the_configured_schema(pg_db): + scoped = _scoped_connection(pg_db, "phase4_regression_test") + Base.metadata.create_all(bind=scoped, checkfirst=True) + + relationship_classification(engine=scoped) + + n_class = scoped.execute( + sa.select(sa.func.count()).select_from(RelationshipClass.__table__) + ).scalar() + n_mapping = scoped.execute( + sa.select(sa.func.count()).select_from(RelationshipMapping.__table__) + ).scalar() + assert n_class and n_class > 0 + assert n_mapping and n_mapping > 0 + + actual_schema = pg_db.connection.execute( + sa.text( + "SELECT table_schema FROM information_schema.tables " + "WHERE table_name = 'relationship_class'" + ) + ).scalar() + assert actual_schema == "phase4_regression_test" + + enum_type = pg_db.connection.execute( + sa.text("SELECT typname FROM pg_type WHERE typname = 'predicatekind'") + ).scalar() + assert enum_type == "predicatekind" + + +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.""" + scoped = _scoped_connection(pg_db, "phase4_idempotent_test") + Base.metadata.create_all(bind=scoped, checkfirst=True) + + relationship_classification(engine=scoped) + relationship_classification(engine=scoped) + + n_class = scoped.execute( + sa.select(sa.func.count()).select_from(RelationshipClass.__table__) + ).scalar() + assert n_class and n_class > 0 diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py new file mode 100644 index 0000000..3741558 --- /dev/null +++ b/tests/test_vocab_split_connection.py @@ -0,0 +1,318 @@ +"""Split-connection vocab routing (Phase 3.2 of the schema_translate_map fix). + +``q_edges``, ``q_predicate_row_with_ancestry``, and ``q_all_predicates_with_ancestry`` +join a vocab-role table (Relationship/Concept_Relationship) against +RelationshipMapping (an omop-graph extension table, not vocab-role). When +``vocab_connection`` names a physically different server than ``connection``, +a single SQL join can't span both. ``KnowledgeGraph`` fetches each side +from its own engine and merges in Python instead (see kg.py's +``_vocab_split``/``_predicate_from_rows``/``_relationship_mapping_lookup``). + +Uses two genuinely distinct, real Postgres connections (``test_cdm``, +``test_orm``) standing in for "primary server" and "vocab server". Each +engine gets its own real, uniquely-named schema via +``oa_configurator.testing.isolated_test_schema()``, since rollback-based +isolation (a single already-open connection) can't stand in for two +genuinely separate physical connections. +""" + +from __future__ import annotations + +from datetime import date +from typing import Iterator, NamedTuple + +import pytest +import sqlalchemy as sa +import sqlalchemy.orm as so + +from oa_configurator.testing import isolated_test_database, isolated_test_schema +from orm_loader.config import OrmLoaderConfig +from orm_loader.helpers import Base + +from omop_alchemy.cdm.model.vocabulary import ( + Concept, + Concept_Class, + Concept_Relationship, + Domain, + Relationship, + Vocabulary, +) + +from omop_graph.config import OmopGraphConfig +from omop_graph.extensions.omop_alchemy import ( + PredicateKind, + RelationshipClass, + RelationshipMapping, +) +from omop_graph.graph.kg import KnowledgeGraph + +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + +META_CONCEPT_ID = 0 +SUBJECT_CONCEPT_ID = 1 +OBJECT_CONCEPT_ID = 2 +_TODAY = date(2020, 1, 1) +_FAR_FUTURE = date(2099, 12, 31) + +_VOCAB_TABLES = ( + Domain.__table__, + Vocabulary.__table__, + Concept_Class.__table__, + Concept.__table__, + Relationship.__table__, + Concept_Relationship.__table__, +) + +# Postgres has no cross-database inline FK (unlike cross-schema, which works +# fine within one database) -- RelationshipMapping's ORM-mapped FK to +# relationship.relationship_id can't be created as DDL when vocab lives on a +# genuinely different database, confirmed empirically while writing this +# test. That FK isn't what's under test here (the Python-side merge is), so +# these shadow tables reproduce RelationshipClass/RelationshipMapping's +# columns without it -- the real ORM classes read/write them identically, +# since a SELECT/INSERT only depends on column shape, not on constraint DDL. +_shadow_metadata = sa.MetaData() +_shadow_relationship_class = sa.Table( + "relationship_class", + _shadow_metadata, + sa.Column( + "predicate_kind", + sa.Enum(PredicateKind, values_callable=lambda obj: [e.value for e in obj]), + primary_key=True, + ), + sa.Column("predicate_subkind", sa.String(20), primary_key=True), + sa.Column("description", sa.String(80), nullable=False), + sa.Column("semantics", sa.String(40), nullable=False), + sa.Column("inference", sa.String(40), nullable=False), +) +_shadow_relationship_mapping = sa.Table( + "relationship_mapping", + _shadow_metadata, + sa.Column("relationship_id", sa.String(20), primary_key=True), + sa.Column( + "predicate_kind", + sa.Enum(PredicateKind, values_callable=lambda obj: [e.value for e in obj]), + primary_key=True, + ), + sa.Column("predicate_subkind", sa.String(20), primary_key=True), +) + + +class _Engines(NamedTuple): + primary: sa.Engine + vocab: sa.Engine + + +@pytest.fixture() +def split_engines() -> Iterator[_Engines]: + """A primary connection (RelationshipMapping/RelationshipClass) and a + genuinely separate physical vocab connection (Relationship/Concept/ + Concept_Relationship).""" + with ( + isolated_test_database(OmopGraphConfig, "test_cdm_db_pg") as primary_db, + isolated_test_database(OrmLoaderConfig, "test_orm_db_pg") as vocab_db, + ): + primary_raw = primary_db.connection.engine + vocab_raw = vocab_db.connection.engine + + with ( + isolated_test_schema(primary_raw) as primary_schema, + isolated_test_schema(vocab_raw) as vocab_schema, + ): + primary_engine = primary_raw.execution_options( + schema_translate_map={None: primary_schema, "vocab": primary_schema, "results": primary_schema} + ) + vocab_engine = vocab_raw.execution_options( + schema_translate_map={None: vocab_schema, "vocab": vocab_schema, "results": vocab_schema} + ) + + _shadow_metadata.create_all(primary_engine) + Base.metadata.create_all(vocab_engine, tables=_VOCAB_TABLES, checkfirst=True) + + # Domain/Vocabulary/Concept_Class/Concept form a genuine bootstrap + # cycle (each reference row's own *_concept_id FK requires a Concept + # row to already exist, and that Concept row's domain_id/ + # vocabulary_id/concept_class_id FKs require the reference rows to + # already exist), the same cycle production bulk-loads handle by + # disabling FK triggers for the load, then re-enabling them. + with vocab_engine.begin() as conn: + for table in _VOCAB_TABLES: + conn.execute(sa.text(f'ALTER TABLE "{vocab_schema}"."{table.name}" DISABLE TRIGGER ALL')) + + _seed(primary_engine, vocab_engine) + + with vocab_engine.begin() as conn: + for table in _VOCAB_TABLES: + conn.execute(sa.text(f'ALTER TABLE "{vocab_schema}"."{table.name}" ENABLE TRIGGER ALL')) + + yield _Engines(primary=primary_engine, vocab=vocab_engine) + + +def _seed(primary_engine: sa.Engine, vocab_engine: sa.Engine) -> None: + with so.Session(vocab_engine) as session: + session.add_all( + [ + Concept( + concept_id=META_CONCEPT_ID, + concept_name="Meta concept", + domain_id="Metadata", + vocabulary_id="OMOP", + concept_class_id="Metadata", + standard_concept="S", + concept_code="META", + valid_start_date=_TODAY, + valid_end_date=_FAR_FUTURE, + ), + Concept( + concept_id=SUBJECT_CONCEPT_ID, + concept_name="Subject concept", + domain_id="Condition", + vocabulary_id="SNOMED", + concept_class_id="Clinical Finding", + standard_concept="S", + concept_code="SUBJ", + valid_start_date=_TODAY, + valid_end_date=_FAR_FUTURE, + ), + Concept( + concept_id=OBJECT_CONCEPT_ID, + concept_name="Object concept", + domain_id="Condition", + vocabulary_id="SNOMED", + concept_class_id="Clinical Finding", + standard_concept="S", + concept_code="OBJ", + valid_start_date=_TODAY, + valid_end_date=_FAR_FUTURE, + ), + Domain(domain_id="Metadata", domain_name="Metadata", domain_concept_id=META_CONCEPT_ID), + Domain(domain_id="Condition", domain_name="Condition", domain_concept_id=META_CONCEPT_ID), + Vocabulary( + vocabulary_id="OMOP", + vocabulary_name="OMOP", + vocabulary_reference="local", + vocabulary_version="test", + vocabulary_concept_id=META_CONCEPT_ID, + ), + Vocabulary( + vocabulary_id="SNOMED", + vocabulary_name="SNOMED", + vocabulary_reference="local", + vocabulary_version="test", + vocabulary_concept_id=META_CONCEPT_ID, + ), + Concept_Class( + concept_class_id="Metadata", + concept_class_name="Metadata", + concept_class_concept_id=META_CONCEPT_ID, + ), + Concept_Class( + concept_class_id="Clinical Finding", + concept_class_name="Clinical Finding", + concept_class_concept_id=META_CONCEPT_ID, + ), + Relationship( + relationship_id="maps to", + relationship_name="Maps to", + is_hierarchical="0", + defines_ancestry="0", + reverse_relationship_id="mapped from", + relationship_concept_id=META_CONCEPT_ID, + ), + Relationship( + relationship_id="mapped from", + relationship_name="Mapped from", + is_hierarchical="0", + defines_ancestry="0", + reverse_relationship_id="maps to", + relationship_concept_id=META_CONCEPT_ID, + ), + Concept_Relationship( + concept_id_1=SUBJECT_CONCEPT_ID, + concept_id_2=OBJECT_CONCEPT_ID, + relationship_id="maps to", + valid_start_date=_TODAY, + valid_end_date=_FAR_FUTURE, + invalid_reason=None, + ), + ] + ) + session.commit() + + with so.Session(primary_engine) as session: + session.add_all( + [ + RelationshipClass( + predicate_kind=PredicateKind.IDENTITY, + predicate_subkind="mapping", + description="Identity mapping", + semantics="identity", + inference="none", + ), + RelationshipMapping( + relationship_id="maps to", + predicate_kind=PredicateKind.IDENTITY, + predicate_subkind="mapping", + ), + RelationshipMapping( + relationship_id="mapped from", + predicate_kind=PredicateKind.IDENTITY, + predicate_subkind="mapping", + ), + ] + ) + session.commit() + + +def _split_kg(engines: _Engines) -> KnowledgeGraph: + return KnowledgeGraph(cdm_engine=engines.primary, vocab_engine=engines.vocab) + + +def test_predicate_merges_across_split_connections(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + predicate = kg.predicate("maps to") + + assert predicate.relationship_id == "maps to" + assert predicate.reverse_id == "mapped from" + assert predicate.predicate_kind == PredicateKind.IDENTITY + assert predicate.predicate_subkind == "mapping" + + +def test_predicates_merges_across_split_connections(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + by_id = {p.relationship_id: p for p in kg.predicates()} + + assert set(by_id) == {"maps to", "mapped from"} + assert by_id["maps to"].predicate_kind == PredicateKind.IDENTITY + assert by_id["maps to"].predicate_subkind == "mapping" + + +def test_edges_merges_across_split_connections(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + edges = kg.edges( + concept_ids=SUBJECT_CONCEPT_ID, + direction="out", + active_only=False, + within_domain=False, + ) + + assert len(edges) == 1 + edge = edges[0] + assert edge.subject_id == SUBJECT_CONCEPT_ID + assert edge.object_id == OBJECT_CONCEPT_ID + assert edge.predicate_id == "maps to" + assert edge.predicate_kind == PredicateKind.IDENTITY + assert edge.predicate_subkind == "mapping" + + +def test_edges_predicate_kinds_filter_applies_after_merge(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + edges = kg.edges( + concept_ids=SUBJECT_CONCEPT_ID, + direction="out", + active_only=False, + within_domain=False, + predicate_kinds=frozenset({PredicateKind.HIERARCHY}), + ) + + assert edges == () From ce1d65264cccb6368059590d294b43238ddbf34d Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 1 Sep 2026 05:32:09 +0000 Subject: [PATCH 2/5] Update docstring of oaklib interface --- .../oaklib_interface/omop_implementation.py | 62 ++++++++++--------- 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/src/omop_graph/oaklib_interface/omop_implementation.py b/src/omop_graph/oaklib_interface/omop_implementation.py index 30714dc..6333c5d 100644 --- a/src/omop_graph/oaklib_interface/omop_implementation.py +++ b/src/omop_graph/oaklib_interface/omop_implementation.py @@ -851,23 +851,30 @@ class OMOPAlchemyImplementation( # type: ignore[override] Parameters ---------- engine_string : str | URL | None, optional - The database connection string. Required unless ``resource`` is given. + The database connection string. Ignored when ``kg`` is given + directly; required otherwise, unless ``resource`` is given. resource : OMOPOntologyResource | None, optional An existing resource object. Takes precedence over ``engine_string`` when - both are supplied. To use the oa-configurator-configured default, - resolve it explicitly via ``omop_resource()`` and pass it here. + both are supplied. Ignored when ``kg`` is given directly. To use the + oa-configurator-configured default, resolve it explicitly via + ``omop_resource()`` and pass it here. kg : KnowledgeGraph | None, optional - An existing Knowledge Graph instance. If None, one is created from - ``engine_string`` / ``resource``. + An existing Knowledge Graph instance. Takes this class's own engine + construction out of the picture entirely -- the caller already built + (and is responsible for) whatever engine ``kg`` wraps, so + ``engine_string``/``resource`` are neither required nor consulted. + If None, a ``KnowledgeGraph`` is created from ``engine_string`` / + ``resource`` instead. kg_emb_config : KnowledgeGraphEmbeddingConfiguration | None, optional Embedding configuration forwarded to the ``KnowledgeGraph`` constructor. Required to enable embedding-based similarity. See :class:`~omop_graph.graph.kg.KnowledgeGraphEmbeddingConfiguration`. + Ignored when ``kg`` is given directly. Raises ------ ValueError - If neither ``engine_string`` nor ``resource`` is given. + If ``kg`` is not given and neither ``engine_string`` nor ``resource`` is. """ def __init__( @@ -878,33 +885,28 @@ def __init__( kg_emb_config: Optional[KnowledgeGraphEmbeddingConfiguration] = None, **kwargs, ): - if engine_string is not None: - self.engine_string = engine_string - self.resource = resource or omop_resource(url=self.engine_string) - elif resource is not None: - self.resource = resource - self.engine_string = self.resource.url - else: - raise ValueError( - "OMOPAlchemyImplementation requires either 'engine_string' or " - "'resource'. To use the oa-configurator-configured default, " - "resolve it explicitly first, e.g. " - "OMOPAlchemyImplementation(resource=omop_resource())." - ) - - assert self.engine_string is not None, ( - "No database URL provided for OMOPAlchemyImplementation" - ) - - engine = make_engine( - self.engine_string, - engine_kwargs={"echo": False, "future": True}, - execution_options=self.resource.execution_options, - ) - self._connection = None if kg is None: + if engine_string is not None: + self.engine_string = engine_string + self.resource = resource or omop_resource(url=self.engine_string) + elif resource is not None: + self.resource = resource + self.engine_string = self.resource.url + else: + raise ValueError( + "OMOPAlchemyImplementation requires 'kg', or one of " + "'engine_string'/'resource'. To use the " + "oa-configurator-configured default, resolve it explicitly " + "first, e.g. OMOPAlchemyImplementation(resource=omop_resource())." + ) + + engine = make_engine( + self.engine_string, + engine_kwargs={"echo": False, "future": True}, + execution_options=self.resource.execution_options, + ) kg = KnowledgeGraph(emb_config=kg_emb_config, cdm_engine=engine) bind_default_renderers(kg) From 08a69c10bcde706c770ac5eed3ecfcb56ef925d2 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 1 Sep 2026 06:11:57 +0000 Subject: [PATCH 3/5] Update CI --- .github/workflows/ci.yml | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1cf941a..0e34f48 100644 --- a/.github/workflows/ci.yml +++ b/.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: + uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres.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 From 2a88c335a1a689ef0241f7e910aa6be1a89ea701 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 4 Sep 2026 06:09:19 +0000 Subject: [PATCH 4/5] CI Fix, convenience methods --- .github/workflows/ci.yml | 4 +- src/omop_graph/cli.py | 54 ++++++++----- src/omop_graph/db/session.py | 22 +++++- tests/test_schema_provenance_guard.py | 105 ++++++++++++++++++++++++++ 4 files changed, 162 insertions(+), 23 deletions(-) create mode 100644 tests/test_schema_provenance_guard.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0e34f48..95f5c08 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: omop_graph_test setup-commands: | diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index 06b32de..10d0c8e 100644 --- a/src/omop_graph/cli.py +++ b/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,7 +11,7 @@ import typer from sqlalchemy.orm import sessionmaker -from oa_configurator import ensure_schema, schema_of +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 @@ -17,7 +19,7 @@ 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 @@ -59,10 +61,24 @@ def packaged_predicate_csv_dir() -> Path: return Path(str(resources.files("omop_graph") / "data")) +@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: Optional[str] = None, *, engine: sa.Engine | sa.Connection | None = None, + resolved: ResolvedCDMDatabase | None = None, ) -> None: """Load pre-classified predicates into the database. @@ -74,7 +90,12 @@ def relationship_classification( omop-graph. engine : sqlalchemy.Engine or sqlalchemy.Connection, optional Bindable to run against. Defaults to the active oa-configurator - config's resolved CDM engine. + 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() @@ -145,7 +166,8 @@ def relationship_classification( ) if engine is None: - engine = make_engine() + resolved = resolve_cdm_database() + engine = resolved.create_engine() db_schema = schema_of(engine) ensure_schema(engine, db_schema) ensure_schema(engine, STAGING_SCHEMA) @@ -164,27 +186,25 @@ def relationship_classification( f"{loader_backend.qualified_staging_name(RelationshipClass.__tablename__)} CASCADE" ), ) - if isinstance(engine, sa.Engine): - with engine.begin() as conn: - for stmt in drop_staging_sql: - conn.execute(stmt) - else: + with _open_connection(engine) as connection: for stmt in drop_staging_sql: - engine.execute(stmt) + 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", meaning this line never matched anything, - # with IF EXISTS silently no-op'ing every run. drop_all(tables=[...]) already - # drops a shared Enum type exactly once, correctly deduped, once every - # table using it is in the same tables= list (true here, both tables - # always move together), so no manual DROP TYPE is needed at all. + # 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( diff --git a/src/omop_graph/db/session.py b/src/omop_graph/db/session.py index 387bf8c..246a9b1 100644 --- a/src/omop_graph/db/session.py +++ b/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, *, @@ -42,9 +58,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 diff --git a/tests/test_schema_provenance_guard.py b/tests/test_schema_provenance_guard.py new file mode 100644 index 0000000..e4a61c4 --- /dev/null +++ b/tests/test_schema_provenance_guard.py @@ -0,0 +1,105 @@ +"""schema-provenance guard wired into relationship_classification()'s +production/CLI path (engine=None). + +Monkeypatches omop_graph.cli.resolve_cdm_database (the one call site) rather +than the whole oa-configurator config chain, to point at a real, isolated +Postgres schema without touching the active on-disk config. + +Only the "fires on a genuinely reconfigured schema" case is covered here. +The guard's own agree/no-op/test_only semantics are already exhaustively +covered at the primitive level in oa-configurator's own test suite; what's +worth proving per consuming repo is that this call site is actually wired +to it, and a wiring mistake would show up here too. + +resolved.create_engine() below is a real, committing engine (not the +rollback-protected pg_db.connection), so every provenance row this test +writes is a genuine commit. cleanup_after_test deletes this test's own +schema_provenance rows at teardown (see Phase 10.12 in the plan). +""" + +from __future__ import annotations + +import dataclasses +import uuid + +import pytest +import sqlalchemy as sa +from oa_configurator import Role, SchemaDriftError +from oa_configurator.domains.resources.sql import ( + SCHEMA_PROVENANCE_SCHEMA, + _schema_provenance_table, + record_schema_provenance, +) +from oa_configurator.testing import delete_rows_on_cleanup, isolated_test_schema + +from orm_loader.helpers import Base + +from omop_graph import cli as omop_graph_cli +from omop_graph.extensions.omop_alchemy import RelationshipClass + +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + + +def _resolved(pg_db, *, database_name: str, schema: str): + """pg_db.resolved with a unique name (the guard's own key includes it), + all three schemas pointed at schema, and connection.test_only forced + False so the guard doesn't no-op against pg_db's own test-only marking. + """ + return dataclasses.replace( + pg_db.resolved, + name=database_name, + schema_name=schema, + vocab_schema=schema, + results_schema=schema, + connection=dataclasses.replace(pg_db.resolved.connection, test_only=False), + ) + + +def test_relationship_classification_guard_fires_on_reconfigured_schema( + pg_db, monkeypatch, cleanup_after_test +): + database_name = f"graph_guard_db_{uuid.uuid4().hex[:8]}" + pg_engine = pg_db.connection.engine + table = _schema_provenance_table(SCHEMA_PROVENANCE_SCHEMA) + delete_rows_on_cleanup( + cleanup_after_test, pg_engine, table, table.c.database_name == database_name + ) + with ( + isolated_test_schema(pg_engine, prefix="graph_guard_a") as schema_a, + isolated_test_schema(pg_engine, prefix="graph_guard_b") as schema_b, + ): + resolved_a = _resolved(pg_db, database_name=database_name, schema=schema_a) + engine_a = resolved_a.create_engine() + try: + Base.metadata.create_all(bind=engine_a, checkfirst=True) + # The line above populates the schema outside the guard's own + # view, so an explicit baseline is needed first, mirroring what + # a real deployment retrofitting provenance onto an + # already-populated database would have to do. + with engine_a.begin() as conn: + record_schema_provenance( + conn, resolved_a, role=Role.PRIMARY, new_schema=schema_a, reason="test setup baseline" + ) + monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: resolved_a) + omop_graph_cli.relationship_classification() + finally: + engine_a.dispose() + + resolved_b = _resolved(pg_db, database_name=database_name, schema=schema_b) + engine_b = resolved_b.create_engine() + try: + Base.metadata.create_all(bind=engine_b, checkfirst=True) + monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: resolved_b) + with pytest.raises(SchemaDriftError): + omop_graph_cli.relationship_classification() + + # This test's own setup step above already created every CDM + # table, including these two, so the real proof the guard fired + # before any write is that they're still empty. + with engine_b.connect() as conn: + count = conn.execute( + sa.select(sa.func.count()).select_from(RelationshipClass.__table__) + ).scalar() + assert count == 0 + finally: + engine_b.dispose() From 4bc8faf86a77eb2fbb73883b4d1b52a38516ac0b Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Mon, 7 Sep 2026 05:25:27 +0000 Subject: [PATCH 5/5] Correct vocab resolving split + execution options --- docs/oaklib/interface.md | 4 +- src/omop_graph/db/session.py | 8 +- .../oaklib_interface/omop_factory.py | 11 +- .../oaklib_interface/omop_implementation.py | 14 ++- .../oaklib_interface/omop_resource.py | 12 ++ tests/test_oaklib_schema_awareness.py | 112 +++++++++++++++++- 6 files changed, 151 insertions(+), 10 deletions(-) diff --git a/docs/oaklib/interface.md b/docs/oaklib/interface.md index 918cb9a..45e05d2 100644 --- a/docs/oaklib/interface.md +++ b/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. --- diff --git a/src/omop_graph/db/session.py b/src/omop_graph/db/session.py index 246a9b1..204b52a 100644 --- a/src/omop_graph/db/session.py +++ b/src/omop_graph/db/session.py @@ -47,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 ------- diff --git a/src/omop_graph/oaklib_interface/omop_factory.py b/src/omop_graph/oaklib_interface/omop_factory.py index e1cdbe4..b5962be 100644 --- a/src/omop_graph/oaklib_interface/omop_factory.py +++ b/src/omop_graph/oaklib_interface/omop_factory.py @@ -7,7 +7,7 @@ from sqlalchemy.engine import URL from .omop_resource import OMOPOntologyResource -from oa_configurator import ResolvedCDMDatabase, Resolver +from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, ResolvedCDMDatabase, Resolver from omop_graph.config import OmopGraphConfig @@ -34,6 +34,8 @@ def omop_resource( OMOPOntologyResource """ execution_options = None + vocab_url = None + vocab_execution_options = None if url is None: resolver = Resolver.from_active_config() db_name = resolver.resolve_package_config(OmopGraphConfig).cdm_db @@ -44,10 +46,15 @@ def omop_resource( f"{type(database).__name__}" ) url = database.connection.url - execution_options = {"schema_translate_map": database.schema_translate_map()} + execution_options = {SCHEMA_TRANSLATE_MAP_KEY: database.schema_translate_map()} + if database.connection != database.vocab_connection: + vocab_url = database.vocab_connection.url + vocab_execution_options = execution_options return OMOPOntologyResource( slug=slug, url=url, execution_options=execution_options, + vocab_url=vocab_url, + vocab_execution_options=vocab_execution_options, ) diff --git a/src/omop_graph/oaklib_interface/omop_implementation.py b/src/omop_graph/oaklib_interface/omop_implementation.py index 6333c5d..9840fb5 100644 --- a/src/omop_graph/oaklib_interface/omop_implementation.py +++ b/src/omop_graph/oaklib_interface/omop_implementation.py @@ -857,7 +857,10 @@ class OMOPAlchemyImplementation( # type: ignore[override] An existing resource object. Takes precedence over ``engine_string`` when both are supplied. Ignored when ``kg`` is given directly. To use the oa-configurator-configured default, resolve it explicitly via - ``omop_resource()`` and pass it here. + ``omop_resource()`` and pass it here. When the resolved CDM database + has a genuinely separate ``vocab_connection`` configured, the + resource carries a second URL for it and a real ``vocab_engine`` is + built and passed to ``KnowledgeGraph`` alongside the primary one. kg : KnowledgeGraph | None, optional An existing Knowledge Graph instance. Takes this class's own engine construction out of the picture entirely -- the caller already built @@ -907,7 +910,14 @@ def __init__( engine_kwargs={"echo": False, "future": True}, execution_options=self.resource.execution_options, ) - kg = KnowledgeGraph(emb_config=kg_emb_config, cdm_engine=engine) + vocab_engine = None + if self.resource.vocab_url is not None: + vocab_engine = make_engine( + self.resource.vocab_url, + engine_kwargs={"echo": False, "future": True}, + execution_options=self.resource.vocab_execution_options, + ) + kg = KnowledgeGraph(emb_config=kg_emb_config, cdm_engine=engine, vocab_engine=vocab_engine) bind_default_renderers(kg) super().__init__(kg=kg, **kwargs) diff --git a/src/omop_graph/oaklib_interface/omop_resource.py b/src/omop_graph/oaklib_interface/omop_resource.py index 3fce3e2..57d5b4b 100644 --- a/src/omop_graph/oaklib_interface/omop_resource.py +++ b/src/omop_graph/oaklib_interface/omop_resource.py @@ -32,6 +32,16 @@ class OMOPOntologyResource(OntologyResource): ``schema_translate_map``). Not carried by ``url`` itself, so a caller resolving through oa-configurator needs this to keep the configured schema past this resource object. + vocab_url : str | URL, optional + Connection URL for a genuinely separate vocabulary server. ``None`` + when the CDM database has no configured vocabulary split, or when + this resource wasn't built by ``omop_resource()``'s config-resolving + path (a caller-supplied ``url=`` has no vocabulary split to carry). + vocab_execution_options : dict, optional + Forwarded to the engine built from ``vocab_url``. The same + ``schema_translate_map`` as ``execution_options``: the map itself + doesn't change across roles, only which physical connection it's + applied to. """ url: Optional[Union[str, URL]] = None # type: ignore[assignment] @@ -41,6 +51,8 @@ class OMOPOntologyResource(OntologyResource): in_memory: bool = False # type: ignore[assignment] readonly: bool = True # type: ignore[assignment] execution_options: Optional[dict] = None # type: ignore[assignment] + vocab_url: Optional[Union[str, URL]] = None # type: ignore[assignment] + vocab_execution_options: Optional[dict] = None # type: ignore[assignment] def _parsed_url(self) -> Optional[URL]: """ diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index fad54b5..2b332ae 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -25,7 +25,11 @@ import sqlalchemy as sa -from oa_configurator import qualified +from oa_configurator import ( + ResolvedCDMDatabase, + ResolvedConnection, + qualified, +) from oa_configurator.testing import isolated_test_schema from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary from orm_loader.helpers import Base @@ -36,6 +40,7 @@ from omop_graph.oaklib_interface.omop_factory import omop_resource from omop_graph.oaklib_interface.omop_implementation import OMOPAlchemyImplementation from omop_graph.oaklib_interface.omop_resource import OMOPOntologyResource +from omop_graph.config import OmopGraphConfig _META_CONCEPT_ID = 0 _CONCEPT_ID = 1001 @@ -141,6 +146,111 @@ def test_omop_resource_execution_options_carry_the_configured_schema() -> None: ) +def test_omop_resource_carries_a_configured_split_vocabulary_target(monkeypatch) -> None: + primary = ResolvedConnection( + name="primary", + url="sqlite:///primary.db", + safe_url="sqlite:///primary.db", + _engine_url=sa.make_url("sqlite:///primary.db"), + ) + vocabulary = ResolvedConnection( + name="vocabulary", + url="sqlite:///vocabulary.db", + safe_url="sqlite:///vocabulary.db", + _engine_url=sa.make_url("sqlite:///vocabulary.db"), + ) + resolved = ResolvedCDMDatabase( + name="split", + connection=primary, + schema_name=None, + vocab_connection=vocabulary, + vocab_schema=None, + results_schema=None, + ) + + class FakeResolver: + def resolve_package_config(self, config_type): + assert config_type is OmopGraphConfig + return OmopGraphConfig(cdm_db="split") + + def resolve_database(self, name): + assert name == "split" + return resolved + + monkeypatch.setattr( + "omop_graph.oaklib_interface.omop_factory.Resolver.from_active_config", + lambda: FakeResolver(), + ) + + resource = omop_resource() + + assert resource.url == primary.url + assert resource.vocab_url == vocabulary.url + assert resource.vocab_execution_options == resource.execution_options + + +def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_resource( + pg_db, +) -> None: + """The other half of the split-vocabulary wiring: omop_resource() deriving + vocab_url/vocab_execution_options is only useful if OMOPAlchemyImplementation + actually consumes them. KnowledgeGraph.__init__ eagerly queries via + cdm_engine (loading relationship-mapping data), so cdm_engine needs a + real, committed, populated schema; vocab_engine is never queried at + construction time here, so a syntactically valid but unpopulated URL is + enough to prove the wiring without a second real database.""" + with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_split") as schema: + engine = pg_db.connection.engine.execution_options( + schema_translate_map={None: schema, "vocab": schema, "results": schema} + ) + Base.metadata.create_all(bind=engine, checkfirst=True) + _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Split-wiring concept") + relationship_classification(engine=engine) + + resource = OMOPOntologyResource( + url=pg_db.connection.engine.url.render_as_string(hide_password=False), + execution_options={ + "schema_translate_map": {None: schema, "vocab": schema, "results": schema} + }, + vocab_url="sqlite:///:memory:", + vocab_execution_options={ + "schema_translate_map": {None: None, "vocab": None, "results": None} + }, + ) + + adapter = OMOPAlchemyImplementation(resource=resource) + + assert adapter.kg.vocab_engine is not adapter.kg.cdm_engine + assert str(adapter.kg.vocab_engine.url) == "sqlite:///:memory:" + assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Split-wiring concept" + + +def test_omop_alchemy_implementation_reuses_one_engine_when_no_split_is_configured( + pg_db, +) -> None: + """A resource with no vocab_url (the common case) must not build a second + engine at all, confirming the new branch is additive, not a regression + for every construction that isn't split.""" + with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_nosplit") as schema: + engine = pg_db.connection.engine.execution_options( + schema_translate_map={None: schema, "vocab": schema, "results": schema} + ) + Base.metadata.create_all(bind=engine, checkfirst=True) + _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="No-split concept") + relationship_classification(engine=engine) + + resource = OMOPOntologyResource( + url=pg_db.connection.engine.url.render_as_string(hide_password=False), + execution_options={ + "schema_translate_map": {None: schema, "vocab": schema, "results": schema} + }, + ) + + adapter = OMOPAlchemyImplementation(resource=resource) + + assert adapter.kg.vocab_engine is adapter.kg.cdm_engine + + def test_kg_injection_path_resolves_against_the_configured_schema(pg_db) -> None: """The kg= injection path this stack's own production code should prefer: build a schema-aware engine externally, wrap it, pass kg=.