Single vocab table loading

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. |

omop_alchemy/maintenance/cli_vocab.py

@@ -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 RuntimeError(
+                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: list[str] | None = 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,

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_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(RuntimeError, 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"])