add upload image

src/rethink/application.py

@@ -2,10 +2,10 @@
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, FileResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 
-from rethink import const
+from rethink import const, config
 from rethink.logger import logger
 from .models import database
 from .routes import (
@@ -87,3 +87,22 @@ async def index() -> HTMLResponse:
         content=content,
         status_code=200,
     )
+
+
+@app.get("/i/{path}", response_class=FileResponse)
+async def img(
+        path: str,
+) -> FileResponse:
+    if config.is_local_db():
+        prefix = config.get_settings().LOCAL_STORAGE_PATH
+    else:
+        prefix = const.RETHINK_DIR.parent.parent
+    return FileResponse(
+        path=prefix / ".data" / "images" / path,
+        status_code=200,
+    )
+
+    return RedirectResponse(
+        url=f"https://rethink.run/api/files/{path}",
+        status_code=302,
+    )

src/rethink/controllers/files/upload_files.py

@@ -85,3 +85,21 @@ def get_upload_process(
         startAt=datetime2str(start_at),
         running=running,
     )
+
+
+def upload_image(
+        td: TokenDecode,
+        file: UploadFile,
+) -> schemas.files.ImageUploadResponse:
+    if td.code != const.Code.OK:
+        return schemas.files.ImageUploadResponse(
+            code=td.code.value,
+            msg=const.get_msg_by_code(td.code, td.language),
+            data={},
+        )
+    res = models.files.upload_image(uid=td.uid, files=[file])
+    return schemas.files.ImageUploadResponse(
+        code=const.Code.OK.value,
+        msg=const.get_msg_by_code(const.Code.OK, td.language),
+        data=res,
+    )

src/rethink/controllers/schemas/files.py

@@ -1,3 +1,5 @@
+from typing import List, Dict
+
 from pydantic import BaseModel, NonNegativeInt
 
 
@@ -16,3 +18,13 @@ class FileUploadProcessResponse(BaseModel):
     type: str
     startAt: str
     running: bool
+
+
+class ImageUploadResponse(BaseModel):
+    class Data(BaseModel):
+        errFiles: List[str]
+        succMap: Dict[str, str]
+
+    code: NonNegativeInt
+    msg: str
+    data: Data

src/rethink/models/files.py

@@ -1,18 +1,23 @@
 import datetime
+import io
+from pathlib import Path
 from typing import List, Tuple, Optional
 
 import pymongo.errors
+from PIL import Image
 from bson import ObjectId
 from bson.tz_util import utc
 from fastapi import UploadFile
 
 from rethink import const, models
-from rethink.config import is_local_db
+from rethink.config import is_local_db, get_settings
+from rethink.models import utils
 from rethink.models.database import COLL
 from rethink.models.tps import ImportData
 
 MAX_FILE_SIZE = 1024 * 512  # 512 kb
 MAX_FILE_COUNT = 200
+MAX_IMAGE_SIZE = 1024 * 1024 * 10  # 10 mb
 
 
 def update_process(
@@ -39,10 +44,8 @@ def update_process(
 
 
 def upload_obsidian(uid: str, files: List[UploadFile]) -> Tuple[str, const.Code]:
-    first_import = False
     doc = COLL.import_data.find_one({"uid": uid})
     if doc is None:
-        first_import = True
         doc: ImportData = {
             "_id": ObjectId(),
             "uid": uid,
@@ -240,3 +243,47 @@ def get_upload_process(uid: str) -> Tuple[int, str, datetime.datetime, bool]:
             {"$set": {"running": False}}
         )
     return doc["process"], doc["type"], doc["startAt"], running
+
+
+def upload_image(uid: str, files: List[UploadFile]) -> dict:
+    res = {
+        "errFiles": [],
+        "succMap": {},
+    }
+    u, code = models.user.get(uid=uid)
+    if code != const.Code.OK:
+        res["errFiles"] = [file.filename for file in files]
+        return res
+
+    if is_local_db():
+        img_dir = get_settings().LOCAL_STORAGE_PATH / ".data" / "images"
+        img_dir.mkdir(parents=True, exist_ok=True)
+    else:
+        img_dir = const.RETHINK_DIR.parent.parent / ".data" / "images"  # TODO: replace to cos
+        img_dir.mkdir(parents=True, exist_ok=True)
+    for file in files:
+        filename = file.filename
+        if not file.content_type.startswith("image/"):
+            res["errFiles"].append(filename)
+            continue
+        if file.size > MAX_IMAGE_SIZE:
+            res["errFiles"].append(filename)
+            continue
+        fn = Path(filename)
+        ext = fn.suffix
+        hashed = utils.file_hash(file)
+        img_path = img_dir / (hashed + ext)
+        if img_path.exists():
+            # skip the same image
+            res["succMap"][filename] = f"http://127.0.0.1:8000/i/{img_path.name}"
+            continue
+
+        try:
+            image = Image.open(io.BytesIO(file.file.read()))
+            image.save(img_path, quality=50, optimize=True)
+        except Exception:
+            res["errFiles"].append(filename)
+            continue
+
+        res["succMap"][filename] = f"http://127.0.0.1:8000/i/{img_path.name}"
+    return res

src/rethink/models/utils.py

@@ -1,4 +1,5 @@
 import datetime
+import hashlib
 import math
 import re
 import uuid
@@ -7,6 +8,7 @@
 
 import jwt
 import pypinyin
+from fastapi import UploadFile
 from markdown import Markdown
 
 from rethink import config
@@ -151,3 +153,19 @@ def change_link_title(md: str, nid: str, new_title: str) -> str:
         md,
     )
     return new_md
+
+
+md5 = hashlib.md5()
+
+
+# General-purpose solution that can process large files
+def file_hash(file: UploadFile) -> str:
+    # https://stackoverflow.com/questions/22058048/hashing-a-file-in-python
+
+    while True:
+        data = file.file.read(65536)  # arbitrary number to reduce RAM usage
+        if not data:
+            break
+        md5.update(data)
+    file.file.seek(0)
+    return md5.hexdigest()

src/rethink/routes/files.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from fastapi import Depends, APIRouter, UploadFile
+from fastapi import Depends, APIRouter, UploadFile, Request
 from typing_extensions import Annotated
 
 from rethink.controllers import schemas
@@ -59,3 +59,20 @@ async def get_upload_process(
         td=token_decode,
         rid=rid,
     )
+
+
+@router.post(
+    path="/files/image",
+    response_model=schemas.files.ImageUploadResponse,
+)
+@measure_time_spend
+async def upload_image(
+        token_decode: Annotated[TokenDecode, Depends(token2uid)],
+        req: Request,
+) -> schemas.files.ImageUploadResponse:
+    form = await req.form()
+    file = form.get("file[]")
+    return upload_files.upload_image(
+        td=token_decode,
+        file=file,
+    )