from sqlalchemy import create_engine, inspect
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
import asyncio
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

async def check_table_constraints():
    try:
        # 获取数据库连接字符串
        DATABASE_URL = os.getenv("DATABASE_URL", "mysql+asyncmy://root:123456@localhost:3306/th_agenter")
        
        # 创建异步引擎
        engine = create_async_engine(DATABASE_URL, echo=True)
        
        # 创建会话
        async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
        
        async with async_session() as session:
            # 获取数据库连接
            async with session.begin():
                # 使用inspect查看表结构
                inspector = inspect(engine)
                
                # 获取messages表的所有约束
                constraints = await engine.run_sync(inspector.get_table_constraints, 'messages')
                print("Messages表的所有约束:")
                for constraint in constraints:
                    print(f"  约束名称: {constraint['name']}, 类型: {constraint['type']}")
                    if constraint['type'] == 'PRIMARY KEY':
                        print(f"  主键约束列: {constraint['constrained_columns']}")
    
    except Exception as e:
        print(f"检查约束时出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(check_table_constraints())