demand_server/app/strategies/config_store.py

import json
import re
from dataclasses import dataclass
from typing import Any

from sqlalchemy import text

from app.core.config import settings
from app.db.mysql import SessionLocal

IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")


def _safe_identifier(name: str) -> str:
    if not IDENTIFIER_RE.match(name):
        raise ValueError(f"invalid sql identifier: {name}")
    return name


@dataclass(frozen=True)
class StrategyConfigRecord:
    strategy_id: str
    name: str
    version: str
    params: dict[str, Any]
    active: bool
    daily_write_limit: int = 0
    priority: int = 0
    create_time: str | None = None
    updated_time: str | None = None


def _parse_params(raw: object) -> dict[str, Any]:
    if raw is None:
        return {}
    if isinstance(raw, dict):
        return raw
    if isinstance(raw, (bytes, bytearray)):
        raw = raw.decode("utf-8")
    if isinstance(raw, str):
        text_value = raw.strip()
        if not text_value:
            return {}
        return json.loads(text_value)
    raise TypeError(f"unsupported params type: {type(raw)!r}")


def ensure_strategy_tables() -> None:
    config_table = _safe_identifier(settings.strategy_config_table)
    staging_table = _safe_identifier(settings.strategy_staging_table)

    config_sql = f"""
    CREATE TABLE IF NOT EXISTS {config_table}
    (
        id           BIGINT       NOT NULL AUTO_INCREMENT COMMENT '自增id',
        strategy_id  VARCHAR(64)  NOT NULL COMMENT '全局唯一id，英文名称',
        name         VARCHAR(128) NOT NULL COMMENT '策略名称',
        version      VARCHAR(32)  NOT NULL COMMENT '策略版本，用于追踪代码变更',
        params       JSON         NULL COMMENT '策略运行参数，如阈值、窗口期等',
        daily_write_limit INT     NOT NULL DEFAULT 0 COMMENT '实验：每日写入上限，0=不限',
        priority     INT          NOT NULL DEFAULT 0 COMMENT '实验：跨策略去重优先级，越小越优先',
        active       TINYINT(1)   NOT NULL DEFAULT 1 COMMENT '是否启用，0=不参与产出和合并',
        create_time  TIMESTAMP    NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
        updated_time TIMESTAMP    NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
        PRIMARY KEY (id),
        UNIQUE KEY uk_strategy_id (strategy_id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='策略配置表'
    """

    staging_sql = f"""
    CREATE TABLE IF NOT EXISTS {staging_table}
    (
        id                 BIGINT        NOT NULL AUTO_INCREMENT COMMENT '自增id',
        strategy_config_id VARCHAR(64)   NOT NULL COMMENT '来源策略ID，关联 strategy_config.strategy_id',
        strategy           VARCHAR(32)   NULL COMMENT '策略名称',
        demand_id          VARCHAR(64)   NULL COMMENT '需求id',
        demand_name        VARCHAR(512)  NULL COMMENT '需求',
        weight             DOUBLE        NULL COMMENT '需求权重',
        `type`             VARCHAR(64)   NULL COMMENT '需求类型',
        video_count        INT           NULL COMMENT '视频数量',
        video_list         JSON          NULL COMMENT '视频列表',
        extend             JSON          NULL COMMENT '拓展字段',
        batch_date         VARCHAR(32)   NULL COMMENT '归属批次日期，e.g. 20250526',
        create_time        TIMESTAMP     NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
        PRIMARY KEY (id),
        KEY idx_strategy_config_id (strategy_config_id),
        KEY idx_batch_date (batch_date),
        KEY idx_demand_id (demand_id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='策略暂存表'
    """

    with SessionLocal() as session:
        session.execute(text(config_sql))
        session.execute(text(staging_sql))
        _migrate_strategy_config_columns(session, config_table)
        session.commit()


def _migrate_strategy_config_columns(session, config_table: str) -> None:
    rows = session.execute(
        text(
            """
            SELECT COLUMN_NAME
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = DATABASE()
              AND TABLE_NAME = :table_name
            """
        ),
        {"table_name": config_table},
    ).all()
    existing = {str(row[0]) for row in rows}
    if "daily_write_limit" not in existing:
        session.execute(
            text(
                f"""
                ALTER TABLE {config_table}
                ADD COLUMN daily_write_limit INT NOT NULL DEFAULT 0
                COMMENT '实验：每日写入上限，0=不限'
                """
            )
        )
    if "priority" not in existing:
        session.execute(
            text(
                f"""
                ALTER TABLE {config_table}
                ADD COLUMN priority INT NOT NULL DEFAULT 0
                COMMENT '实验：跨策略去重优先级，越小越优先'
                """
            )
        )


def _format_timestamp(value: object) -> str | None:
    if value is None:
        return None
    if hasattr(value, "isoformat"):
        return value.isoformat(sep=" ", timespec="seconds")
    return str(value)


def _row_to_record(row: object) -> StrategyConfigRecord:
    mapping = row._mapping if hasattr(row, "_mapping") else row
    return StrategyConfigRecord(
        strategy_id=str(mapping["strategy_id"]),
        name=str(mapping["name"]),
        version=str(mapping["version"]),
        params=_parse_params(mapping["params"]),
        active=bool(mapping["active"]),
        daily_write_limit=int(mapping.get("daily_write_limit") or 0),
        priority=int(mapping.get("priority") or 0),
        create_time=_format_timestamp(mapping.get("create_time")),
        updated_time=_format_timestamp(mapping.get("updated_time")),
    )


def fetch_all_configs() -> list[StrategyConfigRecord]:
    config_table = _safe_identifier(settings.strategy_config_table)
    query = text(
        f"""
        SELECT strategy_id, name, version, params, active,
               daily_write_limit, priority, create_time, updated_time
        FROM {config_table}
        ORDER BY strategy_id ASC
        """
    )
    with SessionLocal() as session:
        rows = session.execute(query).mappings().all()

    return [_row_to_record(row) for row in rows]


def fetch_config_by_id(strategy_id: str) -> StrategyConfigRecord | None:
    config_table = _safe_identifier(settings.strategy_config_table)
    query = text(
        f"""
        SELECT strategy_id, name, version, params, active,
               daily_write_limit, priority, create_time, updated_time
        FROM {config_table}
        WHERE strategy_id = :strategy_id
        LIMIT 1
        """
    )
    with SessionLocal() as session:
        row = session.execute(query, {"strategy_id": strategy_id}).mappings().first()

    if row is None:
        return None

    return _row_to_record(row)


def insert_config(
    *,
    strategy_id: str,
    name: str,
    version: str,
    params: dict[str, Any],
    active: bool,
    daily_write_limit: int = 0,
    priority: int = 0,
) -> StrategyConfigRecord:
    config_table = _safe_identifier(settings.strategy_config_table)
    query = text(
        f"""
        INSERT INTO {config_table}
        (strategy_id, name, version, params, active, daily_write_limit, priority)
        VALUES
        (:strategy_id, :name, :version, CAST(:params AS JSON), :active,
         :daily_write_limit, :priority)
        """
    )
    with SessionLocal() as session:
        session.execute(
            query,
            {
                "strategy_id": strategy_id,
                "name": name,
                "version": version,
                "params": json.dumps(params, ensure_ascii=False),
                "active": int(active),
                "daily_write_limit": daily_write_limit,
                "priority": priority,
            },
        )
        session.commit()

    record = fetch_config_by_id(strategy_id)
    if record is None:
        raise RuntimeError(f"failed to load inserted strategy config: {strategy_id}")
    return record


def update_config(
    *,
    strategy_id: str,
    params: dict[str, Any] | None = None,
    active: bool | None = None,
    daily_write_limit: int | None = None,
    priority: int | None = None,
) -> StrategyConfigRecord:
    if (
        params is None
        and active is None
        and daily_write_limit is None
        and priority is None
    ):
        raise ValueError("至少需提供一项可更新字段")

    config_table = _safe_identifier(settings.strategy_config_table)
    set_clauses: list[str] = []
    payload: dict[str, object] = {"strategy_id": strategy_id}
    if params is not None:
        set_clauses.append("params = CAST(:params AS JSON)")
        payload["params"] = json.dumps(params, ensure_ascii=False)
    if active is not None:
        set_clauses.append("active = :active")
        payload["active"] = int(active)
    if daily_write_limit is not None:
        set_clauses.append("daily_write_limit = :daily_write_limit")
        payload["daily_write_limit"] = daily_write_limit
    if priority is not None:
        set_clauses.append("priority = :priority")
        payload["priority"] = priority

    query = text(
        f"""
        UPDATE {config_table}
        SET {", ".join(set_clauses)}
        WHERE strategy_id = :strategy_id
        """
    )
    with SessionLocal() as session:
        result = session.execute(query, payload)
        if result.rowcount == 0:
            raise KeyError(f"strategy config not found: {strategy_id}")
        session.commit()

    record = fetch_config_by_id(strategy_id)
    if record is None:
        raise RuntimeError(f"failed to load updated strategy config: {strategy_id}")
    return record