diff --git a/src/aare/daq/aaredb.py b/src/aare/daq/aaredb.py index 763e4ca6..77eaeee7 100644 --- a/src/aare/daq/aaredb.py +++ b/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]):