feat: 初始化数据库

This commit is contained in:
李如威 2026-01-08 15:26:36 +08:00
parent c3c1d3fbcf
commit ef68e3cbd8
9 changed files with 69 additions and 9 deletions

View File

@ -10,4 +10,6 @@ aiofiles
pillow pillow
loguru loguru
httpx httpx
baidusearch baidusearch
celery[redis]
tortoise-orm[asyncmy]

View File

@ -3,12 +3,16 @@ from fastapi.responses import StreamingResponse
from src.pipeline.services.rag import RagService from src.pipeline.services.rag import RagService
from src.pipeline.schemas.schemas import ChatRequest from src.pipeline.schemas.schemas import ChatRequest
from src.pipeline.utils import logger from src.pipeline.utils import logger
from src.pipeline.models import User as TUser
import json import json
router = APIRouter() router = APIRouter()
@router.post("/stream-chat") @router.post("/stream-chat")
async def stream_chat(body: ChatRequest): async def stream_chat(body: ChatRequest):
logger.debug(await TUser.all())
service = RagService() service = RagService()
async def event_generator(): async def event_generator():

View File

@ -1,10 +1,5 @@
from mimetypes import init
from typing import TypedDict
from dotenv import load_dotenv
import os
from pydantic_settings import BaseSettings from pydantic_settings import BaseSettings
from pydantic import MySQLDsn
load_dotenv()
class Settings(BaseSettings): class Settings(BaseSettings):
@ -26,7 +21,7 @@ class Settings(BaseSettings):
es_port: int es_port: int
es_user: str es_user: str
es_password: str es_password: str
mysql_dsn: MySQLDsn
class Config: class Config:
env_file = ".env" env_file = ".env"

View File

@ -186,5 +186,6 @@ async def init_client():
global client global client
client = AsyncES() client = AsyncES()
async def close_client(): async def close_client():
await client.close() await client.close()

View File

@ -0,0 +1,26 @@
from pathlib import Path
from tortoise import Tortoise, generate_config
from tortoise.contrib.fastapi import RegisterTortoise
from fastapi import FastAPI
from src.pipeline.utils import logger
DEFAULT_MODELS = {"models": ["src.pipeline.models"]}
DEFAULT_INIT_SQL = "src/pipeline/db/init.sql"
async def init_tortoise(app: FastAPI, db_url: str, modules: dict = DEFAULT_MODELS, schema_sql: str = DEFAULT_INIT_SQL):
config = generate_config(db_url, app_modules=modules)
await RegisterTortoise(app, config=config)
if schema_sql:
conn = Tortoise.get_connection("default")
existing = await conn.execute_query_dict("SHOW TABLES")
if not existing:
logger.debug(f"数据库初始化文件: {schema_sql}")
sql_file = Path(schema_sql)
if sql_file.exists():
logger.debug("初始化数据库")
sql = sql_file.read_text()
logger.debug(f'sql:\n{sql}')
await conn.execute_script(sql)

9
src/pipeline/db/init.sql Normal file
View File

@ -0,0 +1,9 @@
-- tb_users table
CREATE TABLE IF NOT EXISTS tb_users (
id INT AUTO_INCREMENT PRIMARY KEY,
username VARCHAR(50) NOT NULL UNIQUE,
password VARCHAR(128) NOT NULL,
user_status TINYINT(1) NOT NULL DEFAULT 1, -- 0=禁用, 1=启用
create_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
update_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
);

View File

@ -1,14 +1,16 @@
from fastapi import FastAPI from fastapi import FastAPI
from src.pipeline.api import include_router from src.pipeline.api import include_router
from src.pipeline.config import config from src.pipeline.config import config
from contextlib import asynccontextmanager
from src.pipeline.core import llm, es from src.pipeline.core import llm, es
from src.pipeline.db import init_tortoise
from contextlib import asynccontextmanager
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
await llm.init_client() await llm.init_client()
await es.init_client() await es.init_client()
await init_tortoise(app, db_url=str(config.mysql_dsn))
yield yield

View File

@ -0,0 +1,21 @@
from email.policy import default
from tortoise import fields
from tortoise.models import Model
class BaseModel(Model):
id = fields.IntField(pk=True)
create_at = fields.DatetimeField(auto_now_add=True)
update_at = fields.DatetimeField(auto_now=True)
class Meta:
abstract = True
class User(BaseModel):
username = fields.CharField(max_length=50, unique=True)
password = fields.CharField(max_length=128)
user_status = fields.IntField(default=1) # 0=禁用, 1=启用
class Meta:
table = "tb_users"

View File