database.py 774 B

1234567891011121314151617181920212223242526272829303132333435
  1. """
  2. 数据库连接和会话管理
  3. """
  4. from sqlalchemy import create_engine
  5. from sqlalchemy.ext.declarative import declarative_base
  6. from sqlalchemy.orm import sessionmaker
  7. import os
  8. from pathlib import Path
  9. from backend.config import DATA_DIR
  10. # 确保数据目录存在
  11. os.makedirs(DATA_DIR, exist_ok=True)
  12. # 数据库URL
  13. SQLALCHEMY_DATABASE_URL = f"sqlite:///{DATA_DIR}/lightstar.db"
  14. # 创建引擎
  15. engine = create_engine(
  16. SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
  17. )
  18. # 创建会话
  19. SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
  20. # 基础模型类
  21. Base = declarative_base()
  22. # 依赖函数,获取数据库会话
  23. def get_db():
  24. db = SessionLocal()
  25. try:
  26. yield db
  27. finally:
  28. db.close()