于小丘
c1c36c606f
Introduced pkg modules for environment config, logging, and startup initialization. Replaced direct config and logging setup in main.py with modularized functions. Updated database and migration modules to use environment variables and improved DEBUG handling. Removed tool.py and migrated password utilities to pkg. Cleaned up legacy comments and unused code in models and routes.
74 lines
2.2 KiB
Python
74 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) # 执行迁移脚本 |