shaguabijia-app-server/alembic/versions/limit_policy_global_bundle.py

"""store all global limit values in one complete JSON document

Revision ID: limit_policy_global_bundle
Revises: limit_policy_whitelist
"""
from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

revision: str = "limit_policy_global_bundle"
down_revision: str | Sequence[str] | None = "limit_policy_whitelist"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

_JSON = sa.JSON().with_variant(postgresql.JSONB(), "postgresql")
_BUNDLE_KEY = "limit_policy_global"

# rule_code, old sparse key, default, legacy structured key, legacy JSON field
_RULES: tuple[tuple[str, str, int, str | None, str | None], ...] = (
    ("compare.start.daily", "compare_daily_limit", 100, None, None),
    ("sms.send.hourly", "sms_send_hourly_limit", 5, None, None),
    ("sms.send.daily", "sms_send_daily_limit", 20, None, None),
    ("sms.phone.cooldown", "sms_phone_cooldown_seconds", 60, None, None),
    ("sms.code.failed_attempts", "sms_code_max_failed_attempts", 5, None, None),
    ("sms.login.hourly", "sms_login_hourly_limit", 5, None, None),
    ("wechat.bind.hourly", "wechat_bind_sms_hourly_limit", 5, None, None),
    ("wechat.conflict.hourly", "wechat_conflict_hourly_limit", 5, None, None),
    (
        "ad.reward_video.daily",
        "ad_reward_video_daily_limit",
        500,
        "ad_daily_limit",
        None,
    ),
    (
        "ad.feed.daily",
        "ad_feed_daily_limit",
        500,
        "ad_daily_limit",
        None,
    ),
    ("ad.reward_video.cooldown", "ad_cooldown_sec", 3, None, None),
    (
        "guide.video.lifetime",
        "guide_video_max_plays",
        3,
        "coupon_guide_video",
        "max_plays",
    ),
    ("phone.rebind.days", "phone_rebind_days", 30, None, None),
    ("risk.sms.hourly", "risk_sms_hourly_threshold", 5, None, None),
    (
        "risk.oneclick.daily",
        "risk_oneclick_daily_threshold",
        20,
        None,
        None,
    ),
    (
        "risk.compare.daily",
        "risk_compare_daily_threshold",
        100,
        None,
        None,
    ),
)


def _table() -> sa.TableClause:
    return sa.table(
        "app_config",
        sa.column("key", sa.String(64)),
        sa.column("value", _JSON),
        sa.column("updated_by_admin_id", sa.Integer),
        sa.column("updated_at", sa.DateTime(timezone=True)),
    )


def _row(conn, table, key: str):
    return conn.execute(
        sa.select(
            table.c.value,
            table.c.updated_by_admin_id,
        ).where(table.c.key == key)
    ).mappings().first()


def _int_or_none(value: Any) -> int | None:
    if isinstance(value, bool):
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def upgrade() -> None:
    conn = op.get_bind()
    table = _table()
    values = {rule_code: default for rule_code, _, default, _, _ in _RULES}

    # Old shared/structured values are the lowest-precedence compatibility
    # source.  Dedicated sparse keys override them.
    for rule_code, _, _, legacy_key, legacy_field in _RULES:
        if legacy_key is None:
            continue
        legacy = _row(conn, table, legacy_key)
        if legacy is None:
            continue
        raw = legacy["value"]
        if legacy_field is not None:
            raw = raw.get(legacy_field) if isinstance(raw, dict) else None
        parsed = _int_or_none(raw)
        if parsed is not None:
            values[rule_code] = parsed

    for rule_code, sparse_key, _, _, _ in _RULES:
        sparse = _row(conn, table, sparse_key)
        parsed = _int_or_none(sparse["value"]) if sparse is not None else None
        if parsed is not None:
            values[rule_code] = parsed

    # If a deployment already wrote the new key, preserve it over old keys.
    bundle = _row(conn, table, _BUNDLE_KEY)
    if bundle is not None and isinstance(bundle["value"], dict):
        for rule_code, raw in bundle["value"].items():
            if rule_code not in values:
                continue
            parsed = _int_or_none(raw)
            if parsed is not None:
                values[rule_code] = parsed

    if bundle is None:
        conn.execute(
            table.insert().values(
                key=_BUNDLE_KEY,
                value=values,
                updated_by_admin_id=None,
            )
        )
    else:
        conn.execute(
            table.update()
            .where(table.c.key == _BUNDLE_KEY)
            .values(value=values, updated_at=sa.func.now())
        )

    sparse_keys = [sparse_key for _, sparse_key, _, _, _ in _RULES]
    conn.execute(table.delete().where(table.c.key.in_(sparse_keys)))


def downgrade() -> None:
    conn = op.get_bind()
    table = _table()
    bundle = _row(conn, table, _BUNDLE_KEY)
    values = (
        bundle["value"]
        if bundle is not None and isinstance(bundle["value"], dict)
        else {}
    )
    admin_id = bundle["updated_by_admin_id"] if bundle is not None else None

    for rule_code, sparse_key, default, _, _ in _RULES:
        value = _int_or_none(values.get(rule_code))
        if value is None:
            value = default
        existing = _row(conn, table, sparse_key)
        if existing is None:
            conn.execute(
                table.insert().values(
                    key=sparse_key,
                    value=value,
                    updated_by_admin_id=admin_id,
                )
            )
        else:
            conn.execute(
                table.update()
                .where(table.c.key == sparse_key)
                .values(
                    value=value,
                    updated_by_admin_id=admin_id,
                    updated_at=sa.func.now(),
                )
            )

    conn.execute(table.delete().where(table.c.key == _BUNDLE_KEY))