70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
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@localhost: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()
|