From d93ac6c5db19f2eb2379a54ed13b98ba495c49ef Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 8 Jul 2026 05:31:03 +0000 Subject: [PATCH 1/2] Single vocab table loading --- docs/cli/reference.md | 6 +- omop_alchemy/maintenance/cli_vocab.py | 77 ++++++++++++++++---- tests/test_load_vocab_source.py | 101 ++++++++++++++++++++++++++ 3 files changed, 170 insertions(+), 14 deletions(-) diff --git a/docs/cli/reference.md b/docs/cli/reference.md index 017810a..e8dfad1 100644 --- a/docs/cli/reference.md +++ b/docs/cli/reference.md @@ -63,14 +63,18 @@ Summarise ORM-managed OMOP tables present in the target database. ### `load-vocab-source` -Load all Athena vocabulary CSVs from the configured source path, optionally toggling indexes and FK triggers for speed. +Load Athena vocabulary CSVs from the configured source path into the target database. By default all vocabulary tables are loaded; use `--table` to restrict the load to one or more specific tables. !!! warning "Clustering" `--bulk-mode` (default) drops all secondary indexes before loading and recreates them afterward. Physical table clustering is **not** run as part of the load. Use [`indexes cluster`](#indexes-cluster) as a separate step once the load completes and you have confirmed sufficient free disk space. +!!! tip "Single-table loads" + When targeting a specific table with `--table`, consider passing `--no-bulk-mode` to avoid dropping and rebuilding indexes across all vocabulary tables for a single-file update. + | Flag | Type / Choices | Default | Description | |---|---|---|---| | `--athena-source` | str (optional) | (saved default) | Path to the unzipped Athena vocabulary CSV directory. Falls back to the saved athena-source default. | +| `--table` | str (repeatable, optional) | (none — all tables) | Restrict the load to one or more vocabulary tables by name (e.g. `--table concept --table vocabulary`). Skips the all-required-files preflight; errors if the named CSV is absent from the source directory. | | `--merge-strategy` | `replace` / `upsert` / `insert_if_empty` | `replace` | CSV merge strategy. `replace` keeps the DB in sync with the source. `upsert` is incremental and non-destructive. `insert_if_empty` is the fast path for a fresh empty target. | | `--staging-chunk-size` | int (optional) | `100000` | [Phase 1] Rows per ORM transaction when loading CSV → staging table. Ignored when the PostgreSQL COPY fast-path is active. Pass `0` to disable chunking. | | `--bulk-mode` / `--no-bulk-mode` | bool | `True` | Disable FK triggers and drop indexes globally before loading, then rebuild after. Much faster for a full vocabulary reload. Ignored on backends that do not support it. | diff --git a/omop_alchemy/maintenance/cli_vocab.py b/omop_alchemy/maintenance/cli_vocab.py index 72327dd..bf8356c 100644 --- a/omop_alchemy/maintenance/cli_vocab.py +++ b/omop_alchemy/maintenance/cli_vocab.py @@ -6,7 +6,7 @@ from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from typing import Literal, TypeAlias, cast +from typing import List, Literal, Optional, TypeAlias, cast import sqlalchemy as sa import sqlalchemy.orm as so @@ -118,6 +118,10 @@ def _emit( ), ) +_ALL_VOCAB_TABLE_NAMES: frozenset[str] = frozenset( + m.__tablename__ for m in REQUIRED_VOCAB_MODELS + OPTIONAL_VOCAB_MODELS +) + class VocabularyLoadError(RuntimeError): """Raised when a single Athena vocabulary table load fails.""" @@ -201,13 +205,16 @@ def _find_vocab_csv_path(source_path: Path, table_name: str) -> Path | None: return None -def _missing_required_files(source_path: Path) -> list[str]: - """Return the table names of required vocabulary CSVs that cannot be found under source_path.""" - missing: list[str] = [] - for model in REQUIRED_VOCAB_MODELS: - if _find_vocab_csv_path(source_path, model.__tablename__) is None: - missing.append(model.__tablename__) - return missing +def _missing_required_files( + source_path: Path, + required_models: tuple[VocabularyModel, ...] = REQUIRED_VOCAB_MODELS, +) -> list[str]: + """Return the table names of vocabulary CSVs in required_models that cannot be found under source_path.""" + return [ + model.__tablename__ + for model in required_models + if _find_vocab_csv_path(source_path, model.__tablename__) is None + ] def _create_missing_vocabulary_tables( @@ -247,6 +254,7 @@ def load_vocab_source( engine: sa.Engine, *, source_path: str | Path, + tables: list[str] | None = None, db_schema: str | None = None, dry_run: bool = False, merge_strategy: MergeStrategy = "replace", @@ -255,18 +263,50 @@ def load_vocab_source( merge_batch_size: int | None = None, progress_callback: VocabularyLoadProgressCallback | None = None, ) -> VocabularyLoadReport: - """Load all Athena vocabulary CSVs from source_path. With bulk_mode, indexes and FK triggers are toggled around the load.""" + """ + Load Athena vocabulary CSVs from source_path into the target database. + + When tables is None (default), all vocabulary tables are loaded and the + full set of required CSVs is validated before any DB work begins. When + tables is provided, only those tables are loaded and the preflight check + is scoped to the requested tables only; an absent CSV for any named table + raises RuntimeError before touching the database. + + With bulk_mode (default on PostgreSQL), secondary indexes and FK triggers + are toggled globally around the load for speed. Pass --no-bulk-mode when + loading a single table to avoid the index drop/rebuild overhead. + """ resolved_source_path = Path(source_path).expanduser().resolve() if not resolved_source_path.exists() or not resolved_source_path.is_dir(): raise RuntimeError( f"Athena source directory not found: {resolved_source_path}" ) - missing_required = _missing_required_files(resolved_source_path) - if missing_required: + if tables is not None: + unknown = sorted(set(tables) - _ALL_VOCAB_TABLE_NAMES) + if unknown: + raise ValueError( + f"Unknown vocabulary table(s): {unknown}. " + f"Valid names: {sorted(_ALL_VOCAB_TABLE_NAMES)}" + ) + + all_models = REQUIRED_VOCAB_MODELS + OPTIONAL_VOCAB_MODELS + if tables is not None: + all_models = tuple(m for m in all_models if m.__tablename__ in tables) + + # When a specific table selection is given, check exactly those CSVs are present. + # Otherwise check only the required tables, missing optional CSVs are allowed. + models_to_preflight = all_models if tables is not None else REQUIRED_VOCAB_MODELS + missing = _missing_required_files(resolved_source_path, required_models=models_to_preflight) + if missing: + if tables is not None: + raise RuntimeError( + "Requested vocabulary CSV file(s) not found in source directory: " + + ", ".join(sorted(missing)) + ) raise RuntimeError( "Missing required Athena vocabulary CSV files: " - + ", ".join(sorted(missing_required)) + + ", ".join(sorted(missing)) ) if not dry_run: @@ -290,7 +330,6 @@ def _set_search_path(dbapi_conn, _record): cur.execute(f"SET search_path TO staging, {_quoted_schema}") cur.close() - all_models = REQUIRED_VOCAB_MODELS + OPTIONAL_VOCAB_MODELS table_count = sum( 1 for m in all_models if _find_vocab_csv_path(resolved_source_path, m.__tablename__) is not None @@ -495,6 +534,17 @@ def load_vocab_source_command( None, help="Path to the unzipped Athena vocabulary CSV directory. Falls back to the saved athena-source default.", ), + tables: Optional[List[str]] = typer.Option( + None, + "--table", + help=( + "Restrict the load to one or more vocabulary tables by name " + "(e.g. --table concept --table vocabulary). " + "When omitted, all vocabulary tables are loaded. " + "Skips the all-required-files preflight check; " + "errors if the named CSV is absent from the source directory." + ), + ), merge_strategy: MergeStrategy = typer.Option( "replace", help=( @@ -575,6 +625,7 @@ def _update_progress(event: VocabularyLoadProgress) -> None: report = load_vocab_source( engine, source_path=effective_athena_source, + tables=tables or None, db_schema=conn.db_schema, dry_run=dry_run, merge_strategy=merge_strategy, diff --git a/tests/test_load_vocab_source.py b/tests/test_load_vocab_source.py index 44ebf02..8055047 100644 --- a/tests/test_load_vocab_source.py +++ b/tests/test_load_vocab_source.py @@ -173,6 +173,7 @@ def fake_load_vocab_source( engine: object, *, source_path: str | Path, + tables: list[str] | None = None, # noqa: ARG001 db_schema: str | None = None, dry_run: bool = False, merge_strategy: MergeStrategy = "replace", @@ -512,3 +513,103 @@ def fake_load_vocab_model_csv( f"Expected all tables to use quote_mode='auto', got: {received_quote_modes}" ) assert "literal" not in received_quote_modes + + +def test_load_vocab_source_tables_unknown_name_raises_value_error(tmp_path): + """Unknown table name in tables= is rejected before any DB connection.""" + engine = sa.create_engine(f"sqlite:///{tmp_path / 'tables_unknown.db'}", future=True) + source_path = _build_required_athena_source(tmp_path) + + with pytest.raises(ValueError, match="Unknown vocabulary table"): + load_vocab_source(engine, source_path=source_path, tables=["not_a_table"]) + + +def test_load_vocab_source_tables_single_loads_only_that_table(monkeypatch, tmp_path): + """tables=['concept'] loads only concept and skips every other table.""" + engine = sa.create_engine(f"sqlite:///{tmp_path / 'tables_single.db'}", future=True) + source_path = _build_required_athena_source(tmp_path) + loaded_tables: list[str] = [] + + def fake_load_vocab_model_csv( + session, + *, + model, + csv_path, + merge_strategy, + quote_mode="auto", # noqa: ARG001 + chunksize=None, + index_strategy="auto", + merge_batch_size: int = 1_000_000, + ) -> int: + loaded_tables.append(model.__tablename__) + return 1 + + monkeypatch.setattr( + "omop_alchemy.maintenance.cli_vocab._load_vocab_model_csv", + fake_load_vocab_model_csv, + ) + + report = load_vocab_source(engine, source_path=source_path, tables=["concept"]) + + assert loaded_tables == ["concept"] + result_names = {r.table_name for r in report.results} + assert result_names == {"concept"} + + +def test_load_vocab_source_tables_multiple_loads_exactly_those(monkeypatch, tmp_path): + """tables=['concept', 'vocabulary'] loads exactly those two tables.""" + engine = sa.create_engine(f"sqlite:///{tmp_path / 'tables_multiple.db'}", future=True) + source_path = _build_required_athena_source(tmp_path) + loaded_tables: list[str] = [] + + def fake_load_vocab_model_csv( + session, + *, + model, + csv_path, + merge_strategy, + quote_mode="auto", # noqa: ARG001 + chunksize=None, + index_strategy="auto", + merge_batch_size: int = 1_000_000, + ) -> int: + loaded_tables.append(model.__tablename__) + return 1 + + monkeypatch.setattr( + "omop_alchemy.maintenance.cli_vocab._load_vocab_model_csv", + fake_load_vocab_model_csv, + ) + + load_vocab_source(engine, source_path=source_path, tables=["concept", "vocabulary"]) + + assert set(loaded_tables) == {"concept", "vocabulary"} + + +def test_load_vocab_source_tables_skips_required_files_preflight(tmp_path): + """tables= skips the all-required-files gate even when most CSVs are absent.""" + engine = sa.create_engine(f"sqlite:///{tmp_path / 'tables_preflight.db'}", future=True) + + # Only concept.csv present — would fail the all-required-files check without tables=. + source_path = tmp_path / "sparse" + source_path.mkdir() + _write_athena_csv(source_path, "concept") + + # Should raise RuntimeError for missing concept CSV — but NOT the "Missing required" error. + # Since concept.csv IS present, the load should proceed without hitting the preflight. + report = load_vocab_source(engine, source_path=source_path, tables=["concept"], dry_run=True) + + result_names = {r.table_name for r in report.results} + assert result_names == {"concept"} + + +def test_load_vocab_source_tables_missing_csv_raises_runtime_error(tmp_path): + """Explicitly named table whose CSV is absent raises RuntimeError, not a silent skip.""" + engine = sa.create_engine(f"sqlite:///{tmp_path / 'tables_missing_csv.db'}", future=True) + + source_path = tmp_path / "empty_source" + source_path.mkdir() + # No CSVs at all — concept is in tables= but its file is missing. + + with pytest.raises(RuntimeError, match="concept"): + load_vocab_source(engine, source_path=source_path, tables=["concept"]) From 3770e2625e58db6e94c4db6255e1515b582fac60 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Thu, 16 Jul 2026 04:20:56 +0000 Subject: [PATCH 2/2] Adapt type for CLI list elements, raise RuntimeError to get proper catch --- omop_alchemy/maintenance/cli_vocab.py | 6 +++--- tests/test_load_vocab_source.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/omop_alchemy/maintenance/cli_vocab.py b/omop_alchemy/maintenance/cli_vocab.py index bf8356c..23bab50 100644 --- a/omop_alchemy/maintenance/cli_vocab.py +++ b/omop_alchemy/maintenance/cli_vocab.py @@ -6,7 +6,7 @@ from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from typing import List, Literal, Optional, TypeAlias, cast +from typing import Literal, TypeAlias, cast import sqlalchemy as sa import sqlalchemy.orm as so @@ -285,7 +285,7 @@ def load_vocab_source( if tables is not None: unknown = sorted(set(tables) - _ALL_VOCAB_TABLE_NAMES) if unknown: - raise ValueError( + raise RuntimeError( f"Unknown vocabulary table(s): {unknown}. " f"Valid names: {sorted(_ALL_VOCAB_TABLE_NAMES)}" ) @@ -534,7 +534,7 @@ def load_vocab_source_command( None, help="Path to the unzipped Athena vocabulary CSV directory. Falls back to the saved athena-source default.", ), - tables: Optional[List[str]] = typer.Option( + tables: list[str] | None = typer.Option( None, "--table", help=( diff --git a/tests/test_load_vocab_source.py b/tests/test_load_vocab_source.py index 8055047..10067b1 100644 --- a/tests/test_load_vocab_source.py +++ b/tests/test_load_vocab_source.py @@ -515,12 +515,12 @@ def fake_load_vocab_model_csv( assert "literal" not in received_quote_modes -def test_load_vocab_source_tables_unknown_name_raises_value_error(tmp_path): +def test_load_vocab_source_tables_unknown_name_raises_runtime_error(tmp_path): """Unknown table name in tables= is rejected before any DB connection.""" engine = sa.create_engine(f"sqlite:///{tmp_path / 'tables_unknown.db'}", future=True) source_path = _build_required_athena_source(tmp_path) - with pytest.raises(ValueError, match="Unknown vocabulary table"): + with pytest.raises(RuntimeError, match="Unknown vocabulary table"): load_vocab_source(engine, source_path=source_path, tables=["not_a_table"])