Better download (#628)
* fix: rfc for get_workdir_files was failing * fix: make downloads chunked with progress --------- Co-authored-by: Rafael Uzarowski <uzarowski.rafael@proton.me>
ehl0wr0ld committed
Jul 31, 2025 at 13:51 UTC
65f928c633ef9ad3a5da4603d7b73963d5e72fae
2 files changed
+79
-15
python/api/download_work_dir_file.py
+77
-14
@@ -1,10 +1,77 @@
1
import base64
2
from io import BytesIO
3
+import mimetypes
4
+import os
5
4
-from python.helpers.api import ApiHandler, Input, Output, Request, Response, send_file
6
+from flask import Response
7
+from python.helpers.api import ApiHandler, Input, Output, Request
8
from python.helpers import files, runtime
9
from python.api import file_info
7
-import os
10
+
11
+
12
+def stream_file_download(file_source, download_name, chunk_size=8192):
13
+ """
14
+ Create a streaming response for file downloads that shows progress in browser.
15
+
16
+ Args:
17
+ file_source: Either a file path (str) or BytesIO object
18
+ download_name: Name for the downloaded file
19
+ chunk_size: Size of chunks to stream (default 8192 bytes)
20
+
21
+ Returns:
22
+ Flask Response object with streaming content
23
+ """
24
+ # Calculate file size for Content-Length header
25
+ if isinstance(file_source, str):
26
+ # File path - get size from filesystem
27
+ file_size = os.path.getsize(file_source)
28
+ elif isinstance(file_source, BytesIO):
29
+ # BytesIO object - get size from buffer
30
+ current_pos = file_source.tell()
31
+ file_source.seek(0, 2) # Seek to end
32
+ file_size = file_source.tell()
33
+ file_source.seek(current_pos) # Restore original position
34
+ else:
35
+ raise ValueError(f"Unsupported file source type: {type(file_source)}")
36
+
37
+ def generate():
38
+ if isinstance(file_source, str):
39
+ # File path - open and stream from disk
40
+ with open(file_source, 'rb') as f:
41
+ while True:
42
+ chunk = f.read(chunk_size)
43
+ if not chunk:
44
+ break
45
+ yield chunk
46
+ elif isinstance(file_source, BytesIO):
47
+ # BytesIO object - stream from memory
48
+ file_source.seek(0) # Ensure we're at the beginning
49
+ while True:
50
+ chunk = file_source.read(chunk_size)
51
+ if not chunk:
52
+ break
53
+ yield chunk
54
+
55
+ # Detect content type based on file extension
56
+ content_type, _ = mimetypes.guess_type(download_name)
57
+ if not content_type:
58
+ content_type = 'application/octet-stream'
59
+
60
+ # Create streaming response with proper headers for immediate streaming
61
+ response = Response(
62
+ generate(),
63
+ content_type=content_type,
64
+ direct_passthrough=True, # Prevent Flask from buffering the response
65
+ headers={
66
+ 'Content-Disposition': f'attachment; filename="{download_name}"',
67
+ 'Content-Length': str(file_size), # Critical for browser progress bars
68
+ 'Cache-Control': 'no-cache',
69
+ 'X-Accel-Buffering': 'no', # Disable nginx buffering
70
+ 'Accept-Ranges': 'bytes' # Allow browser to resume downloads
71
+ }
72
+ )
73
+
74
+ return response
75
76
77
class DownloadFile(ApiHandler):
@@ -32,31 +99,27 @@ class DownloadFile(ApiHandler):
99
if runtime.is_development():
100
b64 = await runtime.call_development_function(fetch_file, zip_file)
101
file_data = BytesIO(base64.b64decode(b64))
35
- return send_file(
102
+ return stream_file_download(
103
file_data,
37
- as_attachment=True,
38
- download_name=os.path.basename(zip_file),
104
+ download_name=os.path.basename(zip_file)
105
)
106
else:
41
- return send_file(
107
+ return stream_file_download(
108
zip_file,
43
- as_attachment=True,
44
- download_name=f"{os.path.basename(file_path)}.zip",
109
+ download_name=f"{os.path.basename(file_path)}.zip"
110
)
111
elif file["is_file"]:
112
if runtime.is_development():
113
b64 = await runtime.call_development_function(fetch_file, file["abs_path"])
114
file_data = BytesIO(base64.b64decode(b64))
50
- return send_file(
115
+ return stream_file_download(
116
file_data,
52
- as_attachment=True,
53
- download_name=os.path.basename(file_path),
117
+ download_name=os.path.basename(file_path)
118
)
119
else:
56
- return send_file(
120
+ return stream_file_download(
121
file["abs_path"],
58
- as_attachment=True,
59
- download_name=os.path.basename(file["file_name"]),
122
+ download_name=os.path.basename(file["file_name"])
123
)
124
raise Exception(f"File {file_path} not found")
125
python/api/get_work_dir_files.py
+2
-1
@@ -1,6 +1,7 @@
1
from python.helpers.api import ApiHandler, Request, Response
2
from python.helpers.file_browser import FileBrowser
3
from python.helpers import runtime
4
+import python.api.get_work_dir_files as get_work_dir_files_module
5
6
7
class GetWorkDirFiles(ApiHandler):
@@ -20,7 +21,7 @@ class GetWorkDirFiles(ApiHandler):
21
22
# browser = FileBrowser()
23
# result = browser.get_files(current_path)
23
- result = await runtime.call_development_function(get_files, current_path)
24
+ result = await runtime.call_development_function(get_work_dir_files_module.get_files, current_path)
25
26
return {"data": result}
27