From 503f9172caa3f6bd65506491b30d83e10f804d05 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 8 Jul 2026 04:48:05 +0000 Subject: [PATCH] Implement changes of orm-loader --- omop_alchemy/maintenance/cli_vocab.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/omop_alchemy/maintenance/cli_vocab.py b/omop_alchemy/maintenance/cli_vocab.py index ca5a711..e8d68a5 100644 --- a/omop_alchemy/maintenance/cli_vocab.py +++ b/omop_alchemy/maintenance/cli_vocab.py @@ -14,6 +14,7 @@ from sqlalchemy.exc import OperationalError import typer from sqlalchemy.pool import NullPool +from orm_loader.backends import STAGING_SCHEMA, resolve_backend from orm_loader.tables.typing import CSVTableProtocol from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn, TimeElapsedColumn @@ -140,9 +141,10 @@ def _is_missing_staging_table_error( exc: Exception, *, model: VocabularyModel, + session: so.Session, ) -> bool: """Return True if the exception is a ProgrammingError caused by the staging table not existing yet.""" - staging_table_name = model.staging_tablename() + staging_table_name = resolve_backend(session).staging_name_for_table(model.__tablename__) message = str(exc).lower() return ( exc.__class__.__name__ == "ProgrammingError" @@ -175,7 +177,7 @@ def _load_vocab_model_csv( try: return int(model.load_csv(session, csv_path, **load_kwargs)) # type: ignore[arg-type] except Exception as exc: - if not _is_missing_staging_table_error(exc, model=model): + if not _is_missing_staging_table_error(exc, model=model, session=session): raise session.rollback() @@ -269,9 +271,15 @@ def load_vocab_source( + ", ".join(sorted(missing_required)) ) + if db_schema == STAGING_SCHEMA: + raise ValueError( + f"db_schema cannot be {STAGING_SCHEMA!r}: that name is reserved " + "for the orm-loader staging schema." + ) + if not dry_run: ensure_schema(engine, db_schema) - ensure_schema(engine, "staging") + ensure_schema(engine, STAGING_SCHEMA) # NullPool: each session/connection is opened fresh and closed immediately after # use. No stale pooled connections survive between tables, which prevents @@ -282,12 +290,14 @@ def load_vocab_source( # SET search_path on the first checkout doesn't survive into the next # checkout (e.g. after create_staging_table's commit). Re-apply on every # new connection via an engine-level connect event so COPY and raw-cursor - # operations always target the right schema. + # operations always target the right schema. The staging schema no longer + # needs to be on the path because orm-loader now qualifies staging + # table identifiers explicitly. _quoted_schema = '"' + db_schema.replace('"', '""') + '"' @sae.listens_for(load_engine, "connect") def _set_search_path(dbapi_conn, _record): cur = dbapi_conn.cursor() - cur.execute(f"SET search_path TO staging, {_quoted_schema}") + cur.execute(f"SET search_path TO {_quoted_schema}") cur.close() all_models = REQUIRED_VOCAB_MODELS + OPTIONAL_VOCAB_MODELS