fhir_repository.py

import os
import time
from dataclasses import dataclass

import boto3
import botocore.exceptions
import simplejson as json
from boto3.dynamodb.conditions import Attr, ConditionBase, Key
from botocore.config import Config
from fhir.resources.R4B.fhirtypes import Id
from fhir.resources.R4B.identifier import Identifier
from fhir.resources.R4B.immunization import Immunization
from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table

from common.models.constants import Constants
from common.models.errors import ResourceNotFoundError
from common.models.immunization_record_metadata import ImmunizationRecordMetadata
from common.models.utils.generic_utils import (
    get_contained_patient,
    get_nhs_number,
)
from common.models.utils.validation_utils import (
    get_vaccine_type,
)
from models.errors import InvalidStoredDataError, UnhandledResponseError


def create_table(table_name=None, endpoint_url=None, region_name="eu-west-2"):
    if not table_name:
        table_name = os.environ["DYNAMODB_TABLE_NAME"]
    config = Config(connect_timeout=1, read_timeout=1, retries={"max_attempts": 1})
    db: DynamoDBServiceResource = boto3.resource(
        "dynamodb", endpoint_url=endpoint_url, region_name=region_name, config=config
    )
    return db.Table(table_name)


def _make_immunization_pk(_id: str):
    return f"Immunization#{_id}"


def _make_patient_pk(_id: str):
    return f"Patient#{_id}"


def _query_identifier(table, index, pk, identifier):
    queryresponse = table.query(IndexName=index, KeyConditionExpression=Key(pk).eq(identifier), Limit=1)
    if queryresponse.get("Count", 0) > 0:
        return queryresponse


def get_fhir_identifier_from_identifier_pk(identifier_pk: str) -> Identifier:
    split_identifier = identifier_pk.split("#", 1)

    if len(split_identifier) != 2:
        raise InvalidStoredDataError(data_type="identifier")

    supplier_code = split_identifier[0]
    supplier_unique_id = split_identifier[1]

    return Identifier(system=supplier_code, value=supplier_unique_id)


@dataclass
class RecordAttributes:
    pk: str
    patient_pk: str
    patient_sk: str
    patient: dict
    vaccine_type: str
    timestamp: int
    identifier: str
    immunization: Immunization

    @classmethod
    def from_immunization(cls, immunization: Immunization, patient: dict | None = None) -> "RecordAttributes":
        """Build DynamoDB attributes from a FHIR Immunization resource."""
        imms_dict = immunization.dict()
        patient_resolved = patient if patient is not None else get_contained_patient(imms_dict)
        nhs_number = get_nhs_number(imms_dict)
        vaccine_type = get_vaccine_type(imms_dict)
        first_identifier = immunization.identifier[0]
        return cls(
            pk=_make_immunization_pk(immunization.id),
            patient_pk=_make_patient_pk(nhs_number),
            patient_sk=f"{vaccine_type}#{immunization.id}",
            patient=patient_resolved,
            vaccine_type=vaccine_type,
            timestamp=int(time.time()),
            identifier=f"{first_identifier.system}#{first_identifier.value}",
            immunization=immunization,
        )


class ImmunizationRepository:
    def __init__(self, table: Table):
        self.table = table

    def get_immunization_by_identifier(
        self, identifier: Identifier
    ) -> tuple[dict | None, ImmunizationRecordMetadata | None]:
        response = self.table.query(
            IndexName="IdentifierGSI",
            KeyConditionExpression=Key("IdentifierPK").eq(self._make_identifier_pk(identifier)),
        )
        item = response.get("Items")

        if not item:
            return None, None

        item = response["Items"][0]
        resource = json.loads(item.get("Resource", {}))

        # VED-915 - investigate and hopefully stop returning deleted items
        imms_record_meta = ImmunizationRecordMetadata(
            identifier,
            int(item.get("Version")),
            is_deleted=self._is_logically_deleted_item(item),
            is_reinstated=self._is_reinstated_item(item),
        )
        return resource, imms_record_meta

    def get_immunization_resource_and_metadata_by_id(
        self, imms_id: str, include_deleted: bool = False
    ) -> tuple[dict | None, ImmunizationRecordMetadata | None]:
        """Retrieves the immunization resource and metadata from the VEDS table"""
        response = self.table.get_item(Key={"PK": _make_immunization_pk(imms_id)})
        item = response.get("Item")

        if not item:
            return None, None

        is_deleted = self._is_logically_deleted_item(item)
        is_reinstated = self._is_reinstated_item(item)

        if is_deleted and not include_deleted:
            return None, None

        # The FHIR Identifier which is returned in the metadata is based on the IdentifierPK from the database because
        # we keep this attribute up to date in case of any changes rather than modifying the JSON resource. For example,
        # when we performed the V2 to V5 data migration as part of issue VED-893.

        identifier_pk = item.get("IdentifierPK")

        if identifier_pk is None:
            raise InvalidStoredDataError(data_type="identifier")

        identifier = get_fhir_identifier_from_identifier_pk(identifier_pk)

        imms_record_meta = ImmunizationRecordMetadata(identifier, int(item.get("Version")), is_deleted, is_reinstated)

        return json.loads(item.get("Resource", {})), imms_record_meta

    def check_immunization_identifier_exists(self, system: str, unique_id: str) -> bool:
        """Checks whether an immunization with the given immunization identifier (system + local ID) exists."""
        response = self.table.query(
            IndexName="IdentifierGSI",
            KeyConditionExpression=Key("IdentifierPK").eq(f"{system}#{unique_id}"),
        )

        if "Items" in response and len(response["Items"]) > 0:
            return True

        return False

    def create_immunization(self, immunization: Immunization, supplier_system: str) -> Id:
        """Creates a new immunization record returning the unique id if successful."""
        attr = RecordAttributes.from_immunization(immunization)

        response = self.table.put_item(
            Item={
                "PK": attr.pk,
                "PatientPK": attr.patient_pk,
                "PatientSK": attr.patient_sk,
                "Resource": immunization.json(use_decimal=True),
                "IdentifierPK": attr.identifier,
                "Operation": "CREATE",
                "Version": 1,
                "SupplierSystem": supplier_system,
            }
        )

        if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
            raise UnhandledResponseError(message="Non-200 response from dynamodb", response=dict(response))

        return immunization.id

    def update_immunization(
        self,
        imms_id: str,
        immunization: Immunization,
        existing_record_meta: ImmunizationRecordMetadata,
        supplier_system: str,
    ) -> int:
        attr = RecordAttributes.from_immunization(immunization)
        reinstate_operation_required = existing_record_meta.is_deleted

        update_exp = self._build_update_expression(is_reinstate=reinstate_operation_required)

        return self._perform_dynamo_update(
            imms_id,
            update_exp,
            attr,
            existing_record_meta.resource_version,
            supplier_system,
            reinstate_operation_required=reinstate_operation_required,
            record_contains_deletion_history=(reinstate_operation_required or existing_record_meta.is_reinstated),
        )

    @staticmethod
    def _build_update_expression(is_reinstate: bool) -> str:
        if is_reinstate:
            return (
                "SET UpdatedAt = :timestamp, PatientPK = :patient_pk, "
                "PatientSK = :patient_sk, #imms_resource = :imms_resource_val, "
                "Operation = :operation, Version = :version, DeletedAt = :respawn, SupplierSystem = :supplier_system "
            )
        else:
            return (
                "SET UpdatedAt = :timestamp, PatientPK = :patient_pk, "
                "PatientSK = :patient_sk, #imms_resource = :imms_resource_val, "
                "Operation = :operation, Version = :version, SupplierSystem = :supplier_system "
            )

    def _perform_dynamo_update(
        self,
        imms_id: str,
        update_exp: str,
        attr: RecordAttributes,
        existing_resource_version: int,
        supplier_system: str,
        reinstate_operation_required: bool,
        record_contains_deletion_history: bool,
    ) -> int:
        updated_version = existing_resource_version + 1
        condition_expression = Attr("PK").eq(attr.pk) & (
            Attr("DeletedAt").exists()
            if record_contains_deletion_history
            else Attr("PK").eq(attr.pk) & Attr("DeletedAt").not_exists()
        )

        expression_attribute_values = {
            ":timestamp": attr.timestamp,
            ":patient_pk": attr.patient_pk,
            ":patient_sk": attr.patient_sk,
            ":imms_resource_val": attr.immunization.json(use_decimal=True),
            ":operation": "UPDATE",
            ":version": updated_version,
            ":supplier_system": supplier_system,
        }
        if reinstate_operation_required:
            expression_attribute_values.update(
                {
                    ":respawn": "reinstated",
                }
            )

        try:
            self.table.update_item(
                Key={"PK": _make_immunization_pk(imms_id)},
                UpdateExpression=update_exp,
                ExpressionAttributeNames={
                    "#imms_resource": "Resource",
                },
                ExpressionAttributeValues=expression_attribute_values,
                ConditionExpression=condition_expression,
            )
        except botocore.exceptions.ClientError as error:
            # Either resource didn't exist or it has already been deleted. See ConditionExpression in the request
            if error.response["Error"]["Code"] == "ConditionalCheckFailedException":
                raise ResourceNotFoundError(resource_type="Immunization", resource_id=imms_id)

            raise error

        return updated_version

    def delete_immunization(self, imms_id: str, supplier_system: str) -> None:
        now_timestamp = int(time.time())

        try:
            self.table.update_item(
                Key={"PK": _make_immunization_pk(imms_id)},
                UpdateExpression=(
                    "SET DeletedAt = :timestamp, Operation = :operation, SupplierSystem = :supplier_system"
                ),
                ExpressionAttributeValues={
                    ":timestamp": now_timestamp,
                    ":operation": "DELETE",
                    ":supplier_system": supplier_system,
                },
                ConditionExpression=(
                    Attr("PK").eq(_make_immunization_pk(imms_id))
                    & (Attr("DeletedAt").not_exists() | Attr("DeletedAt").eq("reinstated"))
                ),
            )
        except botocore.exceptions.ClientError as error:
            if error.response["Error"]["Code"] == "ConditionalCheckFailedException":
                raise ResourceNotFoundError(resource_type="Immunization", resource_id=imms_id)
            else:
                raise error

    def search_immunizations(self, patient_identifier: str, vaccine_types: set) -> list[dict]:
        """Searches for immunisations by patient identifier (NHS Number) and the vaccination type"""
        # vacc_type_condition = self._build_vacc_type_key_condition(vaccine_types)
        is_not_deleted_condition = Attr("DeletedAt").not_exists() | Attr("DeletedAt").eq("reinstated")
        patient_key_condition = Key("PatientPK").eq(_make_patient_pk(patient_identifier))

        ieds_resources = []

        for vacc_type in vaccine_types:
            # Should make these DB keys constants
            key_condition = patient_key_condition & Key("PatientSK").begins_with(vacc_type)
            ieds_resources.extend(self.get_all_items(key_condition, is_not_deleted_condition))

        if len(ieds_resources) == 0:
            return []

        # Return a list of the FHIR immunization resource JSON items
        final_resources = [
            {
                **json.loads(item["Resource"]),
                "meta": {"versionId": int(item.get("Version", 1))},
            }
            for item in ieds_resources
        ]

        return final_resources

    def get_all_items(self, key_condition: ConditionBase, filter_condition: ConditionBase):
        """Query DynamoDB and paginate through all results."""
        all_items = []
        last_evaluated_key = None

        while True:
            query_args = {
                "IndexName": "PatientGSI",
                "KeyConditionExpression": key_condition,
                "FilterExpression": filter_condition,
            }
            if last_evaluated_key:
                query_args["ExclusiveStartKey"] = last_evaluated_key

            response = self.table.query(**query_args)
            if "Items" not in response:
                raise UnhandledResponseError(message="No Items in DynamoDB response", response=response)

            items = response.get("Items", [])
            all_items.extend(items)

            last_evaluated_key = response.get("LastEvaluatedKey")
            if not last_evaluated_key:
                break

        return all_items

    # @staticmethod
    # def _build_vacc_type_key_condition(vacc_types: set) -> ConditionBase:
    #     vacc_type_condition = None
    #
    #     for vacc_type in vacc_types:
    #         key_cond = Key("PatientSK").begins_with(vacc_type)
    #         vacc_type_condition = key_cond if vacc_type_condition is None else vacc_type_condition | key_cond
    #
    #     return vacc_type_condition

    @staticmethod
    def _vaccine_type(patient_sk: str) -> str:
        parsed = [str.strip(s) for s in patient_sk.split("#")]
        return parsed[0]

    @staticmethod
    def _make_identifier_pk(identifier: Identifier) -> str:
        return f"{identifier.system}#{identifier.value}"

    @staticmethod
    def _is_logically_deleted_item(ieds_item: dict) -> bool:
        deleted_at_attr = ieds_item.get("DeletedAt")
        return deleted_at_attr is not None and deleted_at_attr != Constants.REINSTATED_RECORD_STATUS

    @staticmethod
    def _is_reinstated_item(ieds_item: dict) -> bool:
        return ieds_item.get("DeletedAt") == Constants.REINSTATED_RECORD_STATUS