68 lines
1.7 KiB
Python
68 lines
1.7 KiB
Python
"""数据库连接与会话管理模块"""
|
||
from sqlalchemy import create_engine
|
||
from sqlalchemy.orm import sessionmaker, declarative_base
|
||
from app.config import settings
|
||
|
||
# 处理兼容：如果设置了 DATABASE_URL（单数据库模式），则两个都使用它
|
||
if settings.DATABASE_URL:
|
||
inventory_url = settings.DATABASE_URL
|
||
pms_url = settings.DATABASE_URL
|
||
else:
|
||
inventory_url = settings.DATABASE_URL_INVENTORY
|
||
pms_url = settings.DATABASE_URL_PMS
|
||
|
||
# ============ 库存数据库（只读）============
|
||
|
||
engine_inventory = create_engine(
|
||
inventory_url,
|
||
pool_pre_ping=True,
|
||
pool_size=5,
|
||
max_overflow=10,
|
||
# 库存数据库只读，不需要写入优化
|
||
)
|
||
|
||
SessionLocalInventory = sessionmaker(autocommit=False, autoflush=False, bind=engine_inventory)
|
||
|
||
|
||
def get_db_inventory():
|
||
"""获取库存数据库会话的依赖函数（只读）"""
|
||
db = SessionLocalInventory()
|
||
try:
|
||
yield db
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
# ============ PMS 数据库（可读写）============
|
||
|
||
engine_pms = create_engine(
|
||
pms_url,
|
||
pool_pre_ping=True,
|
||
pool_size=10,
|
||
max_overflow=20,
|
||
)
|
||
|
||
SessionLocalPMS = sessionmaker(autocommit=False, autoflush=False, bind=engine_pms)
|
||
|
||
|
||
def get_db_pms():
|
||
"""获取 PMS 数据库会话的依赖函数（可读写）"""
|
||
db = SessionLocalPMS()
|
||
try:
|
||
yield db
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
# ============ 兼容旧接口 ============
|
||
|
||
# 为向后兼容提供默认的 engine 和 SessionLocal
|
||
engine = engine_pms
|
||
SessionLocal = SessionLocalPMS
|
||
Base = declarative_base()
|
||
|
||
|
||
def get_db():
|
||
"""获取数据库会话的依赖函数（兼容旧接口，默认使用 PMS 数据库）"""
|
||
yield from get_db_pms()
|