Address comments in PR #373

* Use logger.error when handling custom bulk exceptions
* Call ctx.exit in bulk_update_embeddings when BulkOperationError is raisd
* Update method to log elapsed time
* Clean up log when refreshing index

Author: jonavellecuerdo

tim/cli.py

@@ -63,14 +63,13 @@ def main(ctx: click.Context, url: str, *, verbose: bool) -> None:
     ctx.obj["CLIENT"] = tim_os.configure_opensearch_client(url)
     logger.info("OpenSearch client configured for endpoint '%s'", url)
 
+    def _log_command_elapsed_time() -> None:
+        elapsed_time = perf_counter() - ctx.obj["start_time"]
+        logger.info(
+            "Total time to complete process: %s", str(timedelta(seconds=elapsed_time))
+        )
 
-@main.result_callback()
-@click.pass_context
-def log_process_time(ctx: click.Context, _result: object, **_kwargs: dict) -> None:
-    elapsed_time = perf_counter() - ctx.obj["START_TIME"]
-    logger.info(
-        "Total time to complete process: %s", str(timedelta(seconds=elapsed_time))
-    )
+    ctx.call_on_close(_log_command_elapsed_time)
 
 
 # Cluster commands
@@ -310,7 +309,7 @@ def bulk_update(
     try:
         index_results.update(tim_os.bulk_index(client, index, records_to_index))
     except BulkIndexingError as exception:
-        logger.info(f"Bulk indexing failed: {exception}")
+        logger.error(f"Bulk indexing failed: {exception}")  # noqa: TRY400
 
     # bulk delete records
     records_to_delete = td.read_dicts_iter(
@@ -343,7 +342,7 @@ def bulk_update(
         "records with embeddings."
     ),
 )
-@click.option("-rid", "--run-id", help="Run ID.")
+@click.option("-rid", "--run-id", required=True, help="Run ID.")
 @click.argument("dataset_path", type=click.Path())
 @click.pass_context
 def bulk_update_embeddings(
@@ -391,7 +390,8 @@ def bulk_update_embeddings(
     try:
         update_results.update(tim_os.bulk_update(client, index, embeddings_to_index))
     except BulkOperationError as exception:
-        logger.info(f"Bulk update with embeddings failed: {exception}")
+        logger.error(f"Bulk update with embeddings failed: {exception}")  # noqa: TRY400
+        ctx.exit(1)
 
     logger.info(f"Bulk update with embeddings complete: {json.dumps(update_results)}")
 
@@ -471,7 +471,3 @@ def reindex_source(
 
     summary_results = {"index": index_results}
     logger.info(f"Reindex source complete: {json.dumps(summary_results)}")
-
-
-if __name__ == "__main__":
-    main()

tim/opensearch.py

@@ -352,7 +352,7 @@ def bulk_delete(
         result["total"] += 1
         if result["total"] % int(os.getenv("STATUS_UPDATE_INTERVAL", "1000")) == 0:
             logger.info("Status update: %s records deleted so far!", result["total"])
-    logger.info("All records deleted, refreshing index.")
+    logger.info("Refreshing index.")
     response = client.indices.refresh(
         index=index,
     )
@@ -412,7 +412,7 @@ def bulk_index(client: OpenSearch, index: str, records: Iterator[dict]) -> dict[
         result["total"] += 1
         if result["total"] % int(os.getenv("STATUS_UPDATE_INTERVAL", "1000")) == 0:
             logger.info("Status update: %s records indexed so far!", result["total"])
-    logger.info("All records ingested, refreshing index.")
+    logger.info("Refreshing index.")
     response = client.indices.refresh(
         index=index,
     )
@@ -467,4 +467,9 @@ def bulk_update(
         result["total"] += 1
         if result["total"] % int(os.getenv("STATUS_UPDATE_INTERVAL", "1000")) == 0:
             logger.info("Status update: %s records updated so far!", result["total"])
+    logger.info("Refreshing index.")
+    response = client.indices.refresh(
+        index=index,
+    )
+    logger.debug(response)
     return result