QuantEngineByItz/tools/initialize_snapshot_admin_db.py

#!/usr/bin/env python3
"""
snapshot_admin.db를 올바른 스키마와 XLSX 데이터로 초기화
"""

import sqlite3
import json
import pandas as pd
from pathlib import Path
from datetime import datetime

def initialize_snapshot_admin_db():
    """snapshot_admin.db 초기화"""

    db_path = Path('src/quant_engine/snapshot_admin.db')
    xlsx_file = Path('GatherTradingData.xlsx')
    json_file = Path('GatherTradingData.json')

    print("="*80)
    print("snapshot_admin.db 초기화")
    print("="*80)

    # JSON 메타데이터 로드
    with open(json_file, encoding='utf-8') as f:
        metadata = json.load(f).get('metadata', {})

    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    # 1. settings 테이블 초기화
    print("\n[1] settings 테이블 초기화")
    cursor.execute("DROP TABLE IF EXISTS settings")
    cursor.execute("""
        CREATE TABLE settings (
            ordinal INTEGER PRIMARY KEY,
            key TEXT NOT NULL,
            value_json TEXT NOT NULL,
            note TEXT DEFAULT '',
            updated_at TEXT NOT NULL
        )
    """)

    # XLSX에서 settings 데이터 로드
    df_settings = pd.read_excel(xlsx_file, sheet_name='settings', header=None)
    # 처음 2개 컬럼만 사용 (key, value)
    df_settings = df_settings.iloc[:, :2]
    df_settings.columns = ['key', 'value']
    timestamp = datetime.now().isoformat()

    print(f"  {len(df_settings)}개 설정 로드 중...")
    for ordinal, (idx, row) in enumerate(df_settings.iterrows(), start=1):
        key = str(row['key'])
        value = str(row['value'])
        value_json = json.dumps(value, ensure_ascii=False)

        cursor.execute("""
            INSERT INTO settings (ordinal, key, value_json, note, updated_at)
            VALUES (?, ?, ?, ?, ?)
        """, (ordinal, key, value_json, "", timestamp))

    cursor.execute("SELECT COUNT(*) FROM settings")
    count = cursor.fetchone()[0]
    print(f"  [OK] {count}개 설정 로드")

    # 2. account_snapshot 테이블 초기화
    print("\n[2] account_snapshot 테이블 초기화")
    cursor.execute("DROP TABLE IF EXISTS account_snapshot")
    cursor.execute("""
        CREATE TABLE account_snapshot (
            ordinal INTEGER PRIMARY KEY,
            row_json TEXT NOT NULL,
            captured_at TEXT NOT NULL DEFAULT '',
            account TEXT NOT NULL DEFAULT '',
            account_type TEXT NOT NULL DEFAULT '',
            ticker TEXT NOT NULL DEFAULT '',
            name TEXT NOT NULL DEFAULT '',
            parse_status TEXT NOT NULL DEFAULT '',
            user_confirmed TEXT NOT NULL DEFAULT '',
            updated_at TEXT NOT NULL
        )
    """)

    # XLSX에서 account_snapshot 데이터 로드
    df_snapshot = pd.read_excel(xlsx_file, sheet_name='account_snapshot', header=1)

    print(f"  {len(df_snapshot)}개 스냅샷 로드 중...")
    for ordinal, (idx, row) in enumerate(df_snapshot.iterrows(), start=1):
        row_dict = row.to_dict()

        # row_json으로 저장
        row_json = json.dumps(row_dict, default=str, ensure_ascii=False)

        # 핵심 필드 추출
        captured_at = str(row_dict.get('captured_at', ''))
        account = str(row_dict.get('account', ''))
        account_type = str(row_dict.get('account_type', ''))
        ticker = str(row_dict.get('ticker', ''))
        name = str(row_dict.get('name', ''))
        parse_status = str(row_dict.get('parse_status', ''))
        user_confirmed = str(row_dict.get('user_confirmed', ''))

        cursor.execute("""
            INSERT INTO account_snapshot (
                ordinal, row_json, captured_at, account, account_type, ticker, name,
                parse_status, user_confirmed, updated_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            ordinal, row_json, captured_at, account, account_type, ticker, name,
            parse_status, user_confirmed, timestamp
        ))

    cursor.execute("SELECT COUNT(*) FROM account_snapshot")
    count = cursor.fetchone()[0]
    print(f"  [OK] {count}개 스냅샷 로드")

    # 3. 인덱스 생성
    print("\n[3] 인덱스 생성")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_account_snapshot_captured_at ON account_snapshot(captured_at)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_account_snapshot_ticker ON account_snapshot(ticker)")
    print(f"  [OK] 인덱스 생성")

    conn.commit()
    conn.close()

    # 검증
    print("\n[검증]")
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    cursor.execute("SELECT COUNT(*) FROM settings")
    settings_count = cursor.fetchone()[0]

    cursor.execute("SELECT COUNT(*) FROM account_snapshot")
    snapshot_count = cursor.fetchone()[0]

    print(f"  settings: {settings_count}개")
    print(f"  account_snapshot: {snapshot_count}개")

    # 스키마 확인
    cursor.execute("PRAGMA table_info(settings)")
    settings_cols = [col[1] for col in cursor.fetchall()]
    print(f"  settings 컬럼: {settings_cols}")

    cursor.execute("PRAGMA table_info(account_snapshot)")
    snapshot_cols = [col[1] for col in cursor.fetchall()]
    print(f"  account_snapshot 컬럼: {snapshot_cols[:5]}... ({len(snapshot_cols)} 개)")

    conn.close()

    print("\n[OK] snapshot_admin.db 초기화 완료")

if __name__ == "__main__":
    initialize_snapshot_admin_db()