'commit'
This commit is contained in:
HuangHai
182
DataX/scheme/generate_schema.py
Normal file
182
DataX/scheme/generate_schema.py
Normal file
@@ -0,0 +1,182 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user