Files
dataClean/database.py
T

40 lines
1001 B
Python
Raw Normal View History

2026-06-12 16:04:03 +08:00
"""数据库连接与初始化"""
from pathlib import Path
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker, declarative_base
from config import settings
engine = create_engine(
f"sqlite:///{settings.database_path}",
connect_args={"check_same_thread": False},
echo=False,
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
@event.listens_for(engine, "connect")
def _set_sqlite_pragma(dbapi_conn, connection_record):
"""启用 SQLite 外键约束"""
cursor = dbapi_conn.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
def get_db():
"""FastAPI 依赖注入用数据库会话"""
db = SessionLocal()
try:
yield db
finally:
db.close()
def init_db():
"""创建所有数据表,并确保数据库目录存在"""
Path(settings.DATABASE_URL).parent.mkdir(parents=True, exist_ok=True)
Base.metadata.create_all(bind=engine)