From b40851d4f1d321782c17af6ed8f062113bf0ed31 Mon Sep 17 00:00:00 2001 From: xuyonghao <898441624@qq.com> Date: 星期四, 26 十二月 2024 16:38:26 +0800 Subject: [PATCH] 合并excel上传,合并,下载 修改 --- app/api/excel.py | 128 ++++++++++++++++++------------------------ 1 files changed, 56 insertions(+), 72 deletions(-) diff --git a/app/api/excel.py b/app/api/excel.py index 2d6a35f..7b8e792 100644 --- a/app/api/excel.py +++ b/app/api/excel.py @@ -1,13 +1,9 @@ -from fastapi import APIRouter, File, UploadFile, Depends +from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends from fastapi.responses import JSONResponse, FileResponse -from fastapi.exceptions import HTTPException -from sqlalchemy.orm import Session -from starlette.websockets import WebSocket, WebSocketDisconnect -from werkzeug.utils import secure_filename -from app.api import get_current_user_websocket -from app.models.agent_model import AgentModel, AgentType -from app.models.base_model import get_db -from app.models.user_model import UserModel +from starlette.websockets import WebSocket + +from app.api import get_current_user, get_current_user_websocket +from app.models import UserModel from app.utils.excelmerge.conformity import run_conformity import shutil import os @@ -18,17 +14,17 @@ EXCEL_FILES_PATH = 'data/output' SOURCE_FILES_PATH = 'data/source' -def allowed_file(filename): + +def allowed_file(filename: str) -> bool: return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS -def create_dir_if_not_exists(path): +def create_dir_if_not_exists(path: str): if not os.path.exists(path): os.makedirs(path) -# 娓呯悊鍑芥暟 -def clear_directory(path): +def clear_directory(path: str) -> dict: for filename in os.listdir(path): file_path = os.path.join(path, filename) try: @@ -41,23 +37,30 @@ return {"message": "鐩綍宸叉竻绌�"} +def user_file_path(userid: str, path: str) -> str: + return os.path.join(path, userid) + + @router.post('/excel/upload') -async def upload_file(files: list[UploadFile] = File(...)): +async def upload_file(files: list[UploadFile] = File(...), current_user: UserModel = Depends(get_current_user)): + user_id = str(current_user.id) if not any(file.filename for file in files): return JSONResponse(content={"error": "娌℃湁鏂囦欢閮ㄥ垎"}, status_code=400) + if not user_id: + return JSONResponse(content={"error": "缂哄皯鍙傛暟user_id"}, status_code=400) + user_source = user_file_path(user_id, SOURCE_FILES_PATH) + user_excel = EXCEL_FILES_PATH - create_dir_if_not_exists(SOURCE_FILES_PATH) - create_dir_if_not_exists(EXCEL_FILES_PATH) - clear_directory(SOURCE_FILES_PATH) - clear_directory(EXCEL_FILES_PATH) + create_dir_if_not_exists(user_source) + create_dir_if_not_exists(user_excel) + clear_directory(user_source) save_path_list = [] for file in files: if file.filename == '': return JSONResponse(content={"error": "娌℃湁閫夋嫨鏂囦欢"}, status_code=400) if file and allowed_file(file.filename): - filename = secure_filename(file.filename) - save_path = os.path.join(SOURCE_FILES_PATH, filename) + save_path = os.path.join(user_source, file.filename) with open(save_path, 'wb') as buffer: shutil.copyfileobj(file.file, buffer) save_path_list.append(save_path) @@ -68,63 +71,34 @@ # ws://localhost:9201/api/document/ws/excel @router.websocket("/ws/excel") -async def ws_excel(websocket: WebSocket): +async def ws_excel(websocket: WebSocket, current_user: UserModel = Depends(get_current_user_websocket)): await websocket.accept() + user_id = str(current_user.id) - create_dir_if_not_exists(SOURCE_FILES_PATH) - create_dir_if_not_exists(EXCEL_FILES_PATH) - clear_directory(SOURCE_FILES_PATH) - clear_directory(EXCEL_FILES_PATH) + user_source = user_file_path(user_id, SOURCE_FILES_PATH) + user_excel = EXCEL_FILES_PATH + create_dir_if_not_exists(user_source) + create_dir_if_not_exists(user_excel) while True: data = await websocket.receive_text() try: if data == "\"鍚堝苟Excel\"": - clear_directory(EXCEL_FILES_PATH) - output_file_path = run_conformity() - clear_directory(EXCEL_FILES_PATH) - output_file_path = run_conformity() - files = os.listdir(EXCEL_FILES_PATH) - if files: - first_file = files[0] - file_name = os.path.basename(first_file) - file_url = f"./api/document/download/{first_file}" + merge_file = run_conformity(user_source, user_excel) + if merge_file is not None: + await websocket.send_json({ - "step_message": "鏂囨。鍚堝苟鎴愬姛锛�", + "message": "鏂囨。鍚堝苟鎴愬姛锛�", "type": "stream", - "files": [{ - "file_name": file_name, - "file_url": file_url - }] + "file_name": f"{merge_file}.xlsx", + "download_url": f"./api/document/download/{merge_file}.xlsx" }) await websocket.send_json({ - "step_message": "", + "message": "鏂囨。鍚堝苟鎴愬姛锛�", "type": "close", }) else: - await websocket.send_json({"error": "鍚堝苟鎿嶄綔鏈敓鎴愭枃浠�", "type": "stream", "files": []}) - elif data == "\"鏌ヨ鍚堝苟杩涘害\"": - files = os.listdir(EXCEL_FILES_PATH) - if not files: - await websocket.send_json({"step_message": "姝e湪鍚堝苟涓�", "type": "stream", "files": []}) - else: - await websocket.send_json({"step_message": "鏂囨。鍚堝苟鎴愬姛锛�", "type": "stream", "files": []}) - elif data == "\"鑾峰彇鏂囦欢\"": - files = os.listdir(EXCEL_FILES_PATH) - if not files: - await websocket.send_json({"error": "鐩綍涓嬫病鏈夌敓鎴愮殑鏂囦欢", "type": "stream", "files": []}) - else: - first_file = files[0] - file_name = os.path.basename(first_file) - file_url = f"./api/document/download/{first_file}" - await websocket.send_json({ - "step_message": "鏂囨。鍚堝苟鎴愬姛锛�", - "type": "stream", - "files": [{ - "file_name": file_name, - "file_url": file_url - }] - }) + await websocket.send_json({"error": "鍚堝苟澶辫触", "type": "stream", "files": []}) else: print(f"Received data: {data}") await websocket.send_json({"error": "鏈煡鎸囦护", "data": str(data)}) @@ -133,12 +107,22 @@ await websocket.close() -@router.get("/download/{filename}") -async def download_file(filename: str): - try: - return FileResponse(os.path.join(EXCEL_FILES_PATH, filename), filename=filename, - media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet') - except FileNotFoundError: - raise HTTPException(status_code=404, detail="鏂囦欢涓嶅瓨鍦�") - except Exception as e: - raise HTTPException(status_code=500, detail="鏈嶅姟鍣ㄩ敊璇�") \ No newline at end of file +@router.get("/download/{file_full_name}") +async def download_file(background_tasks: BackgroundTasks, file_full_name: str): + file_name = os.path.basename(file_full_name) + user_excel = EXCEL_FILES_PATH + file_path = os.path.join(user_excel, file_full_name) + + if not os.path.exists(file_path): + return JSONResponse(content={"error": "鏂囦欢涓嶅瓨鍦�"}, status_code=404) + + def delete_file(): + try: + os.unlink(file_path) + except OSError as e: + print(f"Deleting file error") + + # 寰呬笅杞藉畬鎴愬悗鍒犻櫎鐢熸垚鐨勬枃浠� + background_tasks.add_task(delete_file) + return FileResponse(path=file_path, filename=file_name, + media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") -- Gitblit v1.8.0