Files
aiData/DataX_Bad/scheme/generate_schema.py
HuangHai 4712967fbe 'commit'
2026-02-05 15:55:17 +08:00

183 lines
6.5 KiB
Python

import pymysql
import sys
import os
import re
# 默认数据库配置
DB_CONFIG = {
"host": "rm-bp1ux6tuk49er80t9xo.mysql.rds.aliyuncs.com",
"port": 3306,
"user": "ylt",
"password": "Ycharge666",
"database": "yltcharge",
"charset": "utf8mb4"
}
# MySQL 到 Doris 类型映射
TYPE_MAPPING = {
"int": "INT",
"tinyint": "TINYINT",
"smallint": "SMALLINT",
"mediumint": "INT",
"bigint": "BIGINT",
"float": "FLOAT",
"double": "DOUBLE",
"decimal": "DECIMAL",
"char": "CHAR",
"varchar": "VARCHAR",
"text": "STRING",
"mediumtext": "STRING",
"longtext": "STRING",
"date": "DATE",
"datetime": "DATETIME",
"timestamp": "DATETIME",
"bit": "BOOLEAN"
}
def get_table_schema(table_name):
connection = pymysql.connect(**DB_CONFIG)
try:
with connection.cursor(pymysql.cursors.DictCursor) as cursor:
# 获取列信息
sql = """
SELECT COLUMN_NAME, DATA_TYPE, COLUMN_TYPE, IS_NULLABLE, COLUMN_DEFAULT, COLUMN_COMMENT, EXTRA
FROM information_schema.COLUMNS
WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
ORDER BY ORDINAL_POSITION
"""
cursor.execute(sql, (DB_CONFIG['database'], table_name))
columns = cursor.fetchall()
# 获取主键信息
cursor.execute(f"SHOW KEYS FROM {table_name} WHERE Key_name = 'PRIMARY'")
pk_info = cursor.fetchall()
pk_columns = [info['Column_name'] for info in pk_info]
return columns, pk_columns
finally:
connection.close()
def generate_mysql_ddl(table_name, columns, pk_columns):
ddl = [f"CREATE TABLE `{table_name}` ("]
col_defs = []
for col in columns:
line = f" `{col['COLUMN_NAME']}` {col['COLUMN_TYPE']}"
if col['IS_NULLABLE'] == 'NO':
line += " NOT NULL"
if col['COLUMN_DEFAULT'] is not None:
if isinstance(col['COLUMN_DEFAULT'], str):
line += f" DEFAULT '{col['COLUMN_DEFAULT']}'"
else:
line += f" DEFAULT {col['COLUMN_DEFAULT']}"
if col['EXTRA']:
line += f" {col['EXTRA']}"
if col['COLUMN_COMMENT']:
line += f" COMMENT '{col['COLUMN_COMMENT']}'"
col_defs.append(line)
if pk_columns:
col_defs.append(f" PRIMARY KEY ({', '.join([f'`{c}`' for c in pk_columns])})")
ddl.append(",\n".join(col_defs))
ddl.append(") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;")
return "\n".join(ddl)
def generate_doris_ddl(table_name, columns, pk_columns):
# Doris 建议主键放在前面
ddl = [f"CREATE TABLE `{table_name}` ("]
col_defs = []
# 确定 Unique Key 列表(如果没有主键,则取第一个字段,但用户要求唯一主键)
keys = pk_columns if pk_columns else [columns[0]['COLUMN_NAME']]
for col in columns:
data_type = col['DATA_TYPE'].lower()
doris_type = TYPE_MAPPING.get(data_type, "STRING")
# 处理长度信息
if doris_type in ["VARCHAR", "CHAR"]:
# 从 COLUMN_TYPE 中提取长度,如 varchar(255)
match = re.search(r'\((\d+)\)', col['COLUMN_TYPE'])
if match:
length = match.group(1)
# Doris VARCHAR 最大支持 65533
if int(length) > 65533:
doris_type = "STRING"
else:
doris_type = f"{doris_type}({length})"
else:
doris_type = f"{doris_type}(255)"
elif doris_type == "DECIMAL":
match = re.search(r'\((\d+,\d+)\)', col['COLUMN_TYPE'])
if match:
doris_type = f"DECIMAL({match.group(1)})"
line = f" `{col['COLUMN_NAME']}` {doris_type}"
if col['IS_NULLABLE'] == 'NO':
line += " NOT NULL"
# Doris 默认值处理较简化,暂时只处理简单字符串和数字
if col['COLUMN_DEFAULT'] is not None:
if isinstance(col['COLUMN_DEFAULT'], str) and '(' not in col['COLUMN_DEFAULT']:
line += f" DEFAULT '{col['COLUMN_DEFAULT']}'"
if col['COLUMN_COMMENT']:
line += f" COMMENT '{col['COLUMN_COMMENT']}'"
col_defs.append(line)
ddl.append(",\n".join(col_defs))
ddl.append(")")
ddl.append("UNIQUE KEY(" + ", ".join([f"`{k}`" for k in keys]) + ")")
ddl.append(f"DISTRIBUTED BY HASH(`{keys[0]}`) BUCKETS 1")
ddl.append('PROPERTIES (')
ddl.append(' "replication_num" = "1"')
ddl.append(');')
return "\n".join(ddl)
def main():
output_dir = os.path.join(os.path.dirname(__file__), "output")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# 自动获取表名列表
if len(sys.argv) < 2:
print("未指定表名,将自动从 json 目录获取...")
json_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "json"))
if not os.path.exists(json_dir):
print(f"错误: 找不到目录 {json_dir}")
sys.exit(1)
tables = [f.replace('.json', '') for f in os.listdir(json_dir) if f.endswith('.json')]
if not tables:
print(f"错误: 在 {json_dir} 中没有找到 .json 文件")
sys.exit(1)
print(f"发现 {len(tables)} 张表: {', '.join(tables)}")
else:
tables = sys.argv[1:]
for table in tables:
print(f"正在获取表 {table} 的结构...")
try:
columns, pk_columns = get_table_schema(table)
mysql_ddl = generate_mysql_ddl(table, columns, pk_columns)
doris_ddl = generate_doris_ddl(table, columns, pk_columns)
mysql_file = os.path.join(output_dir, f"{table}_mysql.sql")
doris_file = os.path.join(output_dir, f"{table}_doris.sql")
with open(mysql_file, "w", encoding="utf-8") as f:
f.write(mysql_ddl)
with open(doris_file, "w", encoding="utf-8") as f:
f.write(doris_ddl)
print(f" - MySQL DDL 已生成: {mysql_file}")
print(f" - Doris DDL 已生成: {doris_file}")
except Exception as e:
print(f" - 处理表 {table} 时出错: {e}")
if __name__ == "__main__":
main()