220 lines
7.3 KiB
Python
220 lines
7.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
将 InfoGenie Go 后端使用的 MySQL 表从「测试库」全量同步到「生产库」。
|
||
同步时在目标库对每张表执行：DROP TABLE IF EXISTS → 按源库 SHOW CREATE TABLE 重建 → 插入源库数据，
|
||
确保列结构与测试库一致（旧生产表结构会被覆盖）。
|
||
|
||
使用前务必备份生产库。会删除并重建生产库中同名表。
|
||
|
||
环境变量（未设置时沿用仓库内 test/test_mysql_connect.py 的约定）：
|
||
源库：DB_SRC_HOST, DB_SRC_PORT, DB_SRC_NAME, DB_SRC_USER, DB_SRC_PASSWORD
|
||
目标：DB_DST_HOST, DB_DST_PORT, DB_DST_NAME, DB_DST_USER, DB_DST_PASSWORD
|
||
|
||
示例（PowerShell）：
|
||
$env:DB_DST_HOST='192.168.1.100'
|
||
$env:DB_DST_NAME='infogenie'
|
||
$env:DB_DST_USER='infogenie'
|
||
$env:DB_DST_PASSWORD='***'
|
||
python scripts/sync_mysql_test_to_prod.py --dry-run
|
||
python scripts/sync_mysql_test_to_prod.py --execute
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import sys
|
||
from typing import Any, List, Sequence, Tuple
|
||
|
||
try:
|
||
import pymysql
|
||
except ImportError:
|
||
print("缺少依赖: pip install pymysql", file=sys.stderr)
|
||
sys.exit(2)
|
||
|
||
# 与 internal/database/mysql.go AutoMigrate 一致
|
||
TABLES: Tuple[str, ...] = (
|
||
"ai_configs",
|
||
"site_60s_disabled",
|
||
"site_ai_runtime",
|
||
"site_60s_upstream",
|
||
"site_ai_model_disabled",
|
||
"site_feature_card_clicks",
|
||
)
|
||
|
||
DEFAULT_SRC = {
|
||
"host": "10.1.1.100",
|
||
"port": 3306,
|
||
"name": "infogenie-test",
|
||
"user": "infogenie-test",
|
||
"password": "infogenie-test",
|
||
} # 仅作文档占位；实际以 DB_SRC_* 为准
|
||
|
||
DEFAULT_DST = {
|
||
"host": "192.168.1.100",
|
||
"port": 3306,
|
||
"name": "infogenie",
|
||
"user": "infogenie",
|
||
"password": "infogenie",
|
||
}
|
||
|
||
|
||
def _env_int(key: str, default: int) -> int:
|
||
raw = os.environ.get(key)
|
||
if raw is None or str(raw).strip() == "":
|
||
return default
|
||
return int(str(raw).strip())
|
||
|
||
|
||
def conn_params_src() -> dict:
|
||
return {
|
||
"host": os.environ.get("DB_SRC_HOST", DEFAULT_SRC["host"]),
|
||
"port": _env_int("DB_SRC_PORT", DEFAULT_SRC["port"]),
|
||
"database": os.environ.get("DB_SRC_NAME", DEFAULT_SRC["name"]),
|
||
"user": os.environ.get("DB_SRC_USER", DEFAULT_SRC["user"]),
|
||
"password": os.environ.get("DB_SRC_PASSWORD", DEFAULT_SRC["password"]),
|
||
}
|
||
|
||
|
||
def conn_params_dst() -> dict:
|
||
return {
|
||
"host": os.environ.get("DB_DST_HOST", DEFAULT_DST["host"]),
|
||
"port": _env_int("DB_DST_PORT", DEFAULT_DST["port"]),
|
||
"database": os.environ.get("DB_DST_NAME", DEFAULT_DST["name"]),
|
||
"user": os.environ.get("DB_DST_USER", DEFAULT_DST["user"]),
|
||
"password": os.environ.get("DB_DST_PASSWORD", DEFAULT_DST["password"]),
|
||
}
|
||
|
||
|
||
def open_conn(**kw: Any):
|
||
return pymysql.connect(
|
||
host=kw["host"],
|
||
port=int(kw["port"]),
|
||
user=kw["user"],
|
||
password=kw["password"],
|
||
database=kw["database"],
|
||
charset="utf8mb4",
|
||
connect_timeout=15,
|
||
cursorclass=pymysql.cursors.Cursor,
|
||
)
|
||
|
||
|
||
def fetch_all_rows(cur, table: str) -> Tuple[List[str], List[tuple]]:
|
||
cur.execute(f"SELECT * FROM `{table}`")
|
||
rows = cur.fetchall()
|
||
cols = [d[0] for d in cur.description] if cur.description else []
|
||
return cols, list(rows)
|
||
|
||
|
||
def table_exists(cur, table: str) -> bool:
|
||
cur.execute("SHOW TABLES LIKE %s", (table,))
|
||
return cur.fetchone() is not None
|
||
|
||
|
||
def recreate_table_from_source(source, target, table: str) -> None:
|
||
"""用源库 DDL 在目标库 DROP 后重建，保证列与测试库一致（覆盖旧结构）。"""
|
||
with source.cursor() as sc:
|
||
sc.execute(f"SHOW CREATE TABLE `{table}`")
|
||
row = sc.fetchone()
|
||
if not row:
|
||
raise RuntimeError(f"源库不存在表: {table}")
|
||
ddl = row[1]
|
||
with target.cursor() as tc:
|
||
tc.execute("SET FOREIGN_KEY_CHECKS=0")
|
||
tc.execute(f"DROP TABLE IF EXISTS `{table}`")
|
||
tc.execute(ddl)
|
||
tc.execute("SET FOREIGN_KEY_CHECKS=1")
|
||
target.commit()
|
||
|
||
|
||
def dry_run(src: dict, dst: dict, tables: Sequence[str]) -> None:
|
||
s = open_conn(**src)
|
||
d = open_conn(**dst)
|
||
try:
|
||
with s.cursor() as sc, d.cursor() as dc:
|
||
for t in tables:
|
||
sc.execute(f"SELECT COUNT(*) FROM `{t}`")
|
||
(n_src,) = sc.fetchone()
|
||
if not table_exists(dc, t):
|
||
print(f"{t}: source_rows={n_src} target_rows=(表不存在)")
|
||
else:
|
||
dc.execute(f"SELECT COUNT(*) FROM `{t}`")
|
||
(n_dst,) = dc.fetchone()
|
||
print(f"{t}: source_rows={n_src} target_rows={n_dst}")
|
||
finally:
|
||
s.close()
|
||
d.close()
|
||
|
||
|
||
def sync_tables(src: dict, dst: dict, tables: Sequence[str]) -> None:
|
||
source = open_conn(**src)
|
||
target = open_conn(**dst)
|
||
try:
|
||
with source.cursor() as sc:
|
||
for table in tables:
|
||
recreate_table_from_source(source, target, table)
|
||
print(f"++ {table}: 目标表已按源库 DDL 重建")
|
||
cols, rows = fetch_all_rows(sc, table)
|
||
if not cols:
|
||
raise RuntimeError(f"无法读取列: {table}")
|
||
placeholders = ",".join(["%s"] * len(cols))
|
||
col_sql = ",".join(f"`{c}`" for c in cols)
|
||
insert_sql = f"INSERT INTO `{table}` ({col_sql}) VALUES ({placeholders})"
|
||
with target.cursor() as tc:
|
||
if rows:
|
||
tc.executemany(insert_sql, rows)
|
||
target.commit()
|
||
print(f"OK {table}: copied {len(rows)} rows")
|
||
finally:
|
||
source.close()
|
||
target.close()
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description="Sync InfoGenie MySQL tables from test to production.")
|
||
ap.add_argument(
|
||
"--execute",
|
||
action="store_true",
|
||
help="执行同步（默认仅打印说明；不加此参数则退出码 1）",
|
||
)
|
||
ap.add_argument("--dry-run", action="store_true", help="仅统计源/目标行数，不写生产库")
|
||
ap.add_argument(
|
||
"--tables",
|
||
type=str,
|
||
default="",
|
||
help="逗号分隔表名子集；默认同步全部 GORM 表",
|
||
)
|
||
args = ap.parse_args()
|
||
|
||
src = conn_params_src()
|
||
dst = conn_params_dst()
|
||
|
||
tables: Tuple[str, ...]
|
||
if args.tables.strip():
|
||
tables = tuple(x.strip() for x in args.tables.split(",") if x.strip())
|
||
bad = [t for t in tables if t not in TABLES]
|
||
if bad:
|
||
print(f"未知表名（不在白名单）: {bad}", file=sys.stderr)
|
||
sys.exit(1)
|
||
else:
|
||
tables = TABLES
|
||
|
||
print(f"SOURCE: {src['user']}@{src['host']}:{src['port']}/{src['database']}")
|
||
print(f"TARGET: {dst['user']}@{dst['host']}:{dst['port']}/{dst['database']}")
|
||
print(f"TABLES: {', '.join(tables)}")
|
||
|
||
if args.dry_run:
|
||
dry_run(src, dst, tables)
|
||
return
|
||
|
||
if not args.execute:
|
||
print("未执行：请加 --dry-run 查看行数，或加 --execute 在确认备份后同步。", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
sync_tables(src, dst, tables)
|
||
print("完成。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|