"""PO persistence: merge-writes to the purchase-orders DynamoDB table. Owns the SET-only merge path (_write_fields / _merge_update) with the sticky "Cancelled" ConditionExpression guard, the collapsed _save_merge behind the save_new_po/save_revision public wrappers, and save_cancellation. Uses a lazily built, cached DynamoDB resource kept under the public name ``dynamodb`` so the tests' setattr(persistence, "dynamodb", fake) patch surface is unchanged. """ import logging import os from datetime import datetime, timezone import boto3 logger = logging.getLogger() logger.setLevel(logging.INFO) PO_TABLE = os.environ.get("PO_TABLE", "purchase-orders") # "Cancelled" is a sticky, authoritative status: once a PO reaches it, a later # new_po/revision may enrich other fields but must never move it back to a # non-cancelled status. CANCELLED_STATUS = "Cancelled" # Lazily-built, cached DynamoDB resource. Public name ``dynamodb`` is preserved # so the monkeypatch attribute is unchanged; building at first CALL (not import) # keeps the moto-before-handler invariant and honors any patched fake. dynamodb = None def _get_dynamodb(): global dynamodb if dynamodb is None: dynamodb = boto3.resource("dynamodb") return dynamodb def _write_fields(po_number: str, fields: dict, *, guard_cancelled: bool): """SET the given non-null fields on a PO record via update_item. Only the fields supplied are written; absent fields are left untouched, so a partial payload can never delete data that an earlier email established. The record is created if it does not exist (DynamoDB update_item upsert). When ``guard_cancelled`` is True the write carries a ConditionExpression that only permits it while the record is not already Cancelled. The condition is evaluated atomically by DynamoDB at write time, so a cancellation that lands first always wins — there is no read-then-write TOCTOU window. A failed guard raises ConditionalCheckFailedException for the caller to handle. """ table = _get_dynamodb().Table(PO_TABLE) set_parts = [] attr_names = {} attr_values = {} for key, value in fields.items(): if value is None or key == "po_number": continue name_ph = f"#{key}" val_ph = f":{key}" attr_names[name_ph] = key attr_values[val_ph] = value set_parts.append(f"{name_ph} = {val_ph}") if not set_parts: return params = { "Key": {"po_number": po_number}, "UpdateExpression": "SET " + ", ".join(set_parts), "ExpressionAttributeNames": attr_names, "ExpressionAttributeValues": attr_values, } if guard_cancelled: params["ExpressionAttributeValues"][":__cancelled_marker"] = CANCELLED_STATUS params["ConditionExpression"] = ( "attribute_not_exists(po_status) OR po_status <> :__cancelled_marker" ) table.update_item(**params) def _merge_update(po_number: str, fields: dict): """Merge (SET-only) the given fields onto a PO record, keeping Cancelled sticky. Only the fields supplied are written; absent fields are left untouched. The record is created if it does not exist (DynamoDB update_item upsert). "Cancelled" is a sticky, authoritative status. When the incoming payload carries a non-cancelled ``po_status``, the write is guarded by a ConditionExpression so the status is only applied while the record is not already Cancelled — enforced atomically at write time, eliminating the read-then-write TOCTOU where a concurrently-landing cancellation could be silently un-cancelled. If the guard fails (the PO is already Cancelled), the same fields are re-written WITHOUT po_status/cancelled_at and unconditionally, so the other fields still merge while the Cancelled status stays intact. A payload with no ``po_status``, or one whose status is already "Cancelled", needs no guard — a plain merge is correct. This is what keeps legitimate status updates (non-cancelled PO) and status-less revisions from ever being dropped: the status is only ever suppressed on a true un-cancel transition. """ incoming_status = fields.get("po_status") if incoming_status is None or incoming_status == CANCELLED_STATUS: _write_fields(po_number, fields, guard_cancelled=False) return try: _write_fields(po_number, fields, guard_cancelled=True) except _get_dynamodb().meta.client.exceptions.ConditionalCheckFailedException: logger.info( f"PO {po_number} is Cancelled; suppressing incoming " f"po_status={incoming_status!r} and merging remaining fields" ) enrich_fields = { k: v for k, v in fields.items() if k not in ("po_status", "cancelled_at") } _write_fields(po_number, enrich_fields, guard_cancelled=False) def _save_merge(parsed: dict, log_verb: str): po_number = parsed["po_number"] fields = {k: v for k, v in parsed.items() if v is not None} _merge_update(po_number, fields) logger.info(f"{log_verb} PO {po_number}") def save_new_po(parsed: dict): """Create a PO, merging into any pre-existing record. Uses a merge update rather than a conditional put so that an out-of-order cancellation (which leaves a Cancelled skeleton) is filled in with the full PO data instead of the new_po being silently dropped. "Cancelled" is a sticky status enforced atomically inside _merge_update: if the PO was already cancelled, the new_po backfills its remaining fields (supplier, line_items, amounts) but never un-cancels it. """ _save_merge(parsed, "Created/merged") def save_revision(parsed: dict): """Merge revised data into an existing PO without deleting omitted fields. A revision email often omits unchanged sections (line_items, supplier). The previous full-overwrite put_item permanently dropped those. This SETs only the fields present in the revision, leaving everything else intact. "Cancelled" is a sticky status: a revision may enrich a cancelled PO's fields but must never move it to a non-cancelled status. That invariant is enforced atomically inside _merge_update and applies ONLY to the un-cancel transition — a revision that carries no status change, or one targeting a non-cancelled PO, updates po_status normally. """ _save_merge(parsed, "Revised") def save_cancellation(parsed: dict): """Mark a PO Cancelled, creating a minimal skeleton if it doesn't exist yet. If the cancellation arrives before the new_po, the skeleton it creates is later backfilled by save_new_po (which preserves this Cancelled status), so no PO data is lost on out-of-order delivery. """ table = _get_dynamodb().Table(PO_TABLE) table.update_item( Key={"po_number": parsed["po_number"]}, UpdateExpression="SET po_status = :status, cancelled_at = :cancelled_at, raw_s3_key = :s3_key", ExpressionAttributeValues={ ":status": "Cancelled", ":cancelled_at": parsed.get( "processed_at", datetime.now(timezone.utc).isoformat() ), ":s3_key": parsed.get("raw_s3_key", ""), }, ) logger.info(f"Cancelled PO {parsed['po_number']}")