[wip]

Author: jonavellecuerdo

tests/test_opensearch.py

@@ -459,7 +459,7 @@ def test_bulk_index_creates_records(
     test_opensearch_client, five_valid_index_libguides_records
 ):
     assert tim_os.bulk_index(
-        test_opensearch_client, "test-index", five_valid_index_libguides_records
+        test_opensearch_client, "test-index", five_valid_index_libguides_records, "index"
     ) == {
         "created": 5,
         "updated": 0,
@@ -474,22 +474,22 @@ def test_bulk_index_updates_records(
 ):
     monkeypatch.setenv("STATUS_UPDATE_INTERVAL", "5")
     assert tim_os.bulk_index(
-        test_opensearch_client, "test-index", five_valid_index_libguides_records
+        test_opensearch_client, "test-index", five_valid_index_libguides_records, "index"
     ) == {
         "created": 0,
         "updated": 5,
         "errors": 0,
         "total": 5,
     }
-    assert "Status update: 5 records indexed so far!" in caplog.text
+    assert "Status update: 5 records processed so far!" in caplog.text
 
 
 @my_vcr.use_cassette("opensearch/bulk_index_record_mapper_parsing_error.yaml")
 def test_bulk_index_logs_mapper_parsing_errors(
     caplog, test_opensearch_client, one_invalid_index_libguides_records
 ):
     assert tim_os.bulk_index(
-        test_opensearch_client, "test-index", one_invalid_index_libguides_records
+        test_opensearch_client, "test-index", one_invalid_index_libguides_records, "index"
     ) == {
         "created": 0,
         "updated": 0,

tim/cli.py

@@ -308,7 +308,9 @@ def bulk_update(
         action="index",
     )
     try:
-        index_results.update(tim_os.bulk_index(client, index, records_to_index))
+        index_results.update(
+            tim_os.bulk_index(client, index, records_to_index, action="index")
+        )
     except BulkIndexingError as exception:
         logger.info(f"Bulk indexing failed: {exception}")
 
@@ -401,9 +403,10 @@ def bulk_update_embeddings(
         ]
     )
     # ==== END TEMPORARY CODE ====
-
     try:
-        update_results.update(tim_os.bulk_update(client, index, records_to_update))
+        update_results.update(
+            tim_os.bulk_index(client, index, records_to_update, action="update")
+        )
     except BulkIndexingError as exception:
         logger.info(f"Bulk update with embeddings failed: {exception}")
 
@@ -479,7 +482,9 @@ def reindex_source(
         action="index",
     )
     try:
-        index_results.update(tim_os.bulk_index(client, index, records_to_index))
+        index_results.update(
+            tim_os.bulk_index(client, index, records_to_index, action="index")
+        )
     except BulkIndexingError as exception:
         logger.info(f"Bulk indexing failed: {exception}")
 

tim/opensearch.py

@@ -359,22 +359,26 @@ def bulk_delete(
     return result
 
 
-def bulk_index(client: OpenSearch, index: str, records: Iterator[dict]) -> dict[str, int]:
+def bulk_index(
+    client: OpenSearch, index: str, records: Iterator[dict], action: str
+) -> dict[str, int]:
     """Indexes records into an existing index using the streaming bulk helper.
 
-    This action function uses the OpenSearch "index" action, which is a
-    combination of create and update: if a record with the same _id exists in the
-    index, it will be updated. If it does not exist, the record will be indexed as a
-    new document.
+    This method uses the OpenSearch "index" and "update" operations.
+    - Setting `action` to "index" will either create or update a record.
+      If a record with the same _id exists in the index, it will be updated;
+      if it does not exist, the record will be added as a new document.
+    - Setting `action` to "update" will update a document only if it exists
+      in the index. Otherwise, an error is raised.
 
-    If an error occurs during record indexing, it will be logged and bulk indexing will
-    continue until all records have been processed.
+    If an error occurs during the operation, it will be logged, and the bulk
+    operation will continue until all records have been processed.
 
     Returns total sums of: records created, records updated, errors, and total records
     processed.
     """
     result = {"created": 0, "updated": 0, "errors": 0, "total": 0}
-    actions = helpers.generate_bulk_actions(index, records, "index")
+    actions = helpers.generate_bulk_actions(index, records, action)
     responses = streaming_bulk(
         client,
         actions,
@@ -400,34 +404,17 @@ def bulk_index(client: OpenSearch, index: str, records: Iterator[dict]) -> dict[
             result["updated"] += 1
         else:
             logger.error(
-                "Something unexpected happened during ingest. Bulk index response: %s",
+                "Something unexpected happened during ingest. "
+                f"Bulk {action} response: %s",
                 json.dumps(response),
             )
             result["errors"] += 1
         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("Status update: %s records processed so far!", result["total"])
     logger.info("All records ingested, refreshing index.")
     response = client.indices.refresh(
         index=index,
     )
     logger.debug(response)
     return result
-
-
-def bulk_update(
-    client: OpenSearch, index: str, records: Iterator[dict]
-) -> dict[str, int]:
-    result = {"updated": 0, "errors": 0, "total": 0}
-    actions = helpers.generate_bulk_actions(index, records, "update")
-    responses = streaming_bulk(
-        client,
-        actions,
-        max_chunk_bytes=REQUEST_CONFIG["OPENSEARCH_BULK_MAX_CHUNK_BYTES"],
-        raise_on_error=False,
-    )
-    for response in responses:
-        # TODO: Parse different responses from bulk operations
-        pass
-    logger.debug(response)
-    return result