Maas-Model-platform/backend/scripts/add_tingwu_comments.py

#!/usr/bin/env python
# coding=utf-8

"""
通义听悟任务表字段注释添加脚本
执行 SQL 迁移文件: 038_add_tingwu_task_comments.sql
"""

import os
import sys
import psycopg2
from pathlib import Path

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

from dotenv import load_dotenv

# 加载环境变量
env_path = project_root / '.env'
load_dotenv(env_path)


def get_db_connection():
    """获取数据库连接"""
    return psycopg2.connect(
        host=os.getenv('DB_HOST', 'localhost'),
        port=int(os.getenv('DB_PORT', 5432)),
        user=os.getenv('DB_USER', 'postgres'),
        password=os.getenv('DB_PASSWORD'),
        database=os.getenv('DB_NAME', 'model_square')
    )


def execute_migration():
    """执行迁移脚本"""
    migration_file = project_root / 'migrations' / '038_add_tingwu_task_comments.sql'
    
    if not migration_file.exists():
        print(f"❌ 迁移文件不存在: {migration_file}")
        return False
    
    print(f"📄 读取迁移文件: {migration_file}")
    with open(migration_file, 'r', encoding='utf-8') as f:
        sql_content = f.read()
    
    print("🔗 连接数据库...")
    conn = get_db_connection()
    cursor = conn.cursor()
    
    try:
        print("🚀 执行迁移...")
        cursor.execute(sql_content)
        conn.commit()
        print("✅ 迁移执行成功！")
        
        # 验证注释是否添加成功
        print("\n📊 验证字段注释...")
        cursor.execute("""
            SELECT 
                col.column_name,
                pgd.description
            FROM pg_catalog.pg_statio_all_tables AS st
            INNER JOIN pg_catalog.pg_description pgd ON (
                pgd.objoid = st.relid
            )
            INNER JOIN information_schema.columns col ON (
                pgd.objsubid = col.ordinal_position
                AND col.table_schema = st.schemaname
                AND col.table_name = st.relname
            )
            WHERE st.schemaname = 'aigcspace'
            AND st.relname = 'tingwu_task'
            ORDER BY col.ordinal_position
            LIMIT 10;
        """)
        
        results = cursor.fetchall()
        if results:
            print("\n前10个字段的注释:")
            for column_name, description in results:
                print(f"  • {column_name}: {description}")
        else:
            print("⚠️  未查询到字段注释，可能需要检查")
        
        return True
        
    except Exception as e:
        conn.rollback()
        print(f"❌ 迁移执行失败: {str(e)}")
        return False
        
    finally:
        cursor.close()
        conn.close()


def main():
    """主函数"""
    print("=" * 60)
    print("通义听悟任务表字段注释添加")
    print("=" * 60)
    print()
    
    success = execute_migration()
    
    print()
    print("=" * 60)
    if success:
        print("✅ 所有操作完成")
        print()
        print("📝 注释说明:")
        print("  - 表注释: aigcspace.tingwu_task")
        print("  - 字段注释: 53个字段")
        print("  - 索引说明: 5个索引")
        print()
        print("🔍 查看注释:")
        print("  psql -d model_square -c \"\\d+ aigcspace.tingwu_task\"")
    else:
        print("❌ 操作失败，请检查错误信息")
    print("=" * 60)


if __name__ == '__main__':
    main()