Files
findreve/model/database.py
于小丘 cd35c6fbed Refactor models and routes for item management
Reorganized model structure by replacing 'object' and 'items' with a unified 'item' model using UUIDs, and moved base model logic into separate files. Updated routes to use the new item model and improved request/response handling. Enhanced user and setting models, added utility functions, and improved error handling throughout the codebase. Also added initial .idea project files and minor admin API improvements.

Co-Authored-By: 砂糖橘 <54745033+Foxerine@users.noreply.github.com>
2025-10-05 18:58:46 +08:00

75 lines
2.2 KiB
Python

# ~/models/database.py
from contextlib import asynccontextmanager
from typing import AsyncGenerator
import os
from dotenv import load_dotenv
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession
from .migration import migration
# 加载环境变量
load_dotenv('.env')
# 获取 DEBUG 配置
DEBUG = os.getenv("DEBUG", "false").lower() in ("true", "1", "yes")
ASYNC_DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///data.db")
engine: AsyncEngine = create_async_engine(
ASYNC_DATABASE_URL,
echo=DEBUG, # 根据 DEBUG 配置决定是否输出 SQL 日志
connect_args={
"check_same_thread": False
} if ASYNC_DATABASE_URL.startswith("sqlite") else {},
future=True,
# pool_size=POOL_SIZE,
# max_overflow=64,
)
_async_session_factory = sessionmaker(
engine, class_=AsyncSession, expire_on_commit=False
)
# 数据库类
class Database:
# Database 初始化方法
def __init__(
self, # self 用于引用类的实例
db_path: str = "data.db", # db_path 数据库文件路径,默认为 data.db
):
self.db_path = db_path
@staticmethod
async def get_session() -> AsyncGenerator[AsyncSession, None]:
"""FastAPI dependency to get a database session."""
async with _async_session_factory() as session:
yield session
@staticmethod
@asynccontextmanager
async def session_context() -> AsyncGenerator[AsyncSession, None]:
"""
提供异步上下文管理器用于直接获取数据库会话
使用示例:
async with Database.session_context() as session:
# 执行数据库操作
pass
"""
async with _async_session_factory() as session:
yield session
async def init_db(self, url: str = ASYNC_DATABASE_URL):
"""创建数据库结构"""
async with engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)
# For internal use, create a temporary context manager
async with self.session_context() as session:
await migration(session) # 执行迁移脚本