forked from xiaohei/taiji-AI-PAD
更新api文档
This commit is contained in:
Ubuntu
@@ -0,0 +1,558 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
将 postgres 数据库的表结构和数据完全覆盖到 taiji 数据库
|
||||
|
||||
使用方法:
|
||||
python scripts/sync_postgres_to_taiji.py
|
||||
|
||||
注意:
|
||||
- 需要安装 psycopg2-binary: pip install psycopg2-binary
|
||||
- 此操作会删除 taiji 库中的所有现有数据！
|
||||
- 建议在执行前备份 taiji 库
|
||||
"""
|
||||
|
||||
import psycopg2
|
||||
from psycopg2.extras import Json, register_default_json, register_default_jsonb
|
||||
import json
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
# 数据库连接配置
|
||||
DB_HOST = "taijipda.postgres.database.azure.com"
|
||||
DB_USER = "taiji"
|
||||
DB_PASSWORD = "By@123456."
|
||||
DB_PORT = 5432
|
||||
|
||||
# 源数据库（新结构）和目标数据库（需要更新）
|
||||
SOURCE_DB = "postgres" # 新的表结构
|
||||
TARGET_DB = "taiji" # 需要更新的旧数据库
|
||||
|
||||
|
||||
def get_connection(database):
|
||||
"""获取数据库连接"""
|
||||
return psycopg2.connect(
|
||||
host=DB_HOST,
|
||||
port=DB_PORT,
|
||||
user=DB_USER,
|
||||
password=DB_PASSWORD,
|
||||
database=database,
|
||||
sslmode="require"
|
||||
)
|
||||
|
||||
|
||||
def get_all_tables(conn):
|
||||
"""获取所有用户表"""
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
SELECT table_name
|
||||
FROM information_schema.tables
|
||||
WHERE table_schema = 'public'
|
||||
AND table_type = 'BASE TABLE'
|
||||
ORDER BY table_name
|
||||
""")
|
||||
tables = [row[0] for row in cursor.fetchall()]
|
||||
cursor.close()
|
||||
return tables
|
||||
|
||||
|
||||
def get_table_ddl(conn, table_name):
|
||||
"""获取表的 DDL 语句"""
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 获取列定义
|
||||
cursor.execute("""
|
||||
SELECT
|
||||
column_name,
|
||||
data_type,
|
||||
character_maximum_length,
|
||||
numeric_precision,
|
||||
numeric_scale,
|
||||
is_nullable,
|
||||
column_default,
|
||||
udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public' AND table_name = %s
|
||||
ORDER BY ordinal_position
|
||||
""", (table_name,))
|
||||
|
||||
columns = cursor.fetchall()
|
||||
|
||||
if not columns:
|
||||
cursor.close()
|
||||
return None
|
||||
|
||||
# 构建列定义
|
||||
column_defs = []
|
||||
for col in columns:
|
||||
col_name, data_type, char_max_len, num_precision, num_scale, is_nullable, col_default, udt_name = col
|
||||
|
||||
# 处理数据类型
|
||||
if data_type == 'character varying':
|
||||
if char_max_len:
|
||||
type_str = f"VARCHAR({char_max_len})"
|
||||
else:
|
||||
type_str = "VARCHAR"
|
||||
elif data_type == 'character':
|
||||
type_str = f"CHAR({char_max_len})" if char_max_len else "CHAR"
|
||||
elif data_type == 'numeric':
|
||||
if num_precision and num_scale:
|
||||
type_str = f"NUMERIC({num_precision},{num_scale})"
|
||||
elif num_precision:
|
||||
type_str = f"NUMERIC({num_precision})"
|
||||
else:
|
||||
type_str = "NUMERIC"
|
||||
elif data_type == 'ARRAY':
|
||||
type_str = f"{udt_name.lstrip('_')}[]"
|
||||
elif data_type == 'USER-DEFINED':
|
||||
type_str = udt_name
|
||||
else:
|
||||
type_str = data_type.upper()
|
||||
|
||||
# 构建列定义
|
||||
col_def = f' "{col_name}" {type_str}'
|
||||
|
||||
if is_nullable == 'NO':
|
||||
col_def += " NOT NULL"
|
||||
|
||||
if col_default:
|
||||
col_def += f" DEFAULT {col_default}"
|
||||
|
||||
column_defs.append(col_def)
|
||||
|
||||
# 获取主键约束
|
||||
cursor.execute("""
|
||||
SELECT kcu.column_name
|
||||
FROM information_schema.table_constraints tc
|
||||
JOIN information_schema.key_column_usage kcu
|
||||
ON tc.constraint_name = kcu.constraint_name
|
||||
AND tc.table_schema = kcu.table_schema
|
||||
WHERE tc.constraint_type = 'PRIMARY KEY'
|
||||
AND tc.table_schema = 'public'
|
||||
AND tc.table_name = %s
|
||||
ORDER BY kcu.ordinal_position
|
||||
""", (table_name,))
|
||||
|
||||
pk_columns = [row[0] for row in cursor.fetchall()]
|
||||
|
||||
if pk_columns:
|
||||
pk_cols_str = '", "'.join(pk_columns)
|
||||
pk_def = f' PRIMARY KEY ("{pk_cols_str}")'
|
||||
column_defs.append(pk_def)
|
||||
|
||||
ddl = f'CREATE TABLE IF NOT EXISTS "{table_name}" (\n'
|
||||
ddl += ",\n".join(column_defs)
|
||||
ddl += "\n);"
|
||||
|
||||
cursor.close()
|
||||
return ddl
|
||||
|
||||
|
||||
def get_indexes(conn, table_name):
|
||||
"""获取表的索引"""
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
SELECT indexdef
|
||||
FROM pg_indexes
|
||||
WHERE schemaname = 'public'
|
||||
AND tablename = %s
|
||||
AND indexname NOT LIKE '%%_pkey'
|
||||
""", (table_name,))
|
||||
|
||||
indexes = [row[0] for row in cursor.fetchall()]
|
||||
cursor.close()
|
||||
return indexes
|
||||
|
||||
|
||||
def get_unique_constraints(conn, table_name):
|
||||
"""获取表的唯一约束"""
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
SELECT
|
||||
tc.constraint_name,
|
||||
string_agg(kcu.column_name, ', ' ORDER BY kcu.ordinal_position) as columns
|
||||
FROM information_schema.table_constraints tc
|
||||
JOIN information_schema.key_column_usage kcu
|
||||
ON tc.constraint_name = kcu.constraint_name
|
||||
AND tc.table_schema = kcu.table_schema
|
||||
WHERE tc.constraint_type = 'UNIQUE'
|
||||
AND tc.table_schema = 'public'
|
||||
AND tc.table_name = %s
|
||||
GROUP BY tc.constraint_name
|
||||
""", (table_name,))
|
||||
|
||||
constraints = cursor.fetchall()
|
||||
cursor.close()
|
||||
return constraints
|
||||
|
||||
|
||||
def get_foreign_keys(conn, table_name):
|
||||
"""获取表的外键约束"""
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
SELECT
|
||||
tc.constraint_name,
|
||||
kcu.column_name,
|
||||
ccu.table_name AS foreign_table_name,
|
||||
ccu.column_name AS foreign_column_name
|
||||
FROM information_schema.table_constraints AS tc
|
||||
JOIN information_schema.key_column_usage AS kcu
|
||||
ON tc.constraint_name = kcu.constraint_name
|
||||
AND tc.table_schema = kcu.table_schema
|
||||
JOIN information_schema.constraint_column_usage AS ccu
|
||||
ON ccu.constraint_name = tc.constraint_name
|
||||
AND ccu.table_schema = tc.table_schema
|
||||
WHERE tc.constraint_type = 'FOREIGN KEY'
|
||||
AND tc.table_schema = 'public'
|
||||
AND tc.table_name = %s
|
||||
""", (table_name,))
|
||||
|
||||
fks = cursor.fetchall()
|
||||
cursor.close()
|
||||
return fks
|
||||
|
||||
|
||||
def get_sequences(conn):
|
||||
"""获取所有序列"""
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
SELECT sequence_name
|
||||
FROM information_schema.sequences
|
||||
WHERE sequence_schema = 'public'
|
||||
""")
|
||||
sequences = [row[0] for row in cursor.fetchall()]
|
||||
cursor.close()
|
||||
return sequences
|
||||
|
||||
|
||||
def get_sequence_value(conn, sequence_name):
|
||||
"""获取序列当前值"""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(f'SELECT last_value FROM "{sequence_name}"')
|
||||
value = cursor.fetchone()[0]
|
||||
except:
|
||||
value = 1
|
||||
cursor.close()
|
||||
return value
|
||||
|
||||
|
||||
def get_json_columns(conn, table_name):
|
||||
"""获取表中的 JSON/JSONB 列"""
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = %s
|
||||
AND data_type IN ('json', 'jsonb')
|
||||
""", (table_name,))
|
||||
json_cols = [row[0] for row in cursor.fetchall()]
|
||||
cursor.close()
|
||||
return json_cols
|
||||
|
||||
|
||||
def copy_table_data(source_conn, target_conn, table_name):
|
||||
"""复制表数据"""
|
||||
source_cursor = source_conn.cursor()
|
||||
target_cursor = target_conn.cursor()
|
||||
|
||||
# 获取列名
|
||||
source_cursor.execute("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = 'public' AND table_name = %s
|
||||
ORDER BY ordinal_position
|
||||
""", (table_name,))
|
||||
|
||||
columns = [row[0] for row in source_cursor.fetchall()]
|
||||
|
||||
if not columns:
|
||||
source_cursor.close()
|
||||
target_cursor.close()
|
||||
return 0
|
||||
|
||||
# 获取 JSON 列
|
||||
json_columns = get_json_columns(source_conn, table_name)
|
||||
json_col_indices = [columns.index(col) for col in json_columns if col in columns]
|
||||
|
||||
# 获取数据
|
||||
columns_str = ', '.join([f'"{c}"' for c in columns])
|
||||
source_cursor.execute(f'SELECT {columns_str} FROM "{table_name}"')
|
||||
rows = source_cursor.fetchall()
|
||||
|
||||
if not rows:
|
||||
source_cursor.close()
|
||||
target_cursor.close()
|
||||
return 0
|
||||
|
||||
# 插入数据
|
||||
placeholders = ', '.join(['%s'] * len(columns))
|
||||
insert_sql = f'INSERT INTO "{table_name}" ({columns_str}) VALUES ({placeholders}) ON CONFLICT DO NOTHING'
|
||||
|
||||
inserted = 0
|
||||
for row in rows:
|
||||
try:
|
||||
# 转换 JSON 列的数据
|
||||
row_list = list(row)
|
||||
for idx in json_col_indices:
|
||||
if row_list[idx] is not None:
|
||||
# 如果是 dict 或 list，转换为 Json 对象
|
||||
if isinstance(row_list[idx], (dict, list)):
|
||||
row_list[idx] = Json(row_list[idx])
|
||||
|
||||
target_cursor.execute(insert_sql, tuple(row_list))
|
||||
inserted += 1
|
||||
except Exception as e:
|
||||
print(f" 警告: 插入数据失败 - {e}")
|
||||
|
||||
target_conn.commit()
|
||||
|
||||
source_cursor.close()
|
||||
target_cursor.close()
|
||||
|
||||
return inserted
|
||||
|
||||
|
||||
def drop_all_tables(conn):
|
||||
"""删除目标数据库中的所有表"""
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 获取所有表
|
||||
cursor.execute("""
|
||||
SELECT table_name
|
||||
FROM information_schema.tables
|
||||
WHERE table_schema = 'public'
|
||||
AND table_type = 'BASE TABLE'
|
||||
""")
|
||||
tables = [row[0] for row in cursor.fetchall()]
|
||||
|
||||
if tables:
|
||||
# 禁用外键检查并删除所有表
|
||||
for table in tables:
|
||||
try:
|
||||
cursor.execute(f'DROP TABLE IF EXISTS "{table}" CASCADE')
|
||||
print(f" ✓ 已删除表: {table}")
|
||||
except Exception as e:
|
||||
print(f" ✗ 删除表 {table} 失败: {e}")
|
||||
conn.commit()
|
||||
|
||||
cursor.close()
|
||||
return len(tables)
|
||||
|
||||
|
||||
def get_table_row_count(conn, table_name):
|
||||
"""获取表的行数"""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(f'SELECT COUNT(*) FROM "{table_name}"')
|
||||
count = cursor.fetchone()[0]
|
||||
except:
|
||||
count = 0
|
||||
cursor.close()
|
||||
return count
|
||||
|
||||
|
||||
def verify_migration(source_conn, target_conn, tables):
|
||||
"""验证迁移结果"""
|
||||
print("\n[验证] 检查迁移结果...")
|
||||
all_ok = True
|
||||
|
||||
for table in tables:
|
||||
source_count = get_table_row_count(source_conn, table)
|
||||
target_count = get_table_row_count(target_conn, table)
|
||||
|
||||
if source_count == target_count:
|
||||
print(f" ✓ {table}: {target_count} 行 (匹配)")
|
||||
else:
|
||||
print(f" ✗ {table}: 源={source_count}, 目标={target_count} (不匹配)")
|
||||
all_ok = False
|
||||
|
||||
return all_ok
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 70)
|
||||
print("PostgreSQL 数据库同步工具")
|
||||
print(f"源数据库: {SOURCE_DB} (新结构)")
|
||||
print(f"目标数据库: {TARGET_DB} (将被覆盖)")
|
||||
print(f"执行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
print("=" * 70)
|
||||
|
||||
# 确认操作
|
||||
print("\n⚠️ 警告: 此操作将删除 taiji 数据库中的所有现有数据！")
|
||||
confirm = input("确认继续? (输入 'yes' 继续): ")
|
||||
if confirm.lower() != 'yes':
|
||||
print("操作已取消")
|
||||
sys.exit(0)
|
||||
|
||||
# 连接源数据库
|
||||
print("\n[1] 连接源数据库 (postgres)...")
|
||||
try:
|
||||
source_conn = get_connection(SOURCE_DB)
|
||||
print(f" ✓ 成功连接到 {SOURCE_DB}")
|
||||
except Exception as e:
|
||||
print(f" ✗ 连接失败: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# 连接目标数据库
|
||||
print("\n[2] 连接目标数据库 (taiji)...")
|
||||
try:
|
||||
target_conn = get_connection(TARGET_DB)
|
||||
print(f" ✓ 成功连接到 {TARGET_DB}")
|
||||
except Exception as e:
|
||||
print(f" ✗ 连接失败: {e}")
|
||||
source_conn.close()
|
||||
sys.exit(1)
|
||||
|
||||
# 获取源数据库表列表
|
||||
print("\n[3] 获取源数据库表列表...")
|
||||
tables = get_all_tables(source_conn)
|
||||
print(f" 找到 {len(tables)} 个表:")
|
||||
for t in tables:
|
||||
count = get_table_row_count(source_conn, t)
|
||||
print(f" - {t} ({count} 行)")
|
||||
|
||||
# 删除目标数据库中的旧表
|
||||
print("\n[4] 清理目标数据库旧表...")
|
||||
dropped_count = drop_all_tables(target_conn)
|
||||
print(f" 共删除 {dropped_count} 个旧表")
|
||||
|
||||
# 复制表结构
|
||||
print("\n[5] 复制表结构...")
|
||||
target_cursor = target_conn.cursor()
|
||||
|
||||
for table in tables:
|
||||
print(f" 处理表: {table}")
|
||||
|
||||
# 获取并执行 DDL
|
||||
ddl = get_table_ddl(source_conn, table)
|
||||
if ddl:
|
||||
try:
|
||||
target_cursor.execute(ddl)
|
||||
target_conn.commit()
|
||||
print(f" ✓ 表结构已创建")
|
||||
except Exception as e:
|
||||
target_conn.rollback()
|
||||
if "already exists" in str(e):
|
||||
print(f" ○ 表已存在，跳过创建")
|
||||
else:
|
||||
print(f" ✗ 创建失败: {e}")
|
||||
|
||||
target_cursor.close()
|
||||
|
||||
# 复制数据（先复制，再创建外键约束）
|
||||
print("\n[6] 复制表数据...")
|
||||
for table in tables:
|
||||
print(f" 复制表: {table}")
|
||||
try:
|
||||
count = copy_table_data(source_conn, target_conn, table)
|
||||
print(f" ✓ 已复制 {count} 行数据")
|
||||
except Exception as e:
|
||||
print(f" ✗ 复制失败: {e}")
|
||||
|
||||
# 创建索引
|
||||
print("\n[7] 创建索引...")
|
||||
target_cursor = target_conn.cursor()
|
||||
|
||||
for table in tables:
|
||||
indexes = get_indexes(source_conn, table)
|
||||
for idx in indexes:
|
||||
try:
|
||||
target_cursor.execute(idx)
|
||||
target_conn.commit()
|
||||
print(f" ✓ 索引已创建: {table}")
|
||||
except Exception as e:
|
||||
target_conn.rollback()
|
||||
if "already exists" in str(e):
|
||||
pass # 静默跳过已存在的索引
|
||||
else:
|
||||
print(f" ✗ 索引创建失败 ({table}): {e}")
|
||||
|
||||
target_cursor.close()
|
||||
|
||||
# 创建唯一约束
|
||||
print("\n[8] 创建唯一约束...")
|
||||
target_cursor = target_conn.cursor()
|
||||
|
||||
for table in tables:
|
||||
constraints = get_unique_constraints(source_conn, table)
|
||||
for constraint_name, columns in constraints:
|
||||
try:
|
||||
sql = f'ALTER TABLE "{table}" ADD CONSTRAINT "{constraint_name}" UNIQUE ({columns})'
|
||||
target_cursor.execute(sql)
|
||||
target_conn.commit()
|
||||
print(f" ✓ 唯一约束已创建: {constraint_name}")
|
||||
except Exception as e:
|
||||
target_conn.rollback()
|
||||
if "already exists" in str(e):
|
||||
pass
|
||||
else:
|
||||
print(f" ✗ 唯一约束创建失败 ({constraint_name}): {e}")
|
||||
|
||||
target_cursor.close()
|
||||
|
||||
# 创建外键约束
|
||||
print("\n[9] 创建外键约束...")
|
||||
target_cursor = target_conn.cursor()
|
||||
|
||||
for table in tables:
|
||||
fks = get_foreign_keys(source_conn, table)
|
||||
for constraint_name, column_name, foreign_table, foreign_column in fks:
|
||||
try:
|
||||
sql = f'''
|
||||
ALTER TABLE "{table}"
|
||||
ADD CONSTRAINT "{constraint_name}"
|
||||
FOREIGN KEY ("{column_name}")
|
||||
REFERENCES "{foreign_table}" ("{foreign_column}")
|
||||
'''
|
||||
target_cursor.execute(sql)
|
||||
target_conn.commit()
|
||||
print(f" ✓ 外键已创建: {constraint_name}")
|
||||
except Exception as e:
|
||||
target_conn.rollback()
|
||||
if "already exists" in str(e):
|
||||
pass
|
||||
else:
|
||||
print(f" ✗ 外键创建失败 ({constraint_name}): {e}")
|
||||
|
||||
target_cursor.close()
|
||||
|
||||
# 更新序列
|
||||
print("\n[10] 同步序列值...")
|
||||
sequences = get_sequences(source_conn)
|
||||
target_cursor = target_conn.cursor()
|
||||
|
||||
for seq in sequences:
|
||||
try:
|
||||
value = get_sequence_value(source_conn, seq)
|
||||
target_cursor.execute(f"SELECT setval('{seq}', {value}, true)")
|
||||
target_conn.commit()
|
||||
print(f" ✓ 序列 {seq} 设置为 {value}")
|
||||
except Exception as e:
|
||||
target_conn.rollback()
|
||||
print(f" ✗ 序列 {seq} 同步失败: {e}")
|
||||
|
||||
target_cursor.close()
|
||||
|
||||
# 验证迁移结果
|
||||
print("\n[11] 验证迁移结果...")
|
||||
# 重新连接以获取最新数据
|
||||
target_conn.close()
|
||||
target_conn = get_connection(TARGET_DB)
|
||||
|
||||
verify_ok = verify_migration(source_conn, target_conn, tables)
|
||||
|
||||
# 关闭连接
|
||||
source_conn.close()
|
||||
target_conn.close()
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
if verify_ok:
|
||||
print("✓ 数据库同步完成！所有数据已成功迁移。")
|
||||
else:
|
||||
print("⚠ 数据库同步完成，但部分数据可能不一致，请检查。")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user