AiDOPWarehouse/doc/db/mdp/qms_add_tenant_id_d4.py

# -*- coding: utf-8 -*-
"""WP0-D4: add tenant_id to all qms_* tables (idempotent).

连接参数取自环境变量：MDP_DB_HOST / MDP_DB_PORT / MDP_DB_USER / MDP_DB_PASSWORD / MDP_DB_NAME
"""
import os

import pymysql

DEFAULT_TENANT = None  # 新环境不再 DEFAULT；存量回填由迁移脚本单独处理

def main():
    c = pymysql.connect(
        host=os.environ.get("MDP_DB_HOST", "127.0.0.1"),
        port=int(os.environ.get("MDP_DB_PORT", "3306")),
        user=os.environ.get("MDP_DB_USER", "root"),
        password=os.environ.get("MDP_DB_PASSWORD", ""),
        database=os.environ.get("MDP_DB_NAME", "aidopdev"),
        charset="utf8mb4",
        connect_timeout=15,
        read_timeout=180,
        autocommit=True,
    )
    cur = c.cursor()
    cur.execute(
        """
        SELECT TABLE_NAME FROM information_schema.TABLES
        WHERE TABLE_SCHEMA='aidopdev' AND TABLE_TYPE='BASE TABLE' AND TABLE_NAME LIKE 'qms\\_%'
        ORDER BY TABLE_NAME
        """
    )
    tables = [r[0] for r in cur.fetchall()]
    added, skipped, errors = [], [], []
    for t in tables:
        cur.execute(
            """
            SELECT COUNT(*) FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA='aidopdev' AND TABLE_NAME=%s AND COLUMN_NAME='tenant_id'
            """,
            (t,),
        )
        if cur.fetchone()[0] > 0:
            skipped.append(t)
            continue
        cur.execute(
            """
            SELECT COLUMN_NAME FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA='aidopdev' AND TABLE_NAME=%s
            ORDER BY ORDINAL_POSITION LIMIT 1
            """,
            (t,),
        )
        first = cur.fetchone()[0]
        after = " AFTER `id`" if first.lower() == "id" else ""
        sql = (
            f"ALTER TABLE `{t}` ADD COLUMN `tenant_id` BIGINT NOT NULL "
            f"COMMENT 'tenant_id WP0-D4'{after}"
        )
        try:
            cur.execute(sql)
        except Exception as e1:
            try:
                cur.execute(
                    f"ALTER TABLE `{t}` ADD COLUMN `tenant_id` BIGINT NOT NULL "
                    f"COMMENT 'tenant_id WP0-D4'"
                )
            except Exception as e2:
                errors.append((t, str(e2) or str(e1)))
                print("ERR", t, e2)
                continue

        if DEFAULT_TENANT is not None:
            cur.execute(
                f"UPDATE `{t}` SET tenant_id=%s WHERE tenant_id IS NULL OR tenant_id=0",
                (DEFAULT_TENANT,),
            )
        idx = f"idx_{t}_tenant"[:64]
        try:
            cur.execute(f"CREATE INDEX `{idx}` ON `{t}` (`tenant_id`)")
        except Exception as ie:
            if "Duplicate" not in str(ie):
                print("IDX warn", t, ie)
        added.append(t)
        print("OK", t)

    cur.execute(
        """
        SELECT COUNT(*) FROM information_schema.COLUMNS
        WHERE TABLE_SCHEMA='aidopdev' AND TABLE_NAME LIKE 'qms\\_%' AND COLUMN_NAME='tenant_id'
        """
    )
    print("done added=", len(added), "skipped=", len(skipped), "errors=", len(errors), "with_tenant=", cur.fetchone()[0])
    c.close()


if __name__ == "__main__":
    main()