|
|
"""
|
|
|
数据库迁移脚本：为合约管理优化添加新字段和品种表
|
|
|
"""
|
|
|
import sys
|
|
|
import os
|
|
|
import re
|
|
|
from datetime import datetime
|
|
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
|
|
from app.database import engine, Base, SessionLocal
|
|
|
from app.models import ProductInfo, ContractInfo
|
|
|
from sqlalchemy import text
|
|
|
|
|
|
|
|
|
# 品种分类映射（根据交易所和品种代码）
|
|
|
CATEGORY_MAP = {
|
|
|
# 能源化工
|
|
|
"rb": "能源化工", "hc": "能源化工", "fu": "能源化工", "bu": "能源化工",
|
|
|
"ru": "能源化工", "nr": "能源化工", "sp": "能源化工", "pg": "能源化工",
|
|
|
"pp": "能源化工", "l": "能源化工", "v": "能源化工", "eg": "能源化工",
|
|
|
"sc": "能源化工", "lu": "能源化工", "low": "能源化工",
|
|
|
# 金属
|
|
|
"cu": "金属", "al": "金属", "zn": "金属", "pb": "金属", "ni": "金属",
|
|
|
"sn": "金属", "au": "金属", "ag": "金属", "ss": "能源化工",
|
|
|
# 农产品
|
|
|
"a": "农产品", "b": "农产品", "c": "农产品", "cs": "农产品",
|
|
|
"m": "农产品", "y": "农产品", "p": "农产品", "rm": "农产品",
|
|
|
"oi": "农产品", "cf": "农产品", "sr": "农产品", "ta": "能源化工",
|
|
|
"ma": "能源化工", "fg": "能源化工", "sa": "农产品", "ur": "农产品",
|
|
|
"sf": "金属", "sm": "金属", "ap": "农产品", "cj": "农产品",
|
|
|
"rr": "农产品", "jd": "农产品", "lh": "农产品",
|
|
|
# 金融
|
|
|
"if": "金融", "ih": "金融", "ic": "金融", "im": "金融",
|
|
|
"t": "金融", "tf": "金融", "ts": "金融", "tl": "金融",
|
|
|
# 其他
|
|
|
"px": "能源化工", "pr": "能源化工", "br": "能源化工",
|
|
|
"lc": "金属", "si": "金属", "ec": "金融",
|
|
|
}
|
|
|
|
|
|
# 品种中文名映射
|
|
|
PRODUCT_NAME_MAP = {
|
|
|
# 能源化工
|
|
|
"rb": "螺纹钢", "hc": "热卷", "fu": "燃料油", "bu": "沥青",
|
|
|
"ru": "橡胶", "nr": "20号胶", "sp": "纸浆", "pg": "液化气",
|
|
|
"pp": "聚丙烯", "l": "塑料", "v": "PVC", "eg": "乙二醇",
|
|
|
"sc": "原油", "lu": "低硫燃料油", "low": "低硫燃料油",
|
|
|
"ta": "PTA", "ma": "甲醇", "fg": "玻璃", "ur": "尿素",
|
|
|
"sa": "纯碱", "px": "对二甲苯", "pr": "丙二醇", "br": "合成橡胶",
|
|
|
# 金属
|
|
|
"cu": "铜", "al": "铝", "zn": "锌", "pb": "铅", "ni": "镍",
|
|
|
"sn": "锡", "au": "黄金", "ag": "白银", "ss": "不锈钢",
|
|
|
"sf": "硅铁", "sm": "锰硅", "lc": "碳酸锂", "si": "工业硅",
|
|
|
# 农产品
|
|
|
"a": "豆一", "b": "豆二", "c": "玉米", "cs": "淀粉",
|
|
|
"m": "豆粕", "y": "豆油", "p": "棕榈油", "rm": "菜粕",
|
|
|
"oi": "菜油", "cf": "棉花", "sr": "白糖", "ap": "苹果",
|
|
|
"cj": "红枣", "rr": "粳米", "jd": "鸡蛋", "lh": "生猪",
|
|
|
# 金融
|
|
|
"if": "沪深300", "ih": "上证50", "ic": "中证500", "im": "中证1000",
|
|
|
"t": "10年国债", "tf": "5年国债", "ts": "2年国债", "tl": "30年国债", "ec": "工业硅",
|
|
|
# 能源
|
|
|
"sc": "原油", "lu": "低硫燃油",
|
|
|
}
|
|
|
|
|
|
# 交易所映射
|
|
|
EXCHANGE_MAP = {
|
|
|
"SHFE": "上海期货交易所",
|
|
|
"DCE": "大连商品交易所",
|
|
|
"CZCE": "郑州商品交易所",
|
|
|
"CFFEX": "中国金融期货交易所",
|
|
|
"INE": "上海国际能源交易中心",
|
|
|
"GFEX": "广州期货交易所",
|
|
|
}
|
|
|
|
|
|
|
|
|
def extract_year_month(symbol: str) -> tuple:
|
|
|
"""从合约代码解析交割年月
|
|
|
|
|
|
示例:
|
|
|
rb2401 -> (2024-01, 1)
|
|
|
cu2312 -> (2023-12, 12)
|
|
|
IF2403 -> (2024-03, 3)
|
|
|
"""
|
|
|
# 匹配末尾的数字（年份+月份）
|
|
|
match = re.search(r'(\d{2})(\d{2})$', symbol)
|
|
|
if not match:
|
|
|
return None, None
|
|
|
|
|
|
year_suffix, month = match.groups()
|
|
|
year_suffix = int(year_suffix)
|
|
|
month = int(month)
|
|
|
|
|
|
# 年份处理：假设是 20xx 年
|
|
|
year = 2000 + year_suffix if year_suffix < 100 else year_suffix
|
|
|
|
|
|
return f"{year}-{month:02d}", month
|
|
|
|
|
|
|
|
|
def extract_product_code(symbol: str) -> str:
|
|
|
"""从合约代码提取品种代码
|
|
|
|
|
|
示例:
|
|
|
rb2401 -> rb
|
|
|
cu2312 -> cu
|
|
|
IF2403 -> IF
|
|
|
"""
|
|
|
# 去掉末尾的数字
|
|
|
return re.sub(r'\d+$', '', symbol)
|
|
|
|
|
|
|
|
|
def migrate():
|
|
|
"""执行迁移"""
|
|
|
print("🚀 开始数据库迁移...")
|
|
|
|
|
|
db = SessionLocal()
|
|
|
try:
|
|
|
# 1. 创建新表（如果不存在）
|
|
|
print("📦 创建 product_info 表...")
|
|
|
ProductInfo.__table__.create(engine, checkfirst=True)
|
|
|
|
|
|
# 2. 添加新字段到 contract_info（SQLite 不支持 ALTER TABLE ADD COLUMN 的某些操作，需要检查）
|
|
|
print("🔧 检查 contract_info 表结构...")
|
|
|
|
|
|
# 检查字段是否存在
|
|
|
inspector_result = db.execute(text("PRAGMA table_info(contract_info)")).fetchall()
|
|
|
existing_columns = [row[1] for row in inspector_result]
|
|
|
|
|
|
new_columns = {
|
|
|
"year_month": "ALTER TABLE contract_info ADD COLUMN year_month VARCHAR(7)",
|
|
|
"delivery_month": "ALTER TABLE contract_info ADD COLUMN delivery_month INTEGER",
|
|
|
"is_main": "ALTER TABLE contract_info ADD COLUMN is_main BOOLEAN DEFAULT 0",
|
|
|
"listing_date": "ALTER TABLE contract_info ADD COLUMN listing_date DATETIME",
|
|
|
"volume": "ALTER TABLE contract_info ADD COLUMN volume BIGINT DEFAULT 0",
|
|
|
"open_interest": "ALTER TABLE contract_info ADD COLUMN open_interest BIGINT DEFAULT 0",
|
|
|
}
|
|
|
|
|
|
for col_name, sql in new_columns.items():
|
|
|
if col_name not in existing_columns:
|
|
|
print(f" 添加字段: {col_name}")
|
|
|
db.execute(text(sql))
|
|
|
else:
|
|
|
print(f" ✅ 字段已存在: {col_name}")
|
|
|
|
|
|
db.commit()
|
|
|
|
|
|
# 3. 数据迁移：填充衍生字段
|
|
|
print("📝 迁移现有合约数据...")
|
|
|
contracts = db.query(ContractInfo).all()
|
|
|
updated_count = 0
|
|
|
|
|
|
for contract in contracts:
|
|
|
# 解析 year_month 和 delivery_month
|
|
|
if not contract.year_month and contract.symbol:
|
|
|
year_month, delivery_month = extract_year_month(contract.symbol)
|
|
|
if year_month:
|
|
|
contract.year_month = year_month
|
|
|
contract.delivery_month = delivery_month
|
|
|
|
|
|
# 提取 product_code
|
|
|
if not contract.product and contract.symbol:
|
|
|
contract.product = extract_product_code(contract.symbol)
|
|
|
|
|
|
updated_count += 1
|
|
|
|
|
|
if updated_count > 0:
|
|
|
db.commit()
|
|
|
print(f" ✅ 更新 {updated_count} 条合约记录")
|
|
|
|
|
|
# 4. 生成品种元数据
|
|
|
print(" 生成品种元数据...")
|
|
|
products = db.query(
|
|
|
ContractInfo.product,
|
|
|
ContractInfo.exchange,
|
|
|
ContractInfo.multiplier,
|
|
|
ContractInfo.price_tick
|
|
|
).distinct().all()
|
|
|
|
|
|
product_count = 0
|
|
|
for prod_code, exchange, multiplier, price_tick in products:
|
|
|
if not prod_code:
|
|
|
continue
|
|
|
|
|
|
# 检查是否已存在
|
|
|
existing = db.query(ProductInfo).filter(
|
|
|
ProductInfo.product_code == prod_code
|
|
|
).first()
|
|
|
|
|
|
if existing:
|
|
|
continue
|
|
|
|
|
|
# 查找中文名
|
|
|
# 优先使用映射表
|
|
|
product_name = PRODUCT_NAME_MAP.get(prod_code.lower(), prod_code)
|
|
|
|
|
|
category = CATEGORY_MAP.get(prod_code.lower(), "其他")
|
|
|
|
|
|
product = ProductInfo(
|
|
|
product_code=prod_code,
|
|
|
product_name=product_name,
|
|
|
exchange=exchange,
|
|
|
multiplier=multiplier or 10,
|
|
|
price_tick=price_tick,
|
|
|
category=category,
|
|
|
is_active=True,
|
|
|
)
|
|
|
db.add(product)
|
|
|
db.commit() # 逐条提交避免批量冲突
|
|
|
product_count += 1
|
|
|
|
|
|
if product_count > 0:
|
|
|
print(f" ✅ 创建 {product_count} 个品种元数据")
|
|
|
|
|
|
# 5. 创建索引
|
|
|
print("📑 创建索引...")
|
|
|
try:
|
|
|
db.execute(text("CREATE INDEX IF NOT EXISTS idx_contract_year_month ON contract_info(year_month)"))
|
|
|
db.execute(text("CREATE INDEX IF NOT EXISTS idx_contract_is_main ON contract_info(is_main)"))
|
|
|
db.execute(text("CREATE INDEX IF NOT EXISTS idx_product_category ON product_info(category)"))
|
|
|
db.execute(text("CREATE INDEX IF NOT EXISTS idx_product_exchange ON product_info(exchange)"))
|
|
|
db.commit()
|
|
|
except Exception as e:
|
|
|
print(f" ⚠️ 索引创建警告: {e}")
|
|
|
|
|
|
print("\n✅ 迁移完成！")
|
|
|
|
|
|
except Exception as e:
|
|
|
db.rollback()
|
|
|
print(f"\n❌ 迁移失败: {e}")
|
|
|
import traceback
|
|
|
traceback.print_exc()
|
|
|
finally:
|
|
|
db.close()
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
migrate()
|
