from contextlib import contextmanager

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base

from log.log_manager import logger

Base = declarative_base()

DATABASE_URL = 'postgresql+psycopg://postgres:K8u3fg0o@47.119.128.161:60001/squirrel'
engine = create_engine(
    DATABASE_URL,
    pool_size=10,
    max_overflow=20,
    pool_timeout=30,
    pool_recycle=1800,  # 防止数据库端连接过期
    connect_args={
        'connect_timeout': 15,
        'keepalives_idle': 60,
        'keepalives_interval': 10,
        'keepalives_count': 5
    }
)

def get_engine():
    return engine

# 不在模块导入时自动创建表；提供显式创建函数
def create_tables():
    # 确保相关 model 模块已被导入并注册到 Base
    Base.metadata.create_all(engine)

@contextmanager
def get_session():
    session = sessionmaker(bind=engine)()
    try:
        yield session
        session.commit()  # 自动提交成功的事务
    except Exception as e:
        session.rollback()  # 异常时回滚
        logger.error(f"Database operation failed: {str(e)}")
        raise  # 重新抛出异常
    finally:
        session.close()  # 确保会话关闭