aaredb: added a timeing debug wrapper to aaredb, to track responses. Removed ToDos!

src/aare/daq/aaredb.py

@@ -1,7 +1,9 @@
 import datetime
+import functools
 import io
 import json
 import os
+import time
 from typing import List, Optional
 
 import aareDB
@@ -38,18 +40,23 @@ from jfjoch_client.models import ScanResult
 
 logger = setup_logger("aareDAQ")
 
-#TODO which additional events to add?
-# - Mounted
-# - Mount Failed
-# - Unmounted
-# - Unmounted Failed
-# - Centered
-# - Collected
-# - Collection Failed
-# - Failed
-# - AXCFailed
-# - ALCFailed
-# - Lost
+def time_db_call(func):
+    """Decorator to time AareDB calls and log the duration."""
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        start = time.perf_counter()
+        logger.debug(f"Starting AareDB call: {func.__name__}")
+        try:
+            result = func(*args, **kwargs)
+            duration = time.perf_counter() - start
+            logger.debug(f"Finished AareDB call: {func.__name__} in {duration:.4f} seconds")
+            return result
+        except Exception as e:
+            duration = time.perf_counter() - start
+            logger.debug(f"AareDB call: {func.__name__} FAILED after {duration:.4f} seconds with error: {e}")
+            raise
+    return wrapper
+
 
 class AareWrapper:
     def __init__(
@@ -84,6 +91,7 @@ class AareWrapper:
         self.__cert_file = configuration.cert_file
         self.__key_file = configuration.key_file
 
+    @time_db_call
     def set_pucks_beamline(self, input_list: List[PuckLoadedInfo]):
         o = []
 
@@ -100,6 +108,7 @@ class AareWrapper:
         )
         logger.debug(ret)
 
+    @time_db_call
     def create_manual_sample(self, s: SampleShortInfo):
         from aareDB.models import ManualSampleCreate
 
@@ -114,6 +123,7 @@ class AareWrapper:
         except Exception as e:
             logger.error(f"Error inserting sample: {e}")
 
+    @time_db_call
     def send_sample_event(
         self,
         s: Optional[SampleShortInfo],
@@ -139,6 +149,7 @@ class AareWrapper:
         except Exception as e:
             logger.error(f"Error sending sample event {event_type!s} to db: {e}")
 
+    @time_db_call
     def upload_image(self, sample_id: int, filename: str, bgr_image: np.ndarray, message: Optional[str] = None):
         _, buffer = cv2.imencode('.jpg', bgr_image)
         jpeg_bytes = io.BytesIO(buffer)
@@ -156,10 +167,10 @@ class AareWrapper:
         }
         if message is not None:
             request_kwargs["data"] = {"comment": message}
-
         response = requests.post(url, **request_kwargs)
         logger.debug(f"Response status code: {response.status_code}")
 
+    @time_db_call
     def upload_jpg(self, sample_id: int, filename: str, jpg_image):
         logger.debug(f"jppg_image of type: {type(jpg_image)}")
         url = f"{self.__host}/protected_router/sample_runner/{sample_id}/upload-images"
@@ -174,6 +185,7 @@ class AareWrapper:
                                  headers=headers)
         logger.debug(f"Response status code: {response.status_code}")
 
+    @time_db_call
     def create_rotation_run(self, s: Optional[SampleShortInfo], r:RotationScanRequest, d:DAQStatusModel):
         if s is None:
             return
@@ -249,6 +261,7 @@ class AareWrapper:
         except Exception as e:
             logger.error(e)
 
+    @time_db_call
     def create_gridscan_run(self, s: Optional[SampleShortInfo], r:RasterGridRequest, d:DAQStatusModel):
         if s is None:
             return
@@ -313,6 +326,7 @@ class AareWrapper:
         except Exception as e:
             logger.debug(e)
 
+    @time_db_call
     def ingest_gridscan(self, sample: Optional[SampleShortInfo], raster_result: ScanResult,
                         raster_request: RasterGridRequest, geom: SampleGeometryModel,
                         com: Optional[CenterOfMassModel], beam_mark_pxl:tuple[float,float]):
@@ -383,6 +397,7 @@ class AareWrapper:
             logger.error(e)
             raise e
 
+    @time_db_call
     def ingest_scan(self, sample: Optional[SampleShortInfo], result: ScanResult,
                         geom: SampleGeometryModel, beam_mark_pxl:tuple[float,float]):