feat(api): enhance file preview handling (#9674)

api/controllers/files/image_preview.py

@@ -1,5 +1,5 @@
 from flask import Response, request
-from flask_restful import Resource
+from flask_restful import Resource, reqparse
 from werkzeug.exceptions import NotFound
 
 import services
@@ -41,24 +41,39 @@ class FilePreviewApi(Resource):
     def get(self, file_id):
         file_id = str(file_id)
 
-        timestamp = request.args.get("timestamp")
+        parser = reqparse.RequestParser()
-        nonce = request.args.get("nonce")
+        parser.add_argument("timestamp", type=str, required=True, location="args")
-        sign = request.args.get("sign")
+        parser.add_argument("nonce", type=str, required=True, location="args")
+        parser.add_argument("sign", type=str, required=True, location="args")
+        parser.add_argument("as_attachment", type=bool, required=False, default=False, location="args")
 
-        if not timestamp or not nonce or not sign:
+        args = parser.parse_args()
+
+        if not args["timestamp"] or not args["nonce"] or not args["sign"]:
             return {"content": "Invalid request."}, 400
 
         try:
-            generator, mimetype = FileService.get_signed_file_preview(
+            generator, upload_file = FileService.get_file_generator_by_file_id(
                 file_id=file_id,
-                timestamp=timestamp,
+                timestamp=args["timestamp"],
-                nonce=nonce,
+                nonce=args["nonce"],
-                sign=sign,
+                sign=args["sign"],
             )
         except services.errors.file.UnsupportedFileTypeError:
             raise UnsupportedFileTypeError()
 
-        return Response(generator, mimetype=mimetype)
+        response = Response(
+            generator,
+            mimetype=upload_file.mime_type,
+            direct_passthrough=True,
+            headers={},
+        )
+        if upload_file.size > 0:
+            response.headers["Content-Length"] = str(upload_file.size)
+        if args["as_attachment"]:
+            response.headers["Content-Disposition"] = f"attachment; filename={upload_file.name}"
+
+        return response
 
 
 class WorkspaceWebappLogoApi(Resource):

api/services/file_service.py

@@ -1,7 +1,6 @@
 import datetime
 import hashlib
 import uuid
-from collections.abc import Generator
 from typing import Literal, Union
 
 from flask_login import current_user
@@ -132,7 +131,7 @@ class FileService:
         return upload_file
 
     @staticmethod
-    def get_file_preview(file_id: str) -> str:
+    def get_file_preview(file_id: str):
         upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()
 
         if not upload_file:
@@ -171,7 +170,7 @@ class FileService:
         return generator, upload_file.mime_type
 
     @staticmethod
-    def get_signed_file_preview(file_id: str, timestamp: str, nonce: str, sign: str):
+    def get_file_generator_by_file_id(file_id: str, timestamp: str, nonce: str, sign: str):
         result = file_helpers.verify_file_signature(upload_file_id=file_id, timestamp=timestamp, nonce=nonce, sign=sign)
         if not result:
             raise NotFound("File not found or signature is invalid")
@@ -183,10 +182,10 @@ class FileService:
 
         generator = storage.load(upload_file.key, stream=True)
 
-        return generator, upload_file.mime_type
+        return generator, upload_file
 
     @staticmethod
-    def get_public_image_preview(file_id: str) -> tuple[Generator, str]:
+    def get_public_image_preview(file_id: str):
         upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()
 
         if not upload_file: