44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
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() # 确保会话关闭 |