修复数据库构建问题
This commit is contained in:
@@ -1,15 +1,12 @@
|
||||
# ~/models/database.py
|
||||
from contextlib import asynccontextmanager
|
||||
import aiosqlite
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from sqlmodel import SQLModel
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from typing import AsyncGenerator
|
||||
|
||||
import warnings
|
||||
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
|
||||
|
||||
ASYNC_DATABASE_URL = "sqlite+aiosqlite:///data.db"
|
||||
@@ -19,161 +16,38 @@ engine: AsyncEngine = create_async_engine(
|
||||
echo=True,
|
||||
connect_args={
|
||||
"check_same_thread": False
|
||||
} if ASYNC_DATABASE_URL.startswith("sqlite") else None,
|
||||
} 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)
|
||||
_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, # self 用于引用类的实例
|
||||
db_path: str = "data.db", # db_path 数据库文件路径，默认为 data.db
|
||||
):
|
||||
self.db_path = db_path
|
||||
|
||||
|
||||
@staticmethod
|
||||
@asynccontextmanager
|
||||
async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
||||
"""FastAPI dependency to get a database session."""
|
||||
async with _async_session_factory() as session:
|
||||
yield session
|
||||
|
||||
async def init_db(
|
||||
self,
|
||||
url: str = ASYNC_DATABASE_URL
|
||||
):
|
||||
async def init_db(self, url: str = ASYNC_DATABASE_URL):
|
||||
"""创建数据库结构"""
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(SQLModel.metadata.create_all)
|
||||
|
||||
async with self.get_session() as session:
|
||||
await migration(session) # 执行迁移脚本
|
||||
|
||||
async def add_object(self, key: str, name: str, icon: str = None, phone: str = None):
|
||||
"""
|
||||
添加新对象
|
||||
|
||||
:param key: 序列号
|
||||
:param name: 名称
|
||||
:param icon: 图标
|
||||
:param phone: 电话
|
||||
"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute("SELECT 1 FROM fr_objects WHERE key = ?", (key,)) as cursor:
|
||||
if await cursor.fetchone():
|
||||
raise ValueError(f"序列号 {key} 已存在")
|
||||
|
||||
now = datetime.now()
|
||||
now = now.strftime("%Y-%m-%d %H:%M:%S")
|
||||
await db.execute(
|
||||
"INSERT INTO fr_objects (key, name, icon, phone, create_at, status) VALUES (?, ?, ?, ?, ?, 'ok')",
|
||||
(key, name, icon, phone, now)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def update_object(
|
||||
self,
|
||||
id: int,
|
||||
key: str = None,
|
||||
name: str = None,
|
||||
icon: str = None,
|
||||
status: str = None,
|
||||
phone: int = None,
|
||||
lost_description: Optional[str] = None,
|
||||
find_ip: Optional[str] = None,
|
||||
lost_time: Optional[str] = None):
|
||||
"""
|
||||
更新对象信息
|
||||
|
||||
:param id: 对象ID
|
||||
:param key: 序列号
|
||||
:param name: 名称
|
||||
:param icon: 图标
|
||||
:param status: 状态
|
||||
:param phone: 电话
|
||||
:param lost_description: 丢失描述
|
||||
:param find_ip: 发现IP
|
||||
:param lost_time: 丢失时间
|
||||
"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute("SELECT 1 FROM fr_objects WHERE id = ?", (id,)) as cursor:
|
||||
if not await cursor.fetchone():
|
||||
raise ValueError(f"ID {id} 不存在")
|
||||
|
||||
async with db.execute("SELECT 1 FROM fr_objects WHERE key = ? AND id != ?", (key, id)) as cursor:
|
||||
if await cursor.fetchone():
|
||||
raise ValueError(f"序列号 {key} 已存在")
|
||||
|
||||
await db.execute(
|
||||
f"UPDATE fr_objects SET "
|
||||
f"key = COALESCE(?, key), "
|
||||
f"name = COALESCE(?, name), "
|
||||
f"icon = COALESCE(?, icon), "
|
||||
f"status = COALESCE(?, status), "
|
||||
f"phone = COALESCE(?, phone), "
|
||||
f"context = COALESCE(?, context), "
|
||||
f"find_ip = COALESCE(?, find_ip), "
|
||||
f"lost_at = COALESCE(?, lost_at) "
|
||||
f"WHERE id = ?",
|
||||
(key, name, icon, status, phone, lost_description, find_ip, lost_time, id)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def get_object(self, id: int = None, key: str = None):
|
||||
"""
|
||||
获取对象
|
||||
|
||||
:param id: 对象ID
|
||||
:param key: 序列号
|
||||
"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
if id is not None or key is not None:
|
||||
async with db.execute(
|
||||
"SELECT * FROM fr_objects WHERE id = ? OR key = ?", (id, key)
|
||||
) as cursor:
|
||||
return await cursor.fetchone()
|
||||
else:
|
||||
async with db.execute("SELECT * FROM fr_objects") as cursor:
|
||||
return await cursor.fetchall()
|
||||
|
||||
async def delete_object(self, id: int):
|
||||
"""
|
||||
删除对象
|
||||
|
||||
:param id: 对象ID
|
||||
"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("DELETE FROM fr_objects WHERE id = ?", (id,))
|
||||
await db.commit()
|
||||
|
||||
async def set_setting(self, name: str, value: str):
|
||||
"""
|
||||
设置配置项
|
||||
|
||||
:param name: 配置项名称
|
||||
:param value: 配置项值
|
||||
"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"INSERT OR REPLACE INTO fr_settings (name, value) VALUES (?, ?)",
|
||||
(name, value)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def get_setting(self, name: str):
|
||||
"""
|
||||
获取配置项
|
||||
|
||||
:param name: 配置项名称
|
||||
"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT value FROM fr_settings WHERE name = ?", (name,)
|
||||
) as cursor:
|
||||
result = await cursor.fetchone()
|
||||
return result[0] if result else None
|
||||
|
||||
# For internal use, create a temporary context manager
|
||||
get_session_cm = asynccontextmanager(self.get_session)
|
||||
async with get_session_cm() as session:
|
||||
await migration(session) # 执行迁移脚本
|
||||
@@ -1,4 +1,4 @@
|
||||
from typing import Sequence
|
||||
from loguru import logger
|
||||
from sqlmodel import select
|
||||
from .setting import Setting
|
||||
import tool
|
||||
@@ -13,9 +13,13 @@ async def migration(session):
|
||||
# 先准备基础配置
|
||||
settings: list[Setting] = default_settings.copy()
|
||||
|
||||
if await Setting.get(session, Setting.name == 'version'):
|
||||
# 已有数据，说明不是第一次运行，直接返回
|
||||
return
|
||||
|
||||
# 生成初始密码与密钥
|
||||
admin_password = tool.generate_password()
|
||||
print(f"密码（请牢记，后续不再显示）: {admin_password}")
|
||||
logger.warning(f"密码（请牢记，后续不再显示）: {admin_password}")
|
||||
|
||||
settings.append(Setting(type='string', name='password', value=tool.hash_password(admin_password)))
|
||||
settings.append(Setting(type='string', name='SECRET_KEY', value=tool.generate_password(64)))
|
||||
|
||||
@@ -14,4 +14,7 @@ class ObjectData(BaseModel):
|
||||
icon: str
|
||||
status: Literal['ok', 'lost']
|
||||
phone: str
|
||||
context: Optional[str] = None
|
||||
context: str | None = None
|
||||
lost_description: str | None = None
|
||||
create_time: str
|
||||
lost_time: str | None = None
|
||||
Reference in New Issue
Block a user