Agent/examples/auto_put_ad_mini/roi_control/repository.py

"""MySQL persistence for versioned ROI runs, snapshots, and action audit."""

from __future__ import annotations

import json
from datetime import date, datetime
from decimal import Decimal
from typing import Any, Iterable

from storage import connect


FINAL_STATUSES = {
    "COMPLETED",
    "PARTIAL",
    "REJECTED",
    "EXPIRED",
}


def _is_expired(expires_at: datetime | None, now: datetime) -> bool:
    if expires_at is None:
        return False
    if expires_at.tzinfo is None and now.tzinfo is not None:
        now = now.replace(tzinfo=None)
    elif expires_at.tzinfo is not None and now.tzinfo is None:
        expires_at = expires_at.replace(tzinfo=None)
    return now > expires_at


def json_dumps(value: Any) -> str:
    def default(item: Any) -> Any:
        if isinstance(item, (datetime, date)):
            return item.isoformat()
        if isinstance(item, Decimal):
            return str(item)
        if hasattr(item, "item"):
            return item.item()
        if hasattr(item, "tolist"):
            return item.tolist()
        raise TypeError(f"Unsupported JSON value: {type(item)!r}")

    return json.dumps(value, ensure_ascii=False, default=default)


def publish_fission_parameter_release(
    release: dict[str, Any],
    values: Iterable[dict[str, Any]],
    *,
    published_by: str,
) -> bool:
    """Publish an immutable parameter version; return False if already present."""

    value_rows = list(values)
    connection = connect()
    connection.autocommit(False)
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT content_sha256 FROM roi_fission_parameter_release "
                "WHERE version=%s FOR UPDATE",
                (release["version"],),
            )
            existing = cursor.fetchone()
            if existing:
                if existing["content_sha256"] != release["content_sha256"]:
                    raise ValueError(
                        "传播裂变系数版本已存在且内容不同，禁止覆盖: "
                        f"{release['version']}"
                    )
                connection.rollback()
                return False

            cursor.execute(
                """
                INSERT INTO roi_fission_parameter_release
                    (version, cohort_date, observation_end_date, horizon_days,
                     run_suffix, content_sha256, metadata_json, published_by)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
                """,
                (
                    release["version"],
                    release["cohort_date"],
                    release["observation_end_date"],
                    release["horizon_days"],
                    release["run_suffix"],
                    release["content_sha256"],
                    json_dumps(release["metadata"]),
                    published_by,
                ),
            )
            cursor.executemany(
                """
                INSERT INTO roi_fission_parameter_value
                    (version, entity_type, match_level, key_primary,
                     key_secondary, multiplier, multiplier_vs_first)
                VALUES (%s,%s,%s,%s,%s,%s,%s)
                """,
                [
                    (
                        release["version"],
                        row["entity_type"],
                        row["match_level"],
                        row["key_primary"],
                        row["key_secondary"],
                        row["multiplier"],
                        row.get("multiplier_vs_first"),
                    )
                    for row in value_rows
                ],
            )
        connection.commit()
        return True
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.close()


def load_fission_parameter_release(
    version: str,
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM roi_fission_parameter_release WHERE version=%s",
                (version,),
            )
            release = cursor.fetchone()
            if not release:
                raise ValueError(f"数据库不存在传播裂变系数版本: {version}")
            cursor.execute(
                """
                SELECT entity_type, match_level, key_primary,
                       key_secondary, multiplier, multiplier_vs_first
                FROM roi_fission_parameter_value
                WHERE version=%s
                ORDER BY entity_type, match_level, key_primary, key_secondary
                """,
                (version,),
            )
            return release, list(cursor.fetchall())
    finally:
        connection.close()


def create_or_load_run(record: dict[str, Any]) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT IGNORE INTO roi_metric_run
                    (run_id, run_key, metric_version, policy_version,
                     fission_parameter_version, fission_cohort_date,
                     start_date, end_date, config_json, status)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,'COMPUTING')
                """,
                (
                    record["run_id"],
                    record["run_key"],
                    record["metric_version"],
                    record["policy_version"],
                    record["fission_parameter_version"],
                    record["fission_cohort_date"],
                    record["start_date"],
                    record["end_date"],
                    json_dumps(record["config"]),
                ),
            )
            cursor.execute(
                "SELECT * FROM roi_metric_run WHERE run_key=%s",
                (record["run_key"],),
            )
            run = cursor.fetchone()
            if run:
                return run
            cursor.execute(
                "SELECT * FROM roi_metric_run WHERE run_id=%s",
                (record["run_id"],),
            )
            collision = cursor.fetchone()
            if collision:
                raise RuntimeError(
                    f"ROI run_id already exists with a different run_key: "
                    f"{record['run_id']}"
                )
            return {}
    finally:
        connection.close()


def load_run(run_id: str) -> dict[str, Any] | None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute("SELECT * FROM roi_metric_run WHERE run_id=%s", (run_id,))
            return cursor.fetchone()
    finally:
        connection.close()


def replace_run_results(
    run_id: str,
    *,
    snapshots: Iterable[dict[str, Any]],
    actions: Iterable[dict[str, Any]],
    thresholds: dict[str, Any],
) -> tuple[int, int]:
    snapshot_rows = list(snapshots)
    action_rows = list(actions)
    connection = connect()
    connection.autocommit(False)
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT status FROM roi_metric_run WHERE run_id=%s FOR UPDATE",
                (run_id,),
            )
            run = cursor.fetchone()
            if not run:
                raise ValueError(f"ROI run not found: {run_id}")
            if run["status"] not in {"COMPUTING", "COMPUTED", "FAILED"}:
                raise RuntimeError(
                    f"ROI run {run_id} cannot be replaced in status {run['status']}"
                )
            cursor.execute("DELETE FROM roi_action_item WHERE run_id=%s", (run_id,))
            cursor.execute("DELETE FROM roi_entity_snapshot WHERE run_id=%s", (run_id,))

            snapshot_columns = [
                "run_id", "entity_hash", "entity_type", "channel",
                "account_id", "account_name", "adgroup_id", "adgroup_name",
                "dynamic_creative_id", "audience_name", "conversion_goal",
                "partner_name", "official_account_name",
                "fission_parameter_version", "fission_cohort_date",
                "fission_multiplier_vs_t0", "fission_multiplier_vs_first",
                "fission_match_level",
                "fission_source", "ad_age",
                "avg_first_uv", "first_uv",
                "t0_fission_count", "t0_fission_rate", "cost",
                "efficiency_revenue", "fission_revenue", "actual_total_revenue",
                "actual_roi", "predicted_tail_revenue",
                "predicted_fission_revenue", "total_revenue", "roi",
                "stop_threshold", "scale_threshold", "recommended_action",
                "action_reason", "execution_mode", "ineligible_reason",
                "daily_metrics_json",
            ]
            if snapshot_rows:
                cursor.executemany(
                    f"INSERT INTO roi_entity_snapshot ({', '.join(snapshot_columns)}) "
                    f"VALUES ({', '.join(['%s'] * len(snapshot_columns))})",
                    [
                        [
                            json_dumps(row[column])
                            if column == "daily_metrics_json"
                            else row.get(column)
                            for column in snapshot_columns
                        ]
                        for row in snapshot_rows
                    ],
                )

            action_columns = [
                "run_id", "idempotency_key", "action_type", "account_id",
                "adgroup_id", "dynamic_creative_id", "execution_status",
            ]
            if action_rows:
                cursor.executemany(
                    f"INSERT INTO roi_action_item ({', '.join(action_columns)}) "
                    f"VALUES ({', '.join(['%s'] * len(action_columns))})",
                    [
                        [row.get(column) for column in action_columns]
                        for row in action_rows
                    ],
                )

            cursor.execute(
                """
                UPDATE roi_metric_run
                SET thresholds_json=%s,
                    status='COMPUTED',
                    entity_count=%s,
                    candidate_count=%s,
                    actionable_count=%s,
                    error_message=NULL
                WHERE run_id=%s
                """,
                (
                    json_dumps(thresholds),
                    len(snapshot_rows),
                    sum(bool(row.get("recommended_action")) for row in snapshot_rows),
                    len(action_rows),
                    run_id,
                ),
            )
        connection.commit()
        return len(snapshot_rows), len(action_rows)
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.close()


def mark_published(
    run_id: str,
    *,
    sheet_token: str,
    sheet_url: str,
    message_id: str,
    expires_at: datetime | None,
    requires_approval: bool,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            affected = cursor.execute(
                """
                UPDATE roi_metric_run
                SET status=%s, sheet_token=%s, sheet_url=%s,
                    message_id=%s, expires_at=%s, error_message=NULL
                WHERE run_id=%s AND status IN ('COMPUTED','PENDING_APPROVAL')
                """,
                (
                    "PENDING_APPROVAL" if requires_approval else "COMPLETED",
                    sheet_token,
                    sheet_url,
                    message_id,
                    expires_at,
                    run_id,
                ),
            )
            if not affected:
                raise RuntimeError(f"ROI run cannot be published: {run_id}")
    finally:
        connection.close()


def mark_failed(run_id: str, error_message: str) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE roi_metric_run
                SET status='FAILED', error_message=%s
                WHERE run_id=%s AND status NOT IN ('COMPLETED','PARTIAL','REJECTED')
                """,
                (error_message[:4000], run_id),
            )
    finally:
        connection.close()


def upsert_agency_delivery(record: dict[str, Any]) -> dict[str, Any]:
    """Create or refresh one idempotent agency delivery record."""

    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO roi_agency_delivery
                    (run_id, agency_name, agency_report_version, file_path,
                     file_sha256, creative_rows, ad_rows,
                     route_fingerprint, status)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,'PENDING')
                ON DUPLICATE KEY UPDATE
                    status=CASE
                        WHEN status='SENT' THEN status
                        WHEN NOT (route_fingerprint <=> VALUES(route_fingerprint))
                            THEN 'PENDING'
                        ELSE status
                    END,
                    error_message=CASE
                        WHEN NOT (route_fingerprint <=> VALUES(route_fingerprint))
                            THEN NULL
                        ELSE error_message
                    END,
                    file_path=IF(status='SENT', file_path, VALUES(file_path)),
                    file_sha256=IF(status='SENT', file_sha256, VALUES(file_sha256)),
                    creative_rows=VALUES(creative_rows),
                    ad_rows=VALUES(ad_rows),
                    route_fingerprint=VALUES(route_fingerprint)
                """,
                (
                    record["run_id"],
                    record["agency_name"],
                    record["agency_report_version"],
                    record["file_path"],
                    record["file_sha256"],
                    int(record.get("creative_rows") or 0),
                    int(record.get("ad_rows") or 0),
                    record.get("route_fingerprint"),
                ),
            )
            cursor.execute(
                """
                SELECT * FROM roi_agency_delivery
                WHERE run_id=%s AND agency_name=%s AND agency_report_version=%s
                """,
                (
                    record["run_id"],
                    record["agency_name"],
                    record["agency_report_version"],
                ),
            )
            return cursor.fetchone() or {}
    finally:
        connection.close()


def load_agency_deliveries(run_id: str) -> list[dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT * FROM roi_agency_delivery
                WHERE run_id=%s
                ORDER BY agency_name
                """,
                (run_id,),
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def update_agency_delivery(
    delivery_id: int,
    *,
    status: str,
    sheet_token: str | None = None,
    sheet_url: str | None = None,
    response_code: str | None = None,
    error_message: str | None = None,
    sent_at: datetime | None = None,
    increment_attempt: bool = False,
) -> None:
    allowed_statuses = {"PENDING", "UPLOADED", "SENT", "SKIPPED", "FAILED"}
    if status not in allowed_statuses:
        raise ValueError(f"Unsupported agency delivery status: {status}")
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE roi_agency_delivery
                SET status=%s,
                    sheet_token=COALESCE(%s, sheet_token),
                    sheet_url=COALESCE(%s, sheet_url),
                    response_code=%s,
                    error_message=%s,
                    sent_at=COALESCE(%s, sent_at),
                    attempt_count=attempt_count+%s
                WHERE id=%s
                """,
                (
                    status,
                    sheet_token,
                    sheet_url,
                    response_code,
                    error_message[:4000] if error_message else None,
                    sent_at,
                    1 if increment_attempt else 0,
                    delivery_id,
                ),
            )
    finally:
        connection.close()


def reject_run(run_id: str, *, sender_open_id: str, now: datetime) -> dict[str, Any]:
    connection = connect()
    connection.autocommit(False)
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM roi_metric_run WHERE run_id=%s FOR UPDATE",
                (run_id,),
            )
            run = cursor.fetchone()
            if not run:
                raise ValueError(f"ROI批次不存在: {run_id}")
            if run["status"] in FINAL_STATUSES:
                connection.commit()
                return run
            if run["status"] != "PENDING_APPROVAL":
                raise RuntimeError(f"ROI批次当前不可拒绝: {run['status']}")
            if _is_expired(run.get("expires_at"), now):
                cursor.execute(
                    "UPDATE roi_metric_run SET status='EXPIRED' WHERE run_id=%s",
                    (run_id,),
                )
            else:
                cursor.execute(
                    """
                    UPDATE roi_metric_run
                    SET status='REJECTED', rejected_by=%s, rejected_at=%s
                    WHERE run_id=%s
                    """,
                    (sender_open_id, now, run_id),
                )
        connection.commit()
        return load_run(run_id) or run
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.close()


def claim_run_for_execution(
    run_id: str,
    *,
    sender_open_id: str,
    now: datetime,
) -> dict[str, Any]:
    connection = connect()
    connection.autocommit(False)
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM roi_metric_run WHERE run_id=%s FOR UPDATE",
                (run_id,),
            )
            run = cursor.fetchone()
            if not run:
                raise ValueError(f"ROI批次不存在: {run_id}")
            if run["status"] in FINAL_STATUSES:
                connection.commit()
                return run
            if run["status"] == "PENDING_APPROVAL":
                if _is_expired(run.get("expires_at"), now):
                    cursor.execute(
                        "UPDATE roi_metric_run SET status='EXPIRED' WHERE run_id=%s",
                        (run_id,),
                    )
                    connection.commit()
                    return load_run(run_id) or run
                cursor.execute(
                    """
                    UPDATE roi_metric_run
                    SET status='EXECUTING', approved_by=%s, approved_at=%s
                    WHERE run_id=%s
                    """,
                    (sender_open_id, now, run_id),
                )
            elif run["status"] != "EXECUTING":
                raise RuntimeError(f"ROI批次当前不可确认: {run['status']}")
        connection.commit()
        return load_run(run_id) or run
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.close()


def load_action_items(run_id: str) -> list[dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM roi_action_item WHERE run_id=%s ORDER BY id",
                (run_id,),
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_pending_sheet_runs(now: datetime) -> list[dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT * FROM roi_metric_run
                WHERE status='PENDING_APPROVAL'
                  AND sheet_token IS NOT NULL
                  AND (expires_at IS NULL OR expires_at>=%s)
                ORDER BY created_at
                """,
                (now,),
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def expire_pending_sheet_runs(now: datetime) -> int:
    connection = connect()
    connection.autocommit(False)
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT run_id FROM roi_metric_run
                WHERE status='PENDING_APPROVAL'
                  AND expires_at IS NOT NULL
                  AND expires_at<%s
                FOR UPDATE
                """,
                (now,),
            )
            run_ids = [row["run_id"] for row in cursor.fetchall()]
            if run_ids:
                placeholders = ",".join(["%s"] * len(run_ids))
                cursor.execute(
                    f"""
                    UPDATE roi_action_item
                    SET approval_status='EXPIRED', execution_status='EXPIRED',
                        executed_at=%s
                    WHERE run_id IN ({placeholders})
                      AND approval_status='PENDING'
                    """,
                    [now, *run_ids],
                )
                cursor.execute(
                    f"""
                    UPDATE roi_metric_run SET status='EXPIRED'
                    WHERE run_id IN ({placeholders})
                    """,
                    run_ids,
                )
        connection.commit()
        return len(run_ids)
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.close()


def record_sheet_decision(
    *,
    run_id: str,
    idempotency_key: str,
    decision: str,
    sheet_row_number: int,
    now: datetime,
) -> dict[str, Any] | None:
    if decision not in {"APPROVED", "REJECTED"}:
        raise ValueError(f"Unsupported ROI sheet decision: {decision}")
    connection = connect()
    connection.autocommit(False)
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT status, expires_at FROM roi_metric_run "
                "WHERE run_id=%s FOR UPDATE",
                (run_id,),
            )
            run = cursor.fetchone()
            if not run or run["status"] != "PENDING_APPROVAL":
                connection.rollback()
                return None
            if _is_expired(run.get("expires_at"), now):
                connection.rollback()
                return None
            cursor.execute(
                """
                SELECT * FROM roi_action_item
                WHERE run_id=%s AND idempotency_key=%s
                FOR UPDATE
                """,
                (run_id, idempotency_key),
            )
            item = cursor.fetchone()
            if not item:
                connection.rollback()
                return None
            if item["approval_status"] != "PENDING":
                connection.commit()
                return {**item, "decision_changed": False}
            if decision == "APPROVED":
                cursor.execute(
                    """
                    UPDATE roi_action_item
                    SET approval_status='APPROVED', approval_source='FEISHU_SHEET',
                        approved_at=%s, sheet_row_number=%s
                    WHERE id=%s AND approval_status='PENDING'
                    """,
                    (now, sheet_row_number, item["id"]),
                )
            else:
                cursor.execute(
                    """
                    UPDATE roi_action_item
                    SET approval_status='REJECTED', approval_source='FEISHU_SHEET',
                        rejected_at=%s, sheet_row_number=%s,
                        execution_status='REJECTED', executed_at=%s
                    WHERE id=%s AND approval_status='PENDING'
                    """,
                    (now, sheet_row_number, now, item["id"]),
                )
        connection.commit()
        return {
            **item,
            "approval_status": decision,
            "sheet_row_number": sheet_row_number,
            "decision_changed": True,
        }
    except Exception:
        connection.rollback()
        raise
    finally:
        connection.close()


def load_actions_by_ids(item_ids: Iterable[int]) -> list[dict[str, Any]]:
    ids = [int(value) for value in item_ids]
    if not ids:
        return []
    connection = connect()
    try:
        with connection.cursor() as cursor:
            placeholders = ",".join(["%s"] * len(ids))
            cursor.execute(
                f"SELECT * FROM roi_action_item WHERE id IN ({placeholders}) ORDER BY id",
                ids,
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_unnotified_action_details(
    run_ids: Iterable[str] | None = None,
) -> list[dict[str, Any]]:
    values = [str(value) for value in run_ids] if run_ids is not None else []
    run_filter = ""
    params: list[Any] = []
    if run_ids is not None:
        if not values:
            return []
        placeholders = ",".join(["%s"] * len(values))
        run_filter = f"AND a.run_id IN ({placeholders})"
        params.extend(values)
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                SELECT a.*, r.sheet_url, s.adgroup_name, s.audience_name,
                       s.cost, s.roi
                FROM roi_action_item a
                JOIN roi_metric_run r ON r.run_id=a.run_id
                LEFT JOIN roi_entity_snapshot s ON s.id=(
                    SELECT MIN(s2.id) FROM roi_entity_snapshot s2
                    WHERE s2.run_id=a.run_id
                      AND s2.account_id=a.account_id
                      AND s2.adgroup_id=a.adgroup_id
                      AND (
                          a.dynamic_creative_id IS NULL
                          OR s2.dynamic_creative_id=a.dynamic_creative_id
                      )
                )
                WHERE a.approval_status IN ('APPROVED','REJECTED')
                  AND a.execution_status NOT IN ('PENDING','PREPARED')
                  AND a.result_notified_at IS NULL
                  {run_filter}
                ORDER BY a.run_id, a.id
                LIMIT 100
                """,
                params,
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def mark_action_notifications(
    item_ids: Iterable[int],
    *,
    notified_at: datetime | None,
    error: str | None,
) -> None:
    ids = [int(value) for value in item_ids]
    if not ids:
        return
    connection = connect()
    try:
        with connection.cursor() as cursor:
            placeholders = ",".join(["%s"] * len(ids))
            cursor.execute(
                f"""
                UPDATE roi_action_item
                SET result_notified_at=%s, notification_error=%s
                WHERE id IN ({placeholders})
                """,
                [notified_at, error[:4000] if error else None, *ids],
            )
    finally:
        connection.close()


def finalize_sheet_run_if_resolved(run_id: str, *, now: datetime) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT
                    SUM(approval_status='PENDING') AS pending_approvals,
                    SUM(execution_status='SUCCESS') AS successes,
                    SUM(execution_status IN ('FAILED','OUTCOME_UNKNOWN','VERIFY_FAILED')) AS failures,
                    SUM(execution_status LIKE 'SKIPPED%%') AS skipped,
                    SUM(approval_status='REJECTED') AS rejected
                FROM roi_action_item WHERE run_id=%s
                """,
                (run_id,),
            )
            counts = cursor.fetchone() or {}
            normalized = {key: int(value or 0) for key, value in counts.items()}
            if normalized["pending_approvals"]:
                return {"status": "PENDING_APPROVAL", **normalized}
            status = "PARTIAL" if normalized["failures"] else "COMPLETED"
            cursor.execute(
                "UPDATE roi_metric_run SET status=%s, executed_at=%s WHERE run_id=%s",
                (status, now, run_id),
            )
            return {"status": status, **normalized}
    finally:
        connection.close()


def update_action_item(item_id: int, **values: Any) -> None:
    allowed = {
        "bid_field", "initial_base_bid_fen", "base_bid_fen", "before_bid_fen",
        "target_bid_fen", "before_status", "target_status", "readback_status",
        "execution_status", "skip_reason", "error_message", "pre_state_json",
        "readback_json", "executed_at", "notification_error",
    }
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported ROI action fields: {sorted(unknown)}")
    assignments: list[str] = []
    params: list[Any] = []
    for column, value in values.items():
        assignments.append(f"{column}=%s")
        params.append(
            json_dumps(value)
            if column in {"pre_state_json", "readback_json"} and value is not None
            else value
        )
    if not assignments:
        return
    params.append(item_id)
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"UPDATE roi_action_item SET {', '.join(assignments)} WHERE id=%s",
                params,
            )
    finally:
        connection.close()


def finalize_run(run_id: str, *, now: datetime) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT
                    SUM(execution_status='SUCCESS') AS successes,
                    SUM(execution_status IN ('FAILED','OUTCOME_UNKNOWN','VERIFY_FAILED')) AS failures,
                    SUM(execution_status LIKE 'SKIPPED%%') AS skipped
                FROM roi_action_item WHERE run_id=%s
                """,
                (run_id,),
            )
            counts = cursor.fetchone() or {}
            failures = int(counts.get("failures") or 0)
            status = "PARTIAL" if failures else "COMPLETED"
            cursor.execute(
                "UPDATE roi_metric_run SET status=%s, executed_at=%s WHERE run_id=%s",
                (status, now, run_id),
            )
        return {"status": status, **{key: int(value or 0) for key, value in counts.items()}}
    finally:
        connection.close()