Files
ai_trading_api/create/create_table.py
2025-12-26 16:40:20 +08:00

110 lines
3.8 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
create.py - 动态生成PostgreSQL建表SQL语句
"""
import os
import re
# 1. 定义表名称变量,方便调整
table_name = "cn_pmi_records" # 可以根据需要修改表名
# 2. 定义SQL模板
sql_template = f"""DO $$
BEGIN
IF NOT EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name = '{table_name}'
) THEN
CREATE TABLE {table_name} (
id UUID DEFAULT gen_random_uuid() PRIMARY KEY, -- id
payload JSONB NOT NULL, -- 数据
deleted BOOLEAN NOT NULL DEFAULT FALSE, -- 删除状态
created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, -- 记录创建时间
updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP -- 记录修改时间
);
-- 3 触发器:自动刷新 updated_at
CREATE TRIGGER trg_{table_name}_at
BEFORE UPDATE ON {table_name}
FOR EACH ROW
EXECUTE FUNCTION moddatetime(updated_at);
RAISE NOTICE '{table_name} 表已创建';
ELSE
RAISE NOTICE '{table_name} 表已存在,跳过';
END IF;
END $$;"""
def normalize_blank_lines(text):
"""规范化空行确保END $$;和DO $$之间只有一个空行"""
# 将多个空行替换为单个空行
text = re.sub(r'\n{3,}', '\n\n', text)
# 确保END $$;后面有一个空行再接DO $$
text = re.sub(r'END \$\$;(\s*)DO \$\$', r'END $$;\n\nDO $$', text)
# 清理开头和结尾的多余空行
text = text.strip() + '\n'
return text
def update_sql_file():
"""将生成的SQL语句追加到03_create_table.sql文件中"""
# 定义文件路径
sql_file_path = os.path.join(
os.path.dirname(os.path.dirname(__file__)),
'infra',
'postgres',
'sql',
'03_create_table.sql'
)
try:
# 读取现有文件内容
with open(sql_file_path, 'r', encoding='utf-8') as f:
content = f.read()
# 查找插入位置(在\"部署完成\"日志前)
insert_pattern = r'(DO \$\$.*?RAISE NOTICE \'🚀============ 数据库表部署开始 ============🚀\'.*?END \$\$;)(.*?)(DO \$\$.*?RAISE NOTICE \'✅============ 数据库表部署完成 ============✅\'.*?END \$\$;)'
match = re.search(insert_pattern, content, re.DOTALL)
if match:
# 分割内容
before_insert = match.group(1)
existing_middle = match.group(2)
after_insert = match.group(3)
# 规范化空行
existing_middle = normalize_blank_lines(existing_middle)
# 组合新内容
new_content = f"""{before_insert}\n\n{sql_template}\n\n{existing_middle}\n{after_insert}"""
# 规范化整个内容的空行
new_content = normalize_blank_lines(new_content)
# 写回文件
with open(sql_file_path, 'w', encoding='utf-8') as f:
f.write(new_content)
print(f"✅ 成功更新 {sql_file_path}")
print(f"📋 生成的表名: {table_name}")
print(f"📝 SQL内容已追加到文件中")
else:
print("❌ 无法找到插入位置,请检查文件格式")
except FileNotFoundError:
print(f"❌ 文件 {sql_file_path} 不存在")
except Exception as e:
print(f"❌ 处理文件时出错: {e}")
if __name__ == "__main__":
print(f"🚀 开始生成表 '{table_name}' 的SQL语句...")
update_sql_file()