Merge pull request #409 from AllenNeuralDynamics/add-replace-flag-for-instruments

Add replace flag for instruments

Author: dougollerenshaw

src/aind_metadata_mapper/fip/make_instrument.py

@@ -36,7 +36,12 @@
 from aind_data_schema.components.devices import Computer
 from aind_data_schema_models.modalities import Modality
 
-from aind_metadata_mapper.utils import check_instrument_id, prompt_for_string, save_instrument
+from aind_metadata_mapper.utils import (
+    check_existing_instrument,
+    check_instrument_id,
+    prompt_for_string,
+    save_instrument,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -449,9 +454,25 @@ def main(
         input_func=input_func,
     )
 
+    # Check if instrument with same ID and date already exists
+    modification_date_str = instrument_model.modification_date.isoformat()
+    record_exists = check_existing_instrument(instrument_id, modification_date_str)
+    replace = False
+    if record_exists:
+        logger.info(
+            f"An instrument with ID '{instrument_id}' and modification_date '{modification_date_str}' already exists."
+        )
+        response = input_func("Do you want to overwrite the existing record? [y/N]: ").strip().lower()
+        if response in ("y", "yes"):
+            replace = True
+            logger.info("Will overwrite existing record.")
+        else:
+            logger.info("Cancelled. Not overwriting existing record.")
+            sys.exit(0)
+
     # Save instrument (includes round-trip validation)
     try:
-        save_instrument(instrument_model)
+        save_instrument(instrument_model, replace=replace)
     except ValueError as e:
         logger.error(f"Failed - {e}")
         sys.exit(1)

src/aind_metadata_mapper/utils.py

@@ -13,7 +13,9 @@
 API_BASE_URL = "http://aind-metadata-service/api/v2/instrument"
 
 
-def get_instrument(instrument_id: str, modification_date: Optional[str] = None) -> Optional[dict]:  # pragma: no cover
+def get_instrument(
+    instrument_id: str, modification_date: Optional[str] = None, suppress_warning: bool = False
+) -> Optional[dict]:  # pragma: no cover
     """Get instrument.
 
     Gets the latest record by default, or a specific record if modification_date is provided.
@@ -24,6 +26,8 @@ def get_instrument(instrument_id: str, modification_date: Optional[str] = None)
         Instrument identifier.
     modification_date : Optional[str]
         Optional modification date (YYYY-MM-DD format). If None, returns latest record.
+    suppress_warning : bool
+        If True, suppress warning when modification_date not found.
 
     Returns
     -------
@@ -48,21 +52,22 @@ def get_instrument(instrument_id: str, modification_date: Optional[str] = None)
         for record in matching_records:
             if record.get("modification_date") == modification_date:
                 return record
-        # No matching date found - log available dates
-        available_dates = sorted(
-            set(r.get("modification_date") for r in matching_records if r.get("modification_date"))
-        )
-        logger.warning(
-            f"No record found for instrument_id '{instrument_id}' with modification_date '{modification_date}'. "
-            f"Available dates: {available_dates}"
-        )
+        # No matching date found - log available dates (unless suppressed)
+        if not suppress_warning:
+            available_dates = sorted(
+                set(r.get("modification_date") for r in matching_records if r.get("modification_date"))
+            )
+            logger.warning(
+                f"No record found for instrument_id '{instrument_id}' with modification_date '{modification_date}'. "
+                f"Available dates: {available_dates}"
+            )
         return None
     else:
         # Return latest record
         return sorted(matching_records, key=lambda record: record["modification_date"])[-1]
 
 
-def save_instrument(instrument_model: instrument.Instrument) -> None:  # pragma: no cover
+def save_instrument(instrument_model: instrument.Instrument, replace: bool = False) -> None:  # pragma: no cover
     """Save instrument and validate round-trip.
 
     Saves the instrument, then retrieves it back and verifies that what we get back
@@ -72,6 +77,8 @@ def save_instrument(instrument_model: instrument.Instrument) -> None:  # pragma:
     ----------
     instrument_model : instrument.Instrument
         Instrument to POST.
+    replace : bool
+        If True, overwrite existing record with same instrument_id and modification_date.
 
     Raises
     ------
@@ -83,7 +90,8 @@ def save_instrument(instrument_model: instrument.Instrument) -> None:  # pragma:
     # Use model_dump_json() and parse to ensure dates are properly serialized
     source_dict = json.loads(instrument_model.model_dump_json())
     logger.info(f"POSTing instrument to {API_BASE_URL}")
-    response = requests.post(API_BASE_URL, json=source_dict)
+    params = {"replace": "true"} if replace else {}
+    response = requests.post(API_BASE_URL, json=source_dict, params=params)
     # POST 400 is always an error (e.g., "Record already exists")
     if response.status_code == 400:
         error_msg = response.json().get("message", response.text)
@@ -108,6 +116,28 @@ def save_instrument(instrument_model: instrument.Instrument) -> None:  # pragma:
         raise ValueError("Round-trip test failed: Source and read-back instruments differ")
 
 
+def check_existing_instrument(
+    instrument_id: str,
+    modification_date: str,
+) -> bool:  # pragma: no cover
+    """Check if an instrument with the same ID and modification_date already exists.
+
+    Parameters
+    ----------
+    instrument_id : str
+        Instrument identifier.
+    modification_date : str
+        Modification date (YYYY-MM-DD format).
+
+    Returns
+    -------
+    bool
+        True if a record with the same instrument_id and modification_date exists.
+    """
+    existing = get_instrument(instrument_id, modification_date=modification_date, suppress_warning=True)
+    return existing is not None
+
+
 def check_instrument_id(
     instrument_id: str,
     skip_confirmation: bool = False,