jiachenlong/backend/app/core/database.py

import os
import time
from sqlalchemy import create_engine, event
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool
from typing import Generator
import logging

logger = logging.getLogger(__name__)

DATABASE_URL = os.getenv(
    "DATABASE_URL",
    "postgresql://postgres:postgres@127.0.0.1:5432/zodiac"
)

# 增强版数据库引擎配置
engine = create_engine(
    DATABASE_URL,
    # 连接池配置
    poolclass=QueuePool,
    pool_size=20,           # 常规连接数
    max_overflow=40,        # 允许超出的连接数（高并发时）
    pool_timeout=30,       # 获取连接超时时间（秒）
    pool_recycle=1800,     # 连接回收时间（30分钟），避免连接过期
    pool_pre_ping=True,    # 每次获取连接前检查连接是否有效
    echo=False,
    # 连接参数优化
    connect_args={
        "connect_timeout": 10,
        "application_name": "zodiac-api",
        "options": "-c statement_timeout=30000"  # 查询超时30秒
    }
)

# 添加连接事件监听器
@event.listens_for(engine, "connect")
def set_connect_timeout(dbapi_conn, connection_record):
    """设置连接参数"""
    cursor = dbapi_conn.cursor()
    cursor.execute("SET statement_timeout = 30000")
    cursor.close()

@event.listens_for(engine, "checkout")
def check_connection(dbapi_conn, connection_record, connection_proxy):
    """检出连接时检查"""
    try:
        cursor = dbapi_conn.cursor()
        cursor.execute("SELECT 1")
        cursor.close()
    except Exception as e:
        logger.warning(f"连接检查失败: {e}")
        raise Exception("数据库连接无效")

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

def get_db() -> Generator:
    """获取数据库会话，带错误处理"""
    db = SessionLocal(expire_on_commit=False)
    try:
        yield db
    except Exception as e:
        logger.error(f"数据库会话错误: {e}")
        db.rollback()
        raise
    finally:
        db.close()