douyin/api/upload.py

import asyncio
import hashlib
import os
import pathlib
import re
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
import jwt
from fastapi import FastAPI,APIRouter,  File, HTTPException, Depends, Request,Header, UploadFile  
from fastapi.responses import JSONResponse 
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials  
import aiofiles
from starlette.status import HTTP_401_UNAUTHORIZED 
from api.swl_jwt import verify_jwt_token,get_current_user
from config import *
from db.docs import DocumentsRepository
from grpc_m.send_data_to_vector import send_to_grpc_vetcor

upload_router = APIRouter()  
security = HTTPBearer()
@upload_router.post('/upload')
async def upload(open_id=Depends(verify_jwt_token),file: UploadFile = File(...)  ):  
    path = await save_to_user_dir(open_id, file)
    if path:
        res = DocumentsRepository().add_document_with_categories(open_id, path)  
        if res:
            doc_model, category_model, doc_categ_model = res
            send_to_grpc_vetcor(category_model.id, doc_model)
            return {"message": "upload success"}
        else:  
            raise HTTPException(status_code=500, detail="Failed to add document to database")
    else:  
        # 保存文件失败，返回 400 Bad Request 或其他适当的错误状态码  
        raise HTTPException(status_code=400, detail="upload fail")

def is_valid_filename(s):  
    """检查字符串是否只包含合法的文件名字符"""  
    # 这里的正则表达式只允许字母、数字、下划线、点、破折号和空格  
    # 你可以根据需要调整这个正则表达式  
    return re.match(r"^[a-zA-Z0-9_.\- ]+$", s) is not None  
  
def get_user_dir(open_id):  
    if is_valid_filename(open_id):  
        # 如果open_id合法，直接使用它作为目录名  
        user_dir = os.path.join(MNT_DOUYIN_DATA, open_id)  
    else:  
        # 否则，计算它的哈希值并用作目录名  
        hash_object = hashlib.md5(open_id.encode())  
        hex_dig = hash_object.hexdigest()
        user_dir = os.path.join(MNT_DOUYIN_DATA, "hash8_" + hex_dig)  
    # 如果目录不存在，创建它  
    if not os.path.exists(user_dir):  
        os.makedirs(user_dir)  
    return user_dir  

def get_user_docs_dir(open_id):
    user_dir = get_user_dir(open_id)
    user_docs_dir = os.path.join(user_dir,"docs")
    if not os.path.exists(user_docs_dir):  
        os.makedirs(user_docs_dir)
    return user_docs_dir

async def save_to_user_dir(open_id, file:UploadFile):
    file_path = os.path.join(get_user_docs_dir(open_id), file.filename)
    async with aiofiles.open(file_path, "wb") as buffer:  
        chunk = await file.read(8192)  
        while chunk:  
            await buffer.write(chunk)  
            chunk = await file.read(8192)  
        logger.info(f"{open_id} save to {file_path}")
        return file_path