update:切片上传

vastsa · vastsa · commit 7fc733d2254e · 2023-02-28T17:22:49.000+08:00
diff --git a/core/storage.py b/core/storage.py
@@ -1,5 +1,7 @@
 import os
 import asyncio
+import time
+import uuid
 from datetime import datetime
 from pathlib import Path
 from typing import BinaryIO
@@ -16,61 +18,6 @@
         import oss2
 
 
-class AliyunFileStorage:
-    def __init__(self):
-        auth = oss2.Auth(settings.KeyId, settings.KeySecret)
-        self.bucket = oss2.Bucket(auth, settings.OSS_ENDPOINT, settings.BUCKET_NAME)
-
-    def upload_file(self, upload_filepath, remote_filepath):
-        self.bucket.put_object_from_file(remote_filepath, upload_filepath)
-
-    async def get_text(self, file: UploadFile, key: str):
-        ext = file.filename.split('.')[-1]
-        now = datetime.now()
-        path = f"FileCodeBox/upload/{now.year}/{now.month}/{now.day}"
-        text = f"{path}/{f'{key}.{ext}'}"
-        return f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/{text}"
-
-    async def get_url(self, info: Codes):
-        text = info.text.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
-        url = self.bucket.sign_url('GET', text, settings.ACCESSTIME, slash_safe=True)
-        return url
-
-    @staticmethod
-    async def get_size(file: UploadFile):
-        f = file.file
-        f.seek(0, os.SEEK_END)
-        size = f.tell()
-        f.seek(0, os.SEEK_SET)
-        return size
-
-    @staticmethod
-    def _save(filepath, file: BinaryIO):
-        with open(filepath, 'wb') as f:
-            chunk_size = 256 * 1024
-            chunk = file.read(chunk_size)
-            while chunk:
-                f.write(chunk)
-                chunk = file.read(chunk_size)
-
-    async def save_file(self, file: UploadFile, remote_filepath: str):
-        now = int(datetime.now().timestamp())
-        upload_filepath = settings.DATA_ROOT + str(now)
-        await asyncio.to_thread(self._save, upload_filepath, file.file)
-        self.upload_file(upload_filepath, remote_filepath)
-        remote_filepath = remote_filepath.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
-        self.upload_file(upload_filepath, remote_filepath)
-        await asyncio.to_thread(os.remove, upload_filepath)
-
-    async def delete_files(self, texts):
-        tasks = [self.delete_file(text) for text in texts]
-        await asyncio.gather(*tasks)
-
-    async def delete_file(self, text: str):
-        text = text.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
-        self.bucket.delete_object(text)
-
-
 class FileSystemStorage:
     def __init__(self):
         self.DATA_ROOT = Path(settings.DATA_ROOT)
@@ -110,6 +57,38 @@ def _save(filepath, file: BinaryIO):
                 f.write(chunk)
                 chunk = file.read(chunk_size)
 
+    @staticmethod
+    def _save_chunk(filepath, file: bytes):
+        with open(filepath, 'wb') as f:
+            f.write(file)
+
+    async def create_upload_file(self):
+        file_key = uuid.uuid4().hex
+        file_path = self.DATA_ROOT / f"temp/{file_key}"
+        if not file_path.exists():
+            file_path.mkdir(parents=True)
+        return file_key
+
+    async def save_chunk_file(self, file_key, file_chunk, chunk_index, chunk_total):
+        file_path = self.DATA_ROOT / f"temp/{file_key}/"
+        await asyncio.to_thread(self._save_chunk, file_path / f"{chunk_total}-{chunk_index}.temp", file_chunk)
+
+    async def merge_chunks(self, file_key, file_name, total_chunks: int):
+        ext = file_name.split('.')[-1]
+        now = datetime.now()
+        path = self.DATA_ROOT / f"upload/{now.year}/{now.month}/{now.day}/"
+        if not path.exists():
+            path.mkdir(parents=True)
+        text = f"{self.STATIC_URL}/{(path / f'{file_key}.{ext}').relative_to(self.DATA_ROOT)}"
+        with open(path / f'{file_key}.{ext}', 'wb') as f:
+            for i in range(1, total_chunks + 1):
+                now_temp = self.DATA_ROOT / f'temp/{file_key}/{total_chunks}-{i}.temp'
+                with open(now_temp, 'rb') as r:
+                    f.write(r.read())
+                    await asyncio.to_thread(os.remove, now_temp)
+        await asyncio.to_thread(os.rmdir, self.DATA_ROOT / f'temp/{file_key}/')
+        return text
+
     async def save_file(self, file: UploadFile, text: str):
         filepath = await self.get_filepath(text)
         await asyncio.to_thread(self._save, filepath, file.file)
@@ -134,6 +113,61 @@ def judge_delete_folder(self, filepath):
                 break
 
 
+class AliyunFileStorage:
+    def __init__(self):
+        auth = oss2.Auth(settings.KeyId, settings.KeySecret)
+        self.bucket = oss2.Bucket(auth, settings.OSS_ENDPOINT, settings.BUCKET_NAME)
+
+    def upload_file(self, upload_filepath, remote_filepath):
+        self.bucket.put_object_from_file(remote_filepath, upload_filepath)
+
+    async def get_text(self, file: UploadFile, key: str):
+        ext = file.filename.split('.')[-1]
+        now = datetime.now()
+        path = f"FileCodeBox/upload/{now.year}/{now.month}/{now.day}"
+        text = f"{path}/{f'{key}.{ext}'}"
+        return f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/{text}"
+
+    async def get_url(self, info: Codes):
+        text = info.text.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
+        url = self.bucket.sign_url('GET', text, settings.ACCESSTIME, slash_safe=True)
+        return url
+
+    @staticmethod
+    async def get_size(file: UploadFile):
+        f = file.file
+        f.seek(0, os.SEEK_END)
+        size = f.tell()
+        f.seek(0, os.SEEK_SET)
+        return size
+
+    @staticmethod
+    def _save(filepath, file: BinaryIO):
+        with open(filepath, 'wb') as f:
+            chunk_size = 256 * 1024
+            chunk = file.read(chunk_size)
+            while chunk:
+                f.write(chunk)
+                chunk = file.read(chunk_size)
+
+    async def save_file(self, file: UploadFile, remote_filepath: str):
+        now = int(datetime.now().timestamp())
+        upload_filepath = settings.DATA_ROOT + str(now)
+        await asyncio.to_thread(self._save, upload_filepath, file.file)
+        self.upload_file(upload_filepath, remote_filepath)
+        remote_filepath = remote_filepath.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
+        self.upload_file(upload_filepath, remote_filepath)
+        await asyncio.to_thread(os.remove, upload_filepath)
+
+    async def delete_files(self, texts):
+        tasks = [self.delete_file(text) for text in texts]
+        await asyncio.gather(*tasks)
+
+    async def delete_file(self, text: str):
+        text = text.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
+        self.bucket.delete_object(text)
+
+
 STORAGE_ENGINE = {
     "filesystem": FileSystemStorage,
     "aliyunsystem": AliyunFileStorage
diff --git a/main.py b/main.py
@@ -9,12 +9,12 @@
 from starlette.staticfiles import StaticFiles
 from core.utils import error_ip_limit, upload_ip_limit, get_code, storage, delete_expire_files, get_token
 from core.depends import admin_required
-from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException, BackgroundTasks
+from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException, BackgroundTasks, Header
 from core.database import init_models, Options, Codes, get_session
 from settings import settings
 
 # 实例化FastAPI
-app = FastAPI(debug=settings.DEBUG, redoc_url=None, docs_url=None, openapi_url=None)
+app = FastAPI(debug=settings.DEBUG, redoc_url=None, )
 
 
 @app.on_event('startup')
@@ -172,6 +172,24 @@ async def get_file(code: str, token: str, ip: str = Depends(error_ip_limit), s:
         return FileResponse(filepath, filename=info.name)
 
 
+@app.post('/file/create/')
+async def create_file():
+    # 生成随机字符串
+    return {'code': 200, 'data': await storage.create_upload_file()}
+
+
+@app.post('/file/upload/{file_key}/')
+async def upload_file(file_key: str, file: bytes = File(...), chunk_index: int = Form(...),
+                      total_chunks: int = Form(...)):
+    await storage.save_chunk_file(file_key, file, chunk_index, total_chunks)
+    return {'code': 200}
+
+
+@app.get('/file/merge/{file_key}')
+async def merge_chunks(file_key: str, file_name: str, total_chunks: int):
+    return {'code': 200, 'data': await storage.merge_chunks(file_key, file_name, total_chunks)}
+
+
 @app.post('/share', dependencies=[Depends(admin_required)], description='分享文件')
 async def share(background_tasks: BackgroundTasks, text: str = Form(default=None),
                 style: str = Form(default='2'), value: int = Form(default=1), file: UploadFile = File(default=None),
diff --git a/templates/index.html b/templates/index.html
@@ -164,6 +164,7 @@
                         drag
                         action="/share"
                         multiple
+                        :http-request="uploadFile"
                         style="margin: 1rem 0;"
                         :data="uploadData"
                         :before-upload="beforeUpload"
@@ -183,7 +184,6 @@
                         placeholder="请输入内容，使用按钮存入"
                         v-model="uploadData.text">
                 </el-input>
-
                 <div class="el-upload__tip">
                     <el-button round @click="pageNum=0">
                         <div class="el-icon-back"></div>
@@ -253,7 +253,7 @@ <h1 @click="copyText(file.code,0)" style="margin: 0;display: inline;cursor: poin
             </el-row>
         </el-card>
     </el-drawer>
-    <div style="text-align: center; margin-top: 1rem;color: #606266">
+    <div style="text-align: center; margin-top: 1rem;color: #606266;visibility: hidden">
         <a style="text-decoration: none;color: #606266" target="_blank" href="https://github.com/vastsa/FileCodeBox">FileCodeBox
             V1.6 Beta</a>
     </div>
@@ -346,18 +346,50 @@ <h1 @click="copyText(file.code,0)" style="margin: 0;display: inline;cursor: poin
                         data: data,
                         ...config
                     }).then(res => {
-                        this.$message({
-                            message: res.data.detail,
-                            type: 'success'
-                        });
+                        if (res.data.detail) {
+                            this.$message({
+                                message: res.data.detail,
+                                type: 'success'
+                            });
+                        }
                         resolve(res.data)
                     }).catch(err => {
-                        this.$message({
-                            message: err.response.data.detail,
-                            type: 'error'
-                        });
+                        if (this.data.detail) {
+                            this.$message({
+                                message: err.response.data.detail,
+                                type: 'error'
+                            });
+                        }
+                    })
+                })
+            },
+            uploadChunk: async function (chunk, file_key, chunk_index, total_chunks) {
+                const formData = new FormData();
+                formData.append('file', chunk);
+                formData.append('file_key', file_key);
+                formData.append('chunk_index', chunk_index);
+                formData.append('total_chunks', total_chunks);
+                await this.http('post', `/file/upload/${file_key}/`, formData)
+            },
+            uploadFile: async function (e) {
+                this.http('post', '/file/create/').then(async res => {
+                    const file = e.file;
+                    let chunk_index = 0;
+                    const shardSize = 1024 * 1024 * 5;
+                    const {name, size} = file;
+                    const total_chunks = Math.ceil(size / shardSize);
+                    while (chunk_index < total_chunks) {
+                        const start = chunk_index * shardSize
+                        const end = Math.min(start + shardSize, size)
+                        chunk_index += 1;
+                        await this.uploadChunk(file.slice(start, end), res.data, chunk_index, total_chunks);
+                    }
+                    this.http('get', `/file/merge/${res.data}/?file_name=${file.name}&total_chunks=${total_chunks}`).then(res => {
+                        console.log(res)
                     })
+                    console.log(chunk_index, name, size, total_chunks);
                 })
+
             },
             copyText: function (value, style = 1) {
                 if (style === 0) {
@@ -439,11 +471,13 @@ <h1 @click="copyText(file.code,0)" style="margin: 0;display: inline;cursor: poin
                 }
             },
             successUpload(response, file, fileList) {
-                this.$message({
-                    message: response.detail,
-                    type: 'success'
-                });
-                this.jiFiles.unshift(response.data);
+                if (response) {
+                    this.$message({
+                        message: response.detail,
+                        type: 'success'
+                    });
+                    this.jiFiles.unshift(response.data);
+                }
             },
             errorUpload(error, file, fileList) {
                 error = JSON.parse(error.message)
