[PRMP-570] Step Function Error Handling (#863)

Author: SWhyteAnswer

lambdas/handlers/migration_dynamodb_handler.py

@@ -45,7 +45,7 @@ def extract_table_info(event):
 
 
 def validate_event_input(event):
-    required_fields = ["segment", "totalSegments", "migrationScript"]
+    required_fields = ["segment", "totalSegments", "migrationScript", "executionId"]
     for field in required_fields:
         if field not in event:
             raise ValueError(f"Missing required field: '{field}' in event")
@@ -67,18 +67,23 @@ def validate_event_input(event):
     if not migration_script:
         raise ValueError("'migrationScript' cannot be empty")
 
-    run_migration = bool(event.get("run_migration", False))
+    run_migration = bool(event.get("runMigration", False))
 
     table_name, environment, region = extract_table_info(event)
 
+    execution_id = event.get("executionId")
+    if not execution_id:
+        raise ValueError("'executionId' cannot be empty")
+
     return (
         segment,
         total_segments,
         table_name,
         environment,
         region,
-        run_migration,
+        run_migration, 
         migration_script,
+        execution_id
     )
 
 
@@ -87,11 +92,12 @@ def lambda_handler(event, context):
         (
             segment,
             total_segments,
-            table_name,
+            table_name, 
             environment,
             region,
             run_migration,
-            migration_script,
+            migration_script, 
+            execution_id
         ) = validate_event_input(event)
 
         logger.info(
@@ -105,13 +111,14 @@ def lambda_handler(event, context):
             environment=environment,
             run_migration=run_migration,
             migration_script=migration_script,
+            execution_id=execution_id
         )
 
         result = service.execute_migration()
         logger.info(
             f"Migration completed: status={result.get('status')}, "
-            f"scanned={result.get('scannedCount')}, updated={result.get('updatedCount')}, "
-            f"errors={result.get('errorCount')}"
+            f"scanned={result.get('scannedCount')}, processed={result.get('processedCount')}, "
+            f"errors={result.get('errorCount')}, skipped={result.get('skippedCount')}"
         )
         return result
 

lambdas/scripts/MigrationBase.py

@@ -1,8 +1,12 @@
 from abc import ABC, abstractmethod
 from typing import Iterable, Callable
+import boto3
+import json
+import os
 
 from services.base.dynamo_service import DynamoDBService
 from utils.audit_logging_setup import LoggingService
+from utils.exceptions import MigrationUnrecoverableException
 
 
 class MigrationBase(ABC):
@@ -15,6 +19,10 @@ def __init__(self, environment: str, table_name: str, run_migration: bool = Fals
         self.table_name = table_name
         self.run_migration = run_migration
         self.target_table = f"{self.table_name}"
+        self.s3_client = boto3.client("s3")
+        self.failed_items_bucket = os.environ.get("MIGRATION_FAILED_ITEMS_STORE_BUCKET_NAME")
+        if not self.failed_items_bucket:
+            raise ValueError("MIGRATION_FAILED_ITEMS_STORE_BUCKET_NAME environment variable is required")
 
         self.logger = LoggingService(self.__class__.__name__)
         self.dynamo_service = DynamoDBService()
@@ -28,36 +36,79 @@ def process_entries(
         label: str,
         entries: Iterable[dict],
         update_fn: Callable[[dict], dict | None],
+        segment: int,
+        execution_id: str
     ):
         """
         Processes a batch of DynamoDB entries for a given migration step.
-
-        :param label: Descriptive name of the step (for logging)
-        :param entries: List of DynamoDB items to process
-        :param update_fn: Function that returns a dict of updated fields or None
         """
         self.logger.info(f"Running '{label}' migration step on {len(entries := list(entries))} items")
+        successful_item_runs = 0
+        skipped_items = 0
+        failed_items = []
 
         for index, entry in enumerate(entries, start=1):
-            item_id = entry.get("ID")
-            self.logger.info(f"[{label}] Processing item {index} (ID: {item_id})")
+            result = self._process_single_entry(label, entry, update_fn, segment)
+            if result == "skipped":
+                skipped_items += 1
+            elif result == "success":
+                successful_item_runs += 1
+            elif isinstance(result, dict):  # failed item
+                failed_items.append(result)
+
+        if failed_items:
+            self._handle_failed_items(label, segment, execution_id, failed_items)
+
+        self.logger.info(f"'{label}' migration step completed.\n")
+        return {
+            "successful_item_runs": successful_item_runs,
+            "failed_items_count": len(failed_items),
+            "skipped_items_count": skipped_items
+        }
 
+    def _process_single_entry(self, label, entry, update_fn, segment):
+        item_id = entry.get("ID")
+        self.logger.info(f"[{label}] Processing item (ID: {item_id})")
+        try:
             updated_fields = update_fn(entry)
             if not updated_fields:
-                self.logger.debug(f"[{label}] Item {item_id} does not require update, skipping.")
-                continue
+                self.logger.info(f"[{label}] Item {item_id} does not require update, skipping.")
+                return "skipped"
+
+            self.logger.info(f"[{label}] Item {item_id} update fields: {updated_fields}")
 
             if self.run_migration:
-                try:
-                    self.dynamo_service.update_item(
-                        table_name=self.target_table,
-                        key_pair={"ID": item_id},
-                        updated_fields=updated_fields,
-                    )
-                    self.logger.info(f"[{label}] Updated item {item_id}: {updated_fields}")
-                except Exception as e:
-                    self.logger.error(f"[{label}] Failed to update item {item_id}: {str(e)}")
+                self.dynamo_service.update_item(
+                    table_name=self.target_table,
+                    key_pair={"ID": item_id},
+                    updated_fields=updated_fields,
+                )
+                self.logger.info(f"[{label}] Updated item {item_id}: {updated_fields}")
+                return "success"
             else:
                 self.logger.info(f"[Dry Run] Would update item {item_id} with {updated_fields}")
+                return "success"
 
-        self.logger.info(f"'{label}' migration step completed.\n")
+        except MigrationUnrecoverableException as e:
+            self.logger.error(f"[{label}] Unrecoverable error for item {item_id} - segment {segment}: {e.message}")
+            return {"item_id": item_id, "error": e.message}
+
+        except Exception as e:
+            self.logger.error(f"[{label}] Error processing item {item_id}: {str(e)}")
+            raise
+
+    def _handle_failed_items(self, label, segment, execution_id, failed_items):
+        self.logger.error(f"'{label}' migration segment: {segment} completed with {len(failed_items)} errors.")
+        error_report_key = f"{execution_id}/{segment}/{label}_errors.json"
+        try:
+            self.s3_client.put_object(
+                Bucket=self.failed_items_bucket,
+                Key=error_report_key,
+                Body=json.dumps(failed_items).encode("utf-8"),
+                ContentType='application/json'
+            )
+            self.logger.error(f"Error report saved to s3://{self.failed_items_bucket}/{error_report_key}")
+        except Exception as s3_error:
+            self.logger.error(f"Failed to save error report to S3: {str(s3_error)}")
+            self.logger.error({"Unlogged failed items": failed_items})
+            raise

lambdas/scripts/dynamodb_migration_prmp_179.py

@@ -1,5 +1,6 @@
 from typing import Callable, Iterable
 
+from utils.exceptions import MigrationUnrecoverableException, MigrationRetryableException
 from scripts.MigrationBase import MigrationBase
 from services.base.dynamo_service import DynamoDBService
 from services.base.s3_service import S3Service
@@ -41,8 +42,10 @@ def main(
         self.logger.info(f"Dry run mode: {not self.run_migration}")
 
         if entries is None:
-            self.logger.error("No entries provided after scanning entire table.")
-            raise ValueError("Entries must be provided to main().")
+            self.logger.error("No entry provided.")
+            raise MigrationRetryableException(
+                message="Entries missing for segment worker", segment_id=None
+            )
 
         return [("S3Metadata", self.get_updated_items)]
 
@@ -51,6 +54,7 @@ def get_updated_items(self, entry: dict) -> dict | None:
         Aggregates updates from all S3 update methods for a single entry.
         Returns a dict of fields to update, or None if no update is needed.
         """
+
         update_items = {}
 
         if s3_metadata_items := self.get_update_s3_metadata_items(entry):
@@ -66,18 +70,18 @@ def get_update_s3_metadata_items(self, entry: dict) -> dict | None:
         file_location = entry.get("FileLocation")
         if not file_location:
             self.logger.warning(f"Missing FileLocation for entry: {entry.get('ID')}")
-            return None
+            raise MigrationUnrecoverableException(message="Missing FileLocation", item_id=entry.get("ID"))
 
         parsed = self.parse_s3_path(file_location)
         if not parsed:
             self.logger.warning(f"Invalid S3 path format: {file_location}")
-            return None
+            raise MigrationUnrecoverableException(message="Invalid S3 path format", item_id=entry.get("ID"))
 
         s3_bucket, s3_key = parsed
         metadata = self.get_s3_metadata(s3_bucket, s3_key)
         if not metadata:
             self.logger.warning(f"Could not retrieve S3 metadata for key: {s3_key}")
-            return None
+            raise MigrationUnrecoverableException(message="Could not retrieve S3 metadata", item_id=entry.get("ID"))
 
         content_length, version_id = metadata
         updated_fields = {}
@@ -113,4 +117,5 @@ def get_s3_metadata(self, bucket: str, key: str) -> tuple[int, str] | None:
                 return s3_head.get("ContentLength"), s3_head.get("VersionId")
         except Exception as e:
             self.logger.error(f"Failed to retrieve S3 metadata for {key}: {e}")
+            raise MigrationUnrecoverableException(message="Failed to retrieve S3 metadata", item_id=key)
         return None

lambdas/scripts/dynamodb_migration_prmp_198.py

@@ -4,6 +4,7 @@
 from scripts.MigrationBase import MigrationBase
 from services.base.dynamo_service import DynamoDBService
 from utils.audit_logging_setup import LoggingService
+from utils.exceptions import MigrationUnrecoverableException, MigrationRetryableException
 
 
 class AuthorMigration(MigrationBase):
@@ -41,7 +42,9 @@ def main(
 
         if entries is None:
             self.logger.error("No entries provided — expected a list of table items.")
-            raise ValueError("Entries must be provided to main().")
+            raise MigrationRetryableException(
+                message="Entries missing for segment worker", segment_id=None
+            )
 
         return [("Author", self.get_update_author_items)]
 
@@ -71,14 +74,18 @@ def get_update_author_items(self, entry: dict) -> dict | None:
             self.logger.warning(
                 f"No completed bulk upload found for NHS number: {nhs_number}"
             )
-            return None
+            raise MigrationUnrecoverableException(
+                message=f"No completed bulk upload found for NHS number: {nhs_number}", item_id=entry.get("ID")
+            )
 
         new_author = bulk_upload_row.get("UploaderOdsCode")
         if not new_author:
             self.logger.warning(
                 f"No uploader ODS code found for NHS number: {nhs_number}"
             )
-            return None
+            raise MigrationUnrecoverableException(
+                message=f"No uploader ODS code found for NHS number: {nhs_number}", item_id=entry.get("ID")
+            )
 
         return {"Author": new_author}
 

lambdas/scripts/dynamodb_migration_prmp_202.py

@@ -5,6 +5,7 @@
 from scripts.MigrationBase import MigrationBase
 from services.base.dynamo_service import DynamoDBService
 from utils.audit_logging_setup import LoggingService
+from utils.exceptions import MigrationUnrecoverableException, MigrationRetryableException
 
 
 class VersionMigration(MigrationBase):
@@ -42,7 +43,9 @@ def main(
 
         if entries is None:
             self.logger.error("No entries provided after scanning entire table.")
-            raise ValueError("Entries must be provided to main().")
+            raise MigrationRetryableException(
+                message="Entries missing for segment worker", segment_id=None
+            )
 
         return [("LGTableValues", self.get_updated_items)]
 
@@ -84,7 +87,9 @@ def get_update_custodian_items(self, entry: dict) -> dict | None:
             self.logger.warning(
                 f"[Custodian] CurrentGpOds is missing for item {entry.get('ID')}"
             )
-            return None
+            raise MigrationUnrecoverableException(
+                message="CurrentGpOds is missing", item_id=entry.get("ID")
+            )
         if current_gp_ods is None or current_gp_ods != custodian:
             return {"Custodian": current_gp_ods}
 
@@ -138,20 +143,23 @@ def get_update_doc_status_items(self, entry: dict) -> dict | None:
 
         inferred_status = document.infer_doc_status()
 
-        if entry.get("uploaded") and entry.get("uploading"):
+        if entry.get("Uploaded") and entry.get("Uploading"):
             self.logger.warning(
                 f"{entry.get('ID')}: Document has a status of uploading and uploaded."
             )
+            raise MigrationUnrecoverableException(
+                message="Document has a status of uploading and uploaded", item_id=entry.get("ID")
+            )
 
         if entry.get("DocStatus", "") == inferred_status:
             return None
 
-        self.logger.warning(f"{entry.get('ID')}: {inferred_status}")
-
         if inferred_status:
             return {"DocStatus": inferred_status}
 
         self.logger.warning(
             f"[DocStatus] Cannot determine status for item {entry.get('ID')}"
         )
-        return None
+        raise MigrationUnrecoverableException(
+            message="Cannot determine document status", item_id=entry.get("ID")
+        )