lhbfx/backend/scripts/init_db.py

from pathlib import Path

from _bootstrap import add_src_to_path

add_src_to_path()

from lhbfx.config import AppConfig, load_config
from lhbfx.db import db_cursor, execute_schema
from lhbfx.seed import seed_traders


def _column_exists(cursor, schema_name: str, table_name: str, column_name: str) -> bool:
    cursor.execute(
        """
        SELECT 1
        FROM information_schema.COLUMNS
        WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s AND COLUMN_NAME = %s
        LIMIT 1
        """,
        (schema_name, table_name, column_name),
    )
    return cursor.fetchone() is not None


def _index_exists(cursor, schema_name: str, table_name: str, index_name: str) -> bool:
    cursor.execute(
        """
        SELECT 1
        FROM information_schema.STATISTICS
        WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s AND INDEX_NAME = %s
        LIMIT 1
        """,
        (schema_name, table_name, index_name),
    )
    return cursor.fetchone() is not None


def apply_incremental_alters(config: AppConfig) -> None:
    schema_name = config.database.database
    with db_cursor(config=config) as (_, cursor):
        if not _column_exists(cursor, schema_name, "lhb_overview", "rid"):
            cursor.execute("ALTER TABLE lhb_overview ADD COLUMN rid VARCHAR(32) NULL AFTER flag")
        if not _column_exists(cursor, schema_name, "lhb_overview", "detail_url"):
            cursor.execute("ALTER TABLE lhb_overview ADD COLUMN detail_url VARCHAR(255) NULL AFTER net_buy")
        if not _index_exists(cursor, schema_name, "lhb_overview", "uniq_lhb_overview_record"):
            cursor.execute("ALTER TABLE lhb_overview ADD UNIQUE KEY uniq_lhb_overview_record (trade_date, stock_code, rid)")

        if not _column_exists(cursor, schema_name, "lhb_detail_seats", "rid"):
            cursor.execute("ALTER TABLE lhb_detail_seats ADD COLUMN rid VARCHAR(32) NULL AFTER stock_name")
        if not _column_exists(cursor, schema_name, "lhb_detail_seats", "section_title"):
            cursor.execute("ALTER TABLE lhb_detail_seats ADD COLUMN section_title VARCHAR(255) NULL AFTER rid")
        if not _column_exists(cursor, schema_name, "lhb_detail_seats", "buy_ratio"):
            cursor.execute("ALTER TABLE lhb_detail_seats ADD COLUMN buy_ratio VARCHAR(32) NULL AFTER buy_amount_wan")
        if not _column_exists(cursor, schema_name, "lhb_detail_seats", "sell_ratio"):
            cursor.execute("ALTER TABLE lhb_detail_seats ADD COLUMN sell_ratio VARCHAR(32) NULL AFTER sell_amount_wan")
        if not _column_exists(cursor, schema_name, "lhb_detail_seats", "detail_url"):
            cursor.execute("ALTER TABLE lhb_detail_seats ADD COLUMN detail_url VARCHAR(255) NULL AFTER matched_seat")
        if not _index_exists(cursor, schema_name, "lhb_detail_seats", "uniq_lhb_detail_record"):
            cursor.execute(
                "ALTER TABLE lhb_detail_seats ADD UNIQUE KEY uniq_lhb_detail_record (trade_date, stock_code, rid, table_title, seat_name)"
            )
        if not _index_exists(cursor, schema_name, "lhb_detail_seats", "idx_lhb_detail_trader_stock_date"):
            cursor.execute(
                "ALTER TABLE lhb_detail_seats ADD KEY idx_lhb_detail_trader_stock_date (matched_trader_name, stock_code, trade_date)"
            )

        if not _index_exists(cursor, schema_name, "warning_events", "idx_warning_events_trader_type_date_code"):
            cursor.execute(
                "ALTER TABLE warning_events ADD KEY idx_warning_events_trader_type_date_code (trader_name, warning_type, trade_date, stock_code)"
            )


def main() -> None:
    root = Path(__file__).resolve().parents[1]
    schema_path = root / "src" / "lhbfx" / "schema.sql"
    config = load_config()
    execute_schema(schema_path=schema_path, config=config)
    apply_incremental_alters(config)
    seed_traders(config=config)
    print("数据库建表完成，游资配置已写入数据库。")


if __name__ == "__main__":
    main()