"""Add ``apikey_budget_alerts_dedup`` + seed ``apikey.budget_check`` cron. Revision ID: 143 Revises: 142 Create Date: 2026-05-10 Why this migration: * :class:`ApiKeyAssignment` already declares ``monthly_budget_usd`` (``Field(..., ge=20)``, USD) on every API-key request — it's been in the schema since M0. PR #240 (2026-05-10) closed the per-IAM cost-attribution gap so ``cmdb_cost_facts.cost_owner_person_id`` is now denormalised at ingest time for AWS Bedrock / Azure OpenAI / future per-applicant flows. * What was missing: a daily cron that compares MTD spend per applicant to ``monthly_budget_usd`` and notifies when 80% / 100% is crossed. Without it the budget field is decorative — applicants declare a number nobody verifies. Schema: * ``apikey_budget_alerts_dedup`` — UNIQUE on ``(assignment_id, billing_month, threshold_pct)`` so each threshold fires at most once per month per assignment, regardless of how many cron ticks see the same overrun. Mirrors :class:`CmdbBillingAlertsDedup` structurally. Cron seed: * ``apikey.budget_check`` — daily at ``09:M Asia/Shanghai`` (M is a hash-spread minute), handler_key matches the ``@register_scheduled_job(key=...)`` decorator on :py:func:`run_apikey_budget_check`. 09:xx SH sits after every cloud's ``cmdb.cost_facts_sync`` tick (04:xx UTC = 12:xx SH for most clouds; the daily fact-ingest is reliably done by 09:00 SH for the prior day) so the budget check evaluates against fresh allocated facts. APScheduler interprets the cron HOUR field IN the ``timezone`` column (see ``app/worker.py`` ``CronTrigger.from_crontab(..., timezone=row.timezone)``), so ``cron="M 9 * * *"`` + ``timezone='Asia/Shanghai'`` fires at 09:M SH. Single-flight by handler_key (the runner's ``execute_job`` outer lock). Idempotent — uses ``CREATE TABLE IF NOT EXISTS``-equivalent guards and ``INSERT ... NOT EXISTS`` for the cron seed. Down-revision drops the table and the seed row. """ from __future__ import annotations import hashlib import json as _json import uuid import sqlalchemy as sa from alembic import op revision = "143" down_revision = "142" branch_labels = None depends_on = None _CRON_NAME = "apikey.budget_check" def _table_exists(bind: sa.engine.Connection, name: str) -> bool: inspector = sa.inspect(bind) return name in inspector.get_table_names() def _cron_minute_for(name: str) -> int: """Spread cron minutes by hashing the job name. The platform's other crons (cmdb.cost_facts_sync, etc.) hash on name → minute so a fleet of jobs sharing the same hour don't stampede the top-of-hour. Same convention here for parity. Stable across deploys (no salt) so the same job lands on the same minute every time. """ return int(hashlib.sha256(name.encode("utf-8")).hexdigest(), 16) % 60 def upgrade() -> None: bind = op.get_bind() # ── 1. Dedup table ───────────────────────────────────────────── if not _table_exists(bind, "apikey_budget_alerts_dedup"): # ``id`` and ``assignment_id`` are CHAR(36) (not VARCHAR(36)) # to match ``api_key_assignments.id`` exactly. Per CLAUDE.md # memory ``feedback_tidb_fk_strictness`` TiDB rejects # CHAR↔VARCHAR FK pairings that local MySQL silently accepts # — using CHAR(36) on both sides lets the FK + ``ON DELETE # CASCADE`` actually take effect in production, so a hard- # deleted assignment doesn't leave orphan dedup rows. The # ``GUID`` SQLAlchemy type used by ``ApiKeyAssignment.id`` / # ``ApiKeyBudgetAlertsDedup.assignment_id`` resolves to # CHAR(36) on MySQL/TiDB. op.create_table( "apikey_budget_alerts_dedup", sa.Column("id", sa.CHAR(36), primary_key=True), sa.Column( "assignment_id", sa.CHAR(36), nullable=False, ), sa.Column("billing_month", sa.String(7), nullable=False), sa.Column("threshold_pct", sa.Integer, nullable=False), sa.Column( "last_alerted_at", sa.DateTime(timezone=True), nullable=False, ), sa.Column( "last_alerted_amount_usd", sa.Numeric(14, 2), nullable=False, ), sa.Column( "last_alerted_budget_usd", sa.Numeric(14, 2), nullable=False, ), sa.Column( "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now(), ), sa.UniqueConstraint( "assignment_id", "billing_month", "threshold_pct", name="uq_apikey_budget_alerts_dedup_assign_month_pct", ), sa.Index( "ix_apikey_budget_alerts_dedup_month", "billing_month", ), ) # FK SHOULD now succeed on both MySQL and TiDB because both # sides are CHAR(36). The try/except stays as defense-in-depth # in case a future schema change drifts the types — better to # log + leave the orphan-cache caveat than to fail the # migration mid-deploy. ``ON DELETE CASCADE`` keeps the dedup # ledger aligned with the assignment lifecycle: a hard-delete # on an assignment cascades the dedup history (the rows are # meaningless without their parent). try: op.create_foreign_key( "fk_apikey_budget_alerts_dedup_assignment", "apikey_budget_alerts_dedup", "api_key_assignments", ["assignment_id"], ["id"], ondelete="CASCADE", ) except Exception: # noqa: BLE001 pass # ── 2. Seed the cron row ─────────────────────────────────────── if _table_exists(bind, "scheduled_jobs"): minute = _cron_minute_for(_CRON_NAME) config_json = _json.dumps({}) # no per-row config — handler iterates bind.execute( sa.text( "INSERT INTO scheduled_jobs " "(id, name, display_name, description, handler_key, " "trigger_type, cron_expr, timezone, enabled, config, " "run_count, failure_count, created_at, updated_at) " "SELECT :id, :name, :display, :desc, :handler, " "'cron', :cron, 'Asia/Shanghai', 1, :config, " "0, 0, NOW(), NOW() " "FROM DUAL WHERE NOT EXISTS " "(SELECT 1 FROM scheduled_jobs WHERE name = :name)" ), { "id": str(uuid.uuid4()), "name": _CRON_NAME, "display": "API Key 月度预算检查", "desc": ( "每日上午 09:30 Asia/Shanghai 检查 status='active' 的 " "ApiKeyAssignment:从 cmdb_cost_facts 按 cost_owner_" "person_id 聚合本月已花费,对照 monthly_budget_usd 触发 " "80% / 100% 阈值告警(申请人 + 直属上级,in-app + " "Feishu),每 (assignment, month, threshold) 在 " "apikey_budget_alerts_dedup 只落一行。" ), "handler": _CRON_NAME, # APScheduler interprets the cron HOUR field IN the # ``timezone`` column (``Asia/Shanghai``) — see # ``app/worker.py`` ``CronTrigger.from_crontab(..., # timezone=row.timezone)``. So ``"M 9 * * *"`` with # ``timezone='Asia/Shanghai'`` fires at 09:M Asia/ # Shanghai. Earlier this row read ``"M 1 * * *"`` # (matching the description "09:30 SH") which would # actually have fired at 01:M Asia/Shanghai — 8 h # before the day's ``cmdb.cost_facts_sync`` (04:xx # UTC = 12:xx SH), evaluating against stale facts. # Caught in code-review pass 1 before merge. "cron": f"{minute} 9 * * *", "config": config_json, }, ) def downgrade() -> None: bind = op.get_bind() if _table_exists(bind, "scheduled_jobs"): bind.execute( sa.text("DELETE FROM scheduled_jobs WHERE name = :n"), {"n": _CRON_NAME}, ) if _table_exists(bind, "apikey_budget_alerts_dedup"): # Drop the FK first (best-effort — may not exist on TiDB). try: op.drop_constraint( "fk_apikey_budget_alerts_dedup_assignment", "apikey_budget_alerts_dedup", type_="foreignkey", ) except Exception: # noqa: BLE001 pass op.drop_table("apikey_budget_alerts_dedup")