python_market_data_service/scripts/create_split_tables.py

#!/usr/bin/env python3
"""
创建拆分后的新表结构

用法:
    python scripts/create_split_tables.py
    python scripts/create_split_tables.py --drop-existing
"""
import argparse
import sys
from pathlib import Path

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import create_engine, text
from app.repositories.database import Base, engine
from app.core.config import get_config
from app.core.logger import info, error


def create_tables(drop_existing: bool = False):
    """创建所有表"""
    config = get_config()
    
    # 创建数据库连接
    db_url = f"mysql+pymysql://{config.database.user}:{config.database.password}@{config.database.host}:{config.database.port}/{config.database.database}"
    engine_local = create_engine(db_url, echo=False)
    
    if drop_existing:
        info("Dropping existing split tables...")
        tables_to_drop = [
            # 股票分表
            "stock_klines_1d_base", "stock_klines_1d_quote", "stock_klines_1d_finance",
            "stock_realtime_quotes",
            # 期货分表
            "futures_klines_1d_base", "futures_klines_1d_quote",
            "futures_klines_1m_base", "futures_klines_5m_base",
            "futures_realtime_quotes",
        ]
        with engine_local.connect() as conn:
            for table in tables_to_drop:
                try:
                    conn.execute(text(f"DROP TABLE IF EXISTS {table}"))
                    info(f"Dropped table: {table}")
                except Exception as e:
                    error(f"Failed to drop {table}: {e}")
            conn.commit()
    
    info("Creating new split tables...")
    
    # 导入所有模型以确保它们被注册
    from app.repositories.models import (
        # 股票分表
        StockKLine1DBase, StockKLine1DQuote, StockKLine1DFinance,
        StockKLine1M, StockKLine5M, StockKLine15M, StockKLine30M, StockKLine60M,
        StockRealTimeQuote,
        # 期货分表
        FuturesKLine1DBase, FuturesKLine1DQuote,
        FuturesKLine1MBase, FuturesKLine5MBase, FuturesKLine15MBase, FuturesKLine30MBase, FuturesKLine60MBase,
        FuturesRealTimeQuote,
    )
    
    # 创建表
    tables = [
        # 股票表 - 日线分表
        StockKLine1DBase.__table__,
        StockKLine1DQuote.__table__,
        StockKLine1DFinance.__table__,
        # 股票表 - 分钟线
        StockKLine1M.__table__,
        StockKLine5M.__table__,
        StockKLine15M.__table__,
        StockKLine30M.__table__,
        StockKLine60M.__table__,
        StockRealTimeQuote.__table__,
        # 期货表 - 日线分表
        FuturesKLine1DBase.__table__,
        FuturesKLine1DQuote.__table__,
        # 期货表 - 分钟线
        FuturesKLine1MBase.__table__,
        FuturesKLine5MBase.__table__,
        FuturesKLine15MBase.__table__,
        FuturesKLine30MBase.__table__,
        FuturesKLine60MBase.__table__,
        FuturesRealTimeQuote.__table__,
    ]
    
    Base.metadata.create_all(bind=engine_local, tables=tables)
    
    info("Tables created successfully!")
    
    # 显示创建的表
    with engine_local.connect() as conn:
        result = conn.execute(text("""
            SELECT TABLE_NAME, TABLE_COMMENT 
            FROM information_schema.TABLES 
            WHERE TABLE_SCHEMA = DATABASE()
            AND (
                TABLE_NAME LIKE 'stock_klines_%' 
                OR TABLE_NAME LIKE 'futures_klines_%'
                OR TABLE_NAME LIKE '%realtime_quotes'
            )
        """))
        tables = result.fetchall()
        
        info("\nCreated split tables:")
        stock_tables = []
        futures_tables = []
        for table_name, comment in tables:
            if table_name.startswith('stock_'):
                stock_tables.append(f"  - {table_name}")
            else:
                futures_tables.append(f"  - {table_name}")
        
        if stock_tables:
            info("\n[Stock Tables]")
            for t in stock_tables:
                info(t)
        
        if futures_tables:
            info("\n[Futures Tables]")
            for t in futures_tables:
                info(t)


def main():
    parser = argparse.ArgumentParser(description="Create split table structure")
    parser.add_argument(
        "--drop-existing",
        action="store_true",
        help="Drop existing tables before creating (WARNING: data will be lost!)"
    )
    
    args = parser.parse_args()
    
    if args.drop_existing:
        confirm = input("WARNING: This will drop existing tables and ALL DATA will be lost!\nType 'yes' to continue: ")
        if confirm != "yes":
            print("Aborted.")
            return
    
    create_tables(drop_existing=args.drop_existing)


if __name__ == "__main__":
    main()