#!/usr/bin/env python3
"""
修复数据库脚本：为 gsdh_data 表添加缺失的 industry_company 列
"""

import psycopg2
import json
import os

def load_config():
    """加载数据库配置"""
    if os.path.exists("config.json"):
        with open("config.json", "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

def fix_industry_column():
    """添加 industry_company 列到 gsdh_data 表"""
    config = load_config()
    
    print(f"连接到数据库: {config.get('db_host')}:{config.get('db_port')}/{config.get('db_name')}")
    
    try:
        conn = psycopg2.connect(
            host=config.get("db_host", "localhost"),
            port=config.get("db_port", "5432"),
            user=config.get("db_user", "gsdh"),
            password=config.get("db_password", "123gsdh"),
            database=config.get("db_name", "gsdh"),
            connect_timeout=5
        )
        cur = conn.cursor()
        
        # 检查列是否存在
        cur.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'gsdh_data' AND column_name = 'industry_company'
        """)
        
        if cur.fetchone():
            print("✅ industry_company 列已存在，无需修复")
        else:
            # 添加列
            print("🔧 正在添加 industry_company 列...")
            cur.execute("""
                ALTER TABLE gsdh_data 
                ADD COLUMN industry_company VARCHAR(200)
            """)
            conn.commit()
            print("✅ industry_company 列已成功添加")
        
        # 验证列存在
        print("\n当前 gsdh_data 表的列：")
        cur.execute("""
            SELECT column_name, data_type 
            FROM information_schema.columns 
            WHERE table_name = 'gsdh_data'
            ORDER BY ordinal_position
        """)
        columns = cur.fetchall()
        for col in columns:
            print(f"  - {col[0]}: {col[1]}")
        
        cur.close()
        conn.close()
        print("\n✨ 修复完成！")
        
    except Exception as e:
        print(f"\n❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    return True

if __name__ == "__main__":
    fix_industry_column()