From 01d4e74e7409cce4cde6b295aed7d052db05d5d7 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 1 Sep 2026 05:00:32 +0000 Subject: [PATCH 01/20] 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 02/20] 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 03/20] 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 04/20] 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 05/20] 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=. From 64bb0d055b41bcae5982e717116fe633b59882f4 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 01:29:03 +0000 Subject: [PATCH 06/20] Schema resolution, fix handily creating engines, tag tables belonging to schema --- src/omop_graph/cli.py | 18 ++- src/omop_graph/cli_utils/cli_add_test_data.py | 22 +++- src/omop_graph/extensions/omop_alchemy.py | 29 +++-- src/omop_graph/graph/queries.py | 6 +- .../oaklib_interface/omop_factory.py | 24 +++- .../oaklib_interface/omop_implementation.py | 61 +++++----- .../oaklib_interface/omop_resource.py | 7 ++ tests/fixtures/mock_cdm.py | 5 +- tests/test_concept_queries.py | 5 +- tests/test_fulltext_vocab_schema_postgres.py | 86 ++++++++++++++ tests/test_oaklib_schema_awareness.py | 106 +++++++++++------- tests/test_relationship_classification.py | 21 +++- tests/test_schema_provenance_guard.py | 4 +- tests/test_vocab_split_connection.py | 28 ++++- 14 files changed, 326 insertions(+), 96 deletions(-) create mode 100644 tests/test_fulltext_vocab_schema_postgres.py diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index 10d0c8e..1a3a86f 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -19,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, resolve_cdm_database +from omop_graph.db.session import resolve_cdm_database from omop_graph.extensions.omop_alchemy import RelationshipClass, RelationshipMapping from omop_graph.cli_utils import populate_test_data @@ -45,9 +45,11 @@ def _main( @app.command() def populate_with_test_data(): """Populate the database with synthetic test data.""" - engine = make_engine() + resolved = resolve_cdm_database() + engine, vocab_engine = resolved.create_engines() Session = sessionmaker(bind=engine, future=True) - populate_test_data(Session()) + VocabSession = sessionmaker(bind=vocab_engine, future=True) + populate_test_data(Session(), vocab_session=VocabSession()) def packaged_predicate_csv_dir() -> Path: @@ -168,6 +170,16 @@ def relationship_classification( if engine is None: resolved = resolve_cdm_database() engine = resolved.create_engine() + if resolved is not None and resolved.connection != resolved.vocab_connection: + raise RuntimeError( + f"relationship_classification() cannot run against database " + f"{resolved.name!r}: its vocab_connection is a genuinely separate " + "connection from the primary one, and RelationshipMapping's FK to " + "relationship.relationship_id needs both in the same database. " + "Point vocab_connection at the same connection as primary for " + "this command, or provision relationship_class/relationship_mapping " + "manually without the FK constraint." + ) db_schema = schema_of(engine) ensure_schema(engine, db_schema) ensure_schema(engine, STAGING_SCHEMA) diff --git a/src/omop_graph/cli_utils/cli_add_test_data.py b/src/omop_graph/cli_utils/cli_add_test_data.py index 394ffee..3ed3a37 100644 --- a/src/omop_graph/cli_utils/cli_add_test_data.py +++ b/src/omop_graph/cli_utils/cli_add_test_data.py @@ -8,7 +8,7 @@ from omop_alchemy.cdm.model.structural.episode import Episode from omop_alchemy.cdm.model.structural.episode_event import Episode_Event -from omop_alchemy.cdm.model.derived import Observation_Period +from omop_alchemy.cdm.model.clinical import Observation_Period from omop_alchemy.cdm.model.health_system import ( Location, Care_Site, @@ -261,12 +261,22 @@ def populate_conditions_and_modifiers( session.commit() -def populate_test_data(session): - """Brute force addition of test data for development/testing purposes.""" +def populate_test_data(session, vocab_session): + """Brute force addition of test data for development/testing purposes. + + Parameters + ---------- + session + Bound to the primary connection; every write (Person, Visit_Occurrence, + Condition_Occurrence, ...) goes through this session. + vocab_session + Bound to the vocab connection, for the Concept/Concept_Ancestor reads + below. + """ # Data concept_by_domain = pd.DataFrame( - session.query(*Concept.__table__.columns).filter( + vocab_session.query(*Concept.__table__.columns).filter( sa.or_( Concept.domain_id.in_( [ @@ -316,7 +326,7 @@ def populate_test_data(session): ) staging_parents = pd.DataFrame( - session.query(*Concept.__table__.columns) + vocab_session.query(*Concept.__table__.columns) .join( Concept_Ancestor, Concept.concept_id == Concept_Ancestor.descendant_concept_id, @@ -332,7 +342,7 @@ def populate_test_data(session): staging_parents[staging_parents.concept_name.str.contains(axis)].concept_id ) s = pd.DataFrame( - session.query(*Concept.__table__.columns) + vocab_session.query(*Concept.__table__.columns) .join( Concept_Ancestor, Concept.concept_id == Concept_Ancestor.descendant_concept_id, diff --git a/src/omop_graph/extensions/omop_alchemy.py b/src/omop_graph/extensions/omop_alchemy.py index f2a6cba..ee55fad 100644 --- a/src/omop_graph/extensions/omop_alchemy.py +++ b/src/omop_graph/extensions/omop_alchemy.py @@ -2,7 +2,13 @@ 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, role_fk +from omop_alchemy.cdm.base import ( + ReferenceTable, + cdm_table, + CDMTableBase, + merge_table_args, + role_fk, +) from oa_configurator import Role from enum import Enum @@ -25,12 +31,17 @@ class RelationshipClass(ReferenceTable, CDMTableBase, Base): """ __tablename__ = "relationship_class" + # Must match RelationshipMapping's schema, or SQLAlchemy silently builds a second, unlinked Table object. + __table_args__ = merge_table_args({"schema": Role.PRIMARY.value}) predicate_kind: so.Mapped[PredicateKind] = so.mapped_column( sa.Enum( PredicateKind, values_callable=lambda obj: [ e.value for e in obj - ], # Use the value of the enum for storage + ], + # Must match __table_args__'s schema to ensure that schema_translate_map + # routes the enum correctly. + schema=Role.PRIMARY.value, ), primary_key=True, ) @@ -60,7 +71,10 @@ class RelationshipMapping(ReferenceTable, CDMTableBase, Base): PredicateKind, values_callable=lambda obj: [ e.value for e in obj - ], # Use the value of the enum for storage + ], + # Must match __table_args__'s schema to ensure that schema_translate_map + # routes the enum correctly. + schema=Role.PRIMARY.value, ), primary_key=True, ) @@ -68,16 +82,17 @@ class RelationshipMapping(ReferenceTable, CDMTableBase, Base): sa.String(20), primary_key=True ) - # Define the Composite Foreign Key in __table_args__ - __table_args__ = ( + # Must match RelationshipClass's schema, or SQLAlchemy silently builds a second, unlinked Table object. + __table_args__ = merge_table_args( sa.ForeignKeyConstraint( ["predicate_kind", "predicate_subkind"], [ - "relationship_class.predicate_kind", - "relationship_class.predicate_subkind", + role_fk(Role.PRIMARY, "relationship_class.predicate_kind"), + role_fk(Role.PRIMARY, "relationship_class.predicate_subkind"), ], name="fk_rel_mapping_to_rel_class", ), + {"schema": Role.PRIMARY.value}, ) diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 4118918..8c62b0b 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -31,7 +31,7 @@ from sqlalchemy.orm import aliased from sqlalchemy.sql import Select -from oa_configurator import schema_inspect +from oa_configurator import Role, schema_inspect from omop_alchemy.backends import ( CONCEPT_NAME_TSVECTOR_COLUMN, @@ -352,8 +352,8 @@ def q_concept_name_fulltext( name_expr = ( Concept_Synonym.concept_synonym_name if synonym else Concept.concept_name ) - - inspector = schema_inspect(engine) + # Fulltext are in VOCAB schema + inspector = schema_inspect(engine, role=Role.VOCAB) target_table = Concept_Synonym if synonym else Concept target_col = ( CONCEPT_SYNONYM_NAME_TSVECTOR_COLUMN diff --git a/src/omop_graph/oaklib_interface/omop_factory.py b/src/omop_graph/oaklib_interface/omop_factory.py index b5962be..3a9438a 100644 --- a/src/omop_graph/oaklib_interface/omop_factory.py +++ b/src/omop_graph/oaklib_interface/omop_factory.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging from typing import Optional, Union from sqlalchemy.engine import URL @@ -10,6 +11,8 @@ from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, ResolvedCDMDatabase, Resolver from omop_graph.config import OmopGraphConfig +logger = logging.getLogger(__name__) + def omop_resource( *, @@ -25,7 +28,8 @@ def omop_resource( ---------- url : str | URL, optional Explicit database connection URL. When omitted the active oa-configurator - config is used. + config is used. An explicit url bypasses oa-configurator entirely, so no + schema_translate_map is applied. Logs a warning when given for this reason. slug : str, optional Slug identifier for the resource. Defaults to 'omop'. @@ -36,6 +40,7 @@ def omop_resource( execution_options = None vocab_url = None vocab_execution_options = None + resolved: ResolvedCDMDatabase | None = None if url is None: resolver = Resolver.from_active_config() db_name = resolver.resolve_package_config(OmopGraphConfig).cdm_db @@ -45,11 +50,19 @@ def omop_resource( f"OmopGraphConfig.cdm_db must resolve to a CDM database, got " f"{type(database).__name__}" ) - url = database.connection.url - execution_options = {SCHEMA_TRANSLATE_MAP_KEY: database.schema_translate_map()} - if database.connection != database.vocab_connection: - vocab_url = database.vocab_connection.url + resolved = database + url = resolved.connection.url + execution_options = {SCHEMA_TRANSLATE_MAP_KEY: resolved.schema_translate_map()} + if resolved.connection != resolved.vocab_connection: + vocab_url = resolved.vocab_connection.url vocab_execution_options = execution_options + else: + logger.warning( + "omop_resource() was given an explicit url, bypassing oa-configurator " + "entirely: no schema_translate_map is applied, so a configured " + "non-default primary/vocab/results schema is silently not respected. " + "Omit url= to resolve the active oa-configurator config instead." + ) return OMOPOntologyResource( slug=slug, @@ -57,4 +70,5 @@ def omop_resource( execution_options=execution_options, vocab_url=vocab_url, vocab_execution_options=vocab_execution_options, + resolved=resolved, ) diff --git a/src/omop_graph/oaklib_interface/omop_implementation.py b/src/omop_graph/oaklib_interface/omop_implementation.py index 9840fb5..77e4a8b 100644 --- a/src/omop_graph/oaklib_interface/omop_implementation.py +++ b/src/omop_graph/oaklib_interface/omop_implementation.py @@ -743,23 +743,24 @@ def entailed_outgoing_relationships( {self._parse_predicate(p) for p in predicates} if predicates else None ) - for edge in self.kg.iter_edges( - concept_id, direction="out", predicate_kinds=None - ): - if pred_filter and edge.predicate_id not in pred_filter: - continue + with self.kg.session_factory() as session: + for edge in self.kg.iter_edges( + session=session, concept_ids=concept_id, direction="out", predicate_kinds=None + ): + if pred_filter and edge.predicate_id not in pred_filter: + continue - pred_curie = self._predicate_curie(edge.predicate_id) + pred_curie = self._predicate_curie(edge.predicate_id) - # hierarchical entailment - if self.kg.predicate_kind(edge.predicate_id) == PredicateKind.HIERARCHY: - yield pred_curie, self._concept_curie(edge.object_id) + # hierarchical entailment + if self.kg.predicate_kind(edge.predicate_id) == PredicateKind.HIERARCHY: + yield pred_curie, self._concept_curie(edge.object_id) - for parent in self.kg.parents(edge.object_id): - yield pred_curie, self._concept_curie(parent) + for parent in self.kg.parents(edge.object_id): + yield pred_curie, self._concept_curie(parent) - else: - yield pred_curie, self._concept_curie(edge.object_id) + else: + yield pred_curie, self._concept_curie(edge.object_id) def entailed_outputgoing_relationships_by_curie( self, *args, **kwargs @@ -819,9 +820,10 @@ def entailed_relationships_between( obj_id = self._parse_concept(object) # direct relationships - for edge in self.kg.iter_edges(subj_id, direction="out"): - if edge.object_id == obj_id: - yield self._predicate_curie(edge.predicate_id) + with self.kg.session_factory() as session: + for edge in self.kg.iter_edges(session=session, concept_ids=subj_id, direction="out"): + if edge.object_id == obj_id: + yield self._predicate_curie(edge.predicate_id) # hierarchical entailment if obj_id in self.kg.parents(subj_id): @@ -905,18 +907,25 @@ def __init__( "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, - ) - vocab_engine = None - if self.resource.vocab_url is not None: - vocab_engine = make_engine( - self.resource.vocab_url, + if self.resource.resolved is not None: + engine, vocab_engine = self.resource.resolved.create_engines( + echo=False, future=True + ) + if vocab_engine is engine: + vocab_engine = None + else: + engine = make_engine( + self.engine_string, engine_kwargs={"echo": False, "future": True}, - execution_options=self.resource.vocab_execution_options, + execution_options=self.resource.execution_options, ) + 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) diff --git a/src/omop_graph/oaklib_interface/omop_resource.py b/src/omop_graph/oaklib_interface/omop_resource.py index 57d5b4b..ec5876e 100644 --- a/src/omop_graph/oaklib_interface/omop_resource.py +++ b/src/omop_graph/oaklib_interface/omop_resource.py @@ -4,6 +4,8 @@ from oaklib.resource import OntologyResource from sqlalchemy.engine import URL, make_url +from oa_configurator import ResolvedCDMDatabase + @dataclass class OMOPOntologyResource(OntologyResource): @@ -42,6 +44,10 @@ class OMOPOntologyResource(OntologyResource): ``schema_translate_map`` as ``execution_options``: the map itself doesn't change across roles, only which physical connection it's applied to. + resolved : ResolvedCDMDatabase, optional + When present, ``OMOPAlchemyImplementation`` builds its engine(s) via + ``resolved.create_engines()`` directly instead of reconstructing one + from ``url``/``execution_options``. """ url: Optional[Union[str, URL]] = None # type: ignore[assignment] @@ -53,6 +59,7 @@ class OMOPOntologyResource(OntologyResource): 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] + resolved: Optional[ResolvedCDMDatabase] = None # type: ignore[assignment] def _parsed_url(self) -> Optional[URL]: """ diff --git a/tests/fixtures/mock_cdm.py b/tests/fixtures/mock_cdm.py index 7783e2e..df4c447 100644 --- a/tests/fixtures/mock_cdm.py +++ b/tests/fixtures/mock_cdm.py @@ -7,6 +7,7 @@ import sqlalchemy as sa from sqlalchemy.orm import Session, sessionmaker +from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role from oa_configurator.testing import isolated_test_database from orm_loader.helpers import Base from omop_alchemy.cdm.model.vocabulary.concept import Concept @@ -38,7 +39,9 @@ def mock_cdm_engine() -> Iterator[sa.Engine]: "test_cdm_db_sqlite", dialect="sqlite", future=True, - execution_options={"schema_translate_map": {None: None, "vocab": None, "results": None}}, + execution_options={ + SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} + }, ) as db: engine = db.connection.engine _create_mock_cdm_tables(engine) diff --git a/tests/test_concept_queries.py b/tests/test_concept_queries.py index 2a55904..e33198d 100644 --- a/tests/test_concept_queries.py +++ b/tests/test_concept_queries.py @@ -9,6 +9,7 @@ import sqlalchemy as sa from sqlalchemy.orm import Session +from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role from oa_configurator.testing import isolated_test_database from omop_alchemy.cdm.model.vocabulary import Concept @@ -31,7 +32,9 @@ def concept_engine() -> Iterator[sa.Engine]: "test_cdm_db_sqlite", dialect="sqlite", future=True, - execution_options={"schema_translate_map": {None: None, "vocab": None, "results": None}}, + execution_options={ + SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} + }, ) as db: engine = db.connection.engine Concept.__table__.create(engine) diff --git a/tests/test_fulltext_vocab_schema_postgres.py b/tests/test_fulltext_vocab_schema_postgres.py new file mode 100644 index 0000000..f2698bb --- /dev/null +++ b/tests/test_fulltext_vocab_schema_postgres.py @@ -0,0 +1,86 @@ +"""q_concept_name_fulltext() must inspect the VOCAB schema, not primary. + +Concept/Concept_Synonym are VOCAB-role tables. schema_inspect(engine) with +no role= defaults to Role.PRIMARY, so whenever vocab_schema differs from +cdm_schema (a normal same-server split, not just a same-schema deployment), +the old code inspected the wrong schema, found no tsvector column, and +raised a false FullTextError even though the column genuinely exists. The +only prior coverage (test_fulltext_optional.py) uses SQLite, where +supports_schemas() is False and every role folds to None -- masking this +completely. Real, distinct primary/vocab Postgres schemas here instead. + +Schemas are created directly on pg_db's own already-open connection +(CREATE SCHEMA, rolled back automatically at teardown) rather than via +isolated_test_schema(): that opens a second, genuinely separate connection +from the same pool, which can starve waiting on pg_db's own still-open +transaction. +""" + +from __future__ import annotations + +import uuid + +import pytest +import sqlalchemy as sa +from oa_configurator import Role +from omop_alchemy.backends import CONCEPT_NAME_TSVECTOR_COLUMN, FullTextError, resolve_backend +from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary +from orm_loader.helpers import Base + +from omop_graph.graph.queries import q_concept_name_fulltext + +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + +_VOCAB_TABLES = (Domain.__table__, Vocabulary.__table__, Concept_Class.__table__, Concept.__table__) + + +def _scoped(pg_db, *, primary_schema: str, vocab_schema: str) -> sa.Connection: + conn = pg_db.connection + conn.execute(sa.text(f"CREATE SCHEMA {primary_schema}")) + conn.execute(sa.text(f"CREATE SCHEMA {vocab_schema}")) + return conn.execution_options( + schema_translate_map={ + Role.PRIMARY.value: primary_schema, + Role.VOCAB.value: vocab_schema, + Role.RESULTS.value: primary_schema, + } + ) + + +def test_fulltext_query_finds_the_tsvector_column_in_the_vocab_schema(pg_db): + primary_schema = f"fulltext_primary_{uuid.uuid4().hex[:8]}" + vocab_schema = f"fulltext_vocab_{uuid.uuid4().hex[:8]}" + scoped = _scoped(pg_db, primary_schema=primary_schema, vocab_schema=vocab_schema) + Base.metadata.create_all(bind=scoped, tables=_VOCAB_TABLES, checkfirst=True) + + backend = resolve_backend(scoped) + backend.install_fulltext_on_table( + scoped, + table_name="concept", + vector_column_name=CONCEPT_NAME_TSVECTOR_COLUMN, + index_name="idx_concept_name_tsvector_test", + create_indexes=True, + fastupdate=True, + role=Role.VOCAB, + ) + + # Confirmed absent from primary_schema: proves a real vocab/primary + # split, not a lucky same-schema coincidence. + assert not sa.inspect(pg_db.connection).has_table("concept", schema=primary_schema) + columns = { + c["name"] for c in sa.inspect(pg_db.connection).get_columns("concept", schema=vocab_schema) + } + assert CONCEPT_NAME_TSVECTOR_COLUMN in columns + + stmt = q_concept_name_fulltext("kidney cancer", engine=scoped) + assert stmt is not None + + +def test_fulltext_query_still_raises_when_the_column_is_genuinely_absent(pg_db): + primary_schema = f"fulltext_primary_{uuid.uuid4().hex[:8]}" + vocab_schema = f"fulltext_vocab_{uuid.uuid4().hex[:8]}" + scoped = _scoped(pg_db, primary_schema=primary_schema, vocab_schema=vocab_schema) + Base.metadata.create_all(bind=scoped, tables=_VOCAB_TABLES, checkfirst=True) + + with pytest.raises(FullTextError): + q_concept_name_fulltext("kidney cancer", engine=scoped) diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index 2b332ae..aeea76f 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -26,8 +26,10 @@ import sqlalchemy as sa from oa_configurator import ( + SCHEMA_TRANSLATE_MAP_KEY, ResolvedCDMDatabase, ResolvedConnection, + Role, qualified, ) from oa_configurator.testing import isolated_test_schema @@ -50,27 +52,17 @@ 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. + """Minimal vocab bootstrap: Domain/Vocabulary/Concept_Class/Concept form an + FK insert cycle, so triggers are disabled around the insert then + re-enabled. Accepts an Engine or an open Connection; an Engine has no + .execute() of its own, so one short-lived connection is opened 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: @@ -141,8 +133,8 @@ def test_omop_resource_execution_options_carry_the_configured_schema() -> None: assert resource.execution_options is not None assert ( - engine.get_execution_options()["schema_translate_map"] - == resource.execution_options["schema_translate_map"] + engine.get_execution_options()[SCHEMA_TRANSLATE_MAP_KEY] + == resource.execution_options[SCHEMA_TRANSLATE_MAP_KEY] ) @@ -201,7 +193,7 @@ def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_ 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} + schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} ) Base.metadata.create_all(bind=engine, checkfirst=True) _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Split-wiring concept") @@ -210,11 +202,11 @@ def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_ 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} + SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} }, vocab_url="sqlite:///:memory:", vocab_execution_options={ - "schema_translate_map": {None: None, "vocab": None, "results": None} + SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} }, ) @@ -233,7 +225,7 @@ def test_omop_alchemy_implementation_reuses_one_engine_when_no_split_is_configur 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} + schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} ) Base.metadata.create_all(bind=engine, checkfirst=True) _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="No-split concept") @@ -242,7 +234,7 @@ def test_omop_alchemy_implementation_reuses_one_engine_when_no_split_is_configur 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} + SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} }, ) @@ -251,18 +243,63 @@ def test_omop_alchemy_implementation_reuses_one_engine_when_no_split_is_configur assert adapter.kg.vocab_engine is adapter.kg.cdm_engine +def test_omop_alchemy_implementation_builds_its_engine_via_resolved_create_engines( + pg_db, monkeypatch +) -> None: + with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_resolved") as schema: + engine = pg_db.connection.engine.execution_options( + schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + ) + Base.metadata.create_all(bind=engine, checkfirst=True) + _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Resolved-path concept") + relationship_classification(engine=engine) + + url = pg_db.connection.engine.url + connection = ResolvedConnection( + name="resolved_path", + url=url.render_as_string(hide_password=False), + safe_url=url.render_as_string(hide_password=True), + _engine_url=url, + ) + resolved = ResolvedCDMDatabase( + name="resolved_path", + connection=connection, + schema_name=schema, + vocab_connection=connection, + vocab_schema=schema, + results_schema=schema, + ) + + class FakeResolver: + def resolve_package_config(self, config_type): + assert config_type is OmopGraphConfig + return OmopGraphConfig(cdm_db="resolved_path") + + def resolve_database(self, name): + assert name == "resolved_path" + return resolved + + monkeypatch.setattr( + "omop_graph.oaklib_interface.omop_factory.Resolver.from_active_config", + lambda: FakeResolver(), + ) + + resource = omop_resource() + assert resource.resolved is resolved + + adapter = OMOPAlchemyImplementation(resource=resource) + + assert adapter.kg.cdm_engine.pool._pre_ping is True + assert adapter.kg.vocab_engine is adapter.kg.cdm_engine + assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Resolved-path concept" + + 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} + schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} ) Base.metadata.create_all(bind=scoped, checkfirst=True) _seed_one_concept(scoped, concept_id=_CONCEPT_ID, name="Test concept") @@ -275,16 +312,9 @@ def test_kg_injection_path_resolves_against_the_configured_schema(pg_db) -> None 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} + schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} ) Base.metadata.create_all(bind=engine, checkfirst=True) _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Bare-string concept") @@ -296,7 +326,7 @@ def test_bare_engine_string_path_resolves_against_the_configured_schema(pg_db) - # 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} + SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} }, ) adapter = OMOPAlchemyImplementation(resource=resource) diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index 6751b0a..582d4a4 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -11,7 +11,11 @@ from __future__ import annotations +import dataclasses + +import pytest import sqlalchemy as sa +from oa_configurator import Role from orm_loader.helpers import Base @@ -23,7 +27,7 @@ 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} + schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} ) @@ -56,6 +60,21 @@ def test_relationship_classification_respects_the_configured_schema(pg_db): assert enum_type == "predicatekind" +def test_relationship_classification_refuses_a_genuinely_split_vocab_connection(pg_db): + """Postgres has no cross-database inline FK, so RelationshipMapping's FK + to relationship.relationship_id (VOCAB-role) can never be created once + vocab_connection is a genuinely separate connection. + """ + resolved = dataclasses.replace( + pg_db.resolved, + vocab_connection=dataclasses.replace( + pg_db.resolved.connection, name="genuinely_different", safe_url="postgresql://other/db" + ), + ) + with pytest.raises(RuntimeError, match="genuinely separate"): + relationship_classification(engine=pg_db.connection, resolved=resolved) + + 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.""" diff --git a/tests/test_schema_provenance_guard.py b/tests/test_schema_provenance_guard.py index e4a61c4..22026ec 100644 --- a/tests/test_schema_provenance_guard.py +++ b/tests/test_schema_provenance_guard.py @@ -45,13 +45,15 @@ def _resolved(pg_db, *, database_name: str, schema: str): 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. """ + patched_connection = dataclasses.replace(pg_db.resolved.connection, test_only=False) 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), + connection=patched_connection, + vocab_connection=patched_connection, ) diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index 3741558..0db4150 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -25,6 +25,7 @@ import sqlalchemy as sa import sqlalchemy.orm as so +from oa_configurator import Role from oa_configurator.testing import isolated_test_database, isolated_test_schema from orm_loader.config import OrmLoaderConfig from orm_loader.helpers import Base @@ -77,13 +78,19 @@ _shadow_metadata, sa.Column( "predicate_kind", - sa.Enum(PredicateKind, values_callable=lambda obj: [e.value for e in obj]), + sa.Enum( + PredicateKind, + values_callable=lambda obj: [e.value for e in obj], + schema=Role.PRIMARY.value, + ), 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), + # Must match RelationshipClass's schema, or SQLAlchemy silently builds a second, unlinked Table object. + schema=Role.PRIMARY.value, ) _shadow_relationship_mapping = sa.Table( "relationship_mapping", @@ -91,10 +98,15 @@ 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]), + sa.Enum( + PredicateKind, + values_callable=lambda obj: [e.value for e in obj], + schema=Role.PRIMARY.value, + ), primary_key=True, ), sa.Column("predicate_subkind", sa.String(20), primary_key=True), + schema=Role.PRIMARY.value, ) @@ -120,10 +132,18 @@ def split_engines() -> Iterator[_Engines]: 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} + schema_translate_map={ + Role.PRIMARY.value: primary_schema, + Role.VOCAB.value: primary_schema, + Role.RESULTS.value: primary_schema, + } ) vocab_engine = vocab_raw.execution_options( - schema_translate_map={None: vocab_schema, "vocab": vocab_schema, "results": vocab_schema} + schema_translate_map={ + Role.PRIMARY.value: vocab_schema, + Role.VOCAB.value: vocab_schema, + Role.RESULTS.value: vocab_schema, + } ) _shadow_metadata.create_all(primary_engine) From bd45961ca0df38098bb06dbae645e99f90da4df4 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 01:38:31 +0000 Subject: [PATCH 07/20] Fix tests --- src/omop_graph/graph/queries.py | 3 ++- tests/conftest.py | 20 +++++--------------- tests/render/test_render_mmd.py | 7 ------- tests/render/test_render_text.py | 9 --------- tests/test_concept_queries.py | 4 ++-- tests/test_fulltext_vocab_schema_postgres.py | 6 +++++- tests/test_grounding.py | 2 +- tests/test_kg_ancestry.py | 2 +- tests/test_oaklib_schema_awareness.py | 7 ++++++- tests/test_vocab_split_connection.py | 19 +++++++++++-------- 10 files changed, 33 insertions(+), 46 deletions(-) delete mode 100644 tests/render/test_render_mmd.py delete mode 100644 tests/render/test_render_text.py diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 8c62b0b..6d6eccc 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -25,6 +25,7 @@ literal, or_, select, + Connection, Engine, column, ) @@ -324,7 +325,7 @@ def q_concept_name_ilike( def q_concept_name_fulltext( query_concept_name: str, *, - engine: Engine, + engine: Engine | Connection, search_constraint: Optional[ConceptFilter] = None, synonym: bool = False, sort: bool = True, diff --git a/tests/conftest.py b/tests/conftest.py index 887403f..7e42648 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -39,20 +39,10 @@ def filter(self, record): @pytest.fixture(autouse=True, scope="session") def configure_logging_whitelist(): + """Attach the filter to the root logger's handlers, not the logger itself, + so the check runs right before text actually reaches the screen. """ - Attaches the filter to the HANDLERS, not the logger. - """ - # Your allowed list - my_whitelisted_loggers = ["omop_graph", "orm_loader", "omop_spires", "tests"] - - # Instantiate the filter - my_filter = WhitelistFilter(my_whitelisted_loggers) - - # Get the root logger - root_logger = logging.getLogger() - - # --- THE FIX --- - # We iterate over the handlers (Console, File, etc.) and attach the filter there. - # This forces the check to happen right before the text hits the screen. - for handler in root_logger.handlers: + whitelisted_loggers = ["omop_graph", "orm_loader", "omop_spires", "tests"] + my_filter = WhitelistFilter(whitelisted_loggers) + for handler in logging.getLogger().handlers: handler.addFilter(my_filter) diff --git a/tests/render/test_render_mmd.py b/tests/render/test_render_mmd.py deleted file mode 100644 index c0e9a90..0000000 --- a/tests/render/test_render_mmd.py +++ /dev/null @@ -1,7 +0,0 @@ -# def test_path_mermaid_snapshot(kg, example_path): -# from omop_graph.render import render_path -# -# out = render_path(kg, example_path, format="mmd") -# -# assert out.startswith("graph LR") -# assert "-->|Is a|" in out diff --git a/tests/render/test_render_text.py b/tests/render/test_render_text.py deleted file mode 100644 index 06446e7..0000000 --- a/tests/render/test_render_text.py +++ /dev/null @@ -1,9 +0,0 @@ -# def test_path_text_snapshot(kg, example_path): -# from omop_graph.render import render_path -# -# out = render_path(kg, example_path, format="text") -# -# assert out == """\ -# Aspirin --[Is a]--> Antiplatelet agent -# Antiplatelet agent --[Is a]--> Drug -# """ diff --git a/tests/test_concept_queries.py b/tests/test_concept_queries.py index e33198d..dc8b1f8 100644 --- a/tests/test_concept_queries.py +++ b/tests/test_concept_queries.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import date -from typing import Iterator +from typing import Iterator, cast import pytest import sqlalchemy as sa @@ -37,7 +37,7 @@ def concept_engine() -> Iterator[sa.Engine]: }, ) as db: engine = db.connection.engine - Concept.__table__.create(engine) + cast(sa.Table, Concept.__table__).create(engine) valid_from = date(2000, 1, 1) valid_until = date(2099, 12, 31) diff --git a/tests/test_fulltext_vocab_schema_postgres.py b/tests/test_fulltext_vocab_schema_postgres.py index f2698bb..e2e43c8 100644 --- a/tests/test_fulltext_vocab_schema_postgres.py +++ b/tests/test_fulltext_vocab_schema_postgres.py @@ -19,6 +19,7 @@ from __future__ import annotations import uuid +from typing import cast import pytest import sqlalchemy as sa @@ -31,7 +32,10 @@ pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] -_VOCAB_TABLES = (Domain.__table__, Vocabulary.__table__, Concept_Class.__table__, Concept.__table__) +_VOCAB_TABLES = cast( + "tuple[sa.Table, ...]", + (Domain.__table__, Vocabulary.__table__, Concept_Class.__table__, Concept.__table__), +) def _scoped(pg_db, *, primary_schema: str, vocab_schema: str) -> sa.Connection: diff --git a/tests/test_grounding.py b/tests/test_grounding.py index 606b7c3..f91216d 100644 --- a/tests/test_grounding.py +++ b/tests/test_grounding.py @@ -17,7 +17,7 @@ PartialLabelResolver, PartialSynonymResolver, ) -from fixtures.mock_cdm import PARENT_CANCER_ID # type: ignore +from fixtures.mock_cdm import PARENT_CANCER_ID class TestQueryTextWithContext: diff --git a/tests/test_kg_ancestry.py b/tests/test_kg_ancestry.py index 95d1b49..887e7c4 100644 --- a/tests/test_kg_ancestry.py +++ b/tests/test_kg_ancestry.py @@ -4,7 +4,7 @@ from omop_graph.graph.nodes import LabelMatchKind from omop_graph.graph.paths import find_standard_paths from omop_graph.reasoning.resolvers.resolvers import CandidateHit -from fixtures.mock_cdm import PARENT_CANCER_ID, CONCEPT_META_ID # type: ignore +from fixtures.mock_cdm import PARENT_CANCER_ID, CONCEPT_META_ID def test_get_potential_ancestors_batch_returns_only_real_ancestors( diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index aeea76f..2e7d348 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -22,8 +22,10 @@ from __future__ import annotations from datetime import date +from typing import cast import sqlalchemy as sa +import sqlalchemy.orm from oa_configurator import ( SCHEMA_TRANSLATE_MAP_KEY, @@ -48,7 +50,10 @@ _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__) +_VOCAB_TABLES = cast( + "tuple[sa.Table, ...]", + (Domain.__table__, Vocabulary.__table__, Concept_Class.__table__, Concept.__table__), +) def _seed_one_concept(bindable: sa.Engine | sa.Connection, *, concept_id: int, name: str) -> None: diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index 0db4150..04f5ac4 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -19,7 +19,7 @@ from __future__ import annotations from datetime import date -from typing import Iterator, NamedTuple +from typing import Iterator, NamedTuple, cast import pytest import sqlalchemy as sa @@ -55,13 +55,16 @@ _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__, +_VOCAB_TABLES = cast( + "tuple[sa.Table, ...]", + ( + 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 From 5ff3df6e0aa424b0370cb36c1f4d6b46521acf2b Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 04:53:29 +0000 Subject: [PATCH 08/20] Unified/aligned test interface, proper split implementation --- src/omop_graph/graph/kg.py | 36 ++-- tests/fixtures/helpers.py | 91 ++++++++++ tests/fixtures/mock_cdm.py | 32 ++-- tests/test_concept_queries.py | 89 +++++++--- .../test_edges_same_connection_regression.py | 4 + tests/test_embedding_optional.py | 6 + tests/test_fulltext_optional.py | 5 +- tests/test_fulltext_vocab_schema_postgres.py | 19 +- tests/test_grounding.py | 10 ++ tests/test_kg_ancestry.py | 4 + tests/test_oaklib_interface.py | 4 + tests/test_oaklib_schema_awareness.py | 167 +++++++++--------- tests/test_predicate_flags.py | 4 + tests/test_relationship_classification.py | 7 +- tests/test_vocab_split_connection.py | 85 +++++---- 15 files changed, 366 insertions(+), 197 deletions(-) create mode 100644 tests/fixtures/helpers.py diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index 174e685..ed1f611 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -330,7 +330,7 @@ def concept_view(self, concept_id: int) -> ConceptView: ConceptView The immutable view of the concept. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: row = session.execute(q_concept_view(concept_id)).one() return ConceptView.from_row(row) @@ -350,7 +350,7 @@ def concept_views( tuple[ConceptView, ...] A tuple of concept views. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: concept_views = tuple( ConceptView.from_row(row) for row in session.execute(q_concept_views(concept_ids, sort=sort)) @@ -373,7 +373,7 @@ def concept_id_by_code(self, vocabulary_id: str, concept_code: str) -> int: int The resolved OMOP Concept ID. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: concept_id = int( session.execute( q_concept_id_by_code(vocabulary_id, concept_code) @@ -413,7 +413,7 @@ def concept_lookup( elif match_kind == LabelMatchKind.PARTIAL: fn = q_concept_name_ilike elif match_kind == LabelMatchKind.FTS: - fn = functools.partial(q_concept_name_fulltext, engine=self.cdm_engine) + fn = functools.partial(q_concept_name_fulltext, engine=self.vocab_engine) else: raise ValueError(f"Unsupported search mode: {match_kind}") try: @@ -429,7 +429,7 @@ def concept_lookup( return () raise - with self.session_factory() as session: + with self.vocab_session_factory() as session: matches = tuple( LabelMatch( input_query=input_query_term, @@ -449,7 +449,7 @@ def concept_ids_by_label(self, label: str) -> Tuple[int, ...]: Find concept IDs that match the label exactly (case-insensitive). """ label = self._normalise_query_term(label) - with self.session_factory() as session: + with self.vocab_session_factory() as session: rows = session.execute(q_concept_name_match(label)).scalars() return tuple(rows) @@ -489,7 +489,7 @@ def predicate_name(self, relationship_id: str) -> str: Retrieve the human-readable name of a relationship. """ # TODO: Not really necessary. The "ID" is mostly human-readable anyways. - with self.session_factory() as session: + with self.vocab_session_factory() as session: predicate_name = session.execute( q_predicate_name(relationship_id) ).scalar_one() @@ -681,7 +681,7 @@ def parents(self, concept_id: int) -> tuple[int, ...]: """ Retrieve parent Concept IDs of concept using Concept_Ancestor table. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: parents = tuple(session.execute(q_parents(concept_id)).scalars()) return parents @@ -689,7 +689,7 @@ def children(self, concept_id) -> tuple[int, ...]: """ Retrieve children Concept IDs of concept using Concept_Ancestor table. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: children = tuple(session.execute(q_children(concept_id)).scalars()) return children @@ -716,7 +716,7 @@ def roots( """ Retrieve root concepts (no parents). """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: roots = tuple( session.execute( q_roots(domain_id=domain_id, vocabulary_id=vocabulary_id) @@ -730,7 +730,7 @@ def leaves( """ Retrieve leaf concepts (no children). """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: leaves = tuple( session.execute( q_leaves(domain_id=domain_id, vocabulary_id=vocabulary_id) @@ -744,7 +744,7 @@ def singletons( """ Retrieve singleton concepts (no parents and no children). """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: return tuple( session.execute( q_singletons(domain_id=domain_id, vocabulary_id=vocabulary_id) @@ -755,7 +755,7 @@ def synonyms_for_concept(self, concept_id: int) -> tuple[str, ...]: """ Retrieve all synonyms for a concept. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: rows = session.execute(q_concept_synonym_filtered(concept_id)).all() return tuple(row.name for row in rows) @@ -782,13 +782,13 @@ def predicates(self) -> tuple[Predicate, ...]: @functools.cached_property def _valid_domains(self) -> frozenset[str]: - with self.session_factory() as session: + with self.vocab_session_factory() as session: rows = session.execute(q_concept_domain_ids()).all() return frozenset(row.domain_id for row in rows) @functools.cached_property def _valid_vocabularies(self) -> frozenset[str]: - with self.session_factory() as session: + with self.vocab_session_factory() as session: rows = session.execute(q_concept_vocabulary_ids()).all() return frozenset(row.vocabulary_id for row in rows) @@ -799,7 +799,7 @@ def get_potential_ancestor( Check if an ancestry relationship exists between a child and parent. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: row = session.execute( q_concept_potential_ancestor(child_id, parent_id) ).first() @@ -834,7 +834,7 @@ def get_potential_ancestors_batch( if not parent_ids: return {} - with self.session_factory() as session: + with self.vocab_session_factory() as session: rows = session.execute( q_concept_potential_ancestors_batch(child_ids, parent_ids) ).all() @@ -852,7 +852,7 @@ def get_num_ancestors(self, concept_ids: tuple[int, ...]) -> Dict[int, int]: """ Get the count of ancestors for a batch of concepts. """ - with self.session_factory() as session: + with self.vocab_session_factory() as session: rows = session.execute(q_concept_num_ancestors(concept_ids)).all() return {row.concept_id: row.num_ancestors for row in rows} diff --git a/tests/fixtures/helpers.py b/tests/fixtures/helpers.py new file mode 100644 index 0000000..4b405d3 --- /dev/null +++ b/tests/fixtures/helpers.py @@ -0,0 +1,91 @@ +"""Shared test-helper primitives for split/non-split, vocab-schema fixtures. + +Extracted from four independently-duplicated copies across the test suite +(Phase 4.5 of the schema_translate_map fix): a vocab-table tuple, a +schema_translate_map dict builder, and an FK-trigger toggle context manager. +Schema *provisioning* strategy (rollback-based CREATE SCHEMA on an open +connection vs. a genuinely committed isolated_test_schema) stays local to +each fixture/test instead of being forced into one shared helper here -- +that choice genuinely differs per test's needs, unlike the three primitives +below. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Iterator, cast + +import sqlalchemy as sa +from oa_configurator import Role, qualified +from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary + +VOCAB_TABLES = cast( + "tuple[sa.Table, ...]", + (Domain.__table__, Vocabulary.__table__, Concept_Class.__table__, Concept.__table__), +) +"""Domain/Vocabulary/Concept_Class/Concept: the minimal vocab bootstrap set +every split/non-split fixture needs, in FK dependency order. A consumer +needing more tables (e.g. Relationship, Concept_Relationship) extends this +tuple rather than redefining its own copy of the shared core. +""" + + +def schema_translate_map( + primary_schema: str, *, vocab_schema: str | None = None, results_schema: str | None = None +) -> dict[str, str]: + """Build a 3-key schema_translate_map dict, defaulting vocab/results to + primary_schema -- the common "route everything through one schema" shape, + with an explicit override for the genuinely-split-vocab case. + """ + return { + Role.PRIMARY.value: primary_schema, + Role.VOCAB.value: vocab_schema if vocab_schema is not None else primary_schema, + Role.RESULTS.value: results_schema if results_schema is not None else primary_schema, + } + + +@contextmanager +def fk_triggers_disabled( + bindable: "sa.Engine | sa.Connection", tables: "tuple[sa.Table, ...]" +) -> Iterator[None]: + """Disable triggers on *tables* for the duration of the block, matching + production bulk-loads' own FK-toggle pattern. + + Several vocab tables form a genuine FK insert cycle (each reference + row's own *_concept_id FK needs a Concept row that doesn't exist yet, + and Concept's own domain_id/vocabulary_id/concept_class_id FKs need + those reference rows already present), so a plain ordered insert can't + satisfy every constraint. A no-op on any non-Postgres dialect: SQLite + doesn't enforce FK constraints by default, so there's nothing to + disable there, and callers that seed both dialects can call this + unconditionally without their own dialect check. Accepts an Engine or + an open Connection; an Engine has no ``.execute()`` of its own, so one + short-lived connection is opened for each toggle. + """ + if bindable.dialect.name != "postgresql": + yield + return + + opened_here = isinstance(bindable, sa.Engine) + conn = bindable.connect() if opened_here else bindable + try: + for table in tables: + conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name)} DISABLE TRIGGER ALL")) + if opened_here: + conn.commit() + finally: + if opened_here: + conn.close() + + try: + yield + finally: + conn = bindable.connect() if opened_here else bindable + try: + for table in 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() diff --git a/tests/fixtures/mock_cdm.py b/tests/fixtures/mock_cdm.py index df4c447..d0f4a2d 100644 --- a/tests/fixtures/mock_cdm.py +++ b/tests/fixtures/mock_cdm.py @@ -7,8 +7,7 @@ import sqlalchemy as sa from sqlalchemy.orm import Session, sessionmaker -from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role -from oa_configurator.testing import isolated_test_database +from oa_configurator.testing import isolated_test_database, isolated_test_schema 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 @@ -27,6 +26,8 @@ ) from omop_graph.graph.kg import KnowledgeGraph +from .helpers import fk_triggers_disabled, schema_translate_map + PARENT_CANCER_ID = 443392 CONCEPT_META_ID = 0 LANGUAGE_CONCEPT_ID = 1 @@ -34,18 +35,12 @@ @pytest.fixture(scope="module") 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_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} - }, - ) as db: - engine = db.connection.engine - _create_mock_cdm_tables(engine) - yield engine + with isolated_test_database(OmopGraphConfig, "test_cdm_db_pg") as db: + raw_engine = db.connection.engine + with isolated_test_schema(raw_engine, prefix="mock_cdm") as schema: + engine = raw_engine.execution_options(schema_translate_map=schema_translate_map(schema)) + _create_mock_cdm_tables(engine) + yield engine def _create_mock_cdm_tables(engine: sa.Engine) -> None: @@ -67,9 +62,12 @@ def _create_mock_cdm_tables(engine: sa.Engine) -> None: Base.metadata.create_all(engine, tables=tables) - session_local = sessionmaker(bind=engine, future=True) - with session_local() as session: - seed_mock_cdm(session) + # Disable triggers for the whole seed sidesteps ordering entirely, + # since these tables have FK dependencies that form a cycle + with fk_triggers_disabled(engine, tuple(tables)): + session_local = sessionmaker(bind=engine, future=True) + with session_local() as session: + seed_mock_cdm(session) @pytest.fixture() diff --git a/tests/test_concept_queries.py b/tests/test_concept_queries.py index dc8b1f8..bcf5d58 100644 --- a/tests/test_concept_queries.py +++ b/tests/test_concept_queries.py @@ -3,19 +3,18 @@ from __future__ import annotations from datetime import date -from typing import Iterator, cast +from typing import Iterator import pytest import sqlalchemy as sa from sqlalchemy.orm import Session -from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, Role -from oa_configurator.testing import isolated_test_database +from oa_configurator.testing import isolated_test_schema +from orm_loader.helpers import Base -from omop_alchemy.cdm.model.vocabulary import Concept +from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary 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, @@ -24,20 +23,16 @@ q_entities, ) +from fixtures.helpers import VOCAB_TABLES, fk_triggers_disabled, schema_translate_map + +_META_CONCEPT_ID = 0 + @pytest.fixture() -def concept_engine() -> Iterator[sa.Engine]: - with isolated_test_database( - OmopGraphConfig, - "test_cdm_db_sqlite", - dialect="sqlite", - future=True, - execution_options={ - SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} - }, - ) as db: - engine = db.connection.engine - cast(sa.Table, Concept.__table__).create(engine) +def concept_engine(pg_db) -> Iterator[sa.Engine]: + with isolated_test_schema(pg_db.connection.engine, prefix="concept_queries") as schema: + engine = pg_db.connection.engine.execution_options(schema_translate_map=schema_translate_map(schema)) + Base.metadata.create_all(engine, tables=VOCAB_TABLES, checkfirst=True) valid_from = date(2000, 1, 1) valid_until = date(2099, 12, 31) @@ -61,17 +56,55 @@ def concept( 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 fk_triggers_disabled(engine, VOCAB_TABLES): + with Session(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=valid_from, + valid_end_date=valid_until, + ), + 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, + ), + 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() yield engine diff --git a/tests/test_edges_same_connection_regression.py b/tests/test_edges_same_connection_regression.py index 3fc64f0..edc21cc 100644 --- a/tests/test_edges_same_connection_regression.py +++ b/tests/test_edges_same_connection_regression.py @@ -9,9 +9,13 @@ from __future__ import annotations +import pytest + from omop_graph.extensions.omop_alchemy import PredicateKind from omop_graph.graph.kg import KnowledgeGraph +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + def test_edges_use_single_eager_join_when_no_split_is_configured( mock_cdm_kg: KnowledgeGraph, diff --git a/tests/test_embedding_optional.py b/tests/test_embedding_optional.py index b73f243..23f4a1a 100644 --- a/tests/test_embedding_optional.py +++ b/tests/test_embedding_optional.py @@ -86,6 +86,8 @@ def _make_standard_concept(concept_id: int, name: str) -> StandardConcept: ) +@pytest.mark.postgresql +@pytest.mark.db_dialect def test_fallback_flag_true_logs_attempt_when_concepts_missing( mock_cdm_kg: KnowledgeGraph, monkeypatch: pytest.MonkeyPatch, @@ -133,6 +135,8 @@ def test_fallback_flag_true_logs_attempt_when_concepts_missing( assert "Computing missing embeddings on-the-fly" in caplog.text +@pytest.mark.postgresql +@pytest.mark.db_dialect def test_fallback_flag_false_logs_disabled_when_concepts_missing( mock_cdm_kg: KnowledgeGraph, monkeypatch: pytest.MonkeyPatch, @@ -256,6 +260,8 @@ def emb(self): # ── omop-emb#48 split: k threaded explicitly instead of via filter.limit ── +@pytest.mark.postgresql +@pytest.mark.db_dialect def test_semantic_similarity_splits_cdm_and_knn_filters( mock_cdm_kg: KnowledgeGraph, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_fulltext_optional.py b/tests/test_fulltext_optional.py index 6a385f4..cec4570 100644 --- a/tests/test_fulltext_optional.py +++ b/tests/test_fulltext_optional.py @@ -5,6 +5,8 @@ from omop_graph.graph.queries import q_concept_name_fulltext +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + @pytest.mark.parametrize("synonym", [False, True]) def test_fulltext_query_requires_tsvector_columns( @@ -12,7 +14,8 @@ def test_fulltext_query_requires_tsvector_columns( ): """Full-text query raises FullTextError when tsvector columns are absent from the database. - The mock CDM engine is SQLite and never has tsvector columns, so the guard in + The mock CDM's own table creation never adds tsvector columns (that's a + separate, opt-in fulltext-install step), so the guard in q_concept_name_fulltext (which inspects the live DB schema) always fires here. """ with pytest.raises(FullTextError): diff --git a/tests/test_fulltext_vocab_schema_postgres.py b/tests/test_fulltext_vocab_schema_postgres.py index e2e43c8..8b94a40 100644 --- a/tests/test_fulltext_vocab_schema_postgres.py +++ b/tests/test_fulltext_vocab_schema_postgres.py @@ -19,23 +19,18 @@ from __future__ import annotations import uuid -from typing import cast import pytest import sqlalchemy as sa from oa_configurator import Role from omop_alchemy.backends import CONCEPT_NAME_TSVECTOR_COLUMN, FullTextError, resolve_backend -from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary from orm_loader.helpers import Base from omop_graph.graph.queries import q_concept_name_fulltext -pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] +from fixtures.helpers import VOCAB_TABLES, schema_translate_map -_VOCAB_TABLES = cast( - "tuple[sa.Table, ...]", - (Domain.__table__, Vocabulary.__table__, Concept_Class.__table__, Concept.__table__), -) +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] def _scoped(pg_db, *, primary_schema: str, vocab_schema: str) -> sa.Connection: @@ -43,11 +38,7 @@ def _scoped(pg_db, *, primary_schema: str, vocab_schema: str) -> sa.Connection: conn.execute(sa.text(f"CREATE SCHEMA {primary_schema}")) conn.execute(sa.text(f"CREATE SCHEMA {vocab_schema}")) return conn.execution_options( - schema_translate_map={ - Role.PRIMARY.value: primary_schema, - Role.VOCAB.value: vocab_schema, - Role.RESULTS.value: primary_schema, - } + schema_translate_map=schema_translate_map(primary_schema, vocab_schema=vocab_schema) ) @@ -55,7 +46,7 @@ def test_fulltext_query_finds_the_tsvector_column_in_the_vocab_schema(pg_db): primary_schema = f"fulltext_primary_{uuid.uuid4().hex[:8]}" vocab_schema = f"fulltext_vocab_{uuid.uuid4().hex[:8]}" scoped = _scoped(pg_db, primary_schema=primary_schema, vocab_schema=vocab_schema) - Base.metadata.create_all(bind=scoped, tables=_VOCAB_TABLES, checkfirst=True) + Base.metadata.create_all(bind=scoped, tables=VOCAB_TABLES, checkfirst=True) backend = resolve_backend(scoped) backend.install_fulltext_on_table( @@ -84,7 +75,7 @@ def test_fulltext_query_still_raises_when_the_column_is_genuinely_absent(pg_db): primary_schema = f"fulltext_primary_{uuid.uuid4().hex[:8]}" vocab_schema = f"fulltext_vocab_{uuid.uuid4().hex[:8]}" scoped = _scoped(pg_db, primary_schema=primary_schema, vocab_schema=vocab_schema) - Base.metadata.create_all(bind=scoped, tables=_VOCAB_TABLES, checkfirst=True) + Base.metadata.create_all(bind=scoped, tables=VOCAB_TABLES, checkfirst=True) with pytest.raises(FullTextError): q_concept_name_fulltext("kidney cancer", engine=scoped) diff --git a/tests/test_grounding.py b/tests/test_grounding.py index f91216d..d4d33ff 100644 --- a/tests/test_grounding.py +++ b/tests/test_grounding.py @@ -60,6 +60,8 @@ def _unconstrained_constraints() -> GroundingConstraints: ) +@pytest.mark.postgresql +@pytest.mark.db_dialect @pytest.mark.parametrize( "query,expected_concept_id", [ @@ -96,6 +98,8 @@ def test_grounding_resolves_expected_standard_concepts( assert ranked[0].concept_id == expected_concept_id +@pytest.mark.postgresql +@pytest.mark.db_dialect def test_grounding_maps_non_standard_candidate_via_relationships( mock_cdm_kg: KnowledgeGraph, ) -> None: @@ -114,6 +118,8 @@ def test_grounding_maps_non_standard_candidate_via_relationships( assert ranked[0].concept_id == 196653 +@pytest.mark.postgresql +@pytest.mark.db_dialect def test_grounding_rejects_concepts_outside_anchored_hierarchy( mock_cdm_kg: KnowledgeGraph, ) -> None: @@ -131,6 +137,8 @@ def test_grounding_rejects_concepts_outside_anchored_hierarchy( assert ranked == [] +@pytest.mark.postgresql +@pytest.mark.db_dialect def test_grounding_maps_non_standard_candidate_without_parent_ids( mock_cdm_kg: KnowledgeGraph, ) -> None: @@ -150,6 +158,8 @@ def test_grounding_maps_non_standard_candidate_without_parent_ids( assert ranked[0].identity_hops == 1 +@pytest.mark.postgresql +@pytest.mark.db_dialect def test_grounding_standard_candidate_without_parent_ids_is_zero_hop( mock_cdm_kg: KnowledgeGraph, ) -> None: diff --git a/tests/test_kg_ancestry.py b/tests/test_kg_ancestry.py index 887e7c4..084d523 100644 --- a/tests/test_kg_ancestry.py +++ b/tests/test_kg_ancestry.py @@ -1,11 +1,15 @@ from __future__ import annotations +import pytest + from omop_graph.graph.kg import KnowledgeGraph from omop_graph.graph.nodes import LabelMatchKind from omop_graph.graph.paths import find_standard_paths from omop_graph.reasoning.resolvers.resolvers import CandidateHit from fixtures.mock_cdm import PARENT_CANCER_ID, CONCEPT_META_ID +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + def test_get_potential_ancestors_batch_returns_only_real_ancestors( mock_cdm_kg: KnowledgeGraph, diff --git a/tests/test_oaklib_interface.py b/tests/test_oaklib_interface.py index ebfe371..19315c9 100644 --- a/tests/test_oaklib_interface.py +++ b/tests/test_oaklib_interface.py @@ -1,8 +1,12 @@ from __future__ import annotations +import pytest + from omop_graph.graph.kg import KnowledgeGraph from omop_graph.oaklib_interface.omop_implementation import OMOPTextAnnotatorInterface +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + def test_annotate_text_standardizes_non_standard_candidate_without_parent_annotation( mock_cdm_kg: KnowledgeGraph, diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index 2e7d348..033bdc5 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -22,7 +22,6 @@ from __future__ import annotations from datetime import date -from typing import cast import sqlalchemy as sa import sqlalchemy.orm @@ -32,7 +31,6 @@ ResolvedCDMDatabase, ResolvedConnection, Role, - qualified, ) from oa_configurator.testing import isolated_test_schema from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary @@ -46,85 +44,62 @@ from omop_graph.oaklib_interface.omop_resource import OMOPOntologyResource from omop_graph.config import OmopGraphConfig +from fixtures.helpers import VOCAB_TABLES, fk_triggers_disabled, schema_translate_map + _META_CONCEPT_ID = 0 _CONCEPT_ID = 1001 _TODAY = date(2020, 1, 1) _FAR_FUTURE = date(2099, 12, 31) -_VOCAB_TABLES = cast( - "tuple[sa.Table, ...]", - (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 vocab bootstrap: Domain/Vocabulary/Concept_Class/Concept form an FK insert cycle, so triggers are disabled around the insert then - re-enabled. Accepts an Engine or an open Connection; an Engine has no - .execute() of its own, so one short-lived connection is opened for the - trigger toggles. + re-enabled (a no-op on SQLite, which doesn't enforce FK constraints by + default). """ - 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")) - 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() + with fk_triggers_disabled(bindable, VOCAB_TABLES): + 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() def test_omop_resource_execution_options_carry_the_configured_schema() -> None: @@ -187,18 +162,27 @@ def resolve_database(self, name): def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_resource( - pg_db, + pg_db, tmp_path ) -> 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.""" + actually consumes them. Every vocab-role KnowledgeGraph query (Phase 4.2's + retrofit) now genuinely resolves through vocab_engine, not cdm_engine, so + unlike the in-memory-SQLite shortcut this test used before that fix, the + vocab side needs a real, separately seeded database -- a SQLite tempfile + persists across the adapter's own connections, unlike ``:memory:``.""" + vocab_db_path = tmp_path / "vocab.db" + vocab_url = f"sqlite:///{vocab_db_path}" + vocab_engine = sa.create_engine(vocab_url).execution_options( + schema_translate_map={Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} + ) + Base.metadata.create_all(bind=vocab_engine, checkfirst=True) + _seed_one_concept(vocab_engine, concept_id=_CONCEPT_ID, name="Split-wiring concept") + vocab_engine.dispose() + with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_split") as schema: engine = pg_db.connection.engine.execution_options( - schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + schema_translate_map=schema_translate_map(schema) ) Base.metadata.create_all(bind=engine, checkfirst=True) _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Split-wiring concept") @@ -209,7 +193,7 @@ def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_ execution_options={ SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} }, - vocab_url="sqlite:///:memory:", + vocab_url=vocab_url, vocab_execution_options={ SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} }, @@ -218,7 +202,7 @@ def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_ 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 str(adapter.kg.vocab_engine.url) == vocab_url assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Split-wiring concept" @@ -230,7 +214,7 @@ def test_omop_alchemy_implementation_reuses_one_engine_when_no_split_is_configur 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={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + schema_translate_map=schema_translate_map(schema) ) Base.metadata.create_all(bind=engine, checkfirst=True) _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="No-split concept") @@ -253,7 +237,7 @@ def test_omop_alchemy_implementation_builds_its_engine_via_resolved_create_engin ) -> None: with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_resolved") as schema: engine = pg_db.connection.engine.execution_options( - schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + schema_translate_map=schema_translate_map(schema) ) Base.metadata.create_all(bind=engine, checkfirst=True) _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Resolved-path concept") @@ -300,11 +284,19 @@ def resolve_database(self, name): def test_kg_injection_path_resolves_against_the_configured_schema(pg_db) -> None: + """Path 1 from this module's docstring: kg= injection. OMOPAlchemyImplementation + must use the injected KnowledgeGraph as-is rather than rebuilding its own + engine from engine_string -- proved here by pairing a real, schema-aware + kg with a deliberately broken engine_string ("sqlite:///:memory:", never + actually queried). If the implementation ever fell back to building its + own engine instead of using kg=, this would error or return nothing + instead of the seeded concept's name. + """ schema = "phase4_oaklib_kg_injection" conn = pg_db.connection conn.execute(sa.text(f"CREATE SCHEMA {schema}")) scoped = conn.execution_options( - schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + schema_translate_map=schema_translate_map(schema) ) Base.metadata.create_all(bind=scoped, checkfirst=True) _seed_one_concept(scoped, concept_id=_CONCEPT_ID, name="Test concept") @@ -317,9 +309,16 @@ def test_kg_injection_path_resolves_against_the_configured_schema(pg_db) -> None def test_bare_engine_string_path_resolves_against_the_configured_schema(pg_db) -> None: + """Path 2 from this module's docstring: bare engine_string=/resource=-only + construction. OAK-lib's own generic materialize() invocation only ever + supplies a URL string, never a live connection or kg= -- this is the one + path the original bug actually broke, since schema_translate_map has to + be carried through purely via execution_options, with no resolved= + object to lean on. + """ with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_bare") as schema: engine = pg_db.connection.engine.execution_options( - schema_translate_map={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + schema_translate_map=schema_translate_map(schema) ) Base.metadata.create_all(bind=engine, checkfirst=True) _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Bare-string concept") diff --git a/tests/test_predicate_flags.py b/tests/test_predicate_flags.py index 888abe0..d20de02 100644 --- a/tests/test_predicate_flags.py +++ b/tests/test_predicate_flags.py @@ -9,8 +9,12 @@ from __future__ import annotations +import pytest + from omop_graph.graph.kg import KnowledgeGraph +pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] + def test_zero_flag_is_not_read_as_true(mock_cdm_kg: KnowledgeGraph) -> None: """The fixture's relationships are all ``is_hierarchical='0'``. diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index 582d4a4..68364e3 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -15,20 +15,19 @@ import pytest import sqlalchemy as sa -from oa_configurator import Role from orm_loader.helpers import Base from omop_graph.cli import relationship_classification from omop_graph.extensions.omop_alchemy import RelationshipClass, RelationshipMapping +from fixtures.helpers import schema_translate_map + 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={Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} - ) + return conn.execution_options(schema_translate_map=schema_translate_map(schema)) def test_relationship_classification_respects_the_configured_schema(pg_db): diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index 04f5ac4..2f45cd5 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -19,7 +19,7 @@ from __future__ import annotations from datetime import date -from typing import Iterator, NamedTuple, cast +from typing import Iterator, NamedTuple import pytest import sqlalchemy as sa @@ -47,6 +47,8 @@ ) from omop_graph.graph.kg import KnowledgeGraph +from fixtures.helpers import VOCAB_TABLES, fk_triggers_disabled, schema_translate_map + pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] META_CONCEPT_ID = 0 @@ -55,17 +57,10 @@ _TODAY = date(2020, 1, 1) _FAR_FUTURE = date(2099, 12, 31) -_VOCAB_TABLES = cast( - "tuple[sa.Table, ...]", - ( - Domain.__table__, - Vocabulary.__table__, - Concept_Class.__table__, - Concept.__table__, - Relationship.__table__, - Concept_Relationship.__table__, - ), -) +# This test also needs Relationship/Concept_Relationship beyond the shared +# core (they're not part of every consumer's minimal vocab bootstrap, but +# are exactly what the split-connection predicate/edge merge is testing). +_VOCAB_TABLES = VOCAB_TABLES + (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 @@ -135,18 +130,10 @@ def split_engines() -> Iterator[_Engines]: isolated_test_schema(vocab_raw) as vocab_schema, ): primary_engine = primary_raw.execution_options( - schema_translate_map={ - Role.PRIMARY.value: primary_schema, - Role.VOCAB.value: primary_schema, - Role.RESULTS.value: primary_schema, - } + schema_translate_map=schema_translate_map(primary_schema) ) vocab_engine = vocab_raw.execution_options( - schema_translate_map={ - Role.PRIMARY.value: vocab_schema, - Role.VOCAB.value: vocab_schema, - Role.RESULTS.value: vocab_schema, - } + schema_translate_map=schema_translate_map(vocab_schema) ) _shadow_metadata.create_all(primary_engine) @@ -158,15 +145,8 @@ def split_engines() -> Iterator[_Engines]: # 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')) + with fk_triggers_disabled(vocab_engine, _VOCAB_TABLES): + _seed(primary_engine, vocab_engine) yield _Engines(primary=primary_engine, vocab=vocab_engine) @@ -339,3 +319,46 @@ def test_edges_predicate_kinds_filter_applies_after_merge(split_engines: _Engine ) assert edges == () + + +# Phase 4.2: KnowledgeGraph's ~15 concept/ancestor/synonym query methods used +# to always query self.session_factory() (primary), silently wrong once +# vocab lives on a genuinely separate connection -- not detected, not +# refused, and (before these tests) not covered at all. Each of these was +# broken against split_engines before the vocab_session_factory() retrofit, +# confirmed by running them against this same fixture on the pre-fix code. + + +def test_concept_view_resolves_against_the_vocab_connection(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + view = kg.concept_view(SUBJECT_CONCEPT_ID) + + assert view.concept_id == SUBJECT_CONCEPT_ID + assert view.concept_name == "Subject concept" + + +def test_concept_id_by_code_resolves_against_the_vocab_connection(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + + assert kg.concept_id_by_code("SNOMED", "SUBJ") == SUBJECT_CONCEPT_ID + + +def test_concept_ids_by_label_resolves_against_the_vocab_connection(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + + assert kg.concept_ids_by_label("Subject concept") == (SUBJECT_CONCEPT_ID,) + + +def test_predicate_name_resolves_against_the_vocab_connection(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + + assert kg.predicate_name("maps to") == "Maps to" + + +def test_valid_domains_and_vocabularies_resolve_against_the_vocab_connection( + split_engines: _Engines, +) -> None: + kg = _split_kg(split_engines) + + assert {"Metadata", "Condition"} <= kg._valid_domains + assert {"OMOP", "SNOMED"} <= kg._valid_vocabularies From e922452277a17177b230dcf156d6908c64d379b9 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 23:41:51 +0000 Subject: [PATCH 09/20] Complete vocab/cdm split --- src/omop_graph/graph/kg.py | 25 ++++++++-------- .../oaklib_interface/omop_implementation.py | 30 ++++++++----------- tests/test_vocab_split_connection.py | 23 ++++++++++++++ 3 files changed, 48 insertions(+), 30 deletions(-) diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index ed1f611..5145c77 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -514,7 +514,6 @@ def predicate_kinds( def relationships( self, - session: Session, subjects: tuple[int, ...] | None, predicates: tuple[str, ...] | None, objects: tuple[int, ...] | None, @@ -542,7 +541,6 @@ def relationships( """ if invert: for s, p, o in self.relationships( - session=session, subjects=objects, predicates=predicates, objects=subjects, @@ -550,14 +548,15 @@ def relationships( yield o, p, s return - for s, p, o in session.execute( - q_relationships( - subjects=subjects, - predicates=predicates, - objects=objects, - ) - ): - yield s, p, o + with self.vocab_session_factory() as session: + for s, p, o in session.execute( + q_relationships( + subjects=subjects, + predicates=predicates, + objects=objects, + ) + ): + yield s, p, o def reverse_predicate_id(self, relationship_id: str) -> Optional[str]: """ @@ -695,7 +694,6 @@ def children(self, concept_id) -> tuple[int, ...]: def entities( self, - session: Session, domain: str | None = None, standard_only: bool = True, filter_obsoletes: bool = True, @@ -707,8 +705,9 @@ def entities( filter_obsoletes=filter_obsoletes, ) - for row in session.execute(query): - yield int(row.concept_id) + with self.vocab_session_factory() as session: + for row in session.execute(query): + yield int(row.concept_id) def roots( self, domain_id: str | None = None, vocabulary_id: str | None = None diff --git a/src/omop_graph/oaklib_interface/omop_implementation.py b/src/omop_graph/oaklib_interface/omop_implementation.py index 77e4a8b..57a3eed 100644 --- a/src/omop_graph/oaklib_interface/omop_implementation.py +++ b/src/omop_graph/oaklib_interface/omop_implementation.py @@ -538,15 +538,13 @@ def entities( # ty: ignore[invalid-method-override] Concept identifiers. """ - with self.kg.session_factory() as session: - cids = tuple( - self.kg.entities( - session=session, - domain=domain, - standard_only=standard_only, - filter_obsoletes=filter_obsoletes, - ) + cids = tuple( + self.kg.entities( + domain=domain, + standard_only=standard_only, + filter_obsoletes=filter_obsoletes, ) + ) for cid in cids: yield self._concept_curie(cid) @@ -644,16 +642,14 @@ def relationships( # ty: ignore[invalid-method-override] else None ) - with self.kg.session_factory() as session: - relationships = tuple( - self.kg.relationships( - session=session, - subjects=subject_ids, - predicates=predicate_ids, - objects=object_ids, - invert=invert, - ) + relationships = tuple( + self.kg.relationships( + subjects=subject_ids, + predicates=predicate_ids, + objects=object_ids, + invert=invert, ) + ) for s, p, o in relationships: yield ( diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index 2f45cd5..a191738 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -362,3 +362,26 @@ def test_valid_domains_and_vocabularies_resolve_against_the_vocab_connection( assert {"Metadata", "Condition"} <= kg._valid_domains assert {"OMOP", "SNOMED"} <= kg._valid_vocabularies + + +def test_entities_resolves_against_the_vocab_connection(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + ids = tuple(kg.entities(domain="Condition")) + + assert set(ids) == {SUBJECT_CONCEPT_ID, OBJECT_CONCEPT_ID} + + +def test_relationships_resolves_against_the_vocab_connection(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + triples = tuple(kg.relationships(subjects=(SUBJECT_CONCEPT_ID,), predicates=None, objects=None)) + + assert triples == ((SUBJECT_CONCEPT_ID, "maps to", OBJECT_CONCEPT_ID),) + + +def test_relationships_invert_swaps_subjects_and_objects(split_engines: _Engines) -> None: + kg = _split_kg(split_engines) + triples = tuple( + kg.relationships(subjects=(OBJECT_CONCEPT_ID,), predicates=None, objects=None, invert=True) + ) + + assert triples == ((OBJECT_CONCEPT_ID, "maps to", SUBJECT_CONCEPT_ID),) From b5a755575da3e09264e3813a08f25356cd9bb84d Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Tue, 15 Sep 2026 23:42:14 +0000 Subject: [PATCH 10/20] Use helper in test --- tests/test_oaklib_schema_awareness.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index 033bdc5..9aa3d25 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -190,9 +190,7 @@ def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_ resource = OMOPOntologyResource( url=pg_db.connection.engine.url.render_as_string(hide_password=False), - execution_options={ - SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} - }, + execution_options={SCHEMA_TRANSLATE_MAP_KEY: schema_translate_map(schema)}, vocab_url=vocab_url, vocab_execution_options={ SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} @@ -223,7 +221,7 @@ def test_omop_alchemy_implementation_reuses_one_engine_when_no_split_is_configur resource = OMOPOntologyResource( url=pg_db.connection.engine.url.render_as_string(hide_password=False), execution_options={ - SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + SCHEMA_TRANSLATE_MAP_KEY: schema_translate_map(schema) }, ) @@ -330,7 +328,7 @@ def test_bare_engine_string_path_resolves_against_the_configured_schema(pg_db) - # 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_KEY: {Role.PRIMARY.value: schema, Role.VOCAB.value: schema, Role.RESULTS.value: schema} + SCHEMA_TRANSLATE_MAP_KEY: schema_translate_map(schema) }, ) adapter = OMOPAlchemyImplementation(resource=resource) From 89b997ed5a0252242c0caba910dec29d08d96172 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 23 Sep 2026 01:14:55 +0000 Subject: [PATCH 11/20] Cleanup of Role references and docstring comments --- src/omop_graph/cli.py | 13 ++++- src/omop_graph/graph/kg.py | 6 +- src/omop_graph/graph/queries.py | 16 +++--- tests/fixtures/helpers.py | 8 ++- tests/test_fulltext_vocab_schema_postgres.py | 24 +++----- tests/test_oaklib_schema_awareness.py | 11 ++-- tests/test_relationship_classification.py | 2 +- tests/test_schema_provenance_guard.py | 6 +- tests/test_vocab_split_connection.py | 60 +++++++------------- 9 files changed, 67 insertions(+), 79 deletions(-) diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index 1a3a86f..c62119c 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -11,7 +11,13 @@ import typer from sqlalchemy.orm import sessionmaker -from oa_configurator import ResolvedCDMDatabase, Role, ensure_schema, guard_schema_provenance, schema_of +from oa_configurator import ( + ResolvedCDMDatabase, + Role, + ensure_schema, + guard_schema_provenance_for, + schema_of +) from orm_loader.backends import STAGING_SCHEMA, resolve_backend from orm_loader.helpers import bulk_load_context @@ -214,7 +220,10 @@ def relationship_classification( # 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): + guard = guard_schema_provenance_for( + connection, resolved, role=Role.PRIMARY, tables=tables_to_drop # ty: ignore[invalid-argument-type] + ) + with guard: 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 diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index 5145c77..b310bec 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -150,7 +150,7 @@ def provider_type(self) -> str: 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 + RelationshipMapping is an omop-graph extension table, not vocab-tagged, so it never lives on a split ``vocab_engine``. This always runs against the primary connection. """ @@ -197,9 +197,9 @@ class KnowledgeGraph(GraphBackend): ``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 + ``cdm_engine``, the three queries that join a vocab-tagged table (Concept/Concept_Relationship/Relationship) against - RelationshipMapping (not vocab-role, since it's an omop-graph + RelationshipMapping (not vocab-tagged, 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. """ diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 6d6eccc..34458a3 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -22,6 +22,7 @@ case, exists, func, + inspect, literal, or_, select, @@ -32,7 +33,7 @@ from sqlalchemy.orm import aliased from sqlalchemy.sql import Select -from oa_configurator import Role, schema_inspect +from oa_configurator import Role, schema_of from omop_alchemy.backends import ( CONCEPT_NAME_TSVECTOR_COLUMN, @@ -354,7 +355,8 @@ def q_concept_name_fulltext( Concept_Synonym.concept_synonym_name if synonym else Concept.concept_name ) # Fulltext are in VOCAB schema - inspector = schema_inspect(engine, role=Role.VOCAB) + inspector = inspect(engine) + vocab_schema = schema_of(engine, schema_tag=Role.VOCAB) target_table = Concept_Synonym if synonym else Concept target_col = ( CONCEPT_SYNONYM_NAME_TSVECTOR_COLUMN @@ -366,7 +368,7 @@ def q_concept_name_fulltext( tsvector_col = next( ( c["name"] - for c in inspector.get_columns(target_table.__tablename__) + for c in inspector.get_columns(target_table.__tablename__, schema=vocab_schema) if c["name"] == target_col ), None, @@ -438,7 +440,7 @@ def q_predicate_row_with_ancestry( ---------- include_classification : bool, optional Join in RelationshipMapping's predicate_kind/predicate_subkind. Set to - False for a split-connection deployment (Relationship is vocab-role, + False for a split-connection deployment (Relationship is vocab-tagged, 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. @@ -506,7 +508,7 @@ def q_all_predicates_with_ancestry(*, include_classification: bool = True) -> Se 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 + The primary-tagged half of a split-connection predicate lookup, pairing with :func:`q_predicate_row_with_ancestry`'s ``include_classification=False``. """ return select( @@ -519,7 +521,7 @@ def q_relationship_mapping_row(relationship_id: str) -> Select: 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. + The primary-tagged 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. """ @@ -546,7 +548,7 @@ def q_edges( ---------- include_classification : bool, optional Join in RelationshipMapping's predicate_kind/predicate_subkind. - Concept_Relationship is vocab-role, RelationshipMapping is not, so + Concept_Relationship is vocab-tagged, 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 diff --git a/tests/fixtures/helpers.py b/tests/fixtures/helpers.py index 4b405d3..bbd0ea9 100644 --- a/tests/fixtures/helpers.py +++ b/tests/fixtures/helpers.py @@ -16,7 +16,7 @@ from typing import Iterator, cast import sqlalchemy as sa -from oa_configurator import Role, qualified +from oa_configurator import Role, qualified, schema_of, validate_schema_tag from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary VOCAB_TABLES = cast( @@ -70,7 +70,8 @@ def fk_triggers_disabled( conn = bindable.connect() if opened_here else bindable try: for table in tables: - conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name)} DISABLE TRIGGER ALL")) + physical_schema = schema_of(conn, schema_tag=validate_schema_tag(table)) + conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name, physical_schema=physical_schema)} DISABLE TRIGGER ALL")) if opened_here: conn.commit() finally: @@ -83,7 +84,8 @@ def fk_triggers_disabled( conn = bindable.connect() if opened_here else bindable try: for table in tables: - conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name)} ENABLE TRIGGER ALL")) + physical_schema = schema_of(conn, schema_tag=validate_schema_tag(table)) + conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name, physical_schema=physical_schema)} ENABLE TRIGGER ALL")) if opened_here: conn.commit() finally: diff --git a/tests/test_fulltext_vocab_schema_postgres.py b/tests/test_fulltext_vocab_schema_postgres.py index 8b94a40..984250d 100644 --- a/tests/test_fulltext_vocab_schema_postgres.py +++ b/tests/test_fulltext_vocab_schema_postgres.py @@ -1,19 +1,13 @@ """q_concept_name_fulltext() must inspect the VOCAB schema, not primary. -Concept/Concept_Synonym are VOCAB-role tables. schema_inspect(engine) with -no role= defaults to Role.PRIMARY, so whenever vocab_schema differs from -cdm_schema (a normal same-server split, not just a same-schema deployment), -the old code inspected the wrong schema, found no tsvector column, and -raised a false FullTextError even though the column genuinely exists. The -only prior coverage (test_fulltext_optional.py) uses SQLite, where -supports_schemas() is False and every role folds to None -- masking this -completely. Real, distinct primary/vocab Postgres schemas here instead. - -Schemas are created directly on pg_db's own already-open connection -(CREATE SCHEMA, rolled back automatically at teardown) rather than via -isolated_test_schema(): that opens a second, genuinely separate connection -from the same pool, which can starve waiting on pg_db's own still-open -transaction. +Concept/Concept_Synonym are VOCAB-tagged; inspecting without an explicit +schema= defaults to Role.PRIMARY, so a real primary/vocab split (unlike +SQLite, where every schema tag folds to None) would raise a false +FullTextError even though the tsvector column exists. + +Schemas are created directly on pg_db's own connection rather than via +isolated_test_schema(), which opens a second connection that can starve +against pg_db's still-open transaction. """ from __future__ import annotations @@ -56,7 +50,7 @@ def test_fulltext_query_finds_the_tsvector_column_in_the_vocab_schema(pg_db): index_name="idx_concept_name_tsvector_test", create_indexes=True, fastupdate=True, - role=Role.VOCAB, + schema_tag=Role.VOCAB.value, ) # Confirmed absent from primary_schema: proves a real vocab/primary diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index 9aa3d25..1a7f1e5 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -164,13 +164,10 @@ def resolve_database(self, name): def test_omop_alchemy_implementation_builds_a_genuine_vocab_engine_from_a_split_resource( pg_db, tmp_path ) -> 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. Every vocab-role KnowledgeGraph query (Phase 4.2's - retrofit) now genuinely resolves through vocab_engine, not cdm_engine, so - unlike the in-memory-SQLite shortcut this test used before that fix, the - vocab side needs a real, separately seeded database -- a SQLite tempfile - persists across the adapter's own connections, unlike ``:memory:``.""" + """Confirms OMOPAlchemyImplementation actually builds vocab_engine from + omop_resource()'s vocab_url/vocab_execution_options, not just carries them. + Needs a real seeded database, not ``:memory:``, since a SQLite tempfile + persists across the adapter's own connections and ``:memory:`` doesn't.""" vocab_db_path = tmp_path / "vocab.db" vocab_url = f"sqlite:///{vocab_db_path}" vocab_engine = sa.create_engine(vocab_url).execution_options( diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index 68364e3..5a84298 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -61,7 +61,7 @@ def test_relationship_classification_respects_the_configured_schema(pg_db): def test_relationship_classification_refuses_a_genuinely_split_vocab_connection(pg_db): """Postgres has no cross-database inline FK, so RelationshipMapping's FK - to relationship.relationship_id (VOCAB-role) can never be created once + to relationship.relationship_id (VOCAB-tagged) can never be created once vocab_connection is a genuinely separate connection. """ resolved = dataclasses.replace( diff --git a/tests/test_schema_provenance_guard.py b/tests/test_schema_provenance_guard.py index 22026ec..187e214 100644 --- a/tests/test_schema_provenance_guard.py +++ b/tests/test_schema_provenance_guard.py @@ -80,7 +80,11 @@ def test_relationship_classification_guard_fires_on_reconfigured_schema( # 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" + conn, + database_name=resolved_a.name, + schema_tag=Role.PRIMARY, + new_physical_schema=schema_a, + reason="test setup baseline", ) monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: resolved_a) omop_graph_cli.relationship_classification() diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index a191738..a688538 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -1,18 +1,14 @@ -"""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 +"""Split-connection vocab routing. + +q_edges/q_predicate_row_with_ancestry/q_all_predicates_with_ancestry join a +vocab-tagged table against RelationshipMapping. A genuinely separate +vocab_connection can't do it in one SQL join. Instead, KnowledgeGraph fetches each side +from its own engine and merges in Python (see kg.py's +_vocab_split/_predicate_from_rows/_relationship_mapping_lookup). + +Uses two real, distinct Postgres connections (test_cdm, test_orm) standing +in for primary/vocab servers, each with its own schema via +isolated_test_schema(). Note: rollback-based isolation can't stand in for two genuinely separate physical connections. """ @@ -57,19 +53,14 @@ _TODAY = date(2020, 1, 1) _FAR_FUTURE = date(2099, 12, 31) -# This test also needs Relationship/Concept_Relationship beyond the shared -# core (they're not part of every consumer's minimal vocab bootstrap, but -# are exactly what the split-connection predicate/edge merge is testing). +# Relationship/Concept_Relationship aren't in the shared minimal vocab +# bootstrap, but are exactly what the split-connection merge is testing. _VOCAB_TABLES = VOCAB_TABLES + (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. +# Postgres has no cross-database inline FK, so RelationshipMapping's FK to +# relationship.relationship_id can't be created as DDL across genuinely +# separate databases. These shadow tables reproduce the real columns +# without it, since a SELECT/INSERT only depends on column shape. _shadow_metadata = sa.MetaData() _shadow_relationship_class = sa.Table( "relationship_class", @@ -139,12 +130,9 @@ def split_engines() -> Iterator[_Engines]: _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. + # Domain/Vocabulary/Concept_Class/Concept form a genuine FK bootstrap + # cycle, handled the same way production bulk-loads do: disable FK + # triggers for the load, then re-enable them. with fk_triggers_disabled(vocab_engine, _VOCAB_TABLES): _seed(primary_engine, vocab_engine) @@ -321,14 +309,6 @@ def test_edges_predicate_kinds_filter_applies_after_merge(split_engines: _Engine assert edges == () -# Phase 4.2: KnowledgeGraph's ~15 concept/ancestor/synonym query methods used -# to always query self.session_factory() (primary), silently wrong once -# vocab lives on a genuinely separate connection -- not detected, not -# refused, and (before these tests) not covered at all. Each of these was -# broken against split_engines before the vocab_session_factory() retrofit, -# confirmed by running them against this same fixture on the pre-fix code. - - def test_concept_view_resolves_against_the_vocab_connection(split_engines: _Engines) -> None: kg = _split_kg(split_engines) view = kg.concept_view(SUBJECT_CONCEPT_ID) From cbb2e6cd3c6f1d4bda2ac7224ae27deef8c4cc40 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 23 Sep 2026 04:29:39 +0000 Subject: [PATCH 12/20] Follow-up from internal review --- src/omop_graph/cli.py | 30 ++++++----------------- tests/test_relationship_classification.py | 6 ++++- 2 files changed, 13 insertions(+), 23 deletions(-) diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index c62119c..02331f5 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -1,7 +1,5 @@ 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 @@ -12,10 +10,11 @@ from sqlalchemy.orm import sessionmaker from oa_configurator import ( - ResolvedCDMDatabase, - Role, - ensure_schema, - guard_schema_provenance_for, + ResolvedCDMDatabase, + Role, + ensure_schema, + guard_schema_provenance_for, + open_connection, schema_of ) @@ -69,19 +68,6 @@ 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, *, @@ -204,7 +190,7 @@ def relationship_classification( f"{loader_backend.qualified_staging_name(RelationshipClass.__tablename__)} CASCADE" ), ) - with _open_connection(engine) as connection: + with open_connection(engine) as connection: for stmt in drop_staging_sql: connection.execute(stmt) @@ -219,9 +205,9 @@ def relationship_classification( ] # 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 open_connection(engine) as connection: guard = guard_schema_provenance_for( - connection, resolved, role=Role.PRIMARY, tables=tables_to_drop # ty: ignore[invalid-argument-type] + connection, resolved, schema_tag=Role.PRIMARY, tables=tables_to_drop # ty: ignore[invalid-argument-type] ) with guard: Base.metadata.drop_all(bind=connection, tables=tables_to_drop, checkfirst=True) # type: ignore diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index 5a84298..f310b40 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -31,6 +31,9 @@ def _scoped_connection(pg_db, schema: str) -> sa.Connection: def test_relationship_classification_respects_the_configured_schema(pg_db): + """No resolved= is passed, so guard_schema_provenance_for() no-ops here by + design (a bare engine= caller with no resolved config behind it); this test + is about schema routing, not provenance drift protection.""" scoped = _scoped_connection(pg_db, "phase4_regression_test") Base.metadata.create_all(bind=scoped, checkfirst=True) @@ -76,7 +79,8 @@ def test_relationship_classification_refuses_a_genuinely_split_vocab_connection( 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.""" + DROP TABLE/enum-drop cleanup exists for, must not fail. No resolved= here + either, so guard_schema_provenance_for() no-ops by design, same as above.""" scoped = _scoped_connection(pg_db, "phase4_idempotent_test") Base.metadata.create_all(bind=scoped, checkfirst=True) From 7825eac8a5136a6c9b257b08a38fa4d3c07dddec Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 23 Sep 2026 04:48:41 +0000 Subject: [PATCH 13/20] Disambiguate physical schema from schema tag --- src/omop_graph/cli.py | 4 ++-- src/omop_graph/graph/queries.py | 4 ++-- tests/fixtures/helpers.py | 18 ++++++++++++++---- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index 02331f5..ad7bf99 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -15,7 +15,7 @@ ensure_schema, guard_schema_provenance_for, open_connection, - schema_of + physical_schema_of ) from orm_loader.backends import STAGING_SCHEMA, resolve_backend @@ -172,7 +172,7 @@ def relationship_classification( "this command, or provision relationship_class/relationship_mapping " "manually without the FK constraint." ) - db_schema = schema_of(engine) + db_schema = physical_schema_of(engine) ensure_schema(engine, db_schema) ensure_schema(engine, STAGING_SCHEMA) diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 34458a3..50925f1 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -33,7 +33,7 @@ from sqlalchemy.orm import aliased from sqlalchemy.sql import Select -from oa_configurator import Role, schema_of +from oa_configurator import Role, physical_schema_of from omop_alchemy.backends import ( CONCEPT_NAME_TSVECTOR_COLUMN, @@ -356,7 +356,7 @@ def q_concept_name_fulltext( ) # Fulltext are in VOCAB schema inspector = inspect(engine) - vocab_schema = schema_of(engine, schema_tag=Role.VOCAB) + vocab_schema = physical_schema_of(engine, schema_tag=Role.VOCAB) target_table = Concept_Synonym if synonym else Concept target_col = ( CONCEPT_SYNONYM_NAME_TSVECTOR_COLUMN diff --git a/tests/fixtures/helpers.py b/tests/fixtures/helpers.py index bbd0ea9..ca5514a 100644 --- a/tests/fixtures/helpers.py +++ b/tests/fixtures/helpers.py @@ -16,8 +16,18 @@ from typing import Iterator, cast import sqlalchemy as sa -from oa_configurator import Role, qualified, schema_of, validate_schema_tag -from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary +from oa_configurator import ( + Role, + qualified, + physical_schema_of, + validate_schema_tag +) +from omop_alchemy.cdm.model.vocabulary import ( + Concept, + Concept_Class, + Domain, + Vocabulary +) VOCAB_TABLES = cast( "tuple[sa.Table, ...]", @@ -70,7 +80,7 @@ def fk_triggers_disabled( conn = bindable.connect() if opened_here else bindable try: for table in tables: - physical_schema = schema_of(conn, schema_tag=validate_schema_tag(table)) + physical_schema = physical_schema_of(conn, schema_tag=validate_schema_tag(table)) conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name, physical_schema=physical_schema)} DISABLE TRIGGER ALL")) if opened_here: conn.commit() @@ -84,7 +94,7 @@ def fk_triggers_disabled( conn = bindable.connect() if opened_here else bindable try: for table in tables: - physical_schema = schema_of(conn, schema_tag=validate_schema_tag(table)) + physical_schema = physical_schema_of(conn, schema_tag=validate_schema_tag(table)) conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name, physical_schema=physical_schema)} ENABLE TRIGGER ALL")) if opened_here: conn.commit() From f9bc4bc550d643c10bc59314c58cfa04c61e460b Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 23 Sep 2026 05:31:25 +0000 Subject: [PATCH 14/20] Regression test for guard --- tests/test_relationship_classification.py | 44 +++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index f310b40..537886c 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -12,9 +12,11 @@ from __future__ import annotations import dataclasses +import uuid import pytest import sqlalchemy as sa +from oa_configurator import Role, SchemaDriftError, record_schema_provenance from orm_loader.helpers import Base @@ -91,3 +93,45 @@ def test_relationship_classification_is_idempotent(pg_db): sa.select(sa.func.count()).select_from(RelationshipClass.__table__) ).scalar() assert n_class and n_class > 0 + + +def test_relationship_classification_guard_fires_on_reconfigured_schema(pg_db): + """Unlike the two tests above, resolved= is genuinely passed here -- proves + the guard actually detects drift for relationship_classification itself, + not just that it no-ops correctly when a caller opts out.""" + database_name = f"guard_wiring_test_db_{uuid.uuid4().hex[:8]}" + schema_a = "phase4_guard_wiring_a" + schema_b = "phase4_guard_wiring_b" + + # Both connection and vocab_connection must point at the same object: relationship_classification() + # refuses a genuinely split vocab connection, and dataclasses.replace() only overrides fields + # explicitly passed, so leaving vocab_connection untouched would desync it from the new connection. + single_connection = dataclasses.replace(pg_db.resolved.connection, test_only=False) + resolved_a = dataclasses.replace( + pg_db.resolved, + name=database_name, + schema_name=schema_a, + vocab_schema=schema_a, + results_schema=schema_a, + connection=single_connection, + vocab_connection=single_connection, + ) + # Baseline must be recorded before schema_a has any tables, or the guard's own + # first-time-population check trips on Base.metadata.create_all() below. + scoped_a = _scoped_connection(pg_db, schema_a) + record_schema_provenance( + scoped_a, + database_name=database_name, + schema_tag=Role.PRIMARY.value, + new_physical_schema=schema_a, + reason="test setup", + ) + Base.metadata.create_all(bind=scoped_a, checkfirst=True) + relationship_classification(engine=scoped_a, resolved=resolved_a) + + resolved_b = dataclasses.replace( + resolved_a, schema_name=schema_b, vocab_schema=schema_b, results_schema=schema_b + ) + scoped_b = _scoped_connection(pg_db, schema_b) + with pytest.raises(SchemaDriftError): + relationship_classification(engine=scoped_b, resolved=resolved_b) From 527681365996265d909a348e41e78164217669f3 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 24 Sep 2026 00:17:08 +0000 Subject: [PATCH 15/20] Docstring adaptations --- src/omop_graph/db/session.py | 6 +++--- src/omop_graph/graph/queries.py | 4 +++- .../oaklib_interface/omop_implementation.py | 6 ++++++ .../reasoning/concept_handlers/concept_helpers.py | 12 +++++++++--- .../reasoning/phenotypes/phenotype_simplifier.py | 7 ++++++- src/omop_graph/reasoning/resolvers/resolvers.py | 3 +++ 6 files changed, 30 insertions(+), 8 deletions(-) diff --git a/src/omop_graph/db/session.py b/src/omop_graph/db/session.py index 204b52a..f53460b 100644 --- a/src/omop_graph/db/session.py +++ b/src/omop_graph/db/session.py @@ -49,9 +49,9 @@ def make_engine( execution_options : dict, optional 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. + doesn't define, but may not include the resolver's own reserved + keys: those are always set from the resolved config, and + ``create_engine()`` raises ``ValueError`` if any are overridden here. Returns ------- diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 50925f1..48e0bac 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -542,10 +542,12 @@ def q_edges( within_domain: bool = False, include_classification: bool = True, ) -> Select: - """Query outgoing edges for a batch of concept IDs. + """Query edges for a batch of concept IDs, in either direction. Parameters ---------- + direction : {"in", "out"} + Whether to query incoming or outgoing edges for ``concept_ids``. include_classification : bool, optional Join in RelationshipMapping's predicate_kind/predicate_subkind. Concept_Relationship is vocab-tagged, RelationshipMapping is not, so diff --git a/src/omop_graph/oaklib_interface/omop_implementation.py b/src/omop_graph/oaklib_interface/omop_implementation.py index 57a3eed..e744de2 100644 --- a/src/omop_graph/oaklib_interface/omop_implementation.py +++ b/src/omop_graph/oaklib_interface/omop_implementation.py @@ -731,6 +731,9 @@ def entailed_outgoing_relationships( ) -> Iterable[Tuple[PRED_CURIE, CURIE]]: """ Retrieve outgoing relationships, including those implied by the hierarchy. + + Currently disabled: always raises ``NotImplementedError``, per the + raise message below, since a CDM change broke the body's assumptions. """ raise NotImplementedError("Changes to the CDM currently prevents this function") concept_id = self._parse_concept(curie) @@ -807,6 +810,9 @@ def entailed_relationships_between( ) -> Iterable[PRED_CURIE]: """ Find relationships connecting a subject and object, including hierarchical ones. + + Currently disabled: always raises ``NotImplementedError``, per the + raise message below, since a CDM change broke the body's assumptions. """ raise NotImplementedError( "Change in OMOP CDM made this function not work anymore" diff --git a/src/omop_graph/reasoning/concept_handlers/concept_helpers.py b/src/omop_graph/reasoning/concept_handlers/concept_helpers.py index 25c0822..6a5381d 100644 --- a/src/omop_graph/reasoning/concept_handlers/concept_helpers.py +++ b/src/omop_graph/reasoning/concept_handlers/concept_helpers.py @@ -18,9 +18,15 @@ def standardise_ids( """ Map a set of concept IDs to their Standard Concept IDs. - This function attempts to find a 'Maps to' relationship for each input ID. - If a 'Maps to' edge exists, the target ID is used (Standard Concept). - If no such edge exists, the original ID is returned (fallback to self). + Currently disabled: always raises ``NotImplementedError``. The predicate + search this relied on changed and the body below was never updated to + match (it also still calls ``kg.iter_edges(predicate=...)``, a parameter + that no longer exists; see ``predicate_ids``/``predicate_kinds``). + + Intended behavior once reimplemented: attempts to find a 'Maps to' + relationship for each input ID. If a 'Maps to' edge exists, the target ID + is used (Standard Concept). If no such edge exists, the original ID is + returned (fallback to self). Parameters ---------- diff --git a/src/omop_graph/reasoning/phenotypes/phenotype_simplifier.py b/src/omop_graph/reasoning/phenotypes/phenotype_simplifier.py index 2dbd4a8..b5830ef 100644 --- a/src/omop_graph/reasoning/phenotypes/phenotype_simplifier.py +++ b/src/omop_graph/reasoning/phenotypes/phenotype_simplifier.py @@ -22,7 +22,12 @@ def descendants_exhaustive_subsumes( exclude_roots: set[int] | None = None, ) -> set[int]: """ - Exhaustive descendant closure using ONLY 'Subsumes' + Exhaustive descendant closure using ONLY 'Subsumes'. + + Currently disabled: always raises ``NotImplementedError``. The predicate + search this relied on changed and the body below was never updated to + match (it also still calls ``kg.iter_edges(predicate=...)``, a parameter + that no longer exists; see ``predicate_ids``/``predicate_kinds``). """ if exclude_roots is None: diff --git a/src/omop_graph/reasoning/resolvers/resolvers.py b/src/omop_graph/reasoning/resolvers/resolvers.py index a1c60e4..564677a 100644 --- a/src/omop_graph/reasoning/resolvers/resolvers.py +++ b/src/omop_graph/reasoning/resolvers/resolvers.py @@ -45,6 +45,9 @@ class CandidateHit: The kind of match of this hit. matched_concept_label : str The specific text in the database (name or synonym) that matched. + synonym : bool + Whether ``matched_concept_label`` came from a synonym rather than + the concept's own name. """ concept_id: int From 77db9c0100f4c60cd64c376ae01c46e21a32f452 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 24 Sep 2026 02:49:26 +0000 Subject: [PATCH 16/20] Updated Docs --- docs/reasoning/grounding.md | 12 ++++++------ docs/reasoning/resolvers.md | 7 ++++--- docs/usage/cli.md | 4 ++-- docs/usage/testing.md | 18 +++++++++++++++++- src/omop_graph/graph/edges.py | 20 ++++++++++---------- 5 files changed, 39 insertions(+), 22 deletions(-) diff --git a/docs/reasoning/grounding.md b/docs/reasoning/grounding.md index 19d6ba9..591586d 100644 --- a/docs/reasoning/grounding.md +++ b/docs/reasoning/grounding.md @@ -43,7 +43,7 @@ To accelerate the grounding to standard concepts, `omop-graph` makes use of: | `parent_ids` | `tuple[int, ...]` | `None` | Only accept candidates that are descendants of these OMOP concept IDs (hierarchy validation via `concept_ancestor`). | | `search_constraint` | `ConceptFilter` | `None` | Filters applied to the initial resolver query (concept IDs, domain, vocabulary, standard/active flags, and limit). | | `max_depth` | `int` | `6` | Maximum hop distance allowed between a candidate and its standard anchor. | -| `predicate_kinds` | `frozenset[PredicateKind]` | `{IDENTITY}` | Relationship kinds followed when walking from a non-standard candidate to its standard anchor. | +| `predicate_kinds` | `frozenset[PredicateKind]` | `{IDENTITY}` | Relationship kinds followed when walking from a non-standard candidate to its standard anchor. Currently locked: `__post_init__` raises `ValueError` for any value other than exactly `frozenset({PredicateKind.IDENTITY})` — not yet configurable in practice, despite being a normal dataclass field. | ### ConceptFilter @@ -99,15 +99,15 @@ TotalScore = Relevance - ParsimonyPenalty + BroadnessBonus $$ #### 1. Relevance -Relevance represents the initial semantic fit and is computed as **either** embedding similarity **or** textual similarity — not both simultaneously: +Relevance represents the initial semantic fit and is computed as **either** embedding similarity **or** textual similarity, chosen per candidate rather than globally — not both simultaneously for the same candidate: -- **Without embeddings**: textual similarity is used exclusively. -- **With embeddings** (default when `omop-graph[emb]` is installed and configured): embedding cosine similarity **replaces** the textual score entirely. +- A candidate resolved by `EmbeddingResolver` (`match_kind == LabelMatchKind.EMBEDDING`) gets embedding cosine similarity. +- Every other candidate — resolved via exact/partial/full-text matching — always gets textual similarity, whether or not `omop-graph[emb]` is installed. The two scoring modes: -- **Embedding Similarity**: Cosine similarity between the input text embedding and the concept embedding. Requires `omop-graph[emb]` and a configured `KnowledgeGraphEmbeddingConfiguration` — see the [Knowledge Graph docs](../graph/kg.md#embedding-configuration) and the [omop-emb documentation](https://australiancancerdatanetwork.github.io/omop-emb/) for setup. -- **Textual Similarity**: A custom token-overlap score that heavily penalizes missing words from the user's query but allows for "extra" descriptive words in the OMOP concept name. Used as a fallback when no embedding is available. +- **Embedding Similarity**: Cosine similarity between the input text embedding and the concept embedding. Only applies to candidates resolved via `EmbeddingResolver`, which requires `omop-graph[emb]` and a configured `KnowledgeGraphEmbeddingConfiguration` — see the [Knowledge Graph docs](../graph/kg.md#embedding-configuration), [Resolver Pipelines](resolvers.md), and the [omop-emb documentation](https://australiancancerdatanetwork.github.io/omop-emb/) for setup. +- **Textual Similarity**: A custom token-overlap score that heavily penalizes missing words from the user's query but allows for "extra" descriptive words in the OMOP concept name. Used for every candidate not resolved via embeddings. #### 2. Parsimony: Distance Penalty OMOP is a deep hierarchy. A concept that is 1 hop away from your search term is more likely to be correct than one found 5 hops away. diff --git a/docs/reasoning/resolvers.md b/docs/reasoning/resolvers.md index 0703dc4..285c1eb 100644 --- a/docs/reasoning/resolvers.md +++ b/docs/reasoning/resolvers.md @@ -6,14 +6,15 @@ The backbone of the `ResolverPipeline` are specific **resolvers**. `omop-graph` - **`ExactLabelResolver`**: Exact case-insensitive match. Is there anywhere in the the [`concept`](https://ohdsi.github.io/CommonDataModel/cdm54.html#concept) table the exact string `"Hodgkin lymphoma"` - **`ExactSynonymResolver`**: Similar to `ExactLabelResolver` yet searching the [`concept_synonym`](https://ohdsi.github.io/CommonDataModel/cdm54.html#concept_synonym) table -- **`FullTextResolver`**: Matches irrespective of word order. Not as relevant for the example above but relevant for others (e.g., "Kidney Cancer" -> "Cancer of Kidney"). -- **`FullTextSynonymResolver`**: Similar to `FullTextResolver` yet searching the [`concept_synonym`](https://ohdsi.github.io/CommonDataModel/cdm54.html#concept_synonym) table - **`PartialLabelResolver`**: Substring match. Is the search string `"Hodgkin lymphoma"` a partial component of any [`concept`](https://ohdsi.github.io/CommonDataModel/cdm54.html#concept)? - **`PartialSynonymResolver`**: Similar to `PartialLabelResolver` yet searching the [`concept_synonym`](https://ohdsi.github.io/CommonDataModel/cdm54.html#concept_synonym) table +- **`FullTextResolver`**: Matches irrespective of word order. Not as relevant for the example above but relevant for others (e.g., "Kidney Cancer" -> "Cancer of Kidney"). +- **`FullTextSynonymResolver`**: Similar to `FullTextResolver` yet searching the [`concept_synonym`](https://ohdsi.github.io/CommonDataModel/cdm54.html#concept_synonym) table +- **`EmbeddingResolver`**: Vector-similarity match, appended to the pipeline only when `omop-graph[emb]` (`omop-emb`) is installed. !!! tip - Traversing each of the resolvers one by one can be an exhaustive search. The `ResolverPipeline` therefore offers a `stop_after_resolver` option. If set, retrieval from the DB stops after that resolver has concluded. The resolvers are ordered based on their confidence as above (i.e. **`ExactLabelResolver`** >> **`ExactSynonymResolver`** >> etc.) + Traversing each of the resolvers one by one can be an exhaustive search. The `ResolverPipeline` therefore offers a `stop_after_resolver` option. If set, retrieval from the DB stops after that resolver has concluded. `ALL_RESOLVERS`, the default sequence, is ordered exactly as listed above (i.e. **`ExactLabelResolver`** >> **`ExactSynonymResolver`** >> **`PartialLabelResolver`** >> ... >> **`EmbeddingResolver`**). ```python from omop_alchemy.cdm.query import ConceptFilter diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 2080cc5..7b0a198 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -57,5 +57,5 @@ omop-graph relationship-classification --pred-class-dir | Option | Short | Type | Default | Description | | :--- | :--- | :--- | :--- | :--- | -| **`--pred-class-dir`** | | `String` | **Required** | Path to the directory containing the classification CSVs. | -| **`--verbose`** | `-v` | `Count` | `0` | Increase logging verbosity (use `-v` or `-vv`). | +| **`--pred-class-dir`** | | `String` | `None` (bundled CSVs) | Path to a directory of classification CSVs, overriding the bundled defaults. | +| **`--verbose`**{: title="Global option, not specific to this subcommand — see the note above." } | `-v` | `Count` | `0` | Increase logging verbosity (use `-v` or `-vv`). Global option; must precede the subcommand name (see note above). | diff --git a/docs/usage/testing.md b/docs/usage/testing.md index e750ca2..2f0fd09 100644 --- a/docs/usage/testing.md +++ b/docs/usage/testing.md @@ -20,14 +20,30 @@ The full suite runs against an **in-memory SQLite mock CDM** (`tests/fixtures/mo The grounding test suite is structured with parametrized cases so each clinical term is a separate pytest case for easier isolation and debugging. +### PostgreSQL-only integration suite + +A separate suite, tagged with the `db_dialect` marker and excluded from the +default run (`pytest.toml`'s `-m "not db_dialect"`), runs against a real +PostgreSQL database via oa-configurator's test infrastructure. This is where +schema-drift protection and split-connection behavior are covered: +`test_schema_provenance_guard.py`, `test_vocab_split_connection.py`, +`test_oaklib_schema_awareness.py`, `test_fulltext_vocab_schema_postgres.py`, +`test_predicate_flags.py`, `test_relationship_classification.py`. + ## Running Tests -Run all tests: +Run all tests (SQLite suite only, the default): ```bash pytest ``` +Include the PostgreSQL-only suite: + +```bash +pytest -m db_dialect +``` + Run one file: ```bash diff --git a/src/omop_graph/graph/edges.py b/src/omop_graph/graph/edges.py index 96d97a9..e3c7e1a 100644 --- a/src/omop_graph/graph/edges.py +++ b/src/omop_graph/graph/edges.py @@ -6,13 +6,13 @@ It focuses on data definitions and classification logic, not graph traversal algorithms. -Supported Relationships ------------------------ -* **Mapping:** Semantic equivalence (e.g., source code to standard concept). -* **Versioning:** Lifecycle tracking (e.g., 'replaced by', 'is a'). -* **Ontological:** Hierarchical structure (e.g., 'is a', 'subsumes'). +Supported Relationships (``PredicateKind``) +-------------------------------------------- +* **Hierarchy:** Structural parent/child relationships (e.g., 'is a', 'subsumes'). +* **Identity:** Semantic equivalence (e.g., 'maps to', source code to standard concept). +* **Composition:** Part-whole/component relationships. +* **Association:** General, non-hierarchical relationships between concepts. * **Attribute:** Descriptive properties (e.g., 'has dose form'). -* **Metadata:** Administrative or low-semantic value connections. """ from __future__ import annotations @@ -54,9 +54,9 @@ class EdgeView: invalid_reason : str, optional The reason for invalidation (e.g., 'D' for deleted), if applicable. predicate_kind : PredicateKind - The high-level category of this predicate (e.g. Mapping, Versioning, Ontological, Attribute, Metadata). + The high-level category of this predicate (Hierarchy, Identity, Composition, Association, or Attribute). predicate_subkind : str - The more fine-grained subclass of this predicate (e.g. for Mapping: 'standard to non-standard', 'source to standard', etc.) + The more fine-grained subclass of this predicate (e.g. for Identity: 'standard to non-standard', 'source to standard', etc.) """ subject_id: int @@ -119,9 +119,9 @@ class Predicate: anc_down : bool Whether this relationship defines 'defines_ancestry' downwards (deprecated logic). predicate_kind : PredicateKind - The high-level category of this predicate (e.g. Mapping, Versioning, Ontological, Attribute, Metadata). + The high-level category of this predicate (Hierarchy, Identity, Composition, Association, or Attribute). predicate_subkind : str - The more fine-grained subclass of this predicate (e.g. for Mapping: 'standard to non-standard', 'source to standard', etc.) + The more fine-grained subclass of this predicate (e.g. for Identity: 'standard to non-standard', 'source to standard', etc.) """ relationship_id: str From a762440839ef23392823c38c2f08e32aff3e55c9 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 2 Oct 2026 02:51:32 +0000 Subject: [PATCH 17/20] Read-only KG, updated FTS queries, updated tests, remove unnecessary stuff that was not oaklib compliant --- docs/oaklib/interface.md | 29 +- pyproject.toml | 3 + scripts/benchmarks/benchmark.py | 24 +- scripts/benchmarks/enrich_gold_standard.py | 4 +- scripts/benchmarks/trace_example.py | 8 +- src/omop_graph/cli.py | 15 +- src/omop_graph/db/session.py | 89 ++--- src/omop_graph/graph/kg.py | 2 +- src/omop_graph/graph/queries.py | 51 +-- .../oaklib_interface/omop_factory.py | 74 ----- .../oaklib_interface/omop_implementation.py | 127 +++----- .../oaklib_interface/omop_resource.py | 91 +----- tests/fixtures/helpers.py | 90 +---- tests/fixtures/mock_cdm.py | 25 +- tests/test_concept_queries.py | 17 +- tests/test_fulltext_vocab_schema_postgres.py | 78 ++--- tests/test_oaklib_schema_awareness.py | 307 ++++-------------- tests/test_relationship_classification.py | 158 +++------ tests/test_schema_provenance_guard.py | 102 ++---- tests/test_vocab_split_connection.py | 237 +++++++------- 20 files changed, 452 insertions(+), 1079 deletions(-) delete mode 100644 src/omop_graph/oaklib_interface/omop_factory.py diff --git a/docs/oaklib/interface.md b/docs/oaklib/interface.md index 45e05d2..6d3efbe 100644 --- a/docs/oaklib/interface.md +++ b/docs/oaklib/interface.md @@ -29,31 +29,42 @@ The primary adapter class that inherits from multiple OAK interfaces: * **`TextAnnotatorInterface`**: Provides a pipeline to ground raw text spans to OMOP Concept IDs using the internal `KnowledgeGraph` resolvers. ### Resource Management -To initialize a connection, `omop-graph` uses a specialized resource factory: +The adapter resolves its database through oa-configurator; there is no URL-based construction. -* **`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. +* **`OMOPOntologyResource`**: The oaklib resource for the `omop` scheme. Its `slug` names a `[databases.*]` entry; no slug resolves `OmopGraphConfig.cdm_db`. +* **`resolved=`**: Pass an already-resolved `ResolvedCDMDatabase` to build the engines from it directly. A separate `vocab_connection` produces a second, vocabulary-only engine. +* **`kg=`**: Pass an existing `KnowledgeGraph` to use its engines as-is. --- ## Usage Examples ### Initializing the Adapter -You can get an OAK-compliant adapter by providing a SQLAlchemy connection string: +Select the adapter through oaklib with an oa-configurator database name: ```python -from omop_graph.oaklib_interface import OMOPAlchemyImplementation +from oaklib import get_adapter -# Initialize via connection string -adapter = OMOPAlchemyImplementation( - engine_string="postgresql://user:pass@localhost/omop" -) +# Named [databases.*] entry +adapter = get_adapter("omop:cdm_db") + +# Or OmopGraphConfig.cdm_db +adapter = get_adapter("omop:") # Use standard OAK methods label = adapter.label("OMOP:44819488") print(f"Label: {label}") ``` +Or construct it from an already-resolved database: + +```python +from omop_graph.db.session import resolve_cdm_database +from omop_graph.oaklib_interface.omop_implementation import OMOPAlchemyImplementation + +adapter = OMOPAlchemyImplementation(resolved=resolve_cdm_database("cdm_db")) +``` + ### Searching and Traversal Since the adapter implements the `SearchInterface`, you can perform standardized searches: diff --git a/pyproject.toml b/pyproject.toml index 3b6693d..7dab930 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,9 @@ omop-graph = "omop_graph.cli:app" [project.entry-points."omop.config"] omop_graph = "omop_graph.config:OmopGraphConfig" +[project.entry-points."oaklib.plugins"] +omop = "omop_graph.oaklib_interface.omop_implementation:OMOPAlchemyImplementation" + [build-system] requires = ["hatchling", "hatch-vcs"] build-backend = "hatchling.build" diff --git a/scripts/benchmarks/benchmark.py b/scripts/benchmarks/benchmark.py index 819dab2..fe37566 100644 --- a/scripts/benchmarks/benchmark.py +++ b/scripts/benchmarks/benchmark.py @@ -13,9 +13,6 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, cast, Annotated import typer -import sqlalchemy as sa -from sqlalchemy.orm import sessionmaker - import numpy as np from oa_configurator import Resolver, ResolvedModel, ResolvedProvider from omop_emb.config import ( @@ -45,7 +42,7 @@ PartialLabelResolver, PartialSynonymResolver, ) -from omop_graph.db.session import make_engine +from omop_graph.db.session import resolve_cdm_database app = typer.Typer() @@ -163,21 +160,6 @@ def load_cases(path: Path) -> List[BenchmarkCase]: raise TypeError(f"Unsupported benchmark case file shape: {type(payload).__name__}") -def build_session_factory() -> sessionmaker: - """Build a SQLAlchemy session factory via oa-configurator.""" - return sessionmaker(bind=make_engine(), future=True) - - -def build_engine() -> sa.Engine: - """Build a SQLAlchemy engine via oa-configurator.""" - return make_engine() - - -def build_knowledge_graph() -> KnowledgeGraph: - """Create a KnowledgeGraph backed by the live OMOP CDM database.""" - return KnowledgeGraph(cdm_engine=make_engine()) - - def build_embedding_knowledge_graph( embedding_metric: MetricType, resolved_model: ResolvedModel, @@ -185,7 +167,7 @@ def build_embedding_knowledge_graph( ) -> KnowledgeGraph: """Create a KnowledgeGraph with embedding support configured.""" - cdm_engine = make_engine() + cdm_engine = resolve_cdm_database().create_engine() config = KnowledgeGraphEmbeddingConfiguration( metric_type=embedding_metric, backend=backend, @@ -616,7 +598,7 @@ def run_benchmark( for case in cases } - kg = build_knowledge_graph() + kg = KnowledgeGraph(cdm_engine=resolve_cdm_database().create_engine()) configs = build_grounded_configs() errors: Dict[str, str] = {} diff --git a/scripts/benchmarks/enrich_gold_standard.py b/scripts/benchmarks/enrich_gold_standard.py index fd08df8..548666d 100644 --- a/scripts/benchmarks/enrich_gold_standard.py +++ b/scripts/benchmarks/enrich_gold_standard.py @@ -28,7 +28,7 @@ from sqlalchemy import text -from omop_graph.db.session import make_engine +from omop_graph.db.session import resolve_cdm_database DOMAIN_QUERY = text( """ @@ -148,7 +148,7 @@ def main() -> None: levels = [int(x) for x in args.levels.split(",")] payload = json.loads(args.cases_file.read_text()) - engine = make_engine() + engine = resolve_cdm_database().create_engine() total_stats = Stats.empty_for(levels) for bucket_name, cases in payload.items(): diff --git a/scripts/benchmarks/trace_example.py b/scripts/benchmarks/trace_example.py index 56b1568..f1f8387 100644 --- a/scripts/benchmarks/trace_example.py +++ b/scripts/benchmarks/trace_example.py @@ -32,7 +32,7 @@ ) from omop_graph.config import OmopGraphConfig -from omop_graph.db.session import make_engine +from omop_graph.db.session import resolve_cdm_database from omop_graph.extensions.emb import get_embedding_writer_interface from omop_graph.extensions.omop_alchemy import PredicateKind from omop_alchemy.cdm.query import ConceptFilter @@ -101,7 +101,7 @@ def _main( def _build_kg(metric_type=None, embedding_model: Optional[str] = None) -> KnowledgeGraph: """Build a KG with embedding support, resolved from OmopGraphConfig.""" - cdm_engine = make_engine() + cdm_engine = resolve_cdm_database().create_engine() try: from oa_configurator import Resolver from omop_emb.backends import resolve_backend_from_resolved_vector_store @@ -1719,7 +1719,7 @@ def panel_svg( else: selected = cases - kg = KnowledgeGraph(cdm_engine=make_engine()) + kg = KnowledgeGraph(cdm_engine=resolve_cdm_database().create_engine()) rel_class_cache: Dict[int, Dict] = {} out_dir = trace_dir_pl / "plots" / "panel" @@ -1782,7 +1782,7 @@ def graph_svg( else: selected = cases - kg = KnowledgeGraph(cdm_engine=make_engine()) + kg = KnowledgeGraph(cdm_engine=resolve_cdm_database().create_engine()) out_dir = trace_dir_pl / "plots" / "graph" out_dir.mkdir(parents=True, exist_ok=True) diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index ad7bf99..7650c8a 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -18,7 +18,7 @@ physical_schema_of ) -from orm_loader.backends import STAGING_SCHEMA, resolve_backend +from orm_loader.backends import STAGING_SCHEMA, resolve_backend, staging_schema_claim from orm_loader.helpers import bulk_load_context from orm_loader.helpers.metadata import Base from orm_loader.loaders.loader_interface import PandasLoader @@ -83,9 +83,10 @@ def relationship_classification( `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. + Bindable to run against, carrying ``orm_loader.staging_schema_claim()``. + 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 @@ -161,7 +162,7 @@ def relationship_classification( if engine is None: resolved = resolve_cdm_database() - engine = resolved.create_engine() + engine = resolved.create_engine(schema_claims=[staging_schema_claim()]) if resolved is not None and resolved.connection != resolved.vocab_connection: raise RuntimeError( f"relationship_classification() cannot run against database " @@ -206,9 +207,7 @@ def relationship_classification( # 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: - guard = guard_schema_provenance_for( - connection, resolved, schema_tag=Role.PRIMARY, tables=tables_to_drop # ty: ignore[invalid-argument-type] - ) + guard = guard_schema_provenance_for(connection, resolved, schema_tag=Role.PRIMARY) with guard: 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 diff --git a/src/omop_graph/db/session.py b/src/omop_graph/db/session.py index f53460b..fa4bda7 100644 --- a/src/omop_graph/db/session.py +++ b/src/omop_graph/db/session.py @@ -1,83 +1,36 @@ -"""SQLAlchemy engine helper for the OMOP CDM database.""" +"""Resolution helper for the OMOP CDM database.""" from __future__ import annotations -from typing import Optional, Union - -from sqlalchemy import create_engine, URL, Engine -from sqlalchemy.orm import sessionmaker, Session +from typing import Optional 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, - *, - engine_kwargs: Optional[dict] = None, - execution_options: Optional[dict] = None, -) -> Engine: - """Return a SQLAlchemy engine. - - When url is omitted, reads connection details from the active oa-configurator - stack config (schema translate map applied automatically). Pass url explicitly - to override. +def resolve_cdm_database(name: Optional[str] = None) -> ResolvedCDMDatabase: + """Resolve a CDM database from the active oa-configurator config. Parameters ---------- - url : URL or str, optional - SQLAlchemy database URL. If None, resolved from the active oa-configurator config. - engine_kwargs : dict, optional - 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, a ``schema_translate_map`` here may add keys the resolver - doesn't define, but may not include the resolver's own reserved - keys: those are always set from the resolved config, and - ``create_engine()`` raises ``ValueError`` if any are overridden here. + name : str, optional + Name of the ``[databases.*]`` entry. Defaults to ``OmopGraphConfig.cdm_db``. Returns ------- - Engine - A SQLAlchemy engine instance. - """ - engine_kwargs = engine_kwargs or {} - if url is None: - database = resolve_cdm_database() - return database.create_engine(execution_options=execution_options, **engine_kwargs) - - from sqlalchemy import make_url as _make_url - - if isinstance(url, str): - url = _make_url(url) - engine = create_engine(url, **engine_kwargs) - if execution_options: - engine = engine.execution_options(**execution_options) - return engine + ResolvedCDMDatabase - -def make_session( - url: str, - *, - echo: bool = False, -) -> Session: - engine = make_engine(url, engine_kwargs={"echo": echo}) - SessionLocal = sessionmaker(bind=engine) - return SessionLocal() + Raises + ------ + TypeError + If the entry is not a CDM database. + """ + resolver = Resolver.from_active_config() + if name is None: + name = resolver.resolve_package_config(OmopGraphConfig).cdm_db + resolved = resolver.resolve_database(name) + if not isinstance(resolved, ResolvedCDMDatabase): + raise TypeError( + f"Database {name!r} must resolve to a CDM database, got {type(resolved).__name__}" + ) + return resolved diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index b310bec..466d8ff 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -103,7 +103,7 @@ class KnowledgeGraphEmbeddingConfiguration: similarity computations are consistent. backend : omop_emb.EmbeddingBackend An already-constructed embedding backend, e.g. via - ``omop_emb.backends.resolve_backend_from_resolved``. + ``omop_emb.backends.resolve_backend_from_resolved_vector_store``. resolved_model : oa_configurator.ResolvedModel A model resolved via ``oa_configurator.Resolver.resolve_model()``, carrying real provider connection details. ``model_name``/``provider_type`` (see below) are diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 48e0bac..8468430 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -14,7 +14,7 @@ from __future__ import annotations -from typing import Optional, Tuple, Literal, Union +from typing import Optional, Tuple, Literal, Union, cast from datetime import date from sqlalchemy import ( @@ -22,23 +22,19 @@ case, exists, func, - inspect, literal, or_, select, Connection, Engine, - column, + Table, ) from sqlalchemy.orm import aliased from sqlalchemy.sql import Select -from oa_configurator import Role, physical_schema_of from omop_alchemy.backends import ( - CONCEPT_NAME_TSVECTOR_COLUMN, - CONCEPT_SYNONYM_NAME_TSVECTOR_COLUMN, - FullTextError, + resolve_backend, ) from omop_alchemy.cdm.model.vocabulary import ( Concept, @@ -332,12 +328,11 @@ def q_concept_name_fulltext( sort: bool = True, ) -> Select: """ - Query for concept names using PostgreSQL full-text search via optional + Query for concept names using PostgreSQL full-text search via the pre-computed tsvector columns and GIN indices. - This query only works when the stored tsvector columns have been installed - and registered in the ORM metadata via ``omop-maint fulltext install`` and - ``omop-maint fulltext populate``. If those columns are absent, this raises + Requires the stored tsvector columns from ``omop-alchemy fulltext install`` + and ``omop-alchemy fulltext populate``. If they are absent, this raises ``FullTextError`` instead of falling back to on-demand tsvector generation. @@ -345,42 +340,26 @@ def q_concept_name_fulltext( ---------- query_concept_name : str The concept name to search for. + engine : Engine or Connection + Bindable whose schema_translate_map locates the vocabulary schema. search_constraint : ConceptFilter, optional Additional filters (domain, vocab). synonym : bool, optional Whether to search in synonyms instead of concept names. + sort : bool, optional + Whether to order by match quality. + Raises + ------ + FullTextError + If the stored tsvector column is missing. """ name_expr = ( Concept_Synonym.concept_synonym_name if synonym else Concept.concept_name ) - # Fulltext are in VOCAB schema - inspector = inspect(engine) - vocab_schema = physical_schema_of(engine, schema_tag=Role.VOCAB) target_table = Concept_Synonym if synonym else Concept - target_col = ( - CONCEPT_SYNONYM_NAME_TSVECTOR_COLUMN - if synonym - else CONCEPT_NAME_TSVECTOR_COLUMN - ) stmt = q_concept_synonym() if synonym else q_concept_name() - - tsvector_col = next( - ( - c["name"] - for c in inspector.get_columns(target_table.__tablename__, schema=vocab_schema) - if c["name"] == target_col - ), - None, - ) - - if tsvector_col is None: - raise FullTextError( - f"Full-text search column '{target_col}' not found in table '{target_table.__tablename__}'. " - "Make sure to run 'omop-maint fulltext install' and 'omop-maint fulltext populate' to set up full-text search." - ) - - vector = column(tsvector_col) + vector = resolve_backend(engine).fulltext_vector_column(engine, cast(Table, target_table.__table__)) query = func.plainto_tsquery("english", query_concept_name) stmt = stmt.where(vector.op("@@")(query)) # Hits the GIN index instantly diff --git a/src/omop_graph/oaklib_interface/omop_factory.py b/src/omop_graph/oaklib_interface/omop_factory.py deleted file mode 100644 index 3a9438a..0000000 --- a/src/omop_graph/oaklib_interface/omop_factory.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Factory for creating OAK OMOP ontology resources.""" - -from __future__ import annotations - -import logging -from typing import Optional, Union - -from sqlalchemy.engine import URL - -from .omop_resource import OMOPOntologyResource -from oa_configurator import SCHEMA_TRANSLATE_MAP_KEY, ResolvedCDMDatabase, Resolver -from omop_graph.config import OmopGraphConfig - -logger = logging.getLogger(__name__) - - -def omop_resource( - *, - url: Optional[Union[str, URL]] = None, - slug: Optional[str] = "omop", -) -> OMOPOntologyResource: - """Create an OMOP DatabaseOntologyResource. - - When url is omitted, reads connection details from the active oa-configurator - stack config. Pass url explicitly to override. - - Parameters - ---------- - url : str | URL, optional - Explicit database connection URL. When omitted the active oa-configurator - config is used. An explicit url bypasses oa-configurator entirely, so no - schema_translate_map is applied. Logs a warning when given for this reason. - slug : str, optional - Slug identifier for the resource. Defaults to 'omop'. - - Returns - ------- - OMOPOntologyResource - """ - execution_options = None - vocab_url = None - vocab_execution_options = None - resolved: ResolvedCDMDatabase | None = 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__}" - ) - resolved = database - url = resolved.connection.url - execution_options = {SCHEMA_TRANSLATE_MAP_KEY: resolved.schema_translate_map()} - if resolved.connection != resolved.vocab_connection: - vocab_url = resolved.vocab_connection.url - vocab_execution_options = execution_options - else: - logger.warning( - "omop_resource() was given an explicit url, bypassing oa-configurator " - "entirely: no schema_translate_map is applied, so a configured " - "non-default primary/vocab/results schema is silently not respected. " - "Omit url= to resolve the active oa-configurator config instead." - ) - - return OMOPOntologyResource( - slug=slug, - url=url, - execution_options=execution_options, - vocab_url=vocab_url, - vocab_execution_options=vocab_execution_options, - resolved=resolved, - ) diff --git a/src/omop_graph/oaklib_interface/omop_implementation.py b/src/omop_graph/oaklib_interface/omop_implementation.py index e744de2..e17ffa3 100644 --- a/src/omop_graph/oaklib_interface/omop_implementation.py +++ b/src/omop_graph/oaklib_interface/omop_implementation.py @@ -1,6 +1,6 @@ import logging from collections import defaultdict -from typing import Dict, Iterable, Iterator, List, Optional, Tuple +from typing import ClassVar, Dict, Iterable, Iterator, List, Optional, Tuple import numpy as np from omop_alchemy.cdm.query import ConceptFilter @@ -36,12 +36,10 @@ from omop_graph.reasoning.resolvers.resolver_pipeline import ResolverPipeline from omop_graph.render import bind_default_renderers from omop_graph.oaklib_interface.omop_resource import OMOPOntologyResource -from omop_graph.oaklib_interface.omop_factory import omop_resource +from omop_graph.db.session import resolve_cdm_database - -from sqlalchemy.engine import URL - -from omop_graph.db.session import make_engine +from oa_configurator import ResolvedCDMDatabase, SchemaClaim +from oaklib.resource import OntologyResource logger = logging.getLogger(__name__) @@ -841,97 +839,72 @@ class OMOPAlchemyImplementation( # type: ignore[override] A :class:`OntologyInterface` implementation wrapping a SQL Relational Database conforming to the OMOP CDM. - To connect, either use OMOPAlchemyImplementation directly: + Select it through oaklib with an oa-configurator database name, or omit + the name to use ``OmopGraphConfig.cdm_db``: - >>> from omop_graph.oaklib_interface import OMOPAlchemyImplementation - >>> from omop_graph.oaklib_interface.omop_factory import omop_resource - >>> resource = omop_resource(url='postgresql+psycopg2://uid:pid@host:5432/dbname') - >>> adapter = OMOPAlchemyImplementation(resource=resource) + >>> from oaklib import get_adapter + >>> adapter = get_adapter("omop:cdm_db") - or pass a connection string directly: + or construct it directly from an already-resolved database: - >>> adapter = OMOPAlchemyImplementation(engine_string="sqlite:////path/to/omop.db") + >>> from omop_graph.db.session import resolve_cdm_database + >>> adapter = OMOPAlchemyImplementation(resolved=resolve_cdm_database()) Parameters ---------- - engine_string : str | URL | None, optional - 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. Ignored when ``kg`` is given directly. To use the - oa-configurator-configured default, resolve it explicitly via - ``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. + resource : OntologyResource | None, optional + The oaklib resource. Its ``slug`` names the ``[databases.*]`` entry to + resolve; ``None`` resolves ``OmopGraphConfig.cdm_db``. Ignored when + ``resolved`` or ``kg`` is given. + resolved : ResolvedCDMDatabase | None, optional + An already-resolved CDM database to build the engines from, without + registering schema claims. Ignored when ``kg`` is given. + schema_claims : Iterable[SchemaClaim], optional + Forwarded to ``ResolvedCDMDatabase.create_engines()``. + execution_options : dict | None, optional + Non-schema execution options forwarded to ``ResolvedCDMDatabase.create_engines()``. 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 - (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. + An existing Knowledge Graph, used as-is. The caller owns whatever + engine it wraps. 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 ``kg`` is not given and neither ``engine_string`` nor ``resource`` is. + Ignored when ``kg`` is given. """ + # oaklib's class_resolver keys plugins by class name; this adds the "omop:" scheme. + synonyms: ClassVar[tuple[str, ...]] = ("omop",) + def __init__( self, - engine_string: str | URL | None = None, - resource: OMOPOntologyResource | None = None, + resource: OntologyResource | None = None, + *, + resolved: ResolvedCDMDatabase | None = None, + schema_claims: Iterable[SchemaClaim] = (), + execution_options: dict | None = None, kg: KnowledgeGraph | None = None, kg_emb_config: Optional[KnowledgeGraphEmbeddingConfiguration] = None, **kwargs, ): - self._connection = None - + if resource is None: + resource = OMOPOntologyResource() 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())." - ) - - if self.resource.resolved is not None: - engine, vocab_engine = self.resource.resolved.create_engines( - echo=False, future=True - ) - if vocab_engine is engine: - vocab_engine = None - else: - engine = make_engine( - self.engine_string, - engine_kwargs={"echo": False, "future": True}, - execution_options=self.resource.execution_options, - ) - 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) + if resolved is None: + resolved = resolve_cdm_database(resource.slug) + engine, vocab_engine = resolved.create_engines( + schema_claims=schema_claims, + execution_options=execution_options, + register_claims=False, + echo=False, + future=True, + ) + kg = KnowledgeGraph( + emb_config=kg_emb_config, + cdm_engine=engine, + vocab_engine=None if vocab_engine is engine else vocab_engine, + ) bind_default_renderers(kg) - super().__init__(kg=kg, **kwargs) + super().__init__(resource=resource, kg=kg, **kwargs) # TODO: Implement if necessary! def _all_relationships(self): diff --git a/src/omop_graph/oaklib_interface/omop_resource.py b/src/omop_graph/oaklib_interface/omop_resource.py index ec5876e..e1b1238 100644 --- a/src/omop_graph/oaklib_interface/omop_resource.py +++ b/src/omop_graph/oaklib_interface/omop_resource.py @@ -1,102 +1,40 @@ from dataclasses import dataclass -from typing import Optional, Union +from typing import Optional from oaklib.resource import OntologyResource -from sqlalchemy.engine import URL, make_url - -from oa_configurator import ResolvedCDMDatabase @dataclass class OMOPOntologyResource(OntologyResource): """ - Ontology resource backed by a live SQLAlchemy database. + Ontology resource naming an oa-configurator CDM database. - This class extends the `OntologyResource` to support database URLs specifically - for OMOP backends using SQLAlchemy. + Selected by oaklib as ``get_adapter("omop:")``. Parameters ---------- - url : str | URL, optional - The database connection URL. slug : str, optional - A unique slug/identifier for this resource. + Name of the ``[databases.*]`` entry to resolve. ``None`` resolves + ``OmopGraphConfig.cdm_db``. scheme : str, optional - The scheme identifier. Defaults to 'omop_alchemy'. - local : bool, optional - Whether the resource is local. Defaults to False. - in_memory : bool, optional - Whether the resource is in-memory. Defaults to False. + The oaklib scheme. Defaults to 'omop'. 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. - 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. - resolved : ResolvedCDMDatabase, optional - When present, ``OMOPAlchemyImplementation`` builds its engine(s) via - ``resolved.create_engines()`` directly instead of reconstructing one - from ``url``/``execution_options``. """ - url: Optional[Union[str, URL]] = None # type: ignore[assignment] slug: Optional[str] = None # type: ignore[assignment] - scheme: str = "omop_alchemy" # type: ignore[assignment] - local: bool = False # type: ignore[assignment] - in_memory: bool = False # type: ignore[assignment] + scheme: str = "omop" # 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] - resolved: Optional[ResolvedCDMDatabase] = None # type: ignore[assignment] - - def _parsed_url(self) -> Optional[URL]: - """ - Parse the connection URL into a SQLAlchemy URL object. - - Returns - ------- - URL | None - The parsed URL object, or None if no URL is set. - """ - if not self.url: - return None - return make_url(self.url) if isinstance(self.url, str) else self.url - - @property - def display_slug(self) -> Optional[str]: - """ - Get a safe, redacted identifier for logs / UI. - - Returns - ------- - str | None - The string representation of the parsed URL. - """ - u = self._parsed_url() - return str(u) if u else None def valid(self) -> bool: """ - Check if the database ontology resource is valid. + Always True: a missing slug resolves the configured default database. Returns ------- bool - True if a URL is present, False otherwise. """ - return bool(self.url) + return True @property def local_path(self) -> None: @@ -106,15 +44,6 @@ def local_path(self) -> None: Returns ------- None - Always returns None as database-backed resources have no filesystem path. + Always None as database-backed resources have no filesystem path. """ return None - - def __repr__(self) -> str: - parts = [] - if self.slug: - parts.append(f"slug={self.slug!r}") - if self.url: - parts.append(f"url={str(make_url(self.url))!r}") - parts.append(f"scheme={self.scheme!r}") - return f"{type(self).__name__}({', '.join(parts)})" diff --git a/tests/fixtures/helpers.py b/tests/fixtures/helpers.py index ca5514a..7de6513 100644 --- a/tests/fixtures/helpers.py +++ b/tests/fixtures/helpers.py @@ -1,31 +1,14 @@ -"""Shared test-helper primitives for split/non-split, vocab-schema fixtures. - -Extracted from four independently-duplicated copies across the test suite -(Phase 4.5 of the schema_translate_map fix): a vocab-table tuple, a -schema_translate_map dict builder, and an FK-trigger toggle context manager. -Schema *provisioning* strategy (rollback-based CREATE SCHEMA on an open -connection vs. a genuinely committed isolated_test_schema) stays local to -each fixture/test instead of being forced into one shared helper here -- -that choice genuinely differs per test's needs, unlike the three primitives -below. -""" +"""Shared test-helper primitives for vocab-schema fixtures.""" from __future__ import annotations -from contextlib import contextmanager -from typing import Iterator, cast +from typing import cast import sqlalchemy as sa -from oa_configurator import ( - Role, - qualified, - physical_schema_of, - validate_schema_tag -) from omop_alchemy.cdm.model.vocabulary import ( - Concept, - Concept_Class, - Domain, + Concept, + Concept_Class, + Domain, Vocabulary ) @@ -38,66 +21,3 @@ needing more tables (e.g. Relationship, Concept_Relationship) extends this tuple rather than redefining its own copy of the shared core. """ - - -def schema_translate_map( - primary_schema: str, *, vocab_schema: str | None = None, results_schema: str | None = None -) -> dict[str, str]: - """Build a 3-key schema_translate_map dict, defaulting vocab/results to - primary_schema -- the common "route everything through one schema" shape, - with an explicit override for the genuinely-split-vocab case. - """ - return { - Role.PRIMARY.value: primary_schema, - Role.VOCAB.value: vocab_schema if vocab_schema is not None else primary_schema, - Role.RESULTS.value: results_schema if results_schema is not None else primary_schema, - } - - -@contextmanager -def fk_triggers_disabled( - bindable: "sa.Engine | sa.Connection", tables: "tuple[sa.Table, ...]" -) -> Iterator[None]: - """Disable triggers on *tables* for the duration of the block, matching - production bulk-loads' own FK-toggle pattern. - - Several vocab tables form a genuine FK insert cycle (each reference - row's own *_concept_id FK needs a Concept row that doesn't exist yet, - and Concept's own domain_id/vocabulary_id/concept_class_id FKs need - those reference rows already present), so a plain ordered insert can't - satisfy every constraint. A no-op on any non-Postgres dialect: SQLite - doesn't enforce FK constraints by default, so there's nothing to - disable there, and callers that seed both dialects can call this - unconditionally without their own dialect check. Accepts an Engine or - an open Connection; an Engine has no ``.execute()`` of its own, so one - short-lived connection is opened for each toggle. - """ - if bindable.dialect.name != "postgresql": - yield - return - - opened_here = isinstance(bindable, sa.Engine) - conn = bindable.connect() if opened_here else bindable - try: - for table in tables: - physical_schema = physical_schema_of(conn, schema_tag=validate_schema_tag(table)) - conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name, physical_schema=physical_schema)} DISABLE TRIGGER ALL")) - if opened_here: - conn.commit() - finally: - if opened_here: - conn.close() - - try: - yield - finally: - conn = bindable.connect() if opened_here else bindable - try: - for table in tables: - physical_schema = physical_schema_of(conn, schema_tag=validate_schema_tag(table)) - conn.execute(sa.text(f"ALTER TABLE {qualified(conn, table.name, physical_schema=physical_schema)} ENABLE TRIGGER ALL")) - if opened_here: - conn.commit() - finally: - if opened_here: - conn.close() diff --git a/tests/fixtures/mock_cdm.py b/tests/fixtures/mock_cdm.py index d0f4a2d..b743c82 100644 --- a/tests/fixtures/mock_cdm.py +++ b/tests/fixtures/mock_cdm.py @@ -7,8 +7,8 @@ import sqlalchemy as sa from sqlalchemy.orm import Session, sessionmaker -from oa_configurator.testing import isolated_test_database, isolated_test_schema -from orm_loader.helpers import Base +from oa_configurator.testing import isolated_test_database, scoped_test_schema +from orm_loader.helpers import Base, bulk_load_context from omop_alchemy.cdm.model.vocabulary.concept import Concept from omop_alchemy.cdm.model.vocabulary.concept_ancestor import Concept_Ancestor from omop_alchemy.cdm.model.vocabulary.concept_class import Concept_Class @@ -26,8 +26,6 @@ ) from omop_graph.graph.kg import KnowledgeGraph -from .helpers import fk_triggers_disabled, schema_translate_map - PARENT_CANCER_ID = 443392 CONCEPT_META_ID = 0 LANGUAGE_CONCEPT_ID = 1 @@ -36,11 +34,9 @@ @pytest.fixture(scope="module") def mock_cdm_engine() -> Iterator[sa.Engine]: with isolated_test_database(OmopGraphConfig, "test_cdm_db_pg") as db: - raw_engine = db.connection.engine - with isolated_test_schema(raw_engine, prefix="mock_cdm") as schema: - engine = raw_engine.execution_options(schema_translate_map=schema_translate_map(schema)) - _create_mock_cdm_tables(engine) - yield engine + with scoped_test_schema(db.resolved, prefix="mock_cdm") as scoped: + _create_mock_cdm_tables(scoped.engine) + yield scoped.engine def _create_mock_cdm_tables(engine: sa.Engine) -> None: @@ -62,12 +58,13 @@ def _create_mock_cdm_tables(engine: sa.Engine) -> None: Base.metadata.create_all(engine, tables=tables) - # Disable triggers for the whole seed sidesteps ordering entirely, + # Disabling FK checks for the whole seed sidesteps ordering entirely, # since these tables have FK dependencies that form a cycle - with fk_triggers_disabled(engine, tuple(tables)): - session_local = sessionmaker(bind=engine, future=True) - with session_local() as session: + session_local = sessionmaker(bind=engine, future=True) + with session_local() as session: + with bulk_load_context(session): seed_mock_cdm(session) + session.commit() @pytest.fixture() @@ -366,4 +363,4 @@ def seed_mock_cdm(session: Session) -> None: ) ) - session.commit() + session.flush() diff --git a/tests/test_concept_queries.py b/tests/test_concept_queries.py index bcf5d58..cada934 100644 --- a/tests/test_concept_queries.py +++ b/tests/test_concept_queries.py @@ -9,8 +9,8 @@ import sqlalchemy as sa from sqlalchemy.orm import Session -from oa_configurator.testing import isolated_test_schema -from orm_loader.helpers import Base +from oa_configurator.testing import scoped_test_schema +from orm_loader.helpers import Base, bulk_load_context from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary from omop_alchemy.cdm.query import ConceptFilter @@ -23,15 +23,15 @@ q_entities, ) -from fixtures.helpers import VOCAB_TABLES, fk_triggers_disabled, schema_translate_map +from fixtures.helpers import VOCAB_TABLES _META_CONCEPT_ID = 0 @pytest.fixture() def concept_engine(pg_db) -> Iterator[sa.Engine]: - with isolated_test_schema(pg_db.connection.engine, prefix="concept_queries") as schema: - engine = pg_db.connection.engine.execution_options(schema_translate_map=schema_translate_map(schema)) + with scoped_test_schema(pg_db.resolved, prefix="concept_queries") as scoped: + engine = scoped.engine Base.metadata.create_all(engine, tables=VOCAB_TABLES, checkfirst=True) valid_from = date(2000, 1, 1) @@ -56,8 +56,8 @@ def concept( invalid_reason=invalid_reason, ) - with fk_triggers_disabled(engine, VOCAB_TABLES): - with Session(engine) as session: + with Session(engine) as session: + with bulk_load_context(session): session.add_all( [ Concept( @@ -104,7 +104,8 @@ def concept( concept(5, standard_concept=" ", invalid_reason="X"), ] ) - session.commit() + session.flush() + session.commit() yield engine diff --git a/tests/test_fulltext_vocab_schema_postgres.py b/tests/test_fulltext_vocab_schema_postgres.py index 984250d..b9f0065 100644 --- a/tests/test_fulltext_vocab_schema_postgres.py +++ b/tests/test_fulltext_vocab_schema_postgres.py @@ -4,72 +4,54 @@ schema= defaults to Role.PRIMARY, so a real primary/vocab split (unlike SQLite, where every schema tag folds to None) would raise a false FullTextError even though the tsvector column exists. - -Schemas are created directly on pg_db's own connection rather than via -isolated_test_schema(), which opens a second connection that can starve -against pg_db's still-open transaction. """ from __future__ import annotations -import uuid - import pytest import sqlalchemy as sa from oa_configurator import Role +from oa_configurator.testing import scoped_test_schema from omop_alchemy.backends import CONCEPT_NAME_TSVECTOR_COLUMN, FullTextError, resolve_backend from orm_loader.helpers import Base from omop_graph.graph.queries import q_concept_name_fulltext -from fixtures.helpers import VOCAB_TABLES, schema_translate_map +from fixtures.helpers import VOCAB_TABLES pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] -def _scoped(pg_db, *, primary_schema: str, vocab_schema: str) -> sa.Connection: - conn = pg_db.connection - conn.execute(sa.text(f"CREATE SCHEMA {primary_schema}")) - conn.execute(sa.text(f"CREATE SCHEMA {vocab_schema}")) - return conn.execution_options( - schema_translate_map=schema_translate_map(primary_schema, vocab_schema=vocab_schema) - ) - - def test_fulltext_query_finds_the_tsvector_column_in_the_vocab_schema(pg_db): - primary_schema = f"fulltext_primary_{uuid.uuid4().hex[:8]}" - vocab_schema = f"fulltext_vocab_{uuid.uuid4().hex[:8]}" - scoped = _scoped(pg_db, primary_schema=primary_schema, vocab_schema=vocab_schema) - Base.metadata.create_all(bind=scoped, tables=VOCAB_TABLES, checkfirst=True) - - backend = resolve_backend(scoped) - backend.install_fulltext_on_table( - scoped, - table_name="concept", - vector_column_name=CONCEPT_NAME_TSVECTOR_COLUMN, - index_name="idx_concept_name_tsvector_test", - create_indexes=True, - fastupdate=True, - schema_tag=Role.VOCAB.value, - ) - - # Confirmed absent from primary_schema: proves a real vocab/primary - # split, not a lucky same-schema coincidence. - assert not sa.inspect(pg_db.connection).has_table("concept", schema=primary_schema) - columns = { - c["name"] for c in sa.inspect(pg_db.connection).get_columns("concept", schema=vocab_schema) - } - assert CONCEPT_NAME_TSVECTOR_COLUMN in columns - - stmt = q_concept_name_fulltext("kidney cancer", engine=scoped) - assert stmt is not None + with scoped_test_schema(pg_db.resolved, prefix="fulltext", split_roles=[Role.VOCAB]) as scoped: + engine = scoped.engine + Base.metadata.create_all(bind=engine, tables=VOCAB_TABLES, checkfirst=True) + + with engine.begin() as conn: + resolve_backend(conn).install_fulltext_on_table( + conn, + table_name="concept", + vector_column_name=CONCEPT_NAME_TSVECTOR_COLUMN, + index_name="idx_concept_name_tsvector_test", + create_indexes=True, + fastupdate=True, + schema_tag=Role.VOCAB.value, + ) + + # Confirmed absent from the primary schema: proves a real vocab/primary + # split, not a lucky same-schema coincidence. + inspector = sa.inspect(engine) + assert not inspector.has_table("concept", schema=scoped.schemas[Role.PRIMARY]) + columns = {c["name"] for c in inspector.get_columns("concept", schema=scoped.schemas[Role.VOCAB])} + assert CONCEPT_NAME_TSVECTOR_COLUMN in columns + + stmt = q_concept_name_fulltext("kidney cancer", engine=engine) + assert stmt is not None def test_fulltext_query_still_raises_when_the_column_is_genuinely_absent(pg_db): - primary_schema = f"fulltext_primary_{uuid.uuid4().hex[:8]}" - vocab_schema = f"fulltext_vocab_{uuid.uuid4().hex[:8]}" - scoped = _scoped(pg_db, primary_schema=primary_schema, vocab_schema=vocab_schema) - Base.metadata.create_all(bind=scoped, tables=VOCAB_TABLES, checkfirst=True) + with scoped_test_schema(pg_db.resolved, prefix="fulltext", split_roles=[Role.VOCAB]) as scoped: + Base.metadata.create_all(bind=scoped.engine, tables=VOCAB_TABLES, checkfirst=True) - with pytest.raises(FullTextError): - q_concept_name_fulltext("kidney cancer", engine=scoped) + with pytest.raises(FullTextError): + q_concept_name_fulltext("kidney cancer", engine=scoped.engine) diff --git a/tests/test_oaklib_schema_awareness.py b/tests/test_oaklib_schema_awareness.py index 1a7f1e5..5956623 100644 --- a/tests/test_oaklib_schema_awareness.py +++ b/tests/test_oaklib_schema_awareness.py @@ -1,50 +1,34 @@ -"""OAK-lib adapter schema-awareness gap (Phase 4). +"""OAK-lib adapter construction. -`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: +``OMOPAlchemyImplementation`` builds its engines only through +``ResolvedCDMDatabase.create_engines()``. Three construction paths: -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. +1. ``get_adapter("omop:")``: oaklib passes a resource whose slug names + the oa-configurator database entry. +2. ``resolved=``: an already-resolved CDM database. +3. ``kg=``: an existing ``KnowledgeGraph``, used as-is. """ from __future__ import annotations from datetime import date +import pytest import sqlalchemy as sa import sqlalchemy.orm +from oaklib import get_adapter -from oa_configurator import ( - SCHEMA_TRANSLATE_MAP_KEY, - ResolvedCDMDatabase, - ResolvedConnection, - Role, -) -from oa_configurator.testing import isolated_test_schema +from oa_configurator import Resolver +from oa_configurator.testing import scoped_test_schema from omop_alchemy.cdm.model.vocabulary import Concept, Concept_Class, Domain, Vocabulary -from orm_loader.helpers import Base +from orm_loader.backends import staging_schema_claim +from orm_loader.helpers import Base, bulk_load_context from omop_graph.cli import relationship_classification -from omop_graph.db.session import make_engine +from omop_graph.db.session import resolve_cdm_database from omop_graph.graph.kg import KnowledgeGraph -from omop_graph.oaklib_interface.omop_factory import omop_resource +from omop_graph.oaklib_interface import omop_implementation from omop_graph.oaklib_interface.omop_implementation import OMOPAlchemyImplementation -from omop_graph.oaklib_interface.omop_resource import OMOPOntologyResource -from omop_graph.config import OmopGraphConfig - -from fixtures.helpers import VOCAB_TABLES, fk_triggers_disabled, schema_translate_map _META_CONCEPT_ID = 0 _CONCEPT_ID = 1001 @@ -52,14 +36,12 @@ _FAR_FUTURE = date(2099, 12, 31) -def _seed_one_concept(bindable: sa.Engine | sa.Connection, *, concept_id: int, name: str) -> None: +def _seed_one_concept(engine: sa.Engine, *, concept_id: int, name: str) -> None: """Minimal vocab bootstrap: Domain/Vocabulary/Concept_Class/Concept form an - FK insert cycle, so triggers are disabled around the insert then - re-enabled (a no-op on SQLite, which doesn't enforce FK constraints by - default). + FK insert cycle, so FK checks are disabled around the insert. """ - with fk_triggers_disabled(bindable, VOCAB_TABLES): - with sa.orm.Session(bindable) as session: + with sa.orm.Session(engine) as session: + with bulk_load_context(session): session.add_all( [ Concept( @@ -99,235 +81,74 @@ def _seed_one_concept(bindable: sa.Engine | sa.Connection, *, concept_id: int, n ), ] ) - session.commit() - - -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() + session.flush() + session.commit() - 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_KEY] - == resource.execution_options[SCHEMA_TRANSLATE_MAP_KEY] - ) +def _populate(scoped, *, name: str, guarded: bool = True) -> None: + Base.metadata.create_all(bind=scoped.engine, checkfirst=True) + _seed_one_concept(scoped.engine, concept_id=_CONCEPT_ID, name=name) + relationship_classification(engine=scoped.engine, resolved=scoped.resolved if guarded else 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, tmp_path +@pytest.mark.parametrize(("descriptor", "expected_name"), [("omop:named_db", "named_db"), ("omop:", None)]) +def test_get_adapter_resolves_the_slug_as_the_database_name( + pg_db, monkeypatch, descriptor, expected_name ) -> None: - """Confirms OMOPAlchemyImplementation actually builds vocab_engine from - omop_resource()'s vocab_url/vocab_execution_options, not just carries them. - Needs a real seeded database, not ``:memory:``, since a SQLite tempfile - persists across the adapter's own connections and ``:memory:`` doesn't.""" - vocab_db_path = tmp_path / "vocab.db" - vocab_url = f"sqlite:///{vocab_db_path}" - vocab_engine = sa.create_engine(vocab_url).execution_options( - schema_translate_map={Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} - ) - Base.metadata.create_all(bind=vocab_engine, checkfirst=True) - _seed_one_concept(vocab_engine, concept_id=_CONCEPT_ID, name="Split-wiring concept") - vocab_engine.dispose() + with scoped_test_schema(pg_db.resolved, prefix="oaklib_get_adapter", schema_claims=[staging_schema_claim()]) as scoped: + _populate(scoped, name="Adapter concept") - with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_split") as schema: - engine = pg_db.connection.engine.execution_options( - schema_translate_map=schema_translate_map(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) + def fake_resolve_cdm_database(name): + assert name == expected_name + return scoped.resolved - resource = OMOPOntologyResource( - url=pg_db.connection.engine.url.render_as_string(hide_password=False), - execution_options={SCHEMA_TRANSLATE_MAP_KEY: schema_translate_map(schema)}, - vocab_url=vocab_url, - vocab_execution_options={ - SCHEMA_TRANSLATE_MAP_KEY: {Role.PRIMARY.value: None, Role.VOCAB.value: None, Role.RESULTS.value: None} - }, - ) + monkeypatch.setattr(omop_implementation, "resolve_cdm_database", fake_resolve_cdm_database) + adapter = get_adapter(descriptor) - adapter = OMOPAlchemyImplementation(resource=resource) + assert isinstance(adapter, OMOPAlchemyImplementation) + assert adapter.resource.scheme == "omop" + assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Adapter concept" - assert adapter.kg.vocab_engine is not adapter.kg.cdm_engine - assert str(adapter.kg.vocab_engine.url) == vocab_url - assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Split-wiring concept" +def test_resolve_cdm_database_resolves_a_named_entry(pg_db) -> None: + assert resolve_cdm_database(pg_db.resolved.name).name == pg_db.resolved.name -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=schema_translate_map(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_KEY: schema_translate_map(schema) - }, - ) +def test_resolved_path_reuses_one_engine_without_a_vocab_split(pg_db) -> None: + with scoped_test_schema(pg_db.resolved, prefix="oaklib_resolved", schema_claims=[staging_schema_claim()]) as scoped: + _populate(scoped, name="Resolved-path concept") - adapter = OMOPAlchemyImplementation(resource=resource) + adapter = OMOPAlchemyImplementation(resolved=scoped.resolved) assert adapter.kg.vocab_engine is adapter.kg.cdm_engine - - -def test_omop_alchemy_implementation_builds_its_engine_via_resolved_create_engines( - pg_db, monkeypatch -) -> None: - with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_resolved") as schema: - engine = pg_db.connection.engine.execution_options( - schema_translate_map=schema_translate_map(schema) - ) - Base.metadata.create_all(bind=engine, checkfirst=True) - _seed_one_concept(engine, concept_id=_CONCEPT_ID, name="Resolved-path concept") - relationship_classification(engine=engine) - - url = pg_db.connection.engine.url - connection = ResolvedConnection( - name="resolved_path", - url=url.render_as_string(hide_password=False), - safe_url=url.render_as_string(hide_password=True), - _engine_url=url, - ) - resolved = ResolvedCDMDatabase( - name="resolved_path", - connection=connection, - schema_name=schema, - vocab_connection=connection, - vocab_schema=schema, - results_schema=schema, - ) - - class FakeResolver: - def resolve_package_config(self, config_type): - assert config_type is OmopGraphConfig - return OmopGraphConfig(cdm_db="resolved_path") - - def resolve_database(self, name): - assert name == "resolved_path" - return resolved - - monkeypatch.setattr( - "omop_graph.oaklib_interface.omop_factory.Resolver.from_active_config", - lambda: FakeResolver(), - ) - - resource = omop_resource() - assert resource.resolved is resolved - - adapter = OMOPAlchemyImplementation(resource=resource) - assert adapter.kg.cdm_engine.pool._pre_ping is True - assert adapter.kg.vocab_engine is adapter.kg.cdm_engine assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Resolved-path concept" -def test_kg_injection_path_resolves_against_the_configured_schema(pg_db) -> None: - """Path 1 from this module's docstring: kg= injection. OMOPAlchemyImplementation - must use the injected KnowledgeGraph as-is rather than rebuilding its own - engine from engine_string -- proved here by pairing a real, schema-aware - kg with a deliberately broken engine_string ("sqlite:///:memory:", never - actually queried). If the implementation ever fell back to building its - own engine instead of using kg=, this would error or return nothing - instead of the seeded concept's name. - """ - schema = "phase4_oaklib_kg_injection" - conn = pg_db.connection - conn.execute(sa.text(f"CREATE SCHEMA {schema}")) - scoped = conn.execution_options( - schema_translate_map=schema_translate_map(schema) +def test_resolved_path_builds_a_genuine_vocab_engine_for_a_split_connection(pg_db) -> None: + name = pg_db.resolved.name + resolver = Resolver.from_active_config() + resolver = resolver.with_overrides( + connections={"oaklib_split_vocab": resolver.config.connections[pg_db.resolved.connection.name]}, + databases={name: resolver.config.databases[name].model_copy(update={"vocab_connection": "oaklib_split_vocab"})}, ) - Base.metadata.create_all(bind=scoped, checkfirst=True) - _seed_one_concept(scoped, concept_id=_CONCEPT_ID, name="Test concept") - relationship_classification(engine=scoped) + with scoped_test_schema(resolver.resolve_database(name), prefix="oaklib_split", resolver=resolver, schema_claims=[staging_schema_claim()]) as scoped: + # Both connections share one physical server, so the cross-schema FK + # works; resolved= would make relationship_classification() refuse the split. + _populate(scoped, name="Split-wiring concept", guarded=False) - kg = KnowledgeGraph(cdm_engine=scoped) - adapter = OMOPAlchemyImplementation(engine_string="sqlite:///:memory:", kg=kg) + adapter = OMOPAlchemyImplementation(resolved=scoped.resolved) - assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Test concept" + assert adapter.kg.vocab_engine is not adapter.kg.cdm_engine + assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Split-wiring concept" -def test_bare_engine_string_path_resolves_against_the_configured_schema(pg_db) -> None: - """Path 2 from this module's docstring: bare engine_string=/resource=-only - construction. OAK-lib's own generic materialize() invocation only ever - supplies a URL string, never a live connection or kg= -- this is the one - path the original bug actually broke, since schema_translate_map has to - be carried through purely via execution_options, with no resolved= - object to lean on. - """ - with isolated_test_schema(pg_db.connection.engine, prefix="phase4_oaklib_bare") as schema: - engine = pg_db.connection.engine.execution_options( - schema_translate_map=schema_translate_map(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) +def test_kg_injection_path_uses_the_injected_knowledge_graph(pg_db) -> None: + with scoped_test_schema(pg_db.resolved, prefix="oaklib_kg", schema_claims=[staging_schema_claim()]) as scoped: + _populate(scoped, name="Injected concept") + kg = KnowledgeGraph(cdm_engine=scoped.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_KEY: schema_translate_map(schema) - }, - ) - adapter = OMOPAlchemyImplementation(resource=resource) + adapter = OMOPAlchemyImplementation(kg=kg) - assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Bare-string concept" + assert adapter.kg is kg + assert adapter.label(f"OMOP:{_CONCEPT_ID}") == "Injected concept" diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index 537886c..d56394b 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -1,67 +1,51 @@ """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. +Each test runs against a committed scoped schema from ``scoped_test_schema()``, +dropped on exit. Also covers the DROP TYPE naming-mismatch fix: the enum +column never set an explicit ``name=``, so the real generated type is +``predicatekind``, not ``predicatekindenum``. """ from __future__ import annotations -import dataclasses -import uuid - import pytest import sqlalchemy as sa -from oa_configurator import Role, SchemaDriftError, record_schema_provenance +from oa_configurator import ConnectionConfig, Resolver, Role +from oa_configurator.testing import scoped_test_schema +from orm_loader.backends import staging_schema_claim from orm_loader.helpers import Base from omop_graph.cli import relationship_classification from omop_graph.extensions.omop_alchemy import RelationshipClass, RelationshipMapping -from fixtures.helpers import schema_translate_map - - -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=schema_translate_map(schema)) - def test_relationship_classification_respects_the_configured_schema(pg_db): - """No resolved= is passed, so guard_schema_provenance_for() no-ops here by - design (a bare engine= caller with no resolved config behind it); this test - is about schema routing, not provenance drift protection.""" - 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" + with scoped_test_schema(pg_db.resolved, prefix="relclass_routing", schema_claims=[staging_schema_claim()]) as scoped: + schema = scoped.schemas[Role.PRIMARY] + Base.metadata.create_all(bind=scoped.engine, checkfirst=True) + + relationship_classification(engine=scoped.engine, resolved=scoped.resolved) + + with scoped.engine.connect() as conn: + n_class = conn.execute( + sa.select(sa.func.count()).select_from(RelationshipClass.__table__) + ).scalar() + n_mapping = conn.execute( + sa.select(sa.func.count()).select_from(RelationshipMapping.__table__) + ).scalar() + enum_schema = conn.execute( + sa.text( + "SELECT n.nspname FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace " + "WHERE t.typname = 'predicatekind' AND n.nspname = :schema" + ), + {"schema": schema}, + ).scalar() + assert n_class and n_class > 0 + assert n_mapping and n_mapping > 0 + assert sa.inspect(scoped.engine).has_table("relationship_class", schema=schema) + assert enum_schema == schema def test_relationship_classification_refuses_a_genuinely_split_vocab_connection(pg_db): @@ -69,69 +53,29 @@ def test_relationship_classification_refuses_a_genuinely_split_vocab_connection( to relationship.relationship_id (VOCAB-tagged) can never be created once vocab_connection is a genuinely separate connection. """ - resolved = dataclasses.replace( - pg_db.resolved, - vocab_connection=dataclasses.replace( - pg_db.resolved.connection, name="genuinely_different", safe_url="postgresql://other/db" - ), - ) + name = pg_db.resolved.name + resolver = Resolver.from_active_config() + resolved = resolver.with_overrides( + connections={ + "genuinely_different": ConnectionConfig(dialect="postgresql+psycopg", host="other", database_name="db") + }, + databases={name: resolver.config.databases[name].model_copy(update={"vocab_connection": "genuinely_different"})}, + ).resolve_database(name) with pytest.raises(RuntimeError, match="genuinely separate"): relationship_classification(engine=pg_db.connection, resolved=resolved) 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. No resolved= here - either, so guard_schema_provenance_for() no-ops by design, same as above.""" - 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 - - -def test_relationship_classification_guard_fires_on_reconfigured_schema(pg_db): - """Unlike the two tests above, resolved= is genuinely passed here -- proves - the guard actually detects drift for relationship_classification itself, - not just that it no-ops correctly when a caller opts out.""" - database_name = f"guard_wiring_test_db_{uuid.uuid4().hex[:8]}" - schema_a = "phase4_guard_wiring_a" - schema_b = "phase4_guard_wiring_b" - - # Both connection and vocab_connection must point at the same object: relationship_classification() - # refuses a genuinely split vocab connection, and dataclasses.replace() only overrides fields - # explicitly passed, so leaving vocab_connection untouched would desync it from the new connection. - single_connection = dataclasses.replace(pg_db.resolved.connection, test_only=False) - resolved_a = dataclasses.replace( - pg_db.resolved, - name=database_name, - schema_name=schema_a, - vocab_schema=schema_a, - results_schema=schema_a, - connection=single_connection, - vocab_connection=single_connection, - ) - # Baseline must be recorded before schema_a has any tables, or the guard's own - # first-time-population check trips on Base.metadata.create_all() below. - scoped_a = _scoped_connection(pg_db, schema_a) - record_schema_provenance( - scoped_a, - database_name=database_name, - schema_tag=Role.PRIMARY.value, - new_physical_schema=schema_a, - reason="test setup", - ) - Base.metadata.create_all(bind=scoped_a, checkfirst=True) - relationship_classification(engine=scoped_a, resolved=resolved_a) - - resolved_b = dataclasses.replace( - resolved_a, schema_name=schema_b, vocab_schema=schema_b, results_schema=schema_b - ) - scoped_b = _scoped_connection(pg_db, schema_b) - with pytest.raises(SchemaDriftError): - relationship_classification(engine=scoped_b, resolved=resolved_b) + DROP TABLE/enum-drop cleanup exists for, must not fail.""" + with scoped_test_schema(pg_db.resolved, prefix="relclass_idempotent", schema_claims=[staging_schema_claim()]) as scoped: + Base.metadata.create_all(bind=scoped.engine, checkfirst=True) + + relationship_classification(engine=scoped.engine, resolved=scoped.resolved) + relationship_classification(engine=scoped.engine, resolved=scoped.resolved) + + with scoped.engine.connect() as conn: + n_class = conn.execute( + sa.select(sa.func.count()).select_from(RelationshipClass.__table__) + ).scalar() + assert n_class and n_class > 0 diff --git a/tests/test_schema_provenance_guard.py b/tests/test_schema_provenance_guard.py index 187e214..3d7e9c7 100644 --- a/tests/test_schema_provenance_guard.py +++ b/tests/test_schema_provenance_guard.py @@ -11,26 +11,16 @@ 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). +scoped_test_schema() commits real schemas, so every registry row this test +writes is a genuine commit. The Role-tag rows are reset around the test. """ 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 oa_configurator import SchemaDriftError +from oa_configurator.testing import guarded_resolver, reset_schema_registry_rows, scoped_test_schema from orm_loader.helpers import Base @@ -40,72 +30,28 @@ 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. - """ - patched_connection = dataclasses.replace(pg_db.resolved.connection, test_only=False) - return dataclasses.replace( - pg_db.resolved, - name=database_name, - schema_name=schema, - vocab_schema=schema, - results_schema=schema, - connection=patched_connection, - vocab_connection=patched_connection, - ) - - 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, - database_name=resolved_a.name, - schema_tag=Role.PRIMARY, - new_physical_schema=schema_a, - reason="test setup baseline", - ) - monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: resolved_a) + reset_schema_registry_rows(cleanup_after_test, pg_db.committing_engine) + resolver = guarded_resolver(pg_db.resolved) + resolved = resolver.resolve_database(pg_db.resolved.name) + + with scoped_test_schema(resolved, prefix="graph_guard_a", resolver=resolver) as scoped_a: + Base.metadata.create_all(bind=scoped_a.engine, checkfirst=True) + monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: scoped_a.resolved) + omop_graph_cli.relationship_classification() + + with scoped_test_schema(resolved, prefix="graph_guard_b", resolver=resolver) as scoped_b: + Base.metadata.create_all(bind=scoped_b.engine, checkfirst=True) + monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: scoped_b.resolved) + with pytest.raises(SchemaDriftError): 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() + # The tables already exist, so the proof the guard fired before any + # write is that they're still empty. + with scoped_b.engine.connect() as conn: + count = conn.execute( + sa.select(sa.func.count()).select_from(RelationshipClass.__table__) + ).scalar() + assert count == 0 diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index a688538..d3b5c43 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -21,10 +21,10 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import Role -from oa_configurator.testing import isolated_test_database, isolated_test_schema +from oa_configurator import Resolver, Role +from oa_configurator.testing import isolated_test_database, resolve_with_role_schemas, scoped_test_schema from orm_loader.config import OrmLoaderConfig -from orm_loader.helpers import Base +from orm_loader.helpers import Base, bulk_load_context from omop_alchemy.cdm.model.vocabulary import ( Concept, @@ -43,7 +43,7 @@ ) from omop_graph.graph.kg import KnowledgeGraph -from fixtures.helpers import VOCAB_TABLES, fk_triggers_disabled, schema_translate_map +from fixtures.helpers import VOCAB_TABLES pytestmark = [pytest.mark.postgresql, pytest.mark.db_dialect] @@ -108,126 +108,133 @@ class _Engines(NamedTuple): def split_engines() -> Iterator[_Engines]: """A primary connection (RelationshipMapping/RelationshipClass) and a genuinely separate physical vocab connection (Relationship/Concept/ - Concept_Relationship).""" + Concept_Relationship), built by ``ResolvedCDMDatabase.create_engines()``.""" 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=schema_translate_map(primary_schema) - ) - vocab_engine = vocab_raw.execution_options( - schema_translate_map=schema_translate_map(vocab_schema) + name = primary_db.resolved.name + resolver = Resolver.from_active_config() + resolver = resolver.with_overrides( + databases={ + name: resolver.config.databases[name].model_copy( + update={"vocab_connection": vocab_db.resolved.connection.name} + ) + } + ) + split = resolver.resolve_database(name) + with scoped_test_schema(split, prefix="vocab_split", resolver=resolver) as scoped: + primary_engine, vocab_engine = scoped.resolved.create_engines() + # Vocab tables are created through the vocab server's own resolved: + # create_all on the split vocab engine also emits the PRIMARY-tagged + # predicatekind type, whose schema only exists on the primary server. + vocab_schema = scoped.schemas[Role.VOCAB] + vocab_only = resolve_with_role_schemas( + vocab_db.resolved, {role: vocab_schema for role in vocab_db.resolved.schema_tags()} ) - - _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 FK bootstrap - # cycle, handled the same way production bulk-loads do: disable FK - # triggers for the load, then re-enable them. - with fk_triggers_disabled(vocab_engine, _VOCAB_TABLES): - _seed(primary_engine, vocab_engine) - - yield _Engines(primary=primary_engine, vocab=vocab_engine) + vocab_setup_engine = vocab_only.create_engine() + try: + _shadow_metadata.create_all(primary_engine) + Base.metadata.create_all(vocab_setup_engine, tables=_VOCAB_TABLES, checkfirst=True) + _seed(primary_engine, vocab_setup_engine) + yield _Engines(primary=primary_engine, vocab=vocab_engine) + finally: + vocab_setup_engine.dispose() + primary_engine.dispose() + vocab_engine.dispose() 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, - ), - ] - ) + with bulk_load_context(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.flush() session.commit() with so.Session(primary_engine) as session: From 94e52718e50fadf24e344a4ce408f24cccd2b0c5 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Fri, 2 Oct 2026 04:52:07 +0000 Subject: [PATCH 18/20] Use new interfaces of omop-emb --- docs/graph/kg.md | 4 ++-- scripts/benchmarks/benchmark.py | 4 ++-- scripts/benchmarks/trace_example.py | 4 ++-- src/omop_graph/graph/kg.py | 2 +- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/graph/kg.md b/docs/graph/kg.md index 297add2..13b808c 100644 --- a/docs/graph/kg.md +++ b/docs/graph/kg.md @@ -62,7 +62,7 @@ This requires the optional `omop-emb` package — see the [installation guide](. ```python from sqlalchemy import create_engine from oa_configurator import Resolver -from omop_emb.backends import resolve_backend_from_resolved_vector_store +from omop_emb.backends import open_vector_store_writer from omop_graph.graph.kg import KnowledgeGraph, KnowledgeGraphEmbeddingConfiguration from omop_emb.config import MetricType @@ -71,7 +71,7 @@ engine = create_engine("postgresql://user:pass@localhost/omop") resolver = Resolver.from_active_config() resolved_vector_store = resolver.resolve_vector_store("vector_store") # a [vector_stores.*] entry name resolved_model = resolver.resolve_model("embedding-model") # a [models.*] entry name -backend = resolve_backend_from_resolved_vector_store(resolved_vector_store) +backend = open_vector_store_writer(resolved_vector_store) emb_config = KnowledgeGraphEmbeddingConfiguration( metric_type=MetricType.COSINE, diff --git a/scripts/benchmarks/benchmark.py b/scripts/benchmarks/benchmark.py index fe37566..c5b074b 100644 --- a/scripts/benchmarks/benchmark.py +++ b/scripts/benchmarks/benchmark.py @@ -22,7 +22,7 @@ parse_metric_type, ) from omop_emb.interface import EmbeddingRole -from omop_emb.backends import EmbeddingBackend, resolve_backend_from_resolved_vector_store +from omop_emb.backends import EmbeddingBackend, open_vector_store_writer from omop_emb.backends.index_config import index_config_from_index_type from omop_graph.config import OmopGraphConfig from omop_graph.extensions.emb import get_embedding_writer_interface, MissingExtensionError @@ -574,7 +574,7 @@ def run_benchmark( "via `omop-config configure omop_graph`." ) resolved_vector_store = Resolver.from_active_config().resolve_vector_store(resolved_name) - embedding_backend = resolve_backend_from_resolved_vector_store(resolved_vector_store) + embedding_backend = open_vector_store_writer(resolved_vector_store) embedding_kg = build_embedding_knowledge_graph( embedding_metric=resolved_embedding_metric_type, diff --git a/scripts/benchmarks/trace_example.py b/scripts/benchmarks/trace_example.py index f1f8387..57bb07c 100644 --- a/scripts/benchmarks/trace_example.py +++ b/scripts/benchmarks/trace_example.py @@ -104,7 +104,7 @@ def _build_kg(metric_type=None, embedding_model: Optional[str] = None) -> Knowle cdm_engine = resolve_cdm_database().create_engine() try: from oa_configurator import Resolver - from omop_emb.backends import resolve_backend_from_resolved_vector_store + from omop_emb.backends import open_vector_store_writer from omop_emb.config import MetricType resolved_metric = metric_type if metric_type is not None else MetricType.COSINE @@ -118,7 +118,7 @@ def _build_kg(metric_type=None, embedding_model: Optional[str] = None) -> Knowle resolver = Resolver.from_active_config() resolved_model = resolver.resolve_model(resolved_model_name) resolved_vector_store = resolver.resolve_vector_store(cfg.vector_store_name) - backend = resolve_backend_from_resolved_vector_store(resolved_vector_store) + backend = open_vector_store_writer(resolved_vector_store) emb_config = KnowledgeGraphEmbeddingConfiguration( metric_type=resolved_metric, diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index 466d8ff..cfc9ec6 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -103,7 +103,7 @@ class KnowledgeGraphEmbeddingConfiguration: similarity computations are consistent. backend : omop_emb.EmbeddingBackend An already-constructed embedding backend, e.g. via - ``omop_emb.backends.resolve_backend_from_resolved_vector_store``. + ``omop_emb.backends.open_vector_store_writer``. resolved_model : oa_configurator.ResolvedModel A model resolved via ``oa_configurator.Resolver.resolve_model()``, carrying real provider connection details. ``model_name``/``provider_type`` (see below) are From 406a13cc78591de96a90fa8183320dfd2ffb57c7 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 7 Oct 2026 00:47:50 +0000 Subject: [PATCH 19/20] Update from manual review --- src/omop_graph/cli.py | 24 ++++++++------------ src/omop_graph/graph/kg.py | 19 ++++++++++++---- tests/test_schema_provenance_guard.py | 32 +++++++++------------------ 3 files changed, 35 insertions(+), 40 deletions(-) diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index 7650c8a..6e01289 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -11,11 +11,7 @@ from oa_configurator import ( ResolvedCDMDatabase, - Role, - ensure_schema, - guard_schema_provenance_for, open_connection, - physical_schema_of ) from orm_loader.backends import STAGING_SCHEMA, resolve_backend, staging_schema_claim @@ -51,6 +47,11 @@ def _main( def populate_with_test_data(): """Populate the database with synthetic test data.""" resolved = resolve_cdm_database() + if not resolved.connection.test_only or not resolved.vocab_connection.test_only: + raise RuntimeError( + f"Refusing to populate {resolved.name!r} with synthetic test data: " + "both its primary and vocab connections must be test_only." + ) engine, vocab_engine = resolved.create_engines() Session = sessionmaker(bind=engine, future=True) VocabSession = sessionmaker(bind=vocab_engine, future=True) @@ -173,13 +174,10 @@ def relationship_classification( "this command, or provision relationship_class/relationship_mapping " "manually without the FK constraint." ) - db_schema = physical_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, staging_schema=STAGING_SCHEMA) + loader_backend = resolve_backend(engine, staging_schema_tag=STAGING_SCHEMA) drop_staging_sql = ( sa.text( @@ -204,13 +202,9 @@ def relationship_classification( RelationshipMapping.__table__, RelationshipClass.__table__, ] - # 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: - guard = guard_schema_provenance_for(connection, resolved, schema_tag=Role.PRIMARY) - with guard: - 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 + 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( @@ -230,7 +224,7 @@ def relationship_classification( dedupe=True, merge_strategy="replace", loader=PandasLoader(), - staging_schema=STAGING_SCHEMA, + staging_schema_tag=STAGING_SCHEMA, ) session.commit() diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index cfc9ec6..8ad7b8f 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -31,6 +31,7 @@ if TYPE_CHECKING: from omop_emb import ( EmbeddingBackend, + EmbeddingStoreReader, EmbeddingWriterInterface, EmbeddingReaderInterface, ) @@ -101,9 +102,12 @@ class KnowledgeGraphEmbeddingConfiguration: The similarity/distance metric to use for embedding comparisons (e.g., cosine, euclidean). This is required to ensure that the correct type of index is used in the backend and that similarity computations are consistent. - backend : omop_emb.EmbeddingBackend - An already-constructed embedding backend, e.g. via - ``omop_emb.backends.open_vector_store_writer``. + backend : omop_emb.EmbeddingBackend or omop_emb.EmbeddingStoreReader + An already-constructed embedding backend. A write-capable + ``EmbeddingBackend`` (e.g. via ``omop_emb.backends.open_vector_store_writer``) + is required when ``write=True``; the narrower read-only + ``EmbeddingStoreReader`` (via ``open_vector_store_reader``) suffices + when ``write=False``. resolved_model : oa_configurator.ResolvedModel A model resolved via ``oa_configurator.Resolver.resolve_model()``, carrying real provider connection details. ``model_name``/``provider_type`` (see below) are @@ -123,7 +127,7 @@ class KnowledgeGraphEmbeddingConfiguration: """ metric_type: EmbeddingMetricType - backend: "EmbeddingBackend" + backend: "EmbeddingBackend | EmbeddingStoreReader" resolved_model: ResolvedModel write: bool = False compute_missing_embeddings: bool = False @@ -266,6 +270,13 @@ def emb(self) -> "EmbeddingWriterInterface | EmbeddingReaderInterface": ) if self._emb_config.write: + from omop_emb import EmbeddingBackend + + if not isinstance(self._emb_config.backend, EmbeddingBackend): + raise ValueError( + "write=True requires a write-capable EmbeddingBackend, " + f"got {type(self._emb_config.backend).__name__}." + ) # Write-capable interface: the KG builds its own embedding model backend. self._emb = EmbeddingWriterInterface( backend=self._emb_config.backend, diff --git a/tests/test_schema_provenance_guard.py b/tests/test_schema_provenance_guard.py index 3d7e9c7..adfe3c7 100644 --- a/tests/test_schema_provenance_guard.py +++ b/tests/test_schema_provenance_guard.py @@ -1,15 +1,15 @@ -"""schema-provenance guard wired into relationship_classification()'s +"""Schema-provenance enforcement reaching 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. +Enforcement now happens at create_engine() construction time (see +oa-configurator's architectural note on guard_schema_provenance_for()), not +via a guard wrapped around this call site's own DDL: the SchemaDriftError +below is raised while building scoped_b's own engine, before +relationship_classification() is even reachable. scoped_test_schema() commits real schemas, so every registry row this test writes is a genuine commit. The Role-tag rows are reset around the test. @@ -18,14 +18,12 @@ from __future__ import annotations import pytest -import sqlalchemy as sa from oa_configurator import SchemaDriftError from oa_configurator.testing import guarded_resolver, reset_schema_registry_rows, scoped_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] @@ -42,16 +40,8 @@ def test_relationship_classification_guard_fires_on_reconfigured_schema( monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: scoped_a.resolved) omop_graph_cli.relationship_classification() - with scoped_test_schema(resolved, prefix="graph_guard_b", resolver=resolver) as scoped_b: - Base.metadata.create_all(bind=scoped_b.engine, checkfirst=True) - monkeypatch.setattr(omop_graph_cli, "resolve_cdm_database", lambda: scoped_b.resolved) - with pytest.raises(SchemaDriftError): - omop_graph_cli.relationship_classification() - - # The tables already exist, so the proof the guard fired before any - # write is that they're still empty. - with scoped_b.engine.connect() as conn: - count = conn.execute( - sa.select(sa.func.count()).select_from(RelationshipClass.__table__) - ).scalar() - assert count == 0 + # Drift is now caught at create_engine() construction time, before + # relationship_classification() is even reachable. + with pytest.raises(SchemaDriftError): + with scoped_test_schema(resolved, prefix="graph_guard_b", resolver=resolver): + pass From 631192c246c47a4e5600f4cc2fccd6d108fc8303 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 7 Oct 2026 06:34:28 +0000 Subject: [PATCH 20/20] Address PR comments --- src/omop_graph/cli.py | 52 +++++++++++++++++------ src/omop_graph/extensions/omop_alchemy.py | 37 ++++++++++++++++ src/omop_graph/graph/kg.py | 23 +++------- src/omop_graph/graph/queries.py | 22 ++-------- tests/test_relationship_classification.py | 20 +-------- tests/test_vocab_split_connection.py | 44 ++++++++++++++++++- 6 files changed, 129 insertions(+), 69 deletions(-) diff --git a/src/omop_graph/cli.py b/src/omop_graph/cli.py index 6e01289..f6c3309 100644 --- a/src/omop_graph/cli.py +++ b/src/omop_graph/cli.py @@ -14,6 +14,8 @@ open_connection, ) +from omop_alchemy.cdm.model.vocabulary import Relationship + from orm_loader.backends import STAGING_SCHEMA, resolve_backend, staging_schema_claim from orm_loader.helpers import bulk_load_context from orm_loader.helpers.metadata import Base @@ -21,7 +23,11 @@ from omop_graph.config import OmopGraphConfig from omop_graph.db.session import resolve_cdm_database -from omop_graph.extensions.omop_alchemy import RelationshipClass, RelationshipMapping +from omop_graph.extensions.omop_alchemy import ( + RelationshipClass, + RelationshipMapping, + relationship_mapping_table_without_vocab_fk, +) from omop_graph.cli_utils import populate_test_data app = typer.Typer() @@ -92,6 +98,11 @@ def relationship_classification( 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. + Also determines whether ``vocab_connection`` is a genuinely + separate physical connection: if so, ``RelationshipMapping`` is + created without its FK to ``relationship.relationship_id``, + and relationship IDs absent from the vocab connection's own `relationship` table + are dropped instead of relying on the FK to reject them. """ pred_class_dir_pl = ( Path(pred_class_dir) if pred_class_dir else packaged_predicate_csv_dir() @@ -164,16 +175,29 @@ def relationship_classification( if engine is None: resolved = resolve_cdm_database() engine = resolved.create_engine(schema_claims=[staging_schema_claim()]) - if resolved is not None and resolved.connection != resolved.vocab_connection: - raise RuntimeError( - f"relationship_classification() cannot run against database " - f"{resolved.name!r}: its vocab_connection is a genuinely separate " - "connection from the primary one, and RelationshipMapping's FK to " - "relationship.relationship_id needs both in the same database. " - "Point vocab_connection at the same connection as primary for " - "this command, or provision relationship_class/relationship_mapping " - "manually without the FK constraint." - ) + split_vocab = resolved is not None and resolved.connection != resolved.vocab_connection + + if split_vocab: + if not isinstance(engine, sa.Engine): + raise TypeError( + "relationship_classification() needs an Engine (not a bare " + "Connection) to build the vocab connection's own engine when " + "vocab_connection is genuinely separate." + ) + vocab_engine = resolved.vocab_engine_for(engine) + with open_connection(vocab_engine) as vocab_connection: + known_relationship_ids = { + row.relationship_id + for row in vocab_connection.execute(sa.select(Relationship.relationship_id)) + } + unknown_mask = ~df_rel_mapping_to_export["relationship_id"].isin(known_relationship_ids) + unknown_ids = df_rel_mapping_to_export.loc[unknown_mask, "relationship_id"].unique().tolist() + if unknown_ids: + logger.warning( + f"Dropping {len(unknown_ids)} relationships absent from the vocabulary " + f"connection's relationship table: {unknown_ids}" + ) + df_rel_mapping_to_export = df_rel_mapping_to_export.loc[~unknown_mask] Session = sessionmaker(bind=engine, future=True) session = Session() @@ -204,7 +228,11 @@ def relationship_classification( ] with open_connection(engine) as connection: 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 + if split_vocab: + Base.metadata.create_all(bind=connection, tables=[RelationshipClass.__table__]) # type: ignore + relationship_mapping_table_without_vocab_fk().create(bind=connection, checkfirst=True) + else: + 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/extensions/omop_alchemy.py b/src/omop_graph/extensions/omop_alchemy.py index ee55fad..2f1beb4 100644 --- a/src/omop_graph/extensions/omop_alchemy.py +++ b/src/omop_graph/extensions/omop_alchemy.py @@ -96,6 +96,43 @@ class RelationshipMapping(ReferenceTable, CDMTableBase, Base): ) +def relationship_mapping_table_without_vocab_fk() -> sa.Table: + """RelationshipMapping's DDL without its FK to vocab.relationship_id. + + Postgres has no cross-database inline FK, so when vocab_connection is a + genuinely separate physical connection, CREATE TABLE on the real ORM + table fails there. This reproduces the same columns under the same + name/schema in a throwaway MetaData, so CREATE TABLE succeeds while + RelationshipMapping's own ORM class still reads/writes the same + physical table afterward. + """ + metadata = sa.MetaData() + return sa.Table( + RelationshipMapping.__tablename__, + 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], + schema=Role.PRIMARY.value, + ), + primary_key=True, + ), + sa.Column("predicate_subkind", sa.String(20), primary_key=True), + sa.ForeignKeyConstraint( + ["predicate_kind", "predicate_subkind"], + [ + RelationshipClass.__table__.c.predicate_kind, + RelationshipClass.__table__.c.predicate_subkind, + ], + name="fk_rel_mapping_to_rel_class", + ), + schema=Role.PRIMARY.value, + ) + + @dataclass(frozen=True, slots=True) class RelationshipMappingElement: relationship_id: str diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index 8ad7b8f..3c82f58 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -74,7 +74,6 @@ q_children, q_predicate_name, q_predicate_row_with_ancestry, - q_relationship_mapping_all, q_relationship_mapping_row, q_roots, q_singletons, @@ -151,20 +150,9 @@ 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-tagged, 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: +def _predicate_from_rows( + ancestry_row: Row, mapping_row: "Row | RelationshipMappingElement" +) -> 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 @@ -651,7 +639,7 @@ def iter_edges( include_classification=False, ) ).all() - mapping_by_id = _relationship_mapping_lookup(session) + mapping_by_id = self._relationship_mapping for vrow in vocab_rows: mapping = mapping_by_id.get(vrow.predicate_id) if mapping is None: @@ -778,8 +766,7 @@ def predicates(self) -> tuple[Predicate, ...]: 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) + mapping_by_id = self._relationship_mapping return tuple( _predicate_from_rows(row, mapping_by_id[row.relationship_id]) for row in ancestry_rows diff --git a/src/omop_graph/graph/queries.py b/src/omop_graph/graph/queries.py index 8468430..dd63cab 100644 --- a/src/omop_graph/graph/queries.py +++ b/src/omop_graph/graph/queries.py @@ -497,20 +497,6 @@ def q_relationship_mapping_row(relationship_id: str) -> Select: ).where(RelationshipMapping.relationship_id == relationship_id) -def q_relationship_mapping_all() -> Select: - """Query every RelationshipMapping row, keyed by relationship_id. - - The primary-tagged 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, - ) - - def q_edges( concept_ids: Union[Tuple[int, ...], int], direction: Literal["in", "out"], @@ -531,10 +517,10 @@ def q_edges( Join in RelationshipMapping's predicate_kind/predicate_subkind. Concept_Relationship is vocab-tagged, 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. + ``KnowledgeGraph``'s own cached ``_relationship_mapping`` 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,) diff --git a/tests/test_relationship_classification.py b/tests/test_relationship_classification.py index d56394b..b989d9b 100644 --- a/tests/test_relationship_classification.py +++ b/tests/test_relationship_classification.py @@ -9,9 +9,8 @@ from __future__ import annotations -import pytest import sqlalchemy as sa -from oa_configurator import ConnectionConfig, Resolver, Role +from oa_configurator import Role from oa_configurator.testing import scoped_test_schema from orm_loader.backends import staging_schema_claim @@ -48,23 +47,6 @@ def test_relationship_classification_respects_the_configured_schema(pg_db): assert enum_schema == schema -def test_relationship_classification_refuses_a_genuinely_split_vocab_connection(pg_db): - """Postgres has no cross-database inline FK, so RelationshipMapping's FK - to relationship.relationship_id (VOCAB-tagged) can never be created once - vocab_connection is a genuinely separate connection. - """ - name = pg_db.resolved.name - resolver = Resolver.from_active_config() - resolved = resolver.with_overrides( - connections={ - "genuinely_different": ConnectionConfig(dialect="postgresql+psycopg", host="other", database_name="db") - }, - databases={name: resolver.config.databases[name].model_copy(update={"vocab_connection": "genuinely_different"})}, - ).resolve_database(name) - with pytest.raises(RuntimeError, match="genuinely separate"): - relationship_classification(engine=pg_db.connection, resolved=resolved) - - 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.""" diff --git a/tests/test_vocab_split_connection.py b/tests/test_vocab_split_connection.py index d3b5c43..7ea05be 100644 --- a/tests/test_vocab_split_connection.py +++ b/tests/test_vocab_split_connection.py @@ -21,8 +21,9 @@ import sqlalchemy as sa import sqlalchemy.orm as so -from oa_configurator import Resolver, Role +from oa_configurator import ResolvedCDMDatabase, Resolver, Role from oa_configurator.testing import isolated_test_database, resolve_with_role_schemas, scoped_test_schema +from orm_loader.backends import staging_schema_claim from orm_loader.config import OrmLoaderConfig from orm_loader.helpers import Base, bulk_load_context @@ -35,6 +36,7 @@ Vocabulary, ) +from omop_graph.cli import relationship_classification from omop_graph.config import OmopGraphConfig from omop_graph.extensions.omop_alchemy import ( PredicateKind, @@ -102,6 +104,7 @@ class _Engines(NamedTuple): primary: sa.Engine vocab: sa.Engine + resolved: ResolvedCDMDatabase @pytest.fixture() @@ -137,7 +140,7 @@ def split_engines() -> Iterator[_Engines]: _shadow_metadata.create_all(primary_engine) Base.metadata.create_all(vocab_setup_engine, tables=_VOCAB_TABLES, checkfirst=True) _seed(primary_engine, vocab_setup_engine) - yield _Engines(primary=primary_engine, vocab=vocab_engine) + yield _Engines(primary=primary_engine, vocab=vocab_engine, resolved=scoped.resolved) finally: vocab_setup_engine.dispose() primary_engine.dispose() @@ -372,3 +375,40 @@ def test_relationships_invert_swaps_subjects_and_objects(split_engines: _Engines ) assert triples == ((OBJECT_CONCEPT_ID, "maps to", SUBJECT_CONCEPT_ID),) + + +def test_relationship_classification_succeeds_on_a_genuinely_split_vocab_connection( + split_engines: _Engines, tmp_path +) -> None: + """relationship_classification() creates RelationshipMapping + without that FK in a split connection and validates relationship_id + against the vocab connection's own `relationship` table in Python instead + """ + (tmp_path / "predicate_classification.csv").write_text( + "class,subclass,description,semantics,inference\n" + "Identity,mapping,Identity mapping,identity,none\n" + ) + (tmp_path / "predicate_mapping.csv").write_text( + "class,subclass,r_id,r_name,classification_notes,additional_info,link\n" + "Identity,mapping,maps to,Maps to,,,\n" + "Identity,mapping,totally_unknown_relationship,Unknown,,,\n" + ) + + staging_engine = split_engines.resolved.create_engine(schema_claims=[staging_schema_claim()]) + try: + relationship_classification( + str(tmp_path), engine=staging_engine, resolved=split_engines.resolved + ) + finally: + staging_engine.dispose() + + with so.Session(split_engines.primary) as session: + relationship_ids = set( + session.execute(sa.select(RelationshipMapping.relationship_id)).scalars() + ) + assert relationship_ids == {"maps to"} + + kg = _split_kg(split_engines) + predicate = kg.predicate("maps to") + assert predicate.predicate_kind == PredicateKind.IDENTITY + assert predicate.predicate_subkind == "mapping"