Agent/examples/tencent_realtime_control/storage.py

"""MySQL state and audit storage for real-time control."""

from __future__ import annotations

import os
import json
from contextlib import contextmanager
from datetime import date, datetime
from pathlib import Path
from typing import Any, Iterator

import pymysql
import pymysql.cursors


ROOT = Path(__file__).resolve().parent


def connect() -> pymysql.Connection:
    required = ["DB_HOST", "DB_USER", "DB_NAME"]
    missing = [key for key in required if not os.getenv(key)]
    if missing:
        raise RuntimeError(f"Missing database environment variables: {', '.join(missing)}")
    return pymysql.connect(
        host=os.environ["DB_HOST"],
        port=int(os.getenv("DB_PORT", "3306")),
        user=os.environ["DB_USER"],
        password=os.getenv("DB_PASSWORD", ""),
        database=os.environ["DB_NAME"],
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        autocommit=True,
        connect_timeout=int(os.getenv("DB_CONNECT_TIMEOUT", "10")),
        read_timeout=int(os.getenv("DB_READ_TIMEOUT", "60")),
        write_timeout=int(os.getenv("DB_WRITE_TIMEOUT", "60")),
    )


def initialize_schema() -> None:
    statements = [
        statement.strip()
        for statement in (ROOT / "schema.sql").read_text(encoding="utf-8").split(";")
        if statement.strip()
    ]
    connection = connect()
    try:
        with connection.cursor() as cursor:
            for statement in statements:
                cursor.execute(statement)
            cursor.execute(
                """
                SELECT COLUMN_NAME
                FROM information_schema.COLUMNS
                WHERE TABLE_SCHEMA=%s
                  AND TABLE_NAME='realtime_control_ad_state'
                """,
                (os.environ["DB_NAME"],),
            )
            existing_columns = {row["COLUMN_NAME"] for row in cursor.fetchall()}
            migrations = {
                "operator_pause_mode": "VARCHAR(32) DEFAULT NULL",
                "operator_resume_at": "DATETIME DEFAULT NULL",
                "operator_command_id": "VARCHAR(64) DEFAULT NULL",
                "operator_paused_from_status": "VARCHAR(50) DEFAULT NULL",
                "operator_paused_at": "DATETIME DEFAULT NULL",
            }
            for column, definition in migrations.items():
                if column not in existing_columns:
                    cursor.execute(
                        f"ALTER TABLE realtime_control_ad_state "
                        f"ADD COLUMN {column} {definition}"
                    )
            cursor.execute(
                """
                SELECT INDEX_NAME
                FROM information_schema.STATISTICS
                WHERE TABLE_SCHEMA=%s
                  AND TABLE_NAME='realtime_control_ad_state'
                  AND INDEX_NAME='idx_operator_pause'
                """,
                (os.environ["DB_NAME"],),
            )
            if not cursor.fetchone():
                cursor.execute(
                    """
                    CREATE INDEX idx_operator_pause
                    ON realtime_control_ad_state
                        (operator_pause_mode, operator_resume_at)
                    """
                )
    finally:
        connection.close()


@contextmanager
def advisory_lock(lock_name: str) -> Iterator[bool]:
    connection = connect()
    acquired = False
    try:
        with connection.cursor() as cursor:
            cursor.execute("SELECT GET_LOCK(%s, 0) AS acquired", (lock_name,))
            acquired = bool((cursor.fetchone() or {}).get("acquired"))
        yield acquired
    finally:
        if acquired:
            try:
                with connection.cursor() as cursor:
                    cursor.execute("SELECT RELEASE_LOCK(%s)", (lock_name,))
            except Exception:
                pass
        connection.close()


def load_enabled_accounts() -> list[dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT c.account_id, c.audience_name, c.bid_scene
                FROM ad_creation_account_config c
                JOIN account_whitelist w ON w.account_id = c.account_id
                WHERE c.enabled = TRUE
                  AND w.enabled = TRUE
                ORDER BY c.account_id
                """
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_realtime_accounts() -> list[dict[str, Any]]:
    """All historical automation accounts still enabled by the whitelist."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT c.account_id, c.audience_name, c.bid_scene
                FROM ad_creation_account_config c
                JOIN account_whitelist w ON w.account_id = c.account_id
                WHERE w.enabled = TRUE
                ORDER BY c.account_id
                """
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_managed_accounts() -> list[dict[str, Any]]:
    """Historical automation accounts still allowed by the account whitelist."""
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT c.account_id, c.audience_name, c.bid_scene, c.enabled
                FROM ad_creation_account_config c
                JOIN account_whitelist w ON w.account_id = c.account_id
                WHERE w.enabled = TRUE
                ORDER BY c.account_id
                """
            )
            return list(cursor.fetchall())
    finally:
        connection.close()


def load_daily_state(control_date: date) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM realtime_control_daily_state WHERE control_date=%s",
                (control_date,),
            )
            return cursor.fetchone() or {}
    finally:
        connection.close()


def save_daily_state(control_date: date, **values: Any) -> None:
    allowed = {
        "morning_recovery_done",
        "cutoff_done",
        "last_observed_partition",
        "last_observed_cpm",
        "last_decision",
        "last_inventory_refresh_at",
        "last_evaluated_at",
    }
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported daily-state fields: {sorted(unknown)}")
    columns = ["control_date", *values]
    params = [control_date, *values.values()]
    updates = ", ".join(f"{column}=VALUES({column})" for column in values)
    placeholders = ", ".join(["%s"] * len(columns))
    sql = (
        f"INSERT INTO realtime_control_daily_state ({', '.join(columns)}) "
        f"VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {updates}"
    )
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(sql, params)
    finally:
        connection.close()


def load_ad_states(account_id: int) -> dict[int, dict[str, Any]]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM realtime_control_ad_state WHERE account_id=%s",
                (account_id,),
            )
            return {int(row["adgroup_id"]): row for row in cursor.fetchall()}
    finally:
        connection.close()


def upsert_ad_state(
    *,
    account_id: int,
    adgroup_id: int,
    adgroup_name: str,
    bid_field: str,
    base_bid_fen: int,
    boosted_date: date | None,
    paused_by_strategy: bool,
    pause_reason: str | None,
    last_action: str,
    action_at: datetime,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO realtime_control_ad_state
                    (account_id, adgroup_id, adgroup_name, bid_field, base_bid_fen,
                     boosted_date, paused_by_strategy, pause_reason, last_action,
                     last_action_at, last_seen_at)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                ON DUPLICATE KEY UPDATE
                    adgroup_name=VALUES(adgroup_name),
                    bid_field=VALUES(bid_field),
                    boosted_date=VALUES(boosted_date),
                    paused_by_strategy=VALUES(paused_by_strategy),
                    pause_reason=VALUES(pause_reason),
                    last_action=VALUES(last_action),
                    last_action_at=VALUES(last_action_at),
                    last_seen_at=VALUES(last_seen_at)
                """,
                (
                    account_id,
                    adgroup_id,
                    adgroup_name,
                    bid_field,
                    base_bid_fen,
                    boosted_date,
                    paused_by_strategy,
                    pause_reason,
                    last_action,
                    action_at,
                    action_at,
                ),
            )
    finally:
        connection.close()

def insert_action_log(record: dict[str, Any]) -> None:
    columns = [
        "run_id",
        "control_date",
        "observed_partition",
        "observed_cpm",
        "decision",
        "account_id",
        "adgroup_id",
        "adgroup_name",
        "bid_field",
        "base_bid_fen",
        "before_bid_fen",
        "target_bid_fen",
        "before_status",
        "target_status",
        "apply_mode",
        "execution_status",
        "error_message",
    ]
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO realtime_control_action_log ({', '.join(columns)}) "
                f"VALUES ({', '.join(['%s'] * len(columns))})",
                [record.get(column) for column in columns],
            )
    finally:
        connection.close()


def create_operator_command(record: dict[str, Any]) -> dict[str, Any]:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            try:
                cursor.execute(
                    """
                    INSERT INTO operator_command
                        (command_id, source_message_id, chat_id, sender_open_id,
                         sender_name, action, scope_type, target_account_ids,
                         status, preview_account_count, preview_ad_count, expires_at)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                    """,
                    (
                        record["command_id"],
                        record["source_message_id"],
                        record["chat_id"],
                        record["sender_open_id"],
                        record.get("sender_name"),
                        record["action"],
                        record["scope_type"],
                        json.dumps(record["target_account_ids"]),
                        record["status"],
                        record.get("preview_account_count", 0),
                        record.get("preview_ad_count", 0),
                        record.get("expires_at"),
                    ),
                )
            except pymysql.err.IntegrityError:
                cursor.execute(
                    "SELECT * FROM operator_command WHERE source_message_id=%s",
                    (record["source_message_id"],),
                )
                existing = cursor.fetchone()
                if existing:
                    existing["target_account_ids"] = json.loads(
                        existing.get("target_account_ids") or "[]"
                    )
                    return existing
                raise
        return load_operator_command(record["command_id"]) or {}
    finally:
        connection.close()


def load_operator_command(command_id: str) -> dict[str, Any] | None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                "SELECT * FROM operator_command WHERE command_id=%s",
                (command_id,),
            )
            row = cursor.fetchone()
            if row:
                row["target_account_ids"] = json.loads(
                    row.get("target_account_ids") or "[]"
                )
            return row
    finally:
        connection.close()


def update_operator_command(command_id: str, status: str, **values: Any) -> None:
    allowed = {"confirmed_at", "executed_at", "error_message"}
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported operator-command fields: {sorted(unknown)}")
    assignments = ["status=%s"]
    params: list[Any] = [status]
    for column, value in values.items():
        assignments.append(f"{column}=%s")
        params.append(value)
    params.append(command_id)
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"UPDATE operator_command SET {', '.join(assignments)} "
                "WHERE command_id=%s",
                params,
            )
    finally:
        connection.close()


def transition_operator_command(
    command_id: str,
    *,
    expected_statuses: set[str],
    target_status: str,
    **values: Any,
) -> bool:
    allowed = {"confirmed_at", "executed_at", "error_message"}
    unknown = set(values) - allowed
    if unknown:
        raise ValueError(f"Unsupported operator-command fields: {sorted(unknown)}")
    if not expected_statuses:
        raise ValueError("expected_statuses must not be empty")
    assignments = ["status=%s"]
    params: list[Any] = [target_status]
    for column, value in values.items():
        assignments.append(f"{column}=%s")
        params.append(value)
    placeholders = ", ".join(["%s"] * len(expected_statuses))
    params.extend([command_id, *sorted(expected_statuses)])
    connection = connect()
    try:
        with connection.cursor() as cursor:
            affected = cursor.execute(
                f"""
                UPDATE operator_command
                SET {', '.join(assignments)}
                WHERE command_id=%s
                  AND status IN ({placeholders})
                """,
                params,
            )
            return affected == 1
    finally:
        connection.close()


def insert_operator_command_item(record: dict[str, Any]) -> None:
    columns = [
        "command_id",
        "account_id",
        "audience_name",
        "adgroup_id",
        "adgroup_name",
        "before_status",
        "target_status",
        "readback_status",
        "execution_status",
        "error_message",
    ]
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO operator_command_item ({', '.join(columns)}) "
                f"VALUES ({', '.join(['%s'] * len(columns))})",
                [record.get(column) for column in columns],
            )
    finally:
        connection.close()


def set_operator_pause(
    *,
    account_id: int,
    adgroup_id: int,
    mode: str,
    resume_at: datetime | None,
    command_id: str,
    paused_from_status: str,
    paused_at: datetime,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE realtime_control_ad_state
                SET operator_pause_mode=%s,
                    operator_resume_at=%s,
                    operator_command_id=%s,
                    operator_paused_from_status=%s,
                    operator_paused_at=%s,
                    last_action='OPERATOR_PAUSE',
                    last_action_at=%s,
                    last_seen_at=%s
                WHERE account_id=%s AND adgroup_id=%s
                """,
                (
                    mode,
                    resume_at,
                    command_id,
                    paused_from_status,
                    paused_at,
                    paused_at,
                    paused_at,
                    account_id,
                    adgroup_id,
                ),
            )
    finally:
        connection.close()


def clear_operator_pause(
    account_id: int,
    adgroup_id: int,
    *,
    action: str,
    action_at: datetime,
) -> None:
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                UPDATE realtime_control_ad_state
                SET operator_pause_mode=NULL,
                    operator_resume_at=NULL,
                    operator_command_id=NULL,
                    operator_paused_from_status=NULL,
                    operator_paused_at=NULL,
                    last_action=%s,
                    last_action_at=%s,
                    last_seen_at=%s
                WHERE account_id=%s AND adgroup_id=%s
                """,
                (action, action_at, action_at, account_id, adgroup_id),
            )
    finally:
        connection.close()


def load_operator_pauses(
    account_ids: list[int] | None = None,
) -> list[dict[str, Any]]:
    params: list[Any] = []
    where = "WHERE operator_pause_mode IS NOT NULL"
    if account_ids is not None:
        if not account_ids:
            return []
        where += f" AND account_id IN ({', '.join(['%s'] * len(account_ids))})"
        params.extend(account_ids)
    connection = connect()
    try:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                SELECT *
                FROM realtime_control_ad_state
                {where}
                ORDER BY account_id, adgroup_id
                """,
                params,
            )
            return list(cursor.fetchall())
    finally:
        connection.close()