[wip]

Author: jonavellecuerdo

Pipfile

@@ -21,6 +21,7 @@ mypy = "*"
 pre-commit = "*"
 pytest = "*"
 ruff = "*"
+timdex-dataset-api = { git = "git+https://github.com/MITLibraries/timdex-dataset-api.git"}
 vcrpy = "*"
 
 [requires]

Pipfile.lock

@@ -2,7 +2,9 @@
 import vcr
 from click.testing import CliRunner
 
+from tests.utils import generate_sample_records
 from tim.opensearch import configure_opensearch_client
+from timdex_dataset_api import TIMDEXDataset
 
 EXIT_CODES = {
     "success": 0,
@@ -33,3 +35,30 @@ def test_opensearch_client():
 @pytest.fixture
 def runner():
     return CliRunner()
+
+
+@pytest.fixture
+def local_dataset_location(tmp_path):
+    return str(tmp_path / "fixed_local_dataset/")
+
+
+@pytest.fixture
+def fixed_local_dataset(local_dataset_location):
+    td = TIMDEXDataset(local_dataset_location)
+
+    for source, run_date, run_id, action in [
+        ("alma", "2024-12-01", "abc123", "index"),
+        ("alma", "2024-12-02", "def456", "delete"),
+    ]:
+        td.write(
+            generate_sample_records(
+                num_records=1_000,
+                timdex_record_id_prefix=source,
+                source=source,
+                run_date=run_date,
+                action=action,
+                run_id=run_id,
+            )
+        )
+    td.load()
+    return td

tests/conftest.py

@@ -1,8 +1,10 @@
 import re
+from unittest.mock import patch
 
 from freezegun import freeze_time
 
 from tim.cli import main
+from tim.errors import BulkIndexingError
 
 from .conftest import EXIT_CODES, my_vcr
 
@@ -256,3 +258,85 @@ def test_bulk_delete_with_source_success(caplog, runner):
         "from index 'alma-2022-09-01t00-00-00'" in caplog.text
     )
     assert "Bulk deletion complete!" in caplog.text
+
+
+@patch("tim.helpers.validate_bulk_cli_options")
+@patch("tim.opensearch.bulk_delete")
+@patch("tim.opensearch.bulk_index")
+def test_bulk_update_with_source_success(
+    mock_bulk_index,
+    mock_bulk_delete,
+    mock_validate_bulk_cli_options,
+    caplog,
+    monkeypatch,
+    runner,
+    fixed_local_dataset,
+):
+    monkeypatch.delenv("TIMDEX_OPENSEARCH_ENDPOINT", raising=False)
+    mock_bulk_index.return_value = {
+        "created": 1000,
+        "updated": 0,
+        "errors": 0,
+        "total": 1000,
+    }
+    mock_bulk_delete.return_value = {"deleted": 0, "errors": 0, "total": 0}
+    mock_validate_bulk_cli_options.return_value = "index"
+
+    result = runner.invoke(
+        main,
+        [
+            "bulk-update",
+            "--source",
+            "alma",
+            "--run-date",
+            "2024-12-01",
+            "--run-id",
+            "abc123",
+            fixed_local_dataset.location,
+        ],
+    )
+    assert result.exit_code == EXIT_CODES["success"]
+    assert (
+        "Bulk update complete: {'created': 1000, 'updated': 0, 'deleted': 0, 'total': 1000, 'bulk_index_errors': 0, 'bulk_delete_errors': 0}"
+        in caplog.text
+    )
+
+
+@patch("tim.helpers.validate_bulk_cli_options")
+@patch("tim.opensearch.bulk_delete")
+@patch("tim.opensearch.bulk_index")
+def test_bulk_update_with_source_raise_bulk_indexing_error(
+    mock_bulk_index,
+    mock_bulk_delete,
+    mock_validate_bulk_cli_options,
+    caplog,
+    monkeypatch,
+    runner,
+    fixed_local_dataset,
+):
+    monkeypatch.delenv("TIMDEX_OPENSEARCH_ENDPOINT", raising=False)
+    mock_bulk_index.side_effect = BulkIndexingError(
+        record="alma:0", index="index", error="exception"
+    )
+
+    mock_bulk_delete.return_value = {"deleted": 0, "errors": 0, "total": 0}
+    mock_validate_bulk_cli_options.return_value = "index"
+
+    result = runner.invoke(
+        main,
+        [
+            "bulk-update",
+            "--source",
+            "alma",
+            "--run-date",
+            "2024-12-01",
+            "--run-id",
+            "abc123",
+            fixed_local_dataset.location,
+        ],
+    )
+    assert result.exit_code == EXIT_CODES["success"]
+    assert (
+        "Bulk update complete: {'created': 0, 'updated': 0, 'deleted': 0, 'total': 0, 'bulk_index_errors': 0, 'bulk_delete_errors': 0}"
+        in caplog.text
+    )

tests/test_cli.py

@@ -0,0 +1,30 @@
+import uuid
+from collections.abc import Iterator
+
+from timdex_dataset_api import DatasetRecord
+
+
+def generate_sample_records(
+    num_records: int,
+    timdex_record_id_prefix: str = "alma",
+    source: str | None = "alma",
+    run_date: str | None = "2024-12-01",
+    run_type: str | None = "daily",
+    action: str | None = "index",
+    run_id: str | None = None,
+) -> Iterator[DatasetRecord]:
+    """Generate sample DatasetRecords."""
+    if not run_id:
+        run_id = str(uuid.uuid4())
+
+    for x in range(num_records):
+        yield DatasetRecord(
+            timdex_record_id=f"{timdex_record_id_prefix}:{x}",
+            source_record=b"<record><title>Hello World.</title></record>",
+            transformed_record=b"""{"title":["Hello World."]}""",
+            source=source,
+            run_date=run_date,
+            run_type=run_type,
+            action=action,
+            run_id=run_id,
+        )

tests/utils.py

@@ -1,13 +1,15 @@
 # ruff: noqa: TRY003, EM101
+import json
 import logging
 from datetime import timedelta
 from time import perf_counter
 
 import rich_click as click
-
 from tim import errors, helpers
 from tim import opensearch as tim_os
 from tim.config import PRIMARY_ALIAS, VALID_SOURCES, configure_logger, configure_sentry
+from tim.errors import BulkIndexingError
+from timdex_dataset_api import TIMDEXDataset
 
 logger = logging.getLogger(__name__)
 
@@ -252,6 +254,7 @@ def promote(ctx: click.Context, index: str, alias: list[str]) -> None:
 # Bulk record processing commands
 
 
+# NOTE: FEATURE FLAG: 'bulk_index' may be removed entirely when v2 work done
 @main.command()
 @click.option("-i", "--index", help="Name of the index to bulk index records into.")
 @click.option(
@@ -295,6 +298,7 @@ def bulk_index(ctx: click.Context, index: str, source: str, filepath: str) -> No
     )
 
 
+# NOTE: FEATURE FLAG: 'bulk_index' may be removed entirely when v2 work done
 @main.command()
 @click.option("-i", "--index", help="Name of the index to bulk delete records from.")
 @click.option(
@@ -334,3 +338,62 @@ def bulk_delete(ctx: click.Context, index: str, source: str, filepath: str) -> N
         results["deleted"],
         results["total"],
     )
+
+
+@main.command()
+@click.option("-i", "--index", help="Name of the index to bulk index records into.")
+@click.option(
+    "-s",
+    "--source",
+    type=click.Choice(VALID_SOURCES),
+    help="Source whose primary-aliased index to bulk index records into.",
+)
+@click.option("-d", "--run-date", help="Run date, formatted as YYYY-MM-DD.")
+@click.option("-rid", "--run-id", help="Run ID.")
+@click.argument("dataset_path", type=click.Path())
+@click.pass_context
+def bulk_update(
+    ctx: click.Context,
+    index: str,
+    source: str,
+    run_date: str,
+    run_id: str,
+    dataset_path: str,
+):
+    """Bulk update records for an index.
+
+    Must provide either the name of an existing index in the cluster or a valid source.
+    If source is provided, it will perform indexing and/or deletion of records for
+    the primary-aliased index for the source.
+
+    The method will read transformed records from a TIMDEXDataset
+    located at dataset_path using the 'timdex-dataset-api' library. The dataset
+    is filtered by run date and run ID.
+
+    Logs an error and aborts if the provided index doesn't exist in the cluster.
+    """
+    client = ctx.obj["CLIENT"]
+    index = helpers.validate_bulk_cli_options(index, source, client)
+
+    logger.info(f"Bulk updating records from dataset '{dataset_path}' into '{index}'")
+
+    index_results = {"created": 0, "updated": 0, "errors": 0, "total": 0}
+    delete_results = {"deleted": 0, "errors": 0, "total": 0}
+
+    td = TIMDEXDataset(location=dataset_path)
+    td.load(run_date=run_date, run_id=run_id)
+
+    # bulk index records
+    records_to_index = td.read_transformed_records_iter(action="index")
+    try:
+        index_results.update(tim_os.bulk_index(client, index, records_to_index))
+    except BulkIndexingError as exception:
+        logger.info(f"Bulk indexing failed: {exception}")
+
+    # bulk delete records
+    records_to_delete = td.read_transformed_records_iter(action="delete")
+    delete_results.update(tim_os.bulk_delete(client, index, records_to_delete))
+
+    logger.info(
+        f"Bulk update complete: {helpers.sum_bulk_update_results(index_results, delete_results)}"
+    )

tim/cli.py

@@ -83,6 +83,19 @@ def parse_deleted_records(filepath: str) -> Generator[dict, None, None]:
             yield {"timdex_record_id": item.rstrip()}
 
 
+def sum_bulk_update_results(
+    index_results: dict[str, int], delete_results: dict[str, int]
+):
+    return {
+        "created": index_results["created"],
+        "updated": index_results["updated"],
+        "deleted": delete_results["deleted"],
+        "total": index_results["total"] + delete_results["total"],
+        "bulk_index_errors": index_results["errors"],
+        "bulk_delete_errors": delete_results["errors"],
+    }
+
+
 def validate_bulk_cli_options(
     index: str | None, source: str, client: tim_os.OpenSearch
 ) -> str: