Files
2026-03-13 23:48:49 +08:00

51 lines
1.5 KiB
Python

# database.py
import os
from dotenv import load_dotenv
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker
load_dotenv()
# 数据库连接 URL,可从 .env 配置,默认 SQLite
SQLALCHEMY_DATABASE_URL = os.getenv("SQLALCHEMY_DATABASE_URL", "sqlite:///./data/demo.db")
# 创建数据库引擎
# 增加 timeout=30 允许连接在遇到 locked 时最多等待 30 秒,而不是直接报错
engine = create_engine(
SQLALCHEMY_DATABASE_URL,
connect_args={"check_same_thread": False, "timeout": 30}
)
# ==========================================
# 监听 SQLite 连接,注入 PRAGMA 指令
# ==========================================
@event.listens_for(engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
cursor = dbapi_connection.cursor()
# 1. 开启 WAL 模式:读写分离
cursor.execute("PRAGMA journal_mode=WAL;")
# 2. 优化同步模式:在 WAL 模式下,NORMAL 既能保证不丢数据,又能大幅提升写入速度
cursor.execute("PRAGMA synchronous=NORMAL;")
# 3. 强制开启外键约束:极其重要!SQLite 默认不检查外键,这行能保护你的多态关联不乱套
cursor.execute("PRAGMA foreign_keys=ON;")
cursor.close()
# 创建数据库会话工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# 依赖注入函数:每个请求过来时,给它发一个数据库连接,请求结束时自动关闭
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()