raster: catch diffracton_iamge_get error, but allow automation etc to continue

src/aare/daq/daq.py

@@ -11,6 +11,7 @@ from typing import List, Tuple, Optional, Callable
 import cv2
 import numpy as np
 from aareDB import SampleEventType
+from jfjoch_client.exceptions import NotFoundException
 
 from jfjoch_client import ScanResult, ScanResultImagesInner
 
@@ -1229,6 +1230,56 @@ class AareDAQ:
             centre_of_mass=None,
         )
 
+    @staticmethod
+    def _scan_result_image_count(scan_result: ScanResult | None) -> int | None:
+        if scan_result is None or scan_result.images is None:
+            return None
+        return len(scan_result.images)
+
+    def _upload_raster_diffraction_preview(
+            self,
+            *,
+            sample_id: int,
+            filename: str,
+            image_id: int,
+            scan_result: ScanResult | None,
+            request: RasterGridRequest,
+    ) -> None:
+        image_count = self._scan_result_image_count(scan_result)
+
+        if image_count is not None and not (0 <= image_id < image_count):
+            logger.warning(
+                "Skipping diffraction image upload because image id is outside scan result range",
+                extra=merge_log_context(
+                    sample_log_context(self.sample),
+                    raster_request_log_context(request),
+                    {
+                        "diffraction_image_id": image_id,
+                        "image_count": image_count,
+                    },
+                ),
+            )
+            return
+
+        try:
+            diffraction_image = self.__jfjoch.get_diffraction_image(image_id)
+        except NotFoundException:
+            logger.warning(
+                "JFJoch diffraction preview image was not found after raster; continuing without upload",
+                extra=merge_log_context(
+                    sample_log_context(self.sample),
+                    raster_request_log_context(request),
+                    {
+                        "diffraction_image_id": image_id,
+                        "image_count": image_count,
+                    },
+                ),
+                exc_info=True,
+            )
+            return
+
+        self.__aare.upload_jpg(sample_id, filename, diffraction_image)
+
     @staticmethod
     def _grid_image_id_from_centre_offset(
             *,
@@ -1398,8 +1449,14 @@ class AareDAQ:
                         y_mm=y,
                         request=request,
                     )
-                diffraction_image = self.__jfjoch.get_diffraction_image(diffraction_image_id)
-                self.__aare.upload_jpg(self.sample.db_id, diffraction_image_filename, diffraction_image)
+
+                self._upload_raster_diffraction_preview(
+                    sample_id=self.sample.db_id,
+                    filename=diffraction_image_filename,
+                    image_id=diffraction_image_id,
+                    scan_result=scan_result,
+                    request=request,
+                )
                 self.__aare.send_sample_event(
                     self.sample,
                     event_type=SampleEventType.RASTERED,