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
loguru
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.schemas.schemas import ChatRequest
from src.pipeline.utils import logger
from src.pipeline.models import User as TUser
import json
router = APIRouter()
@router.post("/stream-chat")
async def stream_chat(body: ChatRequest):
logger.debug(await TUser.all())
service = RagService()
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
load_dotenv()
from pydantic import MySQLDsn
class Settings(BaseSettings):
@ -26,7 +21,7 @@ class Settings(BaseSettings):
es_port: int
es_user: str
es_password: str
mysql_dsn: MySQLDsn
class Config:
env_file = ".env"

View File

@ -186,5 +186,6 @@ async def init_client():
global client
client = AsyncES()
async def close_client():
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 src.pipeline.api import include_router
from src.pipeline.config import config
from contextlib import asynccontextmanager
from src.pipeline.core import llm, es
from src.pipeline.db import init_tortoise
from contextlib import asynccontextmanager
@asynccontextmanager
async def lifespan(app: FastAPI):
await llm.init_client()
await es.init_client()
await init_tortoise(app, db_url=str(config.mysql_dsn))
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