import os
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool

COOLBOT_DB_URL = os.getenv(
    "COOLBOT_DB_URL"
)

# 如果未设置环境变量，抛出错误
if not COOLBOT_DB_URL:
    raise ValueError("COOLBOT_DB_URL environment variable is not set")

coolbot_engine = create_engine(
    COOLBOT_DB_URL,
    poolclass=QueuePool,
    pool_size=10,
    max_overflow=20,
    pool_timeout=30,
    pool_recycle=1800,
    pool_pre_ping=True,
    echo=False,
    connect_args={
        "connect_timeout": 10,
        "application_name": "zodiac-coolbot"
    }
)

CoolbotSession = sessionmaker(autocommit=False, autoflush=False, bind=coolbot_engine)

def get_coolbot_db():
    """获取coolbot数据库会话"""
    db = CoolbotSession()
    try:
        yield db
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()