Merge pull request #179 from ral-facilities/DSEGOG-412-Matching-Shot-numbers

Dsegog 412 matching shot numbers

Author: moonraker595

operationsgateway_api/src/records/ingestion/partial_import_checks.py

@@ -40,56 +40,42 @@ def metadata_checks(self):
         ingested_metadata = self.ingested_record.metadata
         stored_metadata = self.stored_record.metadata
 
+        # Match timestamp with or without timezone info
         try:
-            time_match = (ingested_metadata.timestamp).replace(
+            time_match = ingested_metadata.timestamp.replace(
                 tzinfo=None,
             ) == stored_metadata.timestamp.replace(tzinfo=None)
         except Exception:
-            time_match = (ingested_metadata.timestamp) == stored_metadata.timestamp
+            time_match = ingested_metadata.timestamp == stored_metadata.timestamp
 
-        epac_match = (
-            ingested_metadata.epac_ops_data_version
-            == stored_metadata.epac_ops_data_version
-        )
-        shot_match = ingested_metadata.shotnum == stored_metadata.shotnum
         area_match = ingested_metadata.active_area == stored_metadata.active_area
         experiment_match = (
             ingested_metadata.active_experiment == stored_metadata.active_experiment
         )
 
-        if time_match and epac_match and shot_match and area_match and experiment_match:
+        shot_match = ingested_metadata.shotnum == stored_metadata.shotnum
+
+        # Reject if shotnum matches but timestamp doesn't
+        if shot_match and not time_match:
+            raise RejectRecordError("a record with this shotnum already exists")
+
+        # Accept merge if everything matches
+        if shot_match and time_match and area_match and experiment_match:
             log.info("record metadata matches existing record perfectly")
             return "accept_merge"
 
-        elif (
-            time_match
-            and not epac_match
-            and not shot_match
-            and not area_match
-            and not experiment_match
-        ):
-            raise RejectRecordError("timestamp matches, other metadata does not")
-
-        elif (
-            shot_match
-            and not time_match
-            and not epac_match
-            and not area_match
-            and not experiment_match
-        ):
-            raise RejectRecordError("shotnum matches, other metadata does not")
-
-        elif not time_match and not shot_match:
+        # Accept as new record if both timestamp and shotnum are unique
+        if not time_match and not shot_match:
             return "accept_new"
 
-        else:
-            log.error(
-                "Metadata for file being ingested: %s, metadata for file stored in"
-                " database: %s",
-                ingested_metadata,
-                stored_metadata,
-            )
-            raise RejectRecordError("inconsistent metadata")
+        # All other cases are inconsistent
+        log.error(
+            "Metadata for file being ingested: %s, metadata for file "
+            "stored in database: %s",
+            ingested_metadata,
+            stored_metadata,
+        )
+        raise RejectRecordError("inconsistent metadata")
 
     def channel_checks(
         self,

operationsgateway_api/src/records/record.py

@@ -97,6 +97,7 @@ async def update(self) -> None:
 
         for metadata_key, value in self.record.metadata.model_dump(
             exclude_unset=True,
+            exclude={"epac_ops_data_version"},
         ).items():
             await MongoDBInterface.update_one(
                 "records",
@@ -149,6 +150,30 @@ async def find_existing_record(self) -> Union[RecordModel, None]:
         else:
             return None
 
+    async def find_record_by_shotnum(self) -> Union[RecordModel, None]:
+        """
+        Searches the MongoDB 'records' collection for a record with the given
+        shot number. Returns a `RecordModel` if found, otherwise returns `None`.
+        """
+
+        log.debug(
+            "Querying MongoDB to see if a shotnum is already stored in the database",
+        )
+
+        record_dict = await MongoDBInterface.find_one(
+            "records",
+            filter_={"metadata.shotnum": self.record.metadata.shotnum},
+        )
+
+        if record_dict:
+            return RecordModel(
+                _id=record_dict["_id"],
+                metadata=record_dict["metadata"],
+                channels=record_dict["channels"],
+            )
+        else:
+            return None
+
     @staticmethod
     async def find_record(
         conditions: Dict[str, Any],

operationsgateway_api/src/routes/ingest_data.py

@@ -109,7 +109,10 @@ async def submit_hdf(
     ) = await hdf_handler.extract_data()
 
     record_original = Record(record_data)
+    # a record is deemed existing in the db if the timestamp or shotnum exists
     stored_record = await record_original.find_existing_record()
+    if stored_record is None and record_original.record.metadata.shotnum is not None:
+        stored_record = await record_original.find_record_by_shotnum()
 
     file_checker = FileChecks(record_data)
     warning = file_checker.epac_data_version_checks()

test/endpoints/test_submit_hdf.py

@@ -271,37 +271,6 @@ async def test_hdf_rejects(
         assert test_response.json() == expected_response
         assert test_response.status_code == 400
 
-    @pytest.mark.asyncio
-    async def test_partial_import_record_reject(
-        self,
-        reset_record_storage,
-        test_app: TestClient,
-        login_and_get_token,
-    ):
-        _ = await create_test_hdf_file()
-
-        test_file = "test.h5"
-        files = {"file": (test_file, open(test_file, "rb"))}
-
-        test_app.post(
-            "/submit/hdf",
-            headers={"Authorization": f"Bearer {login_and_get_token}"},
-            files=files,
-        )
-
-        _ = await create_test_hdf_file(shotnum=["valid", "missing"])
-
-        files = {"file": (test_file, open(test_file, "rb"))}
-
-        test_response = test_app.post(
-            "/submit/hdf",
-            headers={"Authorization": f"Bearer {login_and_get_token}"},
-            files=files,
-        )
-
-        assert test_response.json() == {"detail": "inconsistent metadata"}
-        assert test_response.status_code == 400
-
     @pytest.mark.asyncio
     async def test_channel_all_fail(
         self,

test/records/ingestion/create_test_hdf.py

@@ -127,8 +127,10 @@ async def create_test_hdf_file(  # noqa: C901
         if shotnum[1] == "exists":
             if shotnum[0] == "invalid":
                 record.attrs.create("shotnum", "string")
-            else:
+            elif shotnum[0] == "valid":
                 record.attrs.create("shotnum", 366272, dtype="u8")
+            elif shotnum[0]:
+                record.attrs.create("shotnum", int(shotnum[0]), dtype="u8")
         if active_area[1] == "exists":
             record.attrs.create("active_area", active_area[0])
         if active_experiment[1] == "exists":

test/records/ingestion/test_channel.py

@@ -403,7 +403,7 @@ async def test_required_attribute_checks(
         with h5py.File("test.h5", "w") as f:
             f.attrs.create("epac_ops_data_version", "1.2")
             record = f["/"]
-            record.attrs.create("timestamp", "2020-04-07 14:28:16")
+            record.attrs.create("timestamp", "2020-04-07T14:28:16Z")
             record.attrs.create("shotnum", 366272, dtype="u8")
             record.attrs.create("active_area", "ea1")
             record.attrs.create("active_experiment", "90097341")
@@ -564,7 +564,7 @@ async def test_optional_attribute_vector(
         with h5py.File("test.h5", "w") as f:
             f.attrs.create("epac_ops_data_version", "1.2")
             record = f["/"]
-            record.attrs.create("timestamp", "2020-04-07 14:28:16")
+            record.attrs.create("timestamp", "2020-04-07T14:28:16Z")
             record.attrs.create("shotnum", 366272, dtype="u8")
             record.attrs.create("active_area", "ea1")
             record.attrs.create("active_experiment", "90097341")
@@ -724,7 +724,7 @@ async def test_dataset_checks_float_image(self, remove_hdf_file: None):
         with h5py.File("test.h5", "w") as f:
             f.attrs.create("epac_ops_data_version", "1.2")
             record = f["/"]
-            record.attrs.create("timestamp", "2020-04-07 14:28:16")
+            record.attrs.create("timestamp", "2020-04-07T14:28:16Z")
             record.attrs.create("shotnum", 366272, dtype="u8")
             record.attrs.create("active_area", "ea1")
             record.attrs.create("active_experiment", "90097341")

test/records/ingestion/test_hdf_handler.py

@@ -2,7 +2,6 @@
 
 import h5py
 import numpy as np
-from pydantic import ValidationError
 import pytest
 
 from operationsgateway_api.src.exceptions import HDFDataExtractionError
@@ -154,7 +153,7 @@ async def test_extract_channel_unexpected_group(
         with h5py.File("test.h5", "w") as f:
             f.attrs.create("epac_ops_data_version", "1.2")
             record = f["/"]
-            record.attrs.create("timestamp", "2020-04-07 14:28:16")
+            record.attrs.create("timestamp", "2020-04-07T14:28:16Z")
             record.attrs.create("shotnum", 366272, dtype="u8")
             record.attrs.create("active_area", "ea1")
             record.attrs.create("active_experiment", "90097341")
@@ -212,7 +211,7 @@ async def test_extract_channel_data_missing(
         with h5py.File("test.h5", "w") as f:
             f.attrs.create("epac_ops_data_version", "1.2")
             record = f["/"]
-            record.attrs.create("timestamp", "2020-04-07 14:28:16")
+            record.attrs.create("timestamp", "2020-04-07T14:28:16Z")
             record.attrs.create("shotnum", 366272, dtype="u8")
             record.attrs.create("active_area", "ea1")
             record.attrs.create("active_experiment", "90097341")