"""Promote Bailian usage rows from daily-aggregate → per-(day × Aliyun member account). Why: Aliyun BSS ``QueryAccountBill(ProductCode='bailian')`` with ``IsGroupByProduct=false`` (the only call that surfaces cross-RD member-account spend on the intl tenant) returns one row per ``(day × OwnerID)``. The previous ingest path SUMMED those rows into one "aggregate" row per day and stuffed the per-account breakdown into ``raw_data.per_account`` JSON — readable but never indexable, joinable, or surfaced on the dashboard. Production-grade chargeback needs the per-account dimension as a first-class column so we can: * JOIN ``cmdb_cloud_scopes.external_id`` to attach the Aliyun member-account display name + scope owner per row * let the AI-spending dashboard's "by user / by department / by scope" breakdowns include Bailian cost mapped to the matching scope owner (today Bailian lands in "(unattributed)") * let the Bailian admin page's 调用量/成本 tab show "this day, this Aliyun account, spent ¥X" instead of "this day, total ¥sum" This migration: 1. Adds ``aliyun_account_id VARCHAR(32) NULL`` (NULL for any rows that aren't BSS-derived — workspace/api-key-scoped rows from a future ingest path will leave it NULL). 2. Replaces the unique constraint ``(integration_id, usage_date, model_name)`` with ``(integration_id, usage_date, model_name, aliyun_account_id)``. MySQL/TiDB index columns can't be NULL-equal, so the new index uses ``COALESCE(aliyun_account_id, '')`` semantics implicitly — we keep the column NULL-able and rely on the sync layer always writing a non-null OwnerID for BSS-derived rows. 3. Explodes existing aggregate rows: for every row whose ``raw_data.per_account`` has entries, emits one new row per entry (with that account's CNY amount, its converted USD via the row's pre-existing rate, and ``aliyun_account_id`` set), then deletes the original aggregate row. The ``raw_data`` field is no longer the source of truth — clear it. 4. Adds a covering index on ``(aliyun_account_id, usage_date)`` for the new per-account drill-down query. Idempotent: re-running on an already-upgraded DB no-ops (the column exists, the unique key is updated, no aggregate rows remain). Downgrade: re-merges per-account rows back into one aggregate row per ``(integration_id, usage_date)`` and restores the JSON breakdown. Lossy on token/request counts (always 0 anyway), exact on costs. """ from __future__ import annotations import json import uuid from collections import defaultdict from decimal import Decimal import sqlalchemy as sa from alembic import op revision = "087" down_revision = "086" branch_labels = None depends_on = None _TABLE = "bailian_usage_daily" _OLD_UNIQUE = "uq_bailian_usage_integ_date_model" _NEW_UNIQUE = "uq_bailian_usage_integ_date_model_account" _NEW_INDEX = "ix_bailian_usage_account_date" def _column_exists(bind, table: str, column: str) -> bool: insp = sa.inspect(bind) return any(c["name"] == column for c in insp.get_columns(table)) def _index_exists(bind, table: str, name: str) -> bool: insp = sa.inspect(bind) for ix in insp.get_indexes(table): if ix["name"] == name: return True for uq in insp.get_unique_constraints(table): if uq["name"] == name: return True return False def upgrade() -> None: bind = op.get_bind() # 1. Add the column (idempotent). if not _column_exists(bind, _TABLE, "aliyun_account_id"): op.add_column( _TABLE, sa.Column("aliyun_account_id", sa.String(32), nullable=True), ) # 2. Index swap: create the NEW unique constraint BEFORE # dropping the OLD one. # Two reasons we have to do it in this order: # a) MySQL refuses to drop ``uq_bailian_usage_integ_date_model`` # while it's the only index satisfying the FK on # ``integration_id`` (error 1553). Both indexes have # ``integration_id`` as their leading column, so the new # one takes over the FK-coverage role and the old one # becomes droppable. # b) The per-account INSERTs we run between these two # constraint operations all share # ``(integration_id, usage_date, model_name='aggregate')`` # and would collide on the OLD constraint. The new # constraint has ``aliyun_account_id`` as a 4th key # column so each (day × OwnerID) row is unique. if not _index_exists(bind, _TABLE, _NEW_UNIQUE): op.create_unique_constraint( _NEW_UNIQUE, _TABLE, ["integration_id", "usage_date", "model_name", "aliyun_account_id"], ) if _index_exists(bind, _TABLE, _OLD_UNIQUE): op.drop_constraint(_OLD_UNIQUE, _TABLE, type_="unique") # 3. Explode existing aggregate rows. # Filter on three independent markers so we ONLY pick up # genuine pre-085 aggregate rows, even on a partially-applied # state where some rows have already been exploded: # * ``aliyun_account_id IS NULL`` — exploded rows have it set # * ``model_name = 'aggregate'`` — pre-085 emitted exactly this # literal model name; future per-(workspace × api-key × model) # ingest paths would write a real model SKU here # * ``raw_data IS NOT NULL`` — only pre-085 rows had the # ``per_account`` JSON breakdown; explosion clears it # Without all three, a partial re-run could re-explode rows # that were already exploded in a prior run (especially under # the new constraint shape that allows NULL aliyun_account_id # rows to coexist with real ones). rows = bind.execute(sa.text( f"SELECT id, integration_id, usage_date, model_name, cost_usd, " f"cost_cny, raw_data, synced_at, created_at " f"FROM `{_TABLE}` " f"WHERE aliyun_account_id IS NULL " f" AND model_name = 'aggregate' " f" AND raw_data IS NOT NULL" )).fetchall() inserts: list[dict] = [] deletes: list[str] = [] for row in rows: raw = row.raw_data if isinstance(raw, str): try: raw = json.loads(raw) except (TypeError, ValueError): continue if not isinstance(raw, dict): continue per_account = raw.get("per_account") or [] if not per_account: continue # Compute the rate the original ingest used so the per-account # USD amounts sum back to the row's USD total. cost_usd / # cost_cny gives the implicit rate. If the original row has # ¥0 or $0 we cannot recover the rate without knowing it out # of band — skip the row entirely rather than emit per-account # rows with a fabricated USD=0 (which would let the dashboard # silently undercount USD totals while CNY stayed correct). # The original aggregate row stays in place for that case; # the next sync run will replace it with first-class # per-account rows that have an authoritative rate. cost_usd_raw = Decimal(str(row.cost_usd or 0)) cost_cny_raw = Decimal(str(row.cost_cny or 0)) if cost_usd_raw <= 0 or cost_cny_raw <= 0: continue rate = cost_usd_raw / cost_cny_raw for entry in per_account: owner = entry.get("owner_id") amount_cny_raw = entry.get("amount_cny") if owner is None or amount_cny_raw in (None, ""): continue try: cny = Decimal(str(amount_cny_raw)) except (ArithmeticError, ValueError): continue if cny <= 0: continue usd = (cny * rate).quantize(Decimal("0.000001")) inserts.append({ "id": str(uuid.uuid4()), "integration_id": row.integration_id, "usage_date": row.usage_date, "model_name": row.model_name, "aliyun_account_id": str(owner), "cost_cny": cny, "cost_usd": usd, "synced_at": row.synced_at, "created_at": row.created_at, }) deletes.append(row.id) # Insert exploded rows. if inserts: bind.execute( sa.text( f"INSERT INTO `{_TABLE}` (" " id, integration_id, workspace_pk, api_key_pk, " " usage_date, model_name, aliyun_account_id, " " input_tokens, output_tokens, request_count, " " cost_usd, cost_cny, synced_at, created_at, updated_at" ") VALUES (" " :id, :integration_id, NULL, NULL, " " :usage_date, :model_name, :aliyun_account_id, " " 0, 0, 0, " " :cost_usd, :cost_cny, :synced_at, :created_at, :created_at" ")" ), inserts, ) # Drop the original aggregate rows. if deletes: bind.execute( sa.text(f"DELETE FROM `{_TABLE}` WHERE id IN :ids").bindparams( sa.bindparam("ids", expanding=True), ), {"ids": deletes}, ) # 4. Per-account drill-down covering index. if not _index_exists(bind, _TABLE, _NEW_INDEX): op.create_index( _NEW_INDEX, _TABLE, ["aliyun_account_id", "usage_date"], ) def downgrade() -> None: bind = op.get_bind() # 1. Re-merge per-account rows back into aggregates so the # pre-085 schema (which only allows one row per # (integration_id, usage_date, model_name)) is satisfied. rows = bind.execute(sa.text( f"SELECT id, integration_id, usage_date, model_name, cost_usd, " f"cost_cny, aliyun_account_id, synced_at, created_at " f"FROM `{_TABLE}` " f"WHERE aliyun_account_id IS NOT NULL" )).fetchall() grouped: dict[tuple, dict] = defaultdict(lambda: { "cost_usd": Decimal("0"), "cost_cny": Decimal("0"), "per_account": [], "synced_at": None, "created_at": None, "ids": [], }) for row in rows: key = (row.integration_id, row.usage_date, row.model_name) g = grouped[key] g["cost_usd"] += Decimal(str(row.cost_usd or 0)) g["cost_cny"] += Decimal(str(row.cost_cny or 0)) g["per_account"].append({ "owner_id": row.aliyun_account_id, "amount_cny": float(row.cost_cny or 0), }) g["synced_at"] = row.synced_at g["created_at"] = row.created_at g["ids"].append(row.id) # Drop the new index + unique key first so we can re-insert. if _index_exists(bind, _TABLE, _NEW_INDEX): op.drop_index(_NEW_INDEX, table_name=_TABLE) if _index_exists(bind, _TABLE, _NEW_UNIQUE): op.drop_constraint(_NEW_UNIQUE, _TABLE, type_="unique") # Delete the per-account rows. all_ids = [i for g in grouped.values() for i in g["ids"]] if all_ids: bind.execute( sa.text(f"DELETE FROM `{_TABLE}` WHERE id IN :ids").bindparams( sa.bindparam("ids", expanding=True), ), {"ids": all_ids}, ) # Re-insert the merged aggregate rows. for (integ, usage_date, model), g in grouped.items(): bind.execute( sa.text( f"INSERT INTO `{_TABLE}` (" " id, integration_id, workspace_pk, api_key_pk, " " usage_date, model_name, " " input_tokens, output_tokens, request_count, " " cost_usd, cost_cny, raw_data, " " synced_at, created_at, updated_at" ") VALUES (" " :id, :integration_id, NULL, NULL, " " :usage_date, :model_name, " " 0, 0, 0, " " :cost_usd, :cost_cny, :raw, " " :synced_at, :created_at, :created_at" ")" ), { "id": str(uuid.uuid4()), "integration_id": integ, "usage_date": usage_date, "model_name": model, "cost_usd": g["cost_usd"], "cost_cny": g["cost_cny"], "raw": json.dumps({"per_account": g["per_account"]}), "synced_at": g["synced_at"], "created_at": g["created_at"], }, ) # Drop the column itself. if _column_exists(bind, _TABLE, "aliyun_account_id"): op.drop_column(_TABLE, "aliyun_account_id") # Restore the original unique constraint. if not _index_exists(bind, _TABLE, _OLD_UNIQUE): op.create_unique_constraint( _OLD_UNIQUE, _TABLE, ["integration_id", "usage_date", "model_name"], )