Move 'MillingStep' record creation into 'do_insert_milling_step' instead, since the 'MillingStepNameId' requires a DB query

Author: tieneupin

src/murfey/server/ispyb.py

@@ -21,6 +21,7 @@
     FoilHole,
     GridSquare,
     MillingStep,
+    MillingStepName,
     ProcessingJob,
     ProcessingJobParameter,
     Proposal,
@@ -711,9 +712,71 @@ def do_update_processing_status(self, record: AutoProcProgram, **kwargs):
             )
             return {"success": False, "return_value": None}
 
-    def do_insert_milling_step(self, record: MillingStep):
+    def do_insert_milling_step(
+        self,
+        # MillingStepName identifiers
+        recipe_name: str,
+        activity_name: str,
+        grid_square_id: int,
+        # Values
+        is_enabled: bool | None = None,
+        status: str | None = None,
+        execution_time: float | None = None,
+        stage_x: float | None = None,
+        stage_y: float | None = None,
+        stage_z: float | None = None,
+        rotation: float | None = None,
+        tilt_alpha: float | None = None,
+        beam_type: str | None = None,
+        beam_voltage: float | None = None,
+        beam_current: float | None = None,
+        milling_angle: float | None = None,
+        depth_correction: float | None = None,
+        lamella_offset: float | None = None,
+        trench_height_front: float | None = None,
+        trench_height_rear: float | None = None,
+        width_overlap_front_left: float | None = None,
+        width_overlap_front_right: float | None = None,
+        width_overlap_rear_left: float | None = None,
+        width_overlap_rear_right: float | None = None,
+    ):
         try:
             with ISPyBSession() as db:
+                # Find the ID of this MillingStep
+                milling_step_name = (
+                    db.query(MillingStepName)
+                    .filter(
+                        MillingStepName.recipe == recipe_name,
+                        MillingStepName.step == activity_name,
+                    )
+                    .one()
+                )
+                record = MillingStep(
+                    # IDs
+                    millingStepNameId=milling_step_name.millingStepNameId,
+                    gridSquareId=grid_square_id,
+                    # Values
+                    isEnabled=is_enabled,
+                    status=status,
+                    executionTime=execution_time,
+                    stageX=stage_x,
+                    stageY=stage_y,
+                    stageZ=stage_z,
+                    rotation=rotation,
+                    alphaTilt=tilt_alpha,
+                    beamType=beam_type,
+                    beamVoltage=beam_voltage,
+                    beamCurrent=beam_current,
+                    millingAngle=milling_angle,
+                    depthCorrection=depth_correction,
+                    lamellaOffset=lamella_offset,
+                    trenchHeightFront=trench_height_front,
+                    trenchHeightRear=trench_height_rear,
+                    widthOverlapFrontLeft=width_overlap_front_left,
+                    widthOverlapFrontRight=width_overlap_front_right,
+                    widthOverlapRearLeft=width_overlap_rear_left,
+                    widthOverlapRearRight=width_overlap_rear_right,
+                )
                 db.add(record)
                 db.commit()
                 log.info(f"Created MillingStep {record.millingStepId}")

src/murfey/workflows/fib/register_milling_progress.py

@@ -3,7 +3,6 @@
 from importlib.metadata import entry_points
 from typing import Any
 
-import ispyb.sqlalchemy._auto_db_schema as ISPyBDB
 from sqlmodel import Session as SQLModelSession, select
 
 import murfey.util.db as MurfeyDB
@@ -202,15 +201,28 @@ def _register_milling_step(
     if site_info.stage_info is None:
         logger.error("No stage info found in current message")
         return None
+    # Check that GridSquare has ID (for type checking)
+    if grid_square.id is None:
+        logger.error("Current GridSquare entry has no ID")
+        return None
 
     # Iteratively go through the LamellaSiteInfo model and insert for each step
     for steps, stage_name in MILLING_STEP_LOOKUP:
-        for step_name, step_id in steps:
+        for step_name, _ in steps:
             step_info: MillingStepInfo | None = site_info.steps.__getattribute__(
                 step_name
             )
+            # Early continues if key information is missing
             if step_info is None:
+                logger.debug(f"No step info found for {step_name}")
+                continue
+            if step_info.recipe_name is None:
+                logger.debug(f"No recipe name found for {step_name}")
                 continue
+            if step_info.step_name is None:
+                logger.debug(f"No step name found for {step_name}")
+                continue
+
             stage_values: StagePositionValues | None = (
                 site_info.stage_info.__getattribute__(stage_name)
             )
@@ -227,33 +239,33 @@ def _register_milling_step(
 
             if milling_step_entry is None:
                 # Create a new ISPyB entry if no Murfey one is found
-                record = ISPyBDB.MillingStep(
+                result = _transport_object.do_insert_milling_step(
                     # IDs
-                    millingStepNameId=step_id,
-                    gridSquareId=grid_square.id,
+                    recipe_name=step_info.recipe_name,
+                    activity_name=step_info.step_name,
+                    grid_square_id=grid_square.id,
                     # Values
-                    isEnabled=step_info.is_enabled,
+                    is_enabled=step_info.is_enabled,
                     status=step_info.status,
-                    executionTime=step_info.execution_time,
-                    stageX=stage_values.x,
-                    stageY=stage_values.y,
-                    stageZ=stage_values.z,
+                    execution_time=step_info.execution_time,
+                    stage_x=stage_values.x,
+                    stage_y=stage_values.y,
+                    stage_z=stage_values.z,
                     rotation=stage_values.rotation,
-                    alphaTilt=stage_values.tilt_alpha,
-                    beamType=step_info.beam_type,
-                    beamVoltage=step_info.voltage,
-                    beamCurrent=step_info.current,
-                    millingAngle=step_info.milling_angle,
-                    depthCorrection=step_info.depth_correction,
-                    lamellaOffset=step_info.lamella_offset,
-                    trenchHeightFront=step_info.trench_height_front,
-                    trenchHeightRear=step_info.trench_height_rear,
-                    widthOverlapFrontLeft=step_info.width_overlap_front_left,
-                    widthOverlapFrontRight=step_info.width_overlap_front_right,
-                    widthOverlapRearLeft=step_info.width_overlap_rear_left,
-                    widthOverlapRearRight=step_info.width_overlap_rear_right,
+                    tilt_alpha=stage_values.tilt_alpha,
+                    beam_type=step_info.beam_type,
+                    beam_voltage=step_info.voltage,
+                    beam_current=step_info.current,
+                    milling_angle=step_info.milling_angle,
+                    depth_correction=step_info.depth_correction,
+                    lamella_offset=step_info.lamella_offset,
+                    trench_height_front=step_info.trench_height_front,
+                    trench_height_rear=step_info.trench_height_rear,
+                    width_overlap_front_left=step_info.width_overlap_front_left,
+                    width_overlap_front_right=step_info.width_overlap_front_right,
+                    width_overlap_rear_left=step_info.width_overlap_rear_left,
+                    width_overlap_rear_right=step_info.width_overlap_rear_right,
                 )
-                result = _transport_object.do_insert_milling_step(record)
                 if result.get("return_value") is None:
                     logger.error(
                         f"No MillingStep entry created for {step_info.step_name}"