"""AI Spending dashboard — schema, seeds, and the Bailian LLM hook-up. Revision ID: 079 Revises: 078 Create Date: 2026-04-28 (Renumbered from 078 → 079 mid-PR because PR #91 ``fix(migrations): forward-fix env column missing on TiDB after 076`` landed ``078_cloud_scope_environment_forward_fix.py`` on main while this branch was still WIP. Same pattern PR #90 used to deconflict 076↔077.) What this migration does ======================== The "AI 支出大盘" feature lives across four data sources Helix already has (software seats / cloud monthly billing / cloud daily billing / Bailian usage) plus three new structures introduced here: 1. **``asset_software_catalog.is_ai_tool``** (BOOLEAN) — flags the software entries that count as AI for the dashboard's SaaS-seat half. Default false; the migration auto-marks any catalog row whose name matches the canonical AI vendor list (OpenAI / Claude / Cursor / GitHub Copilot / Gemini / Bailian / Bedrock / Azure OpenAI / Anthropic / Microsoft Copilot). Admins can flip the flag at any time from ``/admin/assets/software-catalog``. 2. **``api_key_assignments.cloud_scope_id``** (UUID FK → ``cmdb_cloud_scopes.id``, nullable) — lets an Azure OpenAI sub (which is a ``cmdb_cloud_scopes`` row, not an IAM user) be the binding target for an API-key assignment, so chargeback against the sub's monthly billing has a per-employee owner. 3. **``cmdb_cloud_scopes.is_ai_dedicated``** (BOOLEAN) — marks "100 % AI" sub-accounts/subscriptions/projects (``aws-sg`` standalone account, the ``PingCAP Azure OpenAI`` subscription, etc.) so the dashboard can sum the WHOLE billing of that scope into AI without keyword-matching service names. Default false. The migration auto-flips ``true`` for the well-known IDs (``aws-sg`` slug, the ``7bc8ab17-59fb-44c1-92ca-61fd2b814735`` Azure subscription). 4. **``cmdb_ai_resources``** (new table) — one row per AI-capable cloud resource we care about for chargeback. Today: every Azure ``Microsoft.CognitiveServices/accounts`` (kind=OpenAI / Cognitive*) resource discovered under any ``is_ai_dedicated`` Azure scope. Carries an ``owner_person_id`` (admin-assignable, defaults to either the Azure ``owner`` tag → person lookup or the parent scope's ``owner_person_id``) and ``owner_department`` (denormalised for the dashboard's "by department" tab — re-derived on each sync from the owner person's primary department). Soft-delete via ``deleted_at`` so historical cost rows stay intact when a resource disappears from Azure. 5. **``cmdb_ai_resource_cost_daily``** (new table) — per-resource per-day cost in USD + native currency, populated by ``ai_spending.azure_openai_resources_sync``. Uniqueness on ``(ai_resource_id, date)`` so the upsert is idempotent. 6. **``ai_spending_summaries``** (new table) — LLM-generated period summaries cached one row per ``(period, language)``. The auto-job runs on the 1st of each month and generates ``period=YYYY-MM`` bilingual rows; admin "regenerate" replaces the matching row in-place. Stores token counts so we can spot a runaway prompt. 7. **``platform_settings``** seeds — initial defaults for the ``ai_spending.*`` namespace (currency rate, AI service keyword lists per cloud, regen rate-limit). All values are JSON so the admin UI can edit without a migration. 8. **``scheduled_jobs``** seeds — six jobs (three ai_spending.* + three bailian.*) registered **enabled**. Each handler is written to skip cleanly when its preconditions aren't met (no Aliyun integration / no LLM key / no AI-dedicated Azure scope) and record the skip reason in ``scheduled_job_runs``. Auto-enable was the intentional product call: the alternative (disabled by default) means an admin has to remember to flip three rows by hand AND remember which preconditions go with which row, which is strictly worse UX. Trade-off: a re-run of this migration on prod will RE-ENABLE a job that ops manually disabled mid- incident — operators must record any intentional disable somewhere durable (a runbook entry) because migration re-runs reset the state to "as declared in this file": * ``ai_spending.azure_openai_resources_sync`` (daily 04:30 UTC) * ``ai_spending.monthly_summary_generate`` (cron ``5 3 1 * *`` Asia/Shanghai) * ``ai_spending.cny_rate_refresh`` (cron ``0 5 * * 1`` UTC, Mondays) * ``bailian.workspace_sync`` (daily 04:00 UTC) * ``bailian.application_sync`` (daily 04:30 UTC) * ``bailian.usage_sync`` (daily 05:00 UTC) Idempotent: every step checks ``information_schema`` / target rows first OR uses ``ON DUPLICATE KEY UPDATE``, so re-running on a partially-applied DB is safe. """ from __future__ import annotations import json import logging import uuid import sqlalchemy as sa from alembic import op logger = logging.getLogger(__name__) # Sub-strings that we recognise as idempotency races on retry — # safe to log + continue. Anything else (column type mismatch, FK # target missing, charset / collation disagreement, etc.) MUST # surface as a hard error, not get swallowed. # # Round 4 caught two false positives in the previous version of # this list: # * ``"errno: 150"`` matches BOTH "FK already exists" (idempotent) # AND "FK type mismatch / target missing" (structural). Removed — # duplicate-FK case is covered by ``"duplicate foreign key"``. # * ``"doesn't exist"`` substring-matches "Table 'helix.identity_persons' # doesn't exist" — a hard structural failure. Replaced with # more specific phrases that ONLY appear on idempotent paths. # # The discriminator is intentionally narrow: when in doubt, raise. _IDEMPOTENT_RETRY_HINTS = ( "duplicate key name", # MySQL: index/unique already exists "duplicate entry", # MySQL UQ already populated "duplicate foreign key", # MySQL: FK already created "errno: 121", # MySQL: same FK name twice # MySQL 1091 ("Can't DROP 'X'; check that column/key exists") is # the canonical "drop something that's not there" message — the # tail ``check that column/key exists`` is the unambiguous part. # Don't match the bare ``"can't drop"`` prefix because MySQL also # uses that wording for parent-row / dependent-FK structural # failures (1217 etc.) which are NOT idempotent. "check that column/key exists", ) def _try_swallow(label: str, fn): """Run a DDL helper, swallowing ONLY expected idempotent retry patterns. Anything else re-raises so the migration FAILS loudly on real structural problems (FK type mismatch, charset drift, missing target table, etc.). Replaces a previous ``try: ... except Exception: pass`` that swallowed every error silently — exactly the pattern Round 3 review caught: a CHAR-vs-VARCHAR FK rejection on TiDB would otherwise be logged as a warning and quietly drop referential integrity. """ try: fn() except Exception as exc: # noqa: BLE001 msg = str(exc).lower() if any(hint in msg for hint in _IDEMPOTENT_RETRY_HINTS): logger.warning( "migration 079: %s skipped (idempotent retry): %s", label, exc, ) return # Structural failure — log AND re-raise so the deploy fails # loud. Better to fail the deploy than to ship a half-applied # schema with missing FKs / wrong column types. logger.error( "migration 079: %s failed with non-idempotent error: %s", label, exc, ) raise revision = "079" down_revision = "078" branch_labels = None depends_on = None # ───────────────────────────────────────────────────────────────────── # helpers def _column_exists(bind, table: str, column: str) -> bool: return bind.execute( sa.text( "SELECT 1 FROM information_schema.COLUMNS " "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = :t " " AND COLUMN_NAME = :c" ), {"t": table, "c": column}, ).scalar() is not None def _table_exists(bind, name: str) -> bool: return bind.execute( sa.text( "SELECT 1 FROM information_schema.TABLES " "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = :n" ), {"n": name}, ).scalar() is not None def _fk_source_col_type( bind, target_table: str, target_column: str, *, fallback_type: str = "varchar", fallback_collation: str | None = None, length: int = 36, ): """Mirror an FK target's data type AND collation on the source. TiDB rejects FKs whose source column's data_type OR collation differs from the target's (errno 3780). MySQL 8 is more permissive on type but still strict on collation. The pattern was introduced in PR #81/82/83 for migrations 073/074. We re-import it here instead of relying on hand-coded ``sa.CHAR(36)`` / ``sa.String(36)`` because prod TiDB has columns with **mixed collations** — ``cmdb_cloud_scopes.id`` is ``utf8mb4_bin`` while every other FK target is ``utf8mb4_unicode_ci``. Hand-coded types pick up the table-default collation, which is wrong for that one table. """ row = bind.execute( sa.text( "SELECT DATA_TYPE, COLLATION_NAME FROM " "information_schema.COLUMNS " "WHERE TABLE_SCHEMA = DATABASE() " " AND TABLE_NAME = :t AND COLUMN_NAME = :c" ), {"t": target_table, "c": target_column}, ).fetchone() if row: data_type = (row[0] or fallback_type).lower() collation = row[1] or fallback_collation else: data_type = fallback_type.lower() collation = fallback_collation type_cls = sa.CHAR if data_type == "char" else sa.String if collation: return type_cls(length, collation=collation) return type_cls(length) def _column_collation(bind, table: str, column: str) -> str | None: row = bind.execute( sa.text( "SELECT COLLATION_NAME FROM information_schema.COLUMNS " "WHERE TABLE_SCHEMA = DATABASE() " " AND TABLE_NAME = :t AND COLUMN_NAME = :c" ), {"t": table, "c": column}, ).fetchone() return row[0] if row and row[0] else None def _index_exists(bind, table: str, name: str) -> bool: return bind.execute( sa.text( "SELECT 1 FROM information_schema.STATISTICS " "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = :t " " AND INDEX_NAME = :n" ), {"t": table, "n": name}, ).scalar() is not None # Software-catalog name patterns Helix should auto-flag as AI tools on # first run. Match is case-insensitive substring against the # `asset_software_catalog.name` column. Admins can un-flag any row # from the UI; the flag is admin-owned thereafter. _AI_NAME_PATTERNS: list[str] = [ "openai", "chatgpt", "claude", "anthropic", "cursor", "copilot", # covers GitHub Copilot AND Microsoft 365 Copilot "gemini", "bailian", "百炼", "model studio", "bedrock", "azure openai", "perplexity", "midjourney", ] # Subscription/account IDs we know are 100 % AI-dedicated. The # Azure OpenAI subscription GUID was confirmed by the user 2026-04-28; # the AWS SG slug matches the existing autogrant convention. Any miss # (rename, ID change) is recoverable from the admin UI — the seed # only fires on first run if the row exists and isn't already flagged. _AI_DEDICATED_SCOPE_EXTERNAL_IDS: list[str] = [ "7bc8ab17-59fb-44c1-92ca-61fd2b814735", ] _AI_DEDICATED_INTEGRATION_SLUGS: list[str] = [ "aws-sg", ] # Default AI service keyword lists per cloud provider. Used by the # dedup layer's keyword-matching branch (see ai_spending/dedup.py). # Stored in platform_settings so admins can edit without a deploy. _AI_SERVICE_KEYWORDS: dict[str, list[str]] = { "aws": [ "Bedrock", "SageMaker", "Comprehend", "Textract", "Rekognition", "Polly", "Transcribe", "Translate", "Lex", "Kendra", "Q Developer", ], "azure": [ "Azure OpenAI", "Cognitive Services", "Cognitive Search", "Machine Learning", "Bot Service", "Form Recognizer", "Document Intelligence", "Computer Vision", "Speech Services", "Language", ], "gcp": [ "Vertex AI", "Generative Language API", "Gemini", "AI Platform", "Cloud Translation", "Cloud Natural Language", "Cloud Vision API", "Cloud Speech-to-Text", "Cloud Text-to-Speech", "AutoML", "Document AI", ], "aliyun": [ # Aliyun BSS returns ``ProductName`` in Chinese for CN # tenants (the most common deployment): list both the # Chinese full names and the English SKU-anchored variants # so we match on either an English-tenant or a CN-tenant # billing row. PAI is exactly 3 chars and accidentally # collides with unrelated SKUs, so we anchor it with the # full Chinese product name AND the canonical SKU # suffixes that show up in itemised English billing. "机器学习平台PAI", "机器学习PAI", "PAI-EAS", "PAI-DSW", "PAI-DLC", "PAI-Designer", # Bailian / DashScope / Tongyi: ProductName is "百炼" or # "灵积模型服务" or "通义" depending on which surface a # given line item is billed under. "百炼", "灵积", "通义", "DashScope", "Bailian", # "Model Studio" is the English console branding only — # Aliyun BSS does NOT return that string in ProductName. # Removed to keep the keyword set strict; the Bailian path # already covers any cost billed under Bailian/灵积/通义. # NOT "OpenSearch": Aliyun's OpenSearch is a generic # full-text-search product, not AI-specific. Add specific # AI variants here when needed (e.g. "OpenSearch 智能问答") # via /admin/ai-spending/settings — keeping the default # list strict to avoid false positives. ], # Tencent: BSS-equivalent (``BusinessCodeName`` / # ``ProductCodeName``) is Chinese for CN tenants, so list the # Chinese full names. English SKU-style entries are kept as # belt-and-suspenders for the rare overseas-tenant case but # would not match the CN billing rows by themselves. # ``OCR`` / ``ASR`` / ``TTS`` / ``TI`` / ``TBP`` alone are # bad substring keys (collide with unrelated services); the # entries below all ≥ 4 chars or unambiguous full names. "tencent": [ "腾讯混元", "混元大模型", # Removed "文心一言" — that's Baidu's Wenxin Yiyan, not a # Tencent product, never appears in Tencent BSS rows. "智能对话平台", "光学字符识别", # CN ProductCodeName for OCR "通用文字识别", # CN: another OCR variant "智能语音", # CN: ASR/TTS combined "语音合成", # CN: TTS "语音识别", # CN: ASR "TI-ONE", "TI-EMS", "TI 平台", # CN: with space "TI平台", # CN: Tencent BSS sometimes drops the space "腾讯云 TI", ], } # ───────────────────────────────────────────────────────────────────── # upgrade def upgrade() -> None: bind = op.get_bind() # Resolve every FK target's data_type + collation ONCE up front. # Different parent tables have different collations on prod TiDB # (``cmdb_cloud_scopes.id`` is utf8mb4_bin while every other FK # target is utf8mb4_unicode_ci), so hard-coded ``sa.String(36)`` # picks up the table-default collation and TiDB rejects the FK # with errno 3780 ("Referencing column ... and referenced column # ... in foreign key constraint ... are incompatible"). Using # ``_fk_source_col_type`` mirrors the target's actual stored # type+collation. See PR #81/82/83 for the canonical pattern. integrations_id_col = _fk_source_col_type( bind, "integrations", "id", fallback_type="char", fallback_collation="utf8mb4_unicode_ci", ) identity_persons_id_col = _fk_source_col_type( bind, "identity_persons", "id", fallback_type="char", fallback_collation="utf8mb4_unicode_ci", ) cmdb_cloud_scopes_id_col = _fk_source_col_type( bind, "cmdb_cloud_scopes", "id", fallback_type="varchar", fallback_collation="utf8mb4_bin", ) # Bailian workspace FK target type is resolved per-callsite below # (only one callsite uses it, in the workspace_pk ALTER block). # ── 1. asset_software_catalog.is_ai_tool ──────────────────── if _table_exists(bind, "asset_software_catalog") and not _column_exists( bind, "asset_software_catalog", "is_ai_tool", ): op.add_column( "asset_software_catalog", sa.Column( "is_ai_tool", sa.Boolean(), nullable=False, server_default=sa.text("0"), ), ) # Best-effort seed. Patterns are case-insensitive substring on # the canonical software name; misses are admin-fixable. for pat in _AI_NAME_PATTERNS: bind.execute( sa.text( "UPDATE asset_software_catalog " "SET is_ai_tool = 1 " "WHERE LOWER(name) LIKE :pat " " AND is_ai_tool = 0" ), {"pat": f"%{pat.lower()}%"}, ) # ── 2. api_key_assignments.cloud_scope_id ────────────────── if _table_exists(bind, "api_key_assignments"): target_collation = _column_collation( bind, "cmdb_cloud_scopes", "id", ) or "utf8mb4_bin" if not _column_exists( bind, "api_key_assignments", "cloud_scope_id", ): # Add with the parent's exact type+collation so FK creation # below succeeds on TiDB. op.add_column( "api_key_assignments", sa.Column( "cloud_scope_id", cmdb_cloud_scopes_id_col, nullable=True, ), ) else: # Column already exists (deploy of an earlier 079 added it # before this fix landed). Verify collation matches and # ALTER if not — otherwise FK creation below would fail # with errno 3780 like it did on the first prod attempt. current_collation = _column_collation( bind, "api_key_assignments", "cloud_scope_id", ) if current_collation and current_collation != target_collation: op.execute(sa.text( "ALTER TABLE api_key_assignments " "MODIFY COLUMN cloud_scope_id " f"VARCHAR(36) COLLATE {target_collation} NULL" )) # FK creation. Use ``_try_swallow`` so a re-run after a # successful create is idempotent. if _table_exists(bind, "cmdb_cloud_scopes"): _try_swallow( "create FK fk_api_key_assignments_cloud_scope_id", lambda: op.create_foreign_key( "fk_api_key_assignments_cloud_scope_id", "api_key_assignments", "cmdb_cloud_scopes", ["cloud_scope_id"], ["id"], ondelete="SET NULL", ), ) if not _index_exists( bind, "api_key_assignments", "ix_api_key_assignments_cloud_scope_id", ): op.create_index( "ix_api_key_assignments_cloud_scope_id", "api_key_assignments", ["cloud_scope_id"], ) # ── 3. cmdb_cloud_scopes.is_ai_dedicated ─────────────────── if _table_exists(bind, "cmdb_cloud_scopes") and not _column_exists( bind, "cmdb_cloud_scopes", "is_ai_dedicated", ): op.add_column( "cmdb_cloud_scopes", sa.Column( "is_ai_dedicated", sa.Boolean(), nullable=False, server_default=sa.text("0"), ), ) # Auto-flag the Azure OpenAI subscription. ``external_id`` on # cmdb_cloud_scopes holds the provider-side identifier (Azure # subscription GUID, GCP project id, AWS account number, etc.). for ext_id in _AI_DEDICATED_SCOPE_EXTERNAL_IDS: bind.execute( sa.text( "UPDATE cmdb_cloud_scopes " "SET is_ai_dedicated = 1 " "WHERE external_id = :ext " " AND is_ai_dedicated = 0" ), {"ext": ext_id}, ) # Auto-flag every scope under integrations whose slug is in # the curated list (AWS SG single-account = one scope of # scope_type=standalone_iam). Guard on integrations table # existence to keep the migration safe on partial-state DBs # (defensive: should always exist in normal Alembic chain # since 045 created it, but consistency with the rest of # this file's existence checks). if _table_exists(bind, "integrations"): for slug in _AI_DEDICATED_INTEGRATION_SLUGS: bind.execute( sa.text( "UPDATE cmdb_cloud_scopes s " "JOIN integrations i ON i.id = s.integration_id " "SET s.is_ai_dedicated = 1 " "WHERE i.slug = :slug " " AND s.is_ai_dedicated = 0" ), {"slug": slug}, ) # ── 4. cmdb_ai_resources ────────────────────────────────── # # FK source columns mirror parent type+collation via # ``_fk_source_col_type``. Hard-coding ``sa.CHAR(36)`` / # ``sa.String(36)`` would pick up the new table's default # collation (utf8mb4_unicode_ci), which fails on TiDB FK creation # to ``cmdb_cloud_scopes.id`` (utf8mb4_bin). if not _table_exists(bind, "cmdb_ai_resources"): op.create_table( "cmdb_ai_resources", sa.Column("id", sa.CHAR(36), primary_key=True), sa.Column( "integration_id", integrations_id_col, nullable=False, ), sa.Column( "scope_id", cmdb_cloud_scopes_id_col, nullable=True, ), sa.Column("provider", sa.String(32), nullable=False), sa.Column("resource_id", sa.String(512), nullable=False), sa.Column("resource_name", sa.String(256), nullable=False), sa.Column("resource_kind", sa.String(64), nullable=True), sa.Column("resource_location", sa.String(64), nullable=True), sa.Column( "owner_person_id", identity_persons_id_col, nullable=True, ), sa.Column( "owner_department", sa.String(256), nullable=True, ), sa.Column( "auto_owner_source", sa.String(32), nullable=True, # 'tag' | 'scope_owner' | 'manual' | NULL (= unassigned) ), sa.Column("tags", sa.JSON, nullable=True), sa.Column("raw_data", sa.JSON, nullable=True), sa.Column( "first_seen_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "last_seen_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "deleted_at", sa.DateTime(timezone=True), nullable=True, ), sa.Column( "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), onupdate=sa.func.now(), nullable=False, ), sa.UniqueConstraint( "integration_id", "resource_id", name="uq_cmdb_ai_resources_integration_resource", ), sa.Index( "ix_cmdb_ai_resources_owner", "owner_person_id", ), sa.Index( "ix_cmdb_ai_resources_scope", "scope_id", ), ) # FK best-effort; same pattern as integration_id elsewhere. # Avoid hard FK to integrations because TiDB sometimes rejects # multi-cascade compositions if the table was created by a # different migration order. ``_try_swallow`` logs failures # so an unexpected schema issue is visible in the deploy log # rather than silently swallowed. _try_swallow( "fk_cmdb_ai_resources_integration", lambda: op.create_foreign_key( "fk_cmdb_ai_resources_integration", "cmdb_ai_resources", "integrations", ["integration_id"], ["id"], ondelete="CASCADE", ), ) _try_swallow( "fk_cmdb_ai_resources_scope", lambda: op.create_foreign_key( "fk_cmdb_ai_resources_scope", "cmdb_ai_resources", "cmdb_cloud_scopes", ["scope_id"], ["id"], ondelete="SET NULL", ), ) _try_swallow( "fk_cmdb_ai_resources_owner", lambda: op.create_foreign_key( "fk_cmdb_ai_resources_owner", "cmdb_ai_resources", "identity_persons", ["owner_person_id"], ["id"], ondelete="SET NULL", ), ) # ── 5. cmdb_ai_resource_cost_daily ──────────────────────── # # ``cmdb_ai_resources.id`` is CHAR(36) (just declared above); # FK column matches. if not _table_exists(bind, "cmdb_ai_resource_cost_daily"): op.create_table( "cmdb_ai_resource_cost_daily", sa.Column("id", sa.CHAR(36), primary_key=True), sa.Column( "ai_resource_id", sa.CHAR(36), nullable=False, ), sa.Column("date", sa.Date(), nullable=False), sa.Column( "cost_usd", sa.Numeric(14, 6), nullable=False, server_default=sa.text("0"), ), sa.Column( "cost_native", sa.Numeric(14, 6), nullable=False, server_default=sa.text("0"), ), sa.Column( "currency", sa.String(8), nullable=False, server_default=sa.text("'USD'"), ), sa.Column("raw", sa.JSON, nullable=True), sa.Column( "fetched_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), onupdate=sa.func.now(), nullable=False, ), sa.UniqueConstraint( "ai_resource_id", "date", name="uq_cmdb_ai_resource_cost_daily_resource_date", ), sa.Index( "ix_cmdb_ai_resource_cost_daily_date", "date", ), ) _try_swallow( "fk_cmdb_ai_resource_cost_daily_resource", lambda: op.create_foreign_key( "fk_cmdb_ai_resource_cost_daily_resource", "cmdb_ai_resource_cost_daily", "cmdb_ai_resources", ["ai_resource_id"], ["id"], ondelete="CASCADE", ), ) # ── 5b. bailian_usage_daily — schema fix for BSS account-level rows ── # # The original ``bailian_usage_daily`` (migration 074) keyed # uniquely on ``(workspace_pk, api_key_pk, usage_date, model_name)`` # with workspace_pk NOT NULL — assuming Aliyun would expose a # per-workspace usage API. They don't. The canonical billing API # ``QueryAccountBill(ProductCode='bailian')`` is account-scoped, # so every row legitimately has workspace_pk=NULL. To keep the # ledger truthful, we: # 1. Drop the old unique constraint # 2. Make ``workspace_pk`` nullable # 3. Add a new unique constraint on # ``(integration_id, usage_date, model_name)`` # that's the actual grain BSS gives us. Future per-workspace # / per-api-key sources will need to include their own # discriminator in the constraint or write to a different # table. if _table_exists(bind, "bailian_usage_daily"): # Step 1: Drop the old unique constraint (no longer matches # the BSS-account-level row shape). _try_swallow( "drop uq_bailian_usage_day_key_model", lambda: op.drop_constraint( "uq_bailian_usage_day_key_model", "bailian_usage_daily", type_="unique", ), ) # Step 2: Drop the FK on workspace_pk BEFORE altering its # nullability. TiDB rejects nullable changes on FK-source # columns; MySQL is more permissive but accepts the safer # path. We re-create the FK in step 4. # # ``existing_type`` MUST be ``sa.String(36)`` (= VARCHAR(36)) # because migration ``074_bailian_tables.py`` created this # column as ``sa.String(36)`` AND the FK target # ``bailian_workspaces.id`` is also VARCHAR(36) (see same # 074 migration). The ORM in ``bailian/models.py`` declares # ``Mapped[uuid.UUID]`` (which would be CHAR(36)) but the # actual stored column is VARCHAR(36) — the migration is # the source of truth. Using sa.CHAR(36) here would mutate # VARCHAR→CHAR and create a TiDB-strict CHAR↔VARCHAR FK # mismatch on step 4. (See feedback_tidb_fk_strictness memory.) if _column_exists(bind, "bailian_usage_daily", "workspace_pk"): # Drop the FK if it exists (name varies by MySQL version # — ibfk_N or whatever the original migration named). # We look it up via INFORMATION_SCHEMA to be precise. fk_rows = bind.execute(sa.text( "SELECT CONSTRAINT_NAME FROM " "information_schema.KEY_COLUMN_USAGE " "WHERE TABLE_SCHEMA = DATABASE() " " AND TABLE_NAME = 'bailian_usage_daily' " " AND COLUMN_NAME = 'workspace_pk' " " AND REFERENCED_TABLE_NAME IS NOT NULL" )).fetchall() for (fk_name,) in fk_rows: _try_swallow( f"drop FK {fk_name} on bailian_usage_daily.workspace_pk", lambda fk_name=fk_name: op.drop_constraint( fk_name, "bailian_usage_daily", type_="foreignkey", ), ) # Step 3: alter to nullable, preserving VARCHAR(36) + # the target's collation. Use a raw ``ALTER TABLE # MODIFY COLUMN`` because alembic's ``alter_column`` does # not surface the collation knob cleanly — we want to # explicitly mirror ``bailian_workspaces.id``'s collation # (utf8mb4_unicode_ci on prod) so the re-created FK at # step 4 satisfies TiDB strictness. target_collation = _column_collation( bind, "bailian_workspaces", "id", ) or "utf8mb4_unicode_ci" try: op.execute(sa.text( "ALTER TABLE bailian_usage_daily " "MODIFY COLUMN workspace_pk " f"VARCHAR(36) COLLATE {target_collation} NULL" )) except Exception: # Hard failure — the rest of the schema below assumes # workspace_pk is nullable. We'd rather abort the # deploy than ship a half-applied schema where every # subsequent BSS-billing-sourced row insert would # raise NOT NULL violations. logger.error( "migration 079: critical — could not make " "bailian_usage_daily.workspace_pk nullable. " "Without this, the rewritten BSS usage sync " "would insert workspace_pk=NULL rows and fail " "the constraint." ) raise # Step 4: re-create the FK (now on a nullable column, # collation matched). Both source and target are VARCHAR # with the same collation, so TiDB accepts. _try_swallow( "re-create FK on bailian_usage_daily.workspace_pk", lambda: op.create_foreign_key( "fk_bailian_usage_workspace_pk", "bailian_usage_daily", "bailian_workspaces", ["workspace_pk"], ["id"], ondelete="CASCADE", ), ) # Step 5: add the new unique constraint matching the BSS # account-level row grain. _try_swallow( "create uq_bailian_usage_integ_date_model", lambda: op.create_unique_constraint( "uq_bailian_usage_integ_date_model", "bailian_usage_daily", ["integration_id", "usage_date", "model_name"], ), ) # ── 6. ai_spending_summaries ────────────────────────────── # # ``id`` CHAR(36) for consistency with the GUID pattern; # ``generated_by_user_id`` CHAR(36) to match the # ``identity_persons.id`` FK target type (TiDB strictness). if not _table_exists(bind, "ai_spending_summaries"): op.create_table( "ai_spending_summaries", sa.Column("id", sa.CHAR(36), primary_key=True), sa.Column("period", sa.String(7), nullable=False), sa.Column( "language", sa.String(8), nullable=False, # 'en' | 'zh' | 'bilingual' (single payload carrying # both languages — that's how we store the LLM output # by default to save a round-trip) ), sa.Column( "summary_md_en", sa.Text, nullable=True, ), sa.Column( "summary_md_zh", sa.Text, nullable=True, ), sa.Column("findings", sa.JSON, nullable=True), sa.Column("suggestions", sa.JSON, nullable=True), sa.Column("aggregates", sa.JSON, nullable=True), sa.Column( "model_used", sa.String(128), nullable=True, ), sa.Column( "input_token_count", sa.Integer, nullable=True, ), sa.Column( "output_token_count", sa.Integer, nullable=True, ), sa.Column( "generated_by_user_id", identity_persons_id_col, nullable=True, ), sa.Column( "generated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), onupdate=sa.func.now(), nullable=False, ), sa.UniqueConstraint( "period", "language", name="uq_ai_spending_summaries_period_lang", ), ) _try_swallow( "fk_ai_spending_summaries_user", lambda: op.create_foreign_key( "fk_ai_spending_summaries_user", "ai_spending_summaries", "identity_persons", ["generated_by_user_id"], ["id"], ondelete="SET NULL", ), ) # ── 7. platform_settings seeds ──────────────────────────── if _table_exists(bind, "platform_settings"): _seed_setting( bind, "ai_spending.cny_to_usd_rate", 0.14, "CNY → USD conversion rate used by the AI spending " "dashboard. Refreshed every Monday from ECB; manual " "override permitted.", ) _seed_setting( bind, "ai_spending.cny_rate_last_refreshed_at", None, "ISO timestamp of the last automatic CNY/USD refresh " "(NULL = never refreshed; using static default).", ) _seed_setting( bind, "ai_spending.service_keywords", _AI_SERVICE_KEYWORDS, "Per-cloud-provider list of service-name substrings that " "Helix counts as AI for billing aggregation. Match is " "case-insensitive substring on cost_by_service keys. " "Edit to add a new vendor product; no deploy needed.", ) _seed_setting( bind, "ai_spending.regenerate_per_user_per_day", 5, "How many times each admin can hit 'regenerate AI " "summary' per UTC day. Counter resets at 00:00 UTC.", ) _seed_setting( bind, "ai_spending.llm.model", "qwen-max-latest", "DashScope model id used for AI spending summaries. " "qwen-max-latest tracks Aliyun's flagship qwen-max " "release; pin to qwen-max-2025-XX-XX for reproducibility " "if desired. NEVER point this at a non-Bailian model — " "the LLM key is a Bailian-issued sk-xxx, not OpenAI.", ) # The actual LLM api_key is set by the admin from the settings # UI (it's a secret; we don't seed a placeholder so an empty # value clearly means 'not configured' rather than 'broken'). # ── 8. scheduled_jobs seeds ─────────────────────────────── # # All jobs land **enabled=1** by design. Each handler short-circuits # cleanly when its preconditions aren't met (no enabled Aliyun # integration, no ai_dedicated Azure scope, no Bailian sk-xxx LLM # key) and records the skip reason in scheduled_job_runs. Leaving # them disabled-by-default would mean the user has to manually flip # each row after deploy AND remember which prerequisite was missing # — strictly worse UX. The skip-with-reason pattern is the # production-grade equivalent of "no-op" without lying about it. if _table_exists(bind, "scheduled_jobs"): _seed_job( bind, name="ai_spending.azure_openai_resources_sync", display_name="AI 支出 · Azure OpenAI 资源 + 成本同步", description=( "每天遍历所有 is_ai_dedicated=true 的 Azure 子账号," "拉取下面的 Microsoft.CognitiveServices/accounts 资源 " "及其每日成本(GROUP BY ResourceId 调 Cost Management " "API)。资源 owner 优先取 Azure tag 中的 owner/Owner " "字段(按 work_email 匹配 identity_persons),否则继承 " "scope.owner_person_id。管理员可在前端手动覆盖。" ), cron_expr="30 4 * * *", timezone_name="UTC", enabled=True, ) _seed_job( bind, name="ai_spending.monthly_summary_generate", display_name="AI 支出 · 月度总结自动生成", description=( "每月 1 号 03:05 (Asia/Shanghai) 调用百炼 LLM 对上月 " "AI 支出生成中英双语总结 + 优化建议,缓存到 " "ai_spending_summaries。要求 platform_settings " "ai_spending.llm.api_key 已配置;否则任务跳过并在 " "运行历史里留下错误信息。" ), cron_expr="5 3 1 * *", timezone_name="Asia/Shanghai", enabled=True, ) _seed_job( bind, name="ai_spending.cny_rate_refresh", display_name="AI 支出 · 人民币汇率刷新", description=( "每周一 05:00 UTC 拉一次 ECB 公开汇率 " "(api.exchangerate.host),写入 " "ai_spending.cny_to_usd_rate platform_setting。" "失败时不覆盖现有值,下一次重试。" ), cron_expr="0 5 * * 1", timezone_name="UTC", enabled=True, ) # ── Forward-fix: bailian.* jobs that migration 074 declared # in comments but never actually seeded. Without these the # /admin/cloud/bailian page stays empty forever even after # admin enables the Aliyun integration with valid AK/SK, # because no job ever calls ListWorkspaces. Same enabled=1 + # skip-cleanly-on-no-creds posture as the ai_spending block # above. _seed_job( bind, name="bailian.workspace_sync", display_name="阿里云百炼 · 业务空间同步", description=( "每天 04:00 UTC 拉取所有启用的阿里云集成下的百炼业务" "空间到 bailian_workspaces。基于 Bailian OpenAPI " "ListWorkspaces;Aliyun AK 需要 " "AliyunBailianFullAccess 或等价权限。无启用集成时" "任务记录 skip 原因后退出。" ), cron_expr="0 4 * * *", timezone_name="UTC", enabled=True, ) _seed_job( bind, name="bailian.application_sync", display_name="阿里云百炼 · 应用 / 智能体同步", description=( "每天 04:30 UTC 对所有 BailianWorkspace 调用 " "ListApplications,把应用 / 智能体 / RAG 应用同步到 " "bailian_applications。workspace_sync 完成后跑,所以" "比 workspace_sync 晚 30 分钟。" ), cron_expr="30 4 * * *", timezone_name="UTC", enabled=True, ) _seed_job( bind, name="bailian.usage_sync", display_name="阿里云百炼 · Token 调用量同步", description=( "每天 05:00 UTC 通过 BSS QueryUserOmsData 拉取过去 " "7 天的百炼调用量到 bailian_usage_daily,按 (workspace " "× api_key × model × day) 维度。AI 支出大盘读这个表 " "作为百炼支出的权威源。" ), cron_expr="0 5 * * *", timezone_name="UTC", enabled=True, ) def _seed_setting(bind, key: str, value, description: str) -> None: """Idempotent seed of a platform_settings row. Uses ``INSERT ... ON DUPLICATE KEY UPDATE`` so the seed survives: * Re-runs of this migration on a partially-applied DB. * Concurrent Alembic processes (rare in prod, common in CI with parallel runners) — TOCTOU between the prior SELECT-then-INSERT pattern would surface as 1062 and abort the migration. Updating the description on collision keeps the human-readable explanation current with the migration source code, but never overwrites an admin-edited ``value``. """ bind.execute( sa.text( "INSERT INTO platform_settings (`key`, value, description, " "created_at, updated_at) " "VALUES (:k, :v, :d, NOW(), NOW()) " "ON DUPLICATE KEY UPDATE " "description = VALUES(description), updated_at = NOW()" ), { "k": key, "v": json.dumps(value, ensure_ascii=False), "d": description, }, ) def _seed_job( bind, *, name: str, display_name: str, description: str, cron_expr: str, timezone_name: str, enabled: bool = False, ) -> None: """Idempotent seed of a scheduled_jobs row. On re-run / partial-apply: keeps the existing ``id`` / ``run_count`` / ``failure_count`` / ``last_run_at`` (those are operational state that the worker owns) but refreshes the descriptive fields AND ``enabled`` to whatever the migration declares. This means a partially-applied earlier run that inserted a job with ``enabled=0`` will be promoted to the declared state on a subsequent migration run, instead of requiring a manual UI flip. Note: the explicit ``enabled`` column update is the deliberate overwrite the migration owns. If ops have a long-running operational reason to disable a job (e.g. mid-incident), they must record that intent somewhere durable — running this migration would otherwise re-enable on the next deploy. For Helix, this trade-off favours "declared state wins" over "ops state survives migration re-run" because the migration runs only on intentional version bumps, not unattended. """ bind.execute( sa.text( """ INSERT INTO scheduled_jobs (id, name, display_name, description, handler_key, trigger_type, interval_seconds, cron_expr, timezone, enabled, config, run_count, failure_count, created_at, updated_at) VALUES (:id, :name, :display_name, :description, :handler_key, 'cron', NULL, :cron_expr, :tz, :enabled, '{}', 0, 0, NOW(), NOW()) ON DUPLICATE KEY UPDATE display_name = VALUES(display_name), description = VALUES(description), cron_expr = VALUES(cron_expr), timezone = VALUES(timezone), enabled = VALUES(enabled), updated_at = NOW() """ ), { "id": str(uuid.uuid4()), "name": name, "display_name": display_name, "description": description, "handler_key": name, "cron_expr": cron_expr, "tz": timezone_name, "enabled": 1 if enabled else 0, }, ) # ───────────────────────────────────────────────────────────────────── # downgrade — conservative def downgrade() -> None: """Drop only the new tables + scheduled jobs. Leave the new columns (``is_ai_tool``, ``cloud_scope_id``, ``is_ai_dedicated``) in place — they're nullable / boolean-default-false, harmless to older code, and dropping risks losing operator state. Re-applying the migration will simply re-create the dropped tables.""" bind = op.get_bind() if _table_exists(bind, "scheduled_jobs"): bind.execute(sa.text( "DELETE FROM scheduled_jobs WHERE name IN (" "'ai_spending.azure_openai_resources_sync'," "'ai_spending.monthly_summary_generate'," "'ai_spending.cny_rate_refresh'," "'bailian.workspace_sync'," "'bailian.application_sync'," "'bailian.usage_sync')" )) if _table_exists(bind, "platform_settings"): bind.execute(sa.text( "DELETE FROM platform_settings WHERE `key` LIKE 'ai_spending.%'" )) for tbl in ( "ai_spending_summaries", "cmdb_ai_resource_cost_daily", "cmdb_ai_resources", ): if _table_exists(bind, tbl): op.drop_table(tbl)