feat: update record with actual s3 location if mismatch (#146)

* feat: update record with actual s3 location if mismatch

* test: update unit tests

* feat: allow name mismatch with warning

* fix: outdated warning message

Author: helen-m-lin

src/aind_data_asset_indexer/aind_bucket_indexer.py

@@ -8,6 +8,7 @@
 import warnings
 from datetime import datetime, timedelta, timezone
 from typing import Dict, List, Optional
+from uuid import uuid4
 
 import boto3
 import dask.bag as dask_bag
@@ -388,9 +389,8 @@ def _process_docdb_record(
             docdb_record, self.job_settings.s3_bucket
         ):
             logging.warning(
-                f"Record location {docdb_record.get('location')} or name "
-                f"{docdb_record.get('name')} not valid for bucket "
-                f"{self.job_settings.s3_bucket}!"
+                f"Record location {docdb_record.get('location')} not valid "
+                f"for bucket {self.job_settings.s3_bucket}! Skipping."
             )
         else:
             s3_parts = get_s3_bucket_and_prefix(docdb_record["location"])
@@ -587,10 +587,9 @@ def _process_prefix(
             key=object_key,
         )
         if does_metadata_file_exist:
-            # If record not in DocDb, then copy it to DocDb if the location
-            # in the metadata record matches the location the record lives in
-            # Otherwise, log a warning that the metadata record location does
-            # not make sense.
+            # If record not in DocDb, then copy it to DocDb.
+            # If location in the record does not match the actual s3 location,
+            # first log a warning and correct the record.
             s3_full_location = get_s3_location(bucket, object_key)
             if record_id is None:
                 json_contents = download_json_file_from_s3(
@@ -600,40 +599,44 @@ def _process_prefix(
                 )
                 if json_contents:
                     # noinspection PyTypeChecker
-                    if is_record_location_valid(
+                    if not is_record_location_valid(
                         json_contents,
                         expected_bucket=bucket,
                         expected_prefix=s3_prefix,
                     ):
-                        if "_id" in json_contents:
-                            logging.info(
-                                f"Adding record to docdb for: {location}"
-                            )
-                            response = docdb_client.insert_one_docdb_record(
-                                record=json_contents
-                            )
-                            logging.debug(response.json())
-                            cond_copy_then_sync_core_json_files(
-                                metadata_json=json.dumps(
-                                    json_contents, default=str
-                                ),
-                                bucket=bucket,
-                                prefix=s3_prefix,
-                                s3_client=s3_client,
-                                copy_original_md_subdir=(
-                                    self.job_settings.copy_original_md_subdir
-                                ),
-                            )
-                        else:
-                            logging.warning(
-                                f"Metadata record for {location} "
-                                f"does not have an _id field!"
-                            )
-                    else:
                         logging.warning(
                             f"Location field {json_contents.get('location')} "
-                            f"or name field {json_contents.get('name')} does "
-                            f"not match actual location of record {location}!"
+                            "does not match actual location of record "
+                            f"{location}! Updating record with correct "
+                            "location and new id."
+                        )
+                        json_contents.update(
+                            {
+                                "_id": str(uuid4()),
+                                "location": location,
+                            }
+                        )
+                    if "_id" in json_contents:
+                        logging.info(f"Adding record to docdb for: {location}")
+                        response = docdb_client.insert_one_docdb_record(
+                            record=json_contents
+                        )
+                        logging.debug(response.json())
+                        cond_copy_then_sync_core_json_files(
+                            metadata_json=json.dumps(
+                                json_contents, default=str
+                            ),
+                            bucket=bucket,
+                            prefix=s3_prefix,
+                            s3_client=s3_client,
+                            copy_original_md_subdir=(
+                                self.job_settings.copy_original_md_subdir
+                            ),
+                        )
+                    else:
+                        logging.warning(
+                            f"Metadata record for {location} "
+                            f"does not have an _id field!"
                         )
                 else:
                     logging.warning(

src/aind_data_asset_indexer/utils.py

@@ -122,7 +122,9 @@ def is_record_location_valid(
       like 's3://{expected_bucket}/prefix'
       Will return False if there is no s3 scheme, the bucket does not match
       the expected bucket, the prefix contains forward slashes, or the prefix
-      is invalid, or doesn't match the record name or expected prefix.
+      is invalid or does not match the expected prefix.
+      If the record name does not match the prefix, a warning is logged,
+      but the method will still return True.
 
     """
     expected_stripped_prefix = (
@@ -142,15 +144,18 @@ def is_record_location_valid(
                 stripped_prefix == ""
                 or len(stripped_prefix.split("/")) > 1
                 or not is_prefix_valid(stripped_prefix)
-                or record.get("name") != stripped_prefix
                 or (
                     expected_prefix is not None
                     and stripped_prefix != expected_stripped_prefix
                 )
             ):
                 return False
-            else:
-                return True
+            elif record.get("name") != stripped_prefix:
+                logging.warning(
+                    f"Record name {record.get('name')} does not match "
+                    f"prefix {stripped_prefix}."
+                )
+            return True
 
 
 def compute_md5_hash(json_contents: str) -> str:

tests/test_aind_bucket_indexer.py

@@ -759,9 +759,8 @@ def test_process_docdb_record_invalid_location(
             )
         expected_log_messages = [
             "WARNING:root:Record location "
-            "no_s3/bucket/prefix1_2024-01-01_01-01-01 or name "
-            "prefix1_2024-01-01_01-01-01 not valid for bucket "
-            "aind-ephys-data-dev-u5u0i5!"
+            "no_s3/bucket/prefix1_2024-01-01_01-01-01 not valid "
+            "for bucket aind-ephys-data-dev-u5u0i5! Skipping."
         ]
         self.assertEqual(expected_log_messages, captured.output)
         self.assertIsNone(docdb_id_to_delete)
@@ -787,8 +786,8 @@ def test_process_docdb_record_invalid_prefix(
                 },
             )
         expected_log_messages = [
-            "WARNING:root:Record location s3://bucket/prefix1 or name prefix1 "
-            "not valid for bucket aind-ephys-data-dev-u5u0i5!"
+            "WARNING:root:Record location s3://bucket/prefix1 "
+            "not valid for bucket aind-ephys-data-dev-u5u0i5! Skipping."
         ]
         self.assertEqual(expected_log_messages, captured.output)
         self.assertIsNone(docdb_id_to_delete)
@@ -1444,6 +1443,7 @@ def test_process_prefix_no_record_yes_file_good_file_no__id(
         mock_cond_copy_then_sync_core_json_files.assert_not_called()
         mock_upload_metadata_json_str_to_s3.assert_not_called()
 
+    @patch("aind_data_asset_indexer.aind_bucket_indexer.uuid4")
     @patch(
         "aind_data_asset_indexer.aind_bucket_indexer."
         "upload_metadata_json_str_to_s3"
@@ -1467,6 +1467,7 @@ def test_process_prefix_no_record_yes_file_good_file_bad_location(
         mock_download_json_file_from_s3: MagicMock,
         mock_cond_copy_then_sync_core_json_files: MagicMock,
         mock_upload_metadata_json_str_to_s3: MagicMock,
+        mock_uuid: MagicMock,
     ):
         """Tests _process_prefix method when there is no record in DocDb,
         there is and there is metadata.nd.json file in S3, and the file can
@@ -1475,12 +1476,15 @@ def test_process_prefix_no_record_yes_file_good_file_bad_location(
         mock_does_s3_object_exist.return_value = True
         # Test what happens when the location in the record does not match the
         # expected location
-        mocked_downloaded_record = deepcopy(self.example_md_record)
-        mocked_downloaded_record["location"] = (
-            f"s3://{self.basic_job.job_settings.s3_bucket}/"
-            f"ecephys_642478_2020-01-10_10-10-10"
+        mock_download_json_file_from_s3.return_value = {
+            "_id": "5ca4a951-d374-4f4b-8279-d570a35b2286",
+            "name": "ecephys_642478_2023-01-17_13-56-29",
+            "location": "s3://bucket/ecephys_642478_2020-01-10_10-10-10",
+        }
+        mock_uuid.return_value = "a62344ff-1cec-48f4-914e-7482797e6332"
+        mock_docdb_client.insert_one_docdb_record.return_value = MagicMock(
+            status_code=200, json=MagicMock(return_value="inserted")
         )
-        mock_download_json_file_from_s3.return_value = mocked_downloaded_record
 
         location_to_id_map = dict()
         with self.assertLogs(level="DEBUG") as captured:
@@ -1490,16 +1494,34 @@ def test_process_prefix_no_record_yes_file_good_file_bad_location(
                 s3_client=mock_s3_client,
                 location_to_id_map=location_to_id_map,
             )
+        actual_location = (
+            f"s3://{self.basic_job.job_settings.s3_bucket}/"
+            "ecephys_642478_2023-01-17_13-56-29"
+        )
         expected_log_messages = [
-            "WARNING:root:Location field s3://aind-ephys-data-dev-u5u0i5/"
-            "ecephys_642478_2020-01-10_10-10-10 or name field "
-            "ecephys_642478_2023-01-17_13-56-29 does not match actual "
-            "location of record s3://aind-ephys-data-dev-u5u0i5/"
-            "ecephys_642478_2023-01-17_13-56-29!"
+            "WARNING:root:Location field s3://bucket/"
+            "ecephys_642478_2020-01-10_10-10-10 does not match actual "
+            f"location of record {actual_location}! Updating "
+            "record with correct location and new id.",
+            f"INFO:root:Adding record to docdb for: {actual_location}",
+            "DEBUG:root:inserted",
         ]
+        expected_record = {
+            "_id": mock_uuid.return_value,
+            "name": "ecephys_642478_2023-01-17_13-56-29",
+            "location": actual_location,
+        }
         self.assertEqual(expected_log_messages, captured.output)
-        mock_docdb_client.insert_one_docdb_record.assert_not_called()
-        mock_cond_copy_then_sync_core_json_files.assert_not_called()
+        mock_docdb_client.insert_one_docdb_record.assert_called_once_with(
+            record=expected_record
+        )
+        mock_cond_copy_then_sync_core_json_files.assert_called_once_with(
+            metadata_json=json.dumps(expected_record),
+            bucket=self.basic_job.job_settings.s3_bucket,
+            prefix="ecephys_642478_2023-01-17_13-56-29",
+            s3_client=mock_s3_client,
+            copy_original_md_subdir="original_metadata",
+        )
         mock_upload_metadata_json_str_to_s3.assert_not_called()
 
     @patch("aind_data_asset_indexer.aind_bucket_indexer.does_s3_object_exist")

tests/test_utils.py

@@ -289,20 +289,25 @@ def test_is_record_location_valid_false4(self):
             )
         )
 
-    def test_is_record_location_valid_false5(self):
-        """Tests is_record_location_valid returns false when prefix does
+    def test_is_record_location_valid_true2(self):
+        """Tests is_record_location_valid returns true when prefix does
         not match name"""
 
         example_record = {
             "_id": "abc-123",
             "name": "prefix2_2024-01-01_01-01-01",
             "location": "s3://some_bucket/prefix1_2024-01-01_01-01-01",
         }
-        self.assertFalse(
-            is_record_location_valid(
+        with self.assertLogs(level="DEBUG") as captured:
+            result = is_record_location_valid(
                 example_record, expected_bucket="some_bucket"
             )
-        )
+        expected_log_messages = [
+            "WARNING:root:Record name prefix2_2024-01-01_01-01-01 does not "
+            "match prefix prefix1_2024-01-01_01-01-01."
+        ]
+        self.assertEqual(expected_log_messages, captured.output)
+        self.assertTrue(result)
 
     def test_is_record_location_valid_false6(self):
         """Tests is_record_location_valid returns false when prefixes don't