"""数据库配置"""
from sqlalchemy import create_engine, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from .config import DATABASE_URL

engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


def get_db():
    """获取数据库会话"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def _migrate_db():
    """简单的数据库迁移：为旧表添加缺失的列"""
    migrations = [
        ("tasks", "branch", "VARCHAR DEFAULT 'main'"),
    ]
    with engine.connect() as conn:
        for table, column, col_type in migrations:
            try:
                conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {col_type}"))
                conn.commit()
            except Exception:
                pass  # 列已存在则忽略


def init_db():
    """初始化数据库"""
    Base.metadata.create_all(bind=engine)
    _migrate_db()