修复数据库构建问题
This commit is contained in:
33
JWT.py
33
JWT.py
@@ -1,12 +1,35 @@
|
||||
from fastapi.security import OAuth2PasswordBearer
|
||||
from model import database
|
||||
import asyncio
|
||||
from model import Setting
|
||||
from model.database import Database
|
||||
|
||||
oauth2_scheme = OAuth2PasswordBearer(
|
||||
scheme_name='获取 JWT Bearer 令牌',
|
||||
description='用于获取 JWT Bearer 令牌,需要以表单的形式提交',
|
||||
description='用于获取 JWT Bearer 令牌,需要以表单的形式提交',
|
||||
tokenUrl="/api/token"
|
||||
)
|
||||
)
|
||||
|
||||
SECRET_KEY = asyncio.run(database.Database().get_setting('SECRET_KEY'))
|
||||
ALGORITHM = "HS256"
|
||||
|
||||
# 延迟加载 SECRET_KEY
|
||||
_SECRET_KEY_CACHE = None
|
||||
|
||||
async def get_secret_key() -> str:
|
||||
"""
|
||||
获取 JWT 密钥
|
||||
|
||||
:return: JWT 密钥字符串
|
||||
"""
|
||||
global _SECRET_KEY_CACHE
|
||||
|
||||
if _SECRET_KEY_CACHE is None:
|
||||
async with Database.get_session() as session:
|
||||
setting = await Setting.get(
|
||||
session=session,
|
||||
condition=(Setting.name == 'SECRET_KEY')
|
||||
)
|
||||
if setting:
|
||||
_SECRET_KEY_CACHE = setting.value
|
||||
else:
|
||||
raise RuntimeError("SECRET_KEY not found in database")
|
||||
|
||||
return _SECRET_KEY_CACHE
|
||||
@@ -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:
|
||||
# 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) # 执行迁移脚本
|
||||
|
||||
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
|
||||
@@ -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
|
||||
BIN
requirements.txt
BIN
requirements.txt
Binary file not shown.
@@ -1,5 +1,5 @@
|
||||
from fastapi import APIRouter
|
||||
from typing import Annotated, Literal, Optional
|
||||
from typing import Annotated, Literal
|
||||
from fastapi import Depends, Query
|
||||
from fastapi import HTTPException
|
||||
import JWT
|
||||
@@ -8,9 +8,14 @@ from jwt import InvalidTokenError
|
||||
from model import database
|
||||
from model.response import DefaultResponse
|
||||
from model.items import Item
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
from model import Setting
|
||||
|
||||
# 验证是否为管理员
|
||||
async def is_admin(token: Annotated[str, Depends(JWT.oauth2_scheme)]) -> Literal[True]:
|
||||
async def is_admin(
|
||||
token: Annotated[str, Depends(JWT.oauth2_scheme)],
|
||||
session: Annotated[AsyncSession, Depends(database.Database.get_session)]
|
||||
) -> Literal[True]:
|
||||
'''
|
||||
验证是否为管理员。
|
||||
|
||||
@@ -24,9 +29,9 @@ async def is_admin(token: Annotated[str, Depends(JWT.oauth2_scheme)]) -> Literal
|
||||
)
|
||||
|
||||
try:
|
||||
payload = jwt.decode(token, JWT.SECRET_KEY, algorithms=[JWT.ALGORITHM])
|
||||
payload = jwt.decode(token, JWT.get_secret_key(), algorithms=[JWT.ALGORITHM])
|
||||
username = payload.get("sub")
|
||||
if username is None or not await database.Database().get_setting('account') == username:
|
||||
if username is None or not await Setting.get(session, Setting.name == 'account') == username:
|
||||
raise credentials_exception
|
||||
else:
|
||||
return True
|
||||
@@ -64,8 +69,8 @@ async def verity_admin() -> DefaultResponse:
|
||||
response_description='物品信息列表'
|
||||
)
|
||||
async def get_items(
|
||||
id: Optional[int] = Query(default=None, ge=1, description='物品ID'),
|
||||
key: Optional[str] = Query(default=None, description='物品序列号')):
|
||||
id: int | None = Query(default=None, ge=1, description='物品ID'),
|
||||
key: str | None = Query(default=None, description='物品序列号')):
|
||||
'''
|
||||
获得物品信息。
|
||||
|
||||
@@ -80,7 +85,6 @@ async def get_items(
|
||||
items = results
|
||||
item = []
|
||||
for i in items:
|
||||
print(i)
|
||||
item.append(Item(
|
||||
id=i[0],
|
||||
type=i[1],
|
||||
@@ -144,14 +148,15 @@ async def add_items(
|
||||
)
|
||||
async def update_items(
|
||||
id: int = Query(ge=1),
|
||||
key: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
icon: Optional[str] = None,
|
||||
status: Optional[str] = None,
|
||||
phone: Optional[int] = None,
|
||||
lost_description: Optional[str] = None,
|
||||
find_ip: Optional[str] = None,
|
||||
lost_time: Optional[str] = None) -> DefaultResponse:
|
||||
key: str | None = None,
|
||||
name: str | None = None,
|
||||
icon: str | None = None,
|
||||
status: str | None = None,
|
||||
phone: int | None = None,
|
||||
lost_description: str | None = None,
|
||||
find_ip: str | None = None,
|
||||
lost_time: str | None = None
|
||||
) -> DefaultResponse:
|
||||
'''
|
||||
更新物品信息。
|
||||
|
||||
|
||||
@@ -5,31 +5,38 @@ from fastapi import Depends, HTTPException
|
||||
from fastapi.security import OAuth2PasswordRequestForm
|
||||
from fastapi import APIRouter
|
||||
import jwt, JWT
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
from tool import verify_password
|
||||
from loguru import logger
|
||||
|
||||
from model.token import Token
|
||||
from model import Setting, database
|
||||
from tool import verify_password
|
||||
|
||||
Router = APIRouter(tags=["令牌 session"])
|
||||
|
||||
# 创建令牌
|
||||
def create_access_token(data: dict, expires_delta: timedelta | None = None):
|
||||
async def create_access_token(data: dict, expires_delta: timedelta | None = None):
|
||||
to_encode = data.copy()
|
||||
if expires_delta:
|
||||
expire = datetime.now(timezone.utc) + expires_delta
|
||||
else:
|
||||
expire = datetime.now(timezone.utc) + timedelta(minutes=15)
|
||||
to_encode.update({"exp": expire})
|
||||
encoded_jwt = jwt.encode(to_encode, JWT.SECRET_KEY, algorithm='HS256')
|
||||
encoded_jwt = jwt.encode(to_encode, key=await JWT.get_secret_key(), algorithm='HS256')
|
||||
return encoded_jwt
|
||||
|
||||
# 验证账号密码
|
||||
async def authenticate_user(username: str, password: str):
|
||||
async def authenticate_user(session: AsyncSession, username: str, password: str):
|
||||
# 验证账号和密码
|
||||
account = await Setting.get('setting', 'account')
|
||||
stored_password = await Setting.get('setting', 'password')
|
||||
account = await Setting.get(session, Setting.name == 'account')
|
||||
stored_password = await Setting.get(session, Setting.name == 'password')
|
||||
|
||||
if account != username or not verify_password(stored_password, password):
|
||||
if not account or not stored_password:
|
||||
logger.error("Account or password not set in settings.")
|
||||
return False
|
||||
|
||||
if account != username or not verify_password(stored_password.value, password):
|
||||
logger.error("Invalid username or password.")
|
||||
return False
|
||||
|
||||
return {'is_authenticated': True}
|
||||
@@ -44,8 +51,13 @@ async def authenticate_user(username: str, password: str):
|
||||
)
|
||||
async def login_for_access_token(
|
||||
form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
|
||||
session: Annotated[AsyncSession, Depends(database.Database.get_session)],
|
||||
) -> Token:
|
||||
user = await authenticate_user(form_data.username, form_data.password)
|
||||
user = await authenticate_user(
|
||||
session=session,
|
||||
username=form_data.username,
|
||||
password=form_data.password
|
||||
)
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
@@ -53,7 +65,7 @@ async def login_for_access_token(
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
access_token_expires = timedelta(hours=1)
|
||||
access_token = create_access_token(
|
||||
access_token = await create_access_token(
|
||||
data={"sub": form_data.username}, expires_delta=access_token_expires
|
||||
)
|
||||
return Token(access_token=access_token, token_type="bearer")
|
||||
Reference in New Issue
Block a user