"""Backfill an already-approved AWS Bedrock API key application. Records a historical assignment for the investibook project (applicant nizheming@pingcap.com). The approval flow happened outside Helix and the IAM user was provisioned manually, so we just need the paper trail in employee_requests + api_key_assignments. Idempotent on (assigned_to, project_name) — re-running is a no-op once the row exists. Requires the applicant person to exist in identity_persons (HRIS-synced on prod; sync HRIS locally if you need to reproduce there). Usage: # local (DATABASE_URL_SYNC defaults to local docker mysql) uv run python scripts/backfill_apikey_investibook.py # prod (TiDB Cloud) HELIX_PROD_DB_URL_SYNC='mysql+pymysql://...tidbcloud.com:4000/helix?ssl=true' \\ uv run python scripts/backfill_apikey_investibook.py --prod """ from __future__ import annotations import argparse import json import os import sys import uuid from datetime import UTC, datetime from sqlalchemy import create_engine, text from app.core.database_tls import database_engine_options APPLICANT_EMAIL = "nizheming@pingcap.com" PROJECT_NAME = "investibook" MONTHLY_BUDGET = "500.00" USER_ARN = "arn:aws:iam::841162690310:user/investibook" REASON = "investibook 项目接入 AWS Bedrock（流程已审批，补登记录）" LOCAL_DB_URL = "mysql+pymysql://helix:helix@localhost:3306/helix" def _resolve_db_url(use_prod: bool) -> str: if use_prod: url = os.getenv("HELIX_PROD_DB_URL_SYNC") if not url: raise SystemExit( "--prod requires HELIX_PROD_DB_URL_SYNC env var " "(see commented DATABASE_URL_SYNC in repo .env).", ) return url return os.getenv("DATABASE_URL_SYNC") or LOCAL_DB_URL def _find_provider(conn) -> tuple[str, str]: row = conn.execute(text( "SELECT id, name FROM api_key_providers " "WHERE LOWER(name) LIKE '%bedrock%' AND is_active = 1 " "ORDER BY name LIMIT 1", )).fetchone() if not row: raise SystemExit( "No active 'AWS Bedrock' provider found in api_key_providers. " "Create one via /admin/assets/api-keys → 服务商目录 first.", ) return row[0], row[1] def _generate_request_number() -> str: return f"AK-{datetime.now(UTC).strftime('%Y%m%d')}-{uuid.uuid4().hex[:4].upper()}" def backfill(use_prod: bool) -> None: db_url = _resolve_db_url(use_prod) db_url, connect_args = database_engine_options(db_url) engine = create_engine(db_url, future=True, connect_args=connect_args) env = "PROD" if use_prod else "LOCAL" print(f"[{env}] connecting to {db_url.split('@')[-1].split('?')[0]}") with engine.begin() as conn: person = conn.execute(text( "SELECT id, display_name FROM identity_persons " "WHERE work_email = :email", ), {"email": APPLICANT_EMAIL}).fetchone() if not person: raise SystemExit( f"[{env}] {APPLICANT_EMAIL} not found in identity_persons. " "Run HRIS sync (prod) or seed the person row (local) before retrying.", ) person_id, person_name = person provider_id, provider_name = _find_provider(conn) print(f" applicant: {person_name} ({person_id})") print(f" provider: {provider_name} ({provider_id})") existing = conn.execute(text(""" SELECT id, request_number, status FROM api_key_assignments WHERE assigned_to = :person_id AND LOWER(project_name) = LOWER(:project) """), {"person_id": person_id, "project": PROJECT_NAME}).fetchone() if existing: print( f" ! already exists: request_number={existing[1]} " f"status={existing[2]} — skipping insert.", ) return request_number = _generate_request_number() now = datetime.now(UTC).replace(tzinfo=None) valid_from = now.date() feishu_detail = { "template_key": "api_key_request", "approval_code": None, "approval_name": None, "instance_code": None, "instance_status": "APPROVED", "submitted_form": [], "manual_backfill": True, } bedrock_detail = { "granted": True, "in_flight": False, "manual_backfill": True, "user_arn": USER_ARN, "iam_user_name": USER_ARN.rsplit("/", 1)[-1], "region": "ap-southeast-1", "policy_name": "bedrock-invoke", "note": "Pre-existing IAM user provisioned outside Helix; " "record backfilled on " + now.date().isoformat(), } submitted_answers = { "project_name": PROJECT_NAME, "monthly_budget_usd": MONTHLY_BUDGET, "expected_usage": "Bedrock model invocation for investibook", } er_id = str(uuid.uuid4()) er_detail = { "submitted_answers": submitted_answers, "feishu": feishu_detail, "api_key": { "provider_id": provider_id, "provider_name": provider_name, "project_name": PROJECT_NAME, "monthly_budget_usd": MONTHLY_BUDGET, "valid_from": valid_from.isoformat(), "valid_until": None, }, } conn.execute(text(""" INSERT INTO employee_requests (id, request_number, request_type, status, requester_id, reason, detail, created_at, processed_at) VALUES (:id, :req_no, 'api_key_request', 'approved', :requester_id, :reason, :detail, :now, :now) """), { "id": er_id, "req_no": request_number, "requester_id": person_id, "reason": REASON, "detail": json.dumps(er_detail, ensure_ascii=False), "now": now, }) assignment_id = str(uuid.uuid4()) assignment_detail = { "submitted_answers": submitted_answers, "feishu": feishu_detail, "bedrock": bedrock_detail, } conn.execute(text(""" INSERT INTO api_key_assignments (id, provider_id, assigned_to, status, request_number, project_name, reason, monthly_budget_usd, expected_usage, key_value, valid_from, valid_until, extended_count, employee_request_id, approval_request_id, detail, created_at, updated_at) VALUES (:id, :provider_id, :assigned_to, 'active', :req_no, :project, :reason, :budget, :usage, NULL, :valid_from, NULL, 0, :er_id, NULL, :detail, :now, :now) """), { "id": assignment_id, "provider_id": provider_id, "assigned_to": person_id, "req_no": request_number, "project": PROJECT_NAME, "reason": REASON, "budget": MONTHLY_BUDGET, "usage": "Bedrock model invocation for investibook", "valid_from": valid_from, "er_id": er_id, "detail": json.dumps(assignment_detail, ensure_ascii=False), "now": now, }) print(f" + inserted employee_requests id={er_id}") print(f" + inserted api_key_assignments id={assignment_id}") print(f" request_number={request_number} status=active") print(f" user_arn={USER_ARN}") def main() -> int: p = argparse.ArgumentParser() p.add_argument( "--prod", action="store_true", help="Target prod TiDB via HELIX_PROD_DB_URL_SYNC env var " "(default: DATABASE_URL_SYNC env var or local docker mysql)", ) args = p.parse_args() backfill(use_prod=args.prod) return 0 if __name__ == "__main__": sys.exit(main())