375 lines
10 KiB
Python
375 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
将 taiji 数据库的表结构和数据复制到 postgres 数据库
|
|
|
|
使用方法:
|
|
python scripts/copy_database.py
|
|
|
|
注意: 需要安装 psycopg2-binary
|
|
pip install psycopg2-binary
|
|
"""
|
|
|
|
import psycopg2
|
|
from urllib.parse import quote_plus
|
|
import sys
|
|
|
|
# 数据库连接配置
|
|
DB_HOST = "taijipda.postgres.database.azure.com"
|
|
DB_USER = "taiji"
|
|
DB_PASSWORD = "By@123456."
|
|
DB_PORT = 5432
|
|
|
|
# 源数据库和目标数据库
|
|
SOURCE_DB = "taiji"
|
|
TARGET_DB = "postgres"
|
|
|
|
|
|
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_def = f' PRIMARY KEY ("{"\", \"".join(pk_columns)}")'
|
|
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_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 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
|
|
|
|
# 获取数据
|
|
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'
|
|
|
|
for row in rows:
|
|
try:
|
|
target_cursor.execute(insert_sql, row)
|
|
except Exception as e:
|
|
print(f" 警告: 插入数据失败 - {e}")
|
|
|
|
target_conn.commit()
|
|
|
|
source_cursor.close()
|
|
target_cursor.close()
|
|
|
|
return len(rows)
|
|
|
|
|
|
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 main():
|
|
print("=" * 60)
|
|
print("PostgreSQL 数据库复制工具")
|
|
print(f"源数据库: {SOURCE_DB}")
|
|
print(f"目标数据库: {TARGET_DB}")
|
|
print("=" * 60)
|
|
|
|
# 连接源数据库
|
|
print("\n[1] 连接源数据库...")
|
|
try:
|
|
source_conn = get_connection(SOURCE_DB)
|
|
print(f" ✓ 成功连接到 {SOURCE_DB}")
|
|
except Exception as e:
|
|
print(f" ✗ 连接失败: {e}")
|
|
sys.exit(1)
|
|
|
|
# 连接目标数据库
|
|
print("\n[2] 连接目标数据库...")
|
|
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] 清理目标数据库旧表...")
|
|
dropped_count = drop_all_tables(target_conn)
|
|
print(f" 共删除 {dropped_count} 个旧表")
|
|
|
|
# 获取所有表
|
|
print("\n[4] 获取源数据库表列表...")
|
|
tables = get_all_tables(source_conn)
|
|
print(f" 找到 {len(tables)} 个表:")
|
|
for t in tables:
|
|
print(f" - {t}")
|
|
|
|
# 复制表结构
|
|
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}")
|
|
|
|
# 获取并创建索引
|
|
indexes = get_indexes(source_conn, table)
|
|
for idx in indexes:
|
|
try:
|
|
# 修改索引名以避免冲突
|
|
target_cursor.execute(idx)
|
|
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] 同步序列值...")
|
|
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()
|
|
|
|
# 关闭连接
|
|
source_conn.close()
|
|
target_conn.close()
|
|
|
|
print("\n" + "=" * 60)
|
|
print("数据库复制完成!")
|
|
print("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|