[PRMP-202] Migration to Data Standards - LG table values (#793)

Co-authored-by: SWhyteAnswer <[email protected]>
Co-authored-by: Sam Whyte <[email protected]>

Author: PedroSoaresNHS

lambdas/models/document_reference.py

@@ -170,3 +170,12 @@ def set_virus_scanner_result(self, updated_virus_scanner_result) -> None:
 
     def set_uploaded_to_true(self):
         self.uploaded = True
+
+    def infer_doc_status(self) -> str | None:
+        if self.deleted:
+            return "deprecated"
+        if self.uploaded:
+            return "final"
+        if self.uploading:
+            return "preliminary"
+        return None

lambdas/scripts/dynamodb_migration_prmp_202.py

@@ -0,0 +1,210 @@
+import argparse
+from typing import Iterable, Callable
+
+from enums.snomed_codes import SnomedCodes
+from models.document_reference import DocumentReference
+from services.base.dynamo_service import DynamoDBService
+from utils.audit_logging_setup import LoggingService
+
+
+class VersionMigration:
+    def __init__(self, environment: str, table_name: str, run_migration: bool = False):
+        self.environment = environment
+        self.table_name = table_name
+        self.run_migration = run_migration
+        self.logger = LoggingService("CustodianMigration")
+        self.dynamo_service = DynamoDBService()
+
+        self.target_table = f"{self.environment}_{self.table_name}"
+
+    def main(
+            self, entries: Iterable[dict]
+    ) -> list[tuple[str, Callable[[dict], dict | None]]]:
+        """
+        Main entry point for the migration.
+        Returns a list of (label, update function) tuples.
+        Accepts a list of entries for Lambda-based execution, or scans the table if `entries` is None.
+        """
+        self.logger.info("Starting version migration")
+        self.logger.info(f"Target table: {self.target_table}")
+        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().")
+
+        return [
+            ("LGTableValues", self.get_updated_items)
+        ]
+
+    def process_entries(
+            self,
+            label: str,
+            entries: Iterable[dict],
+            update_fn: Callable[[dict], dict | None],
+    ):
+        """
+        Processes a list of entries, applying the update function to each.
+        Logs progress and handles dry-run mode.
+        """
+        self.logger.info(f"Running {label} migration")
+
+        for index, entry in enumerate(entries, start=1):
+            item_id = entry.get("ID")
+            self.logger.info(
+                f"[{label}] Processing item {index} (ID: {item_id})"
+            )
+
+            updated_fields = update_fn(entry)
+            if not updated_fields:
+                self.logger.debug(
+                    f"[{label}] Item {item_id} does not require update, skipping."
+                )
+                continue
+
+            if self.run_migration:
+                self.logger.info(f"Updating item {item_id} with {updated_fields}")
+                try:
+                    self.dynamo_service.update_item(
+                        table_name=self.target_table,
+                        key_pair={"ID": item_id},
+                        updated_fields=updated_fields,
+                    )
+                except Exception as e:
+                    self.logger.error(f"Failed to update item {item_id}: {str(e)}")
+                    continue
+            else:
+                self.logger.info(
+                    f"[Dry Run] Would update item {item_id} with {updated_fields}"
+                )
+
+        self.logger.info(f"{label} migration completed.")  # Moved outside the loop
+
+    def get_updated_items(self, entry: dict) -> dict | None:
+        """
+        Aggregates updates from all update methods for a single entry.
+        Returns a dict of fields to update, or None if no update is needed.
+        """
+        update_items = {}
+
+        if custodian_update_items := self.get_update_custodian_items(entry):
+            update_items.update(custodian_update_items)
+
+        if status_update_items := self.get_update_status_items(entry):
+            update_items.update(status_update_items)
+
+        if snomed_code_update_items := self.get_update_document_snomed_code_type_items(entry):
+            update_items.update(snomed_code_update_items)
+
+        if doc_status_update_items := self.get_update_doc_status_items(entry):
+            update_items.update(doc_status_update_items)
+
+        if version_update_items := self.get_update_version_items(entry):
+            update_items.update(version_update_items)
+
+        return update_items if update_items else None
+
+    def get_update_custodian_items(self, entry: dict) -> dict | None:
+        """
+        Updates the 'Custodian' field if it does not match 'CurrentGpOds'.
+        Returns a dict with the update or None.
+        """
+        current_gp_ods = entry.get("CurrentGpOds")
+        custodian = entry.get("Custodian")
+
+        if current_gp_ods is None:
+            self.logger.warning(f"[Custodian] CurrentGpOds is missing for item {entry.get('ID')}")
+            return None
+        if current_gp_ods is None or current_gp_ods != custodian:
+            return {"Custodian": current_gp_ods}
+
+        return None
+
+    @staticmethod
+    def get_update_status_items(entry: dict) -> dict | None:
+        """
+        Ensures the 'Status' field is set to 'current'.
+        Returns a dict with the update or None.
+        """
+        if entry.get("Status") != "current":
+            return {"Status": "current"}
+        return None
+
+    @staticmethod
+    def get_update_document_snomed_code_type_items(entry: dict) -> dict | None:
+        """
+        Ensures the 'DocumentSnomedCodeType' field matches the expected SNOMED code.
+        Returns a dict with the update or None.
+        """
+        expected_code = SnomedCodes.LLOYD_GEORGE.value.code
+        if entry.get("DocumentSnomedCodeType") != expected_code:
+            return {"DocumentSnomedCodeType": expected_code}
+        return None
+
+    @staticmethod
+    def get_update_version_items(entry: dict) -> dict | None:
+        """
+        Ensures the 'Version' field matches the expected Version code.
+        Returns a dict with the update or None.
+        """
+        expected_version = "1"
+        version_field = "Version"
+        if entry.get(version_field) != expected_version:
+            return {version_field: expected_version}
+        return None
+
+    def get_update_doc_status_items(self, entry: dict) -> dict | None:
+        """
+        Infers and updates the 'DocStatus' field if missing.
+        Returns a dict with the update or None.
+        """
+        try:
+            document = DocumentReference(**entry)
+        except Exception as e:
+            self.logger.warning(f"[DocStatus] Skipping invalid item {entry.get('ID')}: {e}")
+            return None
+
+        inferred_status = document.infer_doc_status()
+
+        if entry.get("uploaded") and entry.get("uploading"):
+            self.logger.warning(f"{entry.get('ID')}: Document has a status of uploading and uploaded.")
+
+        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
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="dynamodb_migration.py",
+        description="Migrate DynamoDB table columns",
+    )
+    parser.add_argument("environment", help="Environment prefix for DynamoDB table")
+    parser.add_argument("table_name", help="DynamoDB table name to migrate")
+    parser.add_argument(
+        "--run-migration",
+        action="store_true",
+        help="Running migration, fields will be updated.",
+    )
+    args = parser.parse_args()
+
+    migration = VersionMigration(
+        environment=args.environment,
+        table_name=args.table_name,
+        run_migration=args.run_migration,
+    )
+
+    entries_to_process = list(
+        migration.dynamo_service.stream_whole_table(migration.target_table)
+    )
+
+    update_functions = migration.main(entries=entries_to_process)
+
+    for label, fn in update_functions:
+        migration.process_entries(label=label, entries=entries_to_process, update_fn=fn)

lambdas/services/base/dynamo_service.py

@@ -1,4 +1,5 @@
 import time
+from typing import Iterator
 from typing import Optional
 
 import boto3
@@ -37,14 +38,14 @@ def get_table(self, table_name):
             raise e
 
     def query_table_by_index(
-        self,
-        table_name,
-        index_name,
-        search_key,
-        search_condition: str,
-        requested_fields: list[str] = None,
-        query_filter: Attr | ConditionBase = None,
-        exclusive_start_key: dict = None,
+            self,
+            table_name,
+            index_name,
+            search_key,
+            search_condition: str,
+            requested_fields: list[str] = None,
+            query_filter: Attr | ConditionBase = None,
+            exclusive_start_key: dict = None,
     ):
         try:
             table = self.get_table(table_name)
@@ -77,7 +78,7 @@ def query_table_by_index(
             raise e
 
     def query_with_pagination(
-        self, table_name: str, search_key: str, search_condition: str
+            self, table_name: str, search_key: str, search_condition: str
     ):
 
         try:
@@ -133,12 +134,12 @@ def create_item(self, table_name, item):
             raise e
 
     def update_item(
-        self,
-        table_name: str,
-        key_pair: dict[str, str],
-        updated_fields: dict,
-        condition_expression: str = None,
-        expression_attribute_values: dict = None,
+            self,
+            table_name: str,
+            key_pair: dict[str, str],
+            updated_fields: dict,
+            condition_expression: str = None,
+            expression_attribute_values: dict = None,
     ):
         table = self.get_table(table_name)
         updated_field_names = list(updated_fields.keys())
@@ -177,10 +178,10 @@ def delete_item(self, table_name: str, key: dict):
             raise e
 
     def scan_table(
-        self,
-        table_name: str,
-        exclusive_start_key: dict = None,
-        filter_expression: str = None,
+            self,
+            table_name: str,
+            exclusive_start_key: dict = None,
+            filter_expression: str = None,
     ):
         try:
             table = self.get_table(table_name)
@@ -199,10 +200,10 @@ def scan_table(
             raise e
 
     def scan_whole_table(
-        self,
-        table_name: str,
-        project_expression: Optional[str] = None,
-        filter_expression: Optional[str] = None,
+            self,
+            table_name: str,
+            project_expression: Optional[str] = None,
+            filter_expression: Optional[str] = None,
     ) -> list[dict]:
         try:
             table = self.get_table(table_name)
@@ -265,7 +266,7 @@ def batch_get_items(self, table_name: str, key_list: list[str]):
                     )
                     request_items = unprocessed_keys
                     retries += 1
-                    time.sleep((2**retries) * 0.1)
+                    time.sleep((2 ** retries) * 0.1)
                 else:
                     break
 
@@ -284,3 +285,38 @@ def get_item(self, table_name: str, key: dict):
                 str(e), {"Result": f"Unable to retrieve item from table: {table_name}"}
             )
             raise e
+
+    def stream_whole_table(
+            self,
+            table_name: str,
+            filter_expression: Optional[str] = None,
+            projection_expression: Optional[str] = None,
+    ) -> Iterator[dict]:
+        """
+        Streams all items from a DynamoDB table using pagination.
+        Yields one item at a time instead of loading everything into memory.
+        """
+        try:
+            table = self.get_table(table_name)
+            scan_kwargs = {}
+
+            if filter_expression:
+                scan_kwargs["FilterExpression"] = filter_expression
+            if projection_expression:
+                scan_kwargs["ProjectionExpression"] = projection_expression
+
+            response = table.scan(**scan_kwargs)
+
+            for item in response.get("Items", []):
+                yield item
+
+            while "LastEvaluatedKey" in response:
+                response = table.scan(
+                    ExclusiveStartKey=response["LastEvaluatedKey"], **scan_kwargs
+                )
+                for item in response.get("Items", []):
+                    yield item
+
+        except ClientError as e:
+            logger.error(str(e), {"Result": f"Unable to stream table: {table_name}"})
+            raise e