Merge branch 'feature/PI-702-prod_bulk_tweak' into release/2024-12-16

Author: jaklinger

infrastructure/terraform/per_workspace/modules/etl/sds/etl-diagram--bulk-transform-and-load.asl.json

@@ -67,6 +67,7 @@
     },
     "Map": {
       "Type": "Map",
+      "MaxConcurrency": 10,
       "ItemProcessor": {
         "ProcessorConfig": {
           "Mode": "INLINE"

scripts/etl/clear_state_inputs.py

@@ -10,8 +10,8 @@
 import boto3
 from etl_utils.constants import CHANGELOG_NUMBER, WorkerKey
 from etl_utils.io import pkl_dumps_lz4
-from sds.epr.bulk_create.bulk_load_fanout import FANOUT
 
+from etl.sds.tests.etl_test_utils.etl_state import _delete_objects_by_prefix
 from test_helpers.aws_session import aws_session
 from test_helpers.terraform import read_terraform_output
 
@@ -38,12 +38,9 @@ def main(changelog_number, workspace):
         s3_client.put_object(
             Bucket=etl_bucket, Key=WorkerKey.LOAD, Body=EMPTY_JSON_DATA
         )
-        for i in range(FANOUT):
-            s3_client.put_object(
-                Bucket=etl_bucket,
-                Key=f"{WorkerKey.LOAD}.{i}",
-                Body=pkl_dumps_lz4(EMPTY_JSON_DATA),
-            )
+        _delete_objects_by_prefix(
+            s3_client=s3_client, bucket=etl_bucket, key_prefix=f"{WorkerKey.LOAD}."
+        )
         s3_client.delete_object(Bucket=etl_bucket, Key=CHANGELOG_NUMBER)
 
     if changelog_number:

src/etl/sds/tests/etl_test_utils/etl_state.py

@@ -10,7 +10,6 @@
 )
 from etl_utils.io import pkl_dumps_lz4
 from mypy_boto3_s3 import S3Client
-from sds.epr.bulk_create.bulk_load_fanout import FANOUT
 
 from test_helpers.terraform import read_terraform_output
 
@@ -40,10 +39,28 @@ def get_etl_config(input_filename: str, etl_type: str = "") -> EtlConfig:
     )
 
 
-def _delete_objects_by_prefix(s3_client: S3Client, bucket: str, key_prefix: str):
-    response = s3_client.list_objects(Bucket=bucket, Prefix=key_prefix)
-    for item in response.get("Contents", []):
-        s3_client.delete_object(Bucket=bucket, Key=item["Key"])
+def _delete_objects_by_prefix(
+    s3_client: S3Client, bucket: str, key_prefix: str, **kwargs
+):
+    # Delete objects if any found
+    response = s3_client.list_objects_v2(Bucket=bucket, Prefix=key_prefix, **kwargs)
+    try:
+        contents = response["Contents"]
+    except KeyError:
+        return
+    else:
+        for item in contents:
+            s3_client.delete_object(Bucket=bucket, Key=item["Key"])
+
+    # Repeat if required
+    continuation_token = response.get("ContinuationToken")
+    if continuation_token:
+        return _delete_objects_by_prefix(
+            s3_client=s3_client,
+            bucket=bucket,
+            key_prefix=key_prefix,
+            ContinuationToken=continuation_token,
+        )
 
 
 def clear_etl_state(s3_client: S3Client, etl_config: EtlConfig):
@@ -61,12 +78,12 @@ def clear_etl_state(s3_client: S3Client, etl_config: EtlConfig):
         Key=WorkerKey.LOAD,
         Body=pkl_dumps_lz4(EMPTY_JSON_DATA),
     )
-    for i in range(FANOUT):
-        s3_client.put_object(
-            Bucket=etl_config.bucket,
-            Key=f"{WorkerKey.LOAD}.{i}",
-            Body=pkl_dumps_lz4(EMPTY_JSON_DATA),
-        )
+
+    # Delete load-fanout files, if they exist
+    _delete_objects_by_prefix(
+        s3_client=s3_client, bucket=etl_config.bucket, key_prefix=f"{WorkerKey.LOAD}."
+    )
+
     s3_client.delete_object(
         Bucket=etl_config.bucket, Key=etl_config.initial_trigger_key
     )

src/etl/sds/worker/bulk/load_bulk/tests/test_load_bulk_worker.py

@@ -24,7 +24,6 @@
 from event.json import json_load
 from moto import mock_aws
 from mypy_boto3_s3 import S3Client
-from sds.epr.bulk_create.bulk_load_fanout import FANOUT
 from sds.epr.constants import AS_DEVICE_SUFFIX, MHS_DEVICE_SUFFIX
 
 from etl.sds.worker.bulk.tests.test_bulk_e2e import PATH_TO_STAGE_DATA
@@ -73,31 +72,33 @@ def test_load_worker_pass(
     from etl.sds.worker.bulk.load_bulk import load_bulk
 
     # Initial state
-    for i in range(FANOUT):
-        with open(PATH_TO_STAGE_DATA / f"3.load_fanout_output.{i}.json") as f:
+    input_paths = sorted(PATH_TO_STAGE_DATA.glob(f"3.load_fanout_output.*.json"))
+    load_fanout_keys = []
+    for i, path in enumerate(input_paths):
+        with open(path) as f:
             input_data = json_load(f)
 
+        fanout_key = f"{WorkerKey.LOAD}.{i}"
         put_object(
-            key=f"{WorkerKey.LOAD}.{i}",
+            key=fanout_key,
             body=pkl_dumps_lz4(deque(input_data)),
         )
+        load_fanout_keys.append(fanout_key)
 
     # Execute the load worker
     with mock_table(TABLE_NAME) as dynamodb_client:
         load_bulk.CACHE.REPOSITORY.client = dynamodb_client
 
         responses = []
-        for i in range(FANOUT):
+        for fanout_key in load_fanout_keys:
             response = load_bulk.handler(
-                event={"s3_input_path": f"s3://{BUCKET_NAME}/{WorkerKey.LOAD}.{i}"},
+                event={"s3_input_path": f"s3://{BUCKET_NAME}/{fanout_key}"},
                 context=None,
             )
             responses.append(response)
 
             # Final state
-            final_unprocessed_data = pkl_loads_lz4(
-                get_object(key=f"{WorkerKey.LOAD}.{i}")
-            )
+            final_unprocessed_data = pkl_loads_lz4(get_object(key=fanout_key))
             assert final_unprocessed_data == deque([])
 
         assert responses == [
@@ -107,7 +108,7 @@ def test_load_worker_pass(
                 "unprocessed_records": 0,
                 "error_message": None,
             }
-        ] * (FANOUT - 1) + [
+        ] * (len(load_fanout_keys) - 1) + [
             {
                 "stage_name": "load",
                 "processed_records": 7,

src/etl/sds/worker/bulk/load_bulk_fanout/load_bulk_fanout.py

@@ -22,7 +22,6 @@
 )
 from event.environment import BaseEnvironment
 from event.step_chain import StepChain
-from sds.epr.bulk_create.bulk_load_fanout import FANOUT, calculate_batch_size
 from sds.epr.bulk_create.bulk_repository import BulkRepository
 
 if TYPE_CHECKING:
@@ -38,6 +37,7 @@ def s3_path(self, key) -> str:
 
 ENVIRONMENT = TransformWorkerEnvironment.build()
 S3_CLIENT = boto3.client("s3")
+EACH_FANOUT_BATCH_SIZE = 10000
 
 
 def execute_step_chain(
@@ -72,11 +72,8 @@ def execute_step_chain(
 
         count_unprocessed_records = len(action_chain.result.unprocessed_records)
 
-        batch_size = calculate_batch_size(
-            sequence=action_chain.result.processed_records, n_batches=FANOUT
-        )
         for i, batch in enumerate(
-            batched(action_chain.result.processed_records, n=batch_size)
+            batched(action_chain.result.processed_records, n=EACH_FANOUT_BATCH_SIZE)
         ):
             count_processed_records = len(batch)
             _action_response = WorkerActionResponse(

src/etl/sds/worker/bulk/load_bulk_fanout/tests/test_load_bulk_fanout_worker.py

@@ -1,6 +1,5 @@
 import os
 from collections import deque
-from itertools import batched
 from pathlib import Path
 from typing import Callable
 from unittest import mock
@@ -17,11 +16,7 @@
 from event.json import json_load
 from moto import mock_aws
 from mypy_boto3_s3 import S3Client
-from sds.epr.bulk_create.bulk_load_fanout import (
-    FANOUT,
-    calculate_batch_size,
-    count_indexes,
-)
+from sds.epr.bulk_create.bulk_load_fanout import count_indexes
 
 from etl.sds.worker.bulk.tests.test_bulk_e2e import PATH_TO_STAGE_DATA
 
@@ -78,30 +73,15 @@ def decompress(obj: dict) -> dict:
     return obj
 
 
-@pytest.mark.parametrize(
-    ("n_batches", "sequence_length", "expected_batch_size", "expected_n_batches"),
-    ((4, 100, 25, 4), (3, 100, 34, 3), (16, 30, 2, 15)),
-)
-def test_calculate_batch_size_general(
-    n_batches: int,
-    sequence_length: int,
-    expected_batch_size: int,
-    expected_n_batches: int,
-):
-    n_batches = n_batches
-    sequence = list(range(sequence_length))
-    batch_size = calculate_batch_size(sequence, n_batches)
-    assert batch_size == expected_batch_size
-
-    batches = list(batched(sequence, batch_size))
-    assert len(batches) == expected_n_batches
-
-
 def test_load_worker_fanout(
-    put_object: Callable[[str], None], get_object: Callable[[str], bytes]
+    put_object: Callable[[str], None],
+    get_object: Callable[[str], bytes],
 ):
+    _EACH_FANOUT_BATCH_SIZE = 10
     from etl.sds.worker.bulk.load_bulk_fanout import load_bulk_fanout
 
+    load_bulk_fanout.EACH_FANOUT_BATCH_SIZE = _EACH_FANOUT_BATCH_SIZE
+
     # Initial state
     with open(PATH_TO_STAGE_DATA / "2.transform_output.json") as f:
         input_data: list[dict[str, dict]] = json_load(f)
@@ -114,31 +94,36 @@ def test_load_worker_fanout(
     # Execute the load worker
     responses = load_bulk_fanout.handler(event={}, context=None)
 
-    assert len(responses) == FANOUT
-    assert responses == [
+    *head_responses, tail_response = responses
+
+    assert len(head_responses) > 1
+
+    expected_head_responses = [
         {
             "stage_name": "load_bulk_fanout",
-            "processed_records": 10,
+            "processed_records": _EACH_FANOUT_BATCH_SIZE,
             "unprocessed_records": 0,
             "s3_input_path": f"s3://my-bucket/input--load/unprocessed.{i}",
             "error_message": None,
         }
-        for i in range(0, FANOUT - 1)
-    ] + [
-        {
-            "stage_name": "load_bulk_fanout",
-            "processed_records": 7,
-            "unprocessed_records": 0,
-            "s3_input_path": f"s3://my-bucket/input--load/unprocessed.{FANOUT - 1}",
-            "error_message": None,
-        },
+        for i in range(len(head_responses))
     ]
+    assert head_responses == expected_head_responses
+
+    tail_processed_records = tail_response.pop("processed_records")
+    assert tail_processed_records <= _EACH_FANOUT_BATCH_SIZE
+    assert tail_response == {
+        "stage_name": "load_bulk_fanout",
+        "unprocessed_records": 0,
+        "s3_input_path": f"s3://my-bucket/input--load/unprocessed.{len(head_responses)}",
+        "error_message": None,
+    }
 
     # Final state
     final_processed_data = pkl_loads_lz4(get_object(key=WorkerKey.LOAD))
     assert final_processed_data == deque([])
     total_size = 0
-    for i in range(10):
+    for i in range(_EACH_FANOUT_BATCH_SIZE):
         final_unprocessed_data = pkl_loads_lz4(get_object(key=f"{WorkerKey.LOAD}.{i}"))
         assert isinstance(final_unprocessed_data, deque)
         total_size += len(final_unprocessed_data)
@@ -164,3 +149,8 @@ def test_load_worker_fanout(
         expected_total_size += count_indexes(obj)
 
     assert total_size == expected_total_size
+
+    total_processed_records_from_response = (
+        tail_processed_records + _EACH_FANOUT_BATCH_SIZE * len(head_responses)
+    )
+    assert total_size == total_processed_records_from_response

src/layers/domain/core/validation.py

@@ -31,7 +31,7 @@ class AccreditedSystem:
         ID_PATTERN = re.compile(rf"^[a-zA-Z-0-9]+$")
 
     class PartyKey:
-        PARTY_KEY_REGEX = rf"^{_ODS_CODE_REGEX}-[0-9]{{6,9}}$"
+        PARTY_KEY_REGEX = rf"^{_ODS_CODE_REGEX}-[0-9]{{5,9}}$"
         ID_PATTERN = re.compile(PARTY_KEY_REGEX)
 
     class CpaId:

src/layers/sds/epr/bulk_create/bulk_create.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from domain.core.aggregate_root import AggregateRoot
 from domain.core.product_team.v1 import ProductTeam
 from domain.core.questionnaire import Questionnaire, QuestionnaireResponse
@@ -91,6 +93,22 @@ def _create_complete_epr_product(
     ]
 
 
+def _impute_manufacturer_org(
+    item: dict[Literal["nhs_mhs_manufacturer_org", "nhs_id_code"], str]
+):
+    """
+    Impute nhs_mhs_manufacturer_org if it is clearly invalid (not alphanumeric)
+    or does not exist by replacing it with the nhs_id_code
+    """
+    manufacturer_org = item["nhs_mhs_manufacturer_org"]
+    item["nhs_mhs_manufacturer_org"] = (
+        manufacturer_org
+        if manufacturer_org is not None and manufacturer_org.isalnum()
+        else item["nhs_id_code"]
+    )
+    return item
+
+
 def create_complete_epr_product(
     party_key_group: list[dict],
     mhs_device_questionnaire: Questionnaire,
@@ -110,6 +128,7 @@ def create_complete_epr_product(
     message_handling_systems: list[dict] = []
     accredited_systems: list[dict] = []
     for item in party_key_group:
+        item = _impute_manufacturer_org(item)
         if item["object_class"].lower() == NhsMhs.OBJECT_CLASS:
             message_handling_systems.append(item)
         else:
@@ -127,7 +146,6 @@ def create_complete_epr_product(
             mhs_device_questionnaire=mhs_device_questionnaire,
             mhs_device_field_mapping=mhs_device_field_mapping,
         )
-
         ods_code = first_mhs["nhs_mhs_manufacturer_org"]
         party_key = first_mhs["nhs_mhs_party_key"]
         product_name = first_mhs["nhs_product_name"] or first_mhs["nhs_mhs_cpa_id"]

src/layers/sds/epr/bulk_create/bulk_load_fanout.py

@@ -1,12 +1,8 @@
-from math import ceil
-
 from domain.core.cpm_product.v1 import CpmProduct
 from domain.core.device.v1 import Device
 from domain.core.device_reference_data.v1 import DeviceReferenceData
 from domain.core.product_team.v1 import ProductTeam
 
-FANOUT = 10
-
 
 def count_indexes(obj: Device | DeviceReferenceData | CpmProduct | ProductTeam):
     count = 1
@@ -15,7 +11,3 @@ def count_indexes(obj: Device | DeviceReferenceData | CpmProduct | ProductTeam):
     if isinstance(obj, (Device)):
         count += len(obj.tags)
     return count
-
-
-def calculate_batch_size(sequence: list, n_batches: int) -> int:
-    return ceil(len(sequence) / (n_batches)) or 1

src/layers/sds/epr/bulk_create/creators.py

@@ -7,7 +7,11 @@
 from domain.core.product_team_key import ProductTeamKey, ProductTeamKeyType
 from domain.core.questionnaire import QuestionnaireResponse
 from domain.core.root import Root
-from sds.epr.constants import EprNameTemplate, SdsDeviceReferenceDataPath
+from sds.epr.constants import (
+    EXCEPTIONAL_ODS_CODES,
+    EprNameTemplate,
+    SdsDeviceReferenceDataPath,
+)
 
 
 def create_epr_product_team(ods_code: str) -> ProductTeam:
@@ -16,6 +20,13 @@ def create_epr_product_team(ods_code: str) -> ProductTeam:
         key_value=EprNameTemplate.PRODUCT_TEAM_KEY.format(ods_code=ods_code),
     )
 
+    if ods_code in EXCEPTIONAL_ODS_CODES:
+        return ProductTeam(
+            name=EprNameTemplate.PRODUCT_TEAM.format(ods_code=ods_code),
+            ods_code=ods_code,
+            keys=[product_team_key],
+        )
+
     org = Root.create_ods_organisation(ods_code=ods_code)
     return org.create_product_team(
         name=EprNameTemplate.PRODUCT_TEAM.format(ods_code=ods_code),