后台更新 服务优化

backend/ali_oss.py

@@ -7,7 +7,9 @@
 import base64
 import io
 import mimetypes
+import os
 import re
+import tempfile
 import threading
 import time
 
@@ -27,6 +29,8 @@ from config import (
 
 _client = None
 _client_lock = threading.Lock()
+_REMOTE_IMAGE_MAX_BYTES = int(os.getenv("OSS_UPLOAD_IMAGE_MAX_BYTES", str(10 * 1024 * 1024)))
+_REMOTE_IMAGE_TIMEOUT = (5, 30)
 
 
 def get_client():
@@ -65,19 +69,72 @@ def get_presigned_url(key: str, expires: int = 7 * 24 * 3600) -> str:
 
 
 
-def upload_file(file_content: bytes, key: str):
-    """上传字节内容到 MinIO，返回可访问链接"""
+class _LimitedReader:
+    """Keep a streamed upload from accepting an unexpectedly huge object."""
+
+    def __init__(self, source, max_bytes: int):
+        self._source = source
+        self._max_bytes = max_bytes
+        self._read_bytes = 0
+
+    def read(self, size=-1):
+        remaining = self._max_bytes - self._read_bytes
+        if remaining < 0:
+            raise ValueError("upload exceeds configured size limit")
+        read_size = remaining + 1 if size is None or size < 0 else min(size, remaining + 1)
+        data = self._source.read(read_size)
+        if not data:
+            return data
+        self._read_bytes += len(data)
+        if self._read_bytes > self._max_bytes:
+            raise ValueError("upload exceeds configured size limit")
+        return data
+
+    def seek(self, offset, whence=0):
+        position = self._source.seek(offset, whence)
+        self._read_bytes = max(0, position)
+        return position
+
+    def tell(self):
+        return self._source.tell()
+
+
+def upload_fileobj(file_obj, key: str, max_bytes: int = 0, content_type: str = ""):
+    """Stream a file-like object to S3/MinIO without materializing it."""
     key = key.lstrip("/")
+    if isinstance(file_obj, (bytes, bytearray)):
+        file_obj = io.BytesIO(file_obj)
+    if not hasattr(file_obj, "read"):
+        raise TypeError("file_obj must be bytes or a readable file-like object")
+    stream_size = None
+    try:
+        file_obj.seek(0)
+        file_obj.seek(0, os.SEEK_END)
+        stream_size = file_obj.tell()
+        file_obj.seek(0)
+    except (AttributeError, OSError, TypeError, ValueError):
+        try:
+            file_obj.seek(0)
+        except (AttributeError, OSError, TypeError, ValueError):
+            pass
+    if max_bytes and max_bytes > 0 and stream_size is not None and stream_size > max_bytes:
+        raise ValueError("upload exceeds configured size limit")
     client = get_client()
-    client.put_object(
-        Bucket=bucket,          # 存储桶名称
-        Key=key,                # 对象名称
-        Body=io.BytesIO(file_content) if isinstance(file_content, (bytes, bytearray)) else file_content,
-        ContentType=_guess_content_type(key),
+    body = _LimitedReader(file_obj, max_bytes) if max_bytes and max_bytes > 0 else file_obj
+    client.upload_fileobj(
+        body,
+        bucket,
+        key,
+        ExtraArgs={"ContentType": content_type or _guess_content_type(key)},
     )
     return file_url_pre + key
 
 
+def upload_file(file_content: bytes, key: str):
+    """上传字节内容到 MinIO，返回可访问链接"""
+    return upload_fileobj(file_content, key)
+
+
 def upload_data_url(data_url: str, prefix: str = "history", key_hint: str = "") -> str:
     """
     将 base64 data URL 上传到对象存储，返回图片链接
@@ -89,14 +146,19 @@ def upload_data_url(data_url: str, prefix: str = "history", key_hint: str = "")
     if not match:
         raise ValueError('无效的 data URL 格式')
     ext = 'png' if match.group(1).lower() in ('png', 'webp') else 'jpg'
-    file_content = base64.b64decode(match.group(2))
+    encoded_payload = match.group(2)
+    if len(encoded_payload) > ((max(_REMOTE_IMAGE_MAX_BYTES, 0) + 2) // 3) * 4:
+        raise ValueError("image exceeds configured upload size limit")
+    file_content = base64.b64decode(encoded_payload, validate=True)
+    if len(file_content) > _REMOTE_IMAGE_MAX_BYTES:
+        raise ValueError("image exceeds configured upload size limit")
     ts = int(time.time() * 1000)
     key = f"{bucket_path}{prefix}/{ts}{key_hint}.{ext}"
-    return upload_file(file_content, key)
+    return upload_fileobj(io.BytesIO(file_content), key, _REMOTE_IMAGE_MAX_BYTES)
 
 
 def upload_data_urls(data_urls: list, prefix: str = "history") -> list:
-    """批量上传 base64 图片到对象存储，返回图片链接列表"""
+    """批量上传图片；每张图���完成上传后立即释放其缓冲区。"""
     urls = []
     ts = int(time.time() * 1000)
     for i, data_url in enumerate(data_urls or []):
@@ -107,12 +169,37 @@ def upload_data_urls(data_urls: list, prefix: str = "history") -> list:
             if not match:
                 continue
             ext = 'png' if match.group(1).lower() in ('png', 'webp') else 'jpg'
-            file_content = base64.b64decode(match.group(2))
-        else:
-            file_content = requests.get(data_url).content
-            ext = "png"
-        key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
-        urls.append(upload_file(file_content, key))
+            encoded_payload = match.group(2)
+            if len(encoded_payload) > ((max(_REMOTE_IMAGE_MAX_BYTES, 0) + 2) // 3) * 4:
+                raise ValueError("image exceeds configured upload size limit")
+            file_content = base64.b64decode(encoded_payload, validate=True)
+            if len(file_content) > _REMOTE_IMAGE_MAX_BYTES:
+                raise ValueError("image exceeds configured upload size limit")
+            key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
+            urls.append(upload_fileobj(io.BytesIO(file_content), key, _REMOTE_IMAGE_MAX_BYTES))
+            del file_content
+            continue
+
+        with requests.get(data_url, stream=True, timeout=_REMOTE_IMAGE_TIMEOUT) as response:
+            response.raise_for_status()
+            content_length = response.headers.get("Content-Length")
+            if content_length and int(content_length) > _REMOTE_IMAGE_MAX_BYTES:
+                raise ValueError("image exceeds configured upload size limit")
+            content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
+            extension = mimetypes.guess_extension(content_type) or ".png"
+            ext = extension.lstrip(".") or "png"
+            key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
+            with tempfile.SpooledTemporaryFile(max_size=2 * 1024 * 1024, mode="w+b") as image_file:
+                total = 0
+                for chunk in response.iter_content(chunk_size=64 * 1024):
+                    if not chunk:
+                        continue
+                    total += len(chunk)
+                    if total > _REMOTE_IMAGE_MAX_BYTES:
+                        raise ValueError("image exceeds configured upload size limit")
+                    image_file.write(chunk)
+                image_file.seek(0)
+                urls.append(upload_fileobj(image_file, key, _REMOTE_IMAGE_MAX_BYTES, content_type))
     return urls
 
 

backend/blueprints/admin_api.py

@@ -34,7 +34,7 @@ from werkzeug.security import generate_password_hash
 from utils.db import get_db
 from utils.auth import admin_required, login_required, get_current_admin_role
 
-from ali_oss import upload_file as oss_upload_file
+from ali_oss import upload_fileobj as oss_upload_fileobj
 
 try:
     from config import bucket_path, backend_java_base_url
@@ -47,6 +47,10 @@ _backend_java_session_local = threading.local()
 _internal_token_lock = threading.Lock()
 IMAGE_VIDEO_DATA_PERMISSION_KEY = 'admin_image_video_task_data'
 SHOP_DATA_CRAWL_DATA_PERMISSION_KEY = 'admin_shop_data_crawl_task_data'
+try:
+    VERSION_UPLOAD_MAX_BYTES = int(os.environ.get('VERSION_UPLOAD_MAX_BYTES', str(512 * 1024 * 1024)))
+except ValueError:
+    VERSION_UPLOAD_MAX_BYTES = 512 * 1024 * 1024
 
 ADMIN_MENU_ACCESS_CONFIG = {
     'dedupe-total-data': {
@@ -2723,13 +2727,28 @@ def upload_version():
         return jsonify({'success': False, 'error': '请选择要上传的 zip 压缩包'})
     if not (file_storage.filename or '').lower().endswith('.zip'):
         return jsonify({'success': False, 'error': '仅支持 .zip 格式'})
+    conn = None
     try:
-        file_content = file_storage.read()
-        if not file_content:
+        if file_storage.content_length and file_storage.content_length > VERSION_UPLOAD_MAX_BYTES:
+            return jsonify({'success': False, 'error': '文件超过允许的大小限制'})
+        file_stream = file_storage.stream
+        stream_size = file_storage.content_length
+        try:
+            file_stream.seek(0)
+            file_stream.seek(0, os.SEEK_END)
+            stream_size = file_stream.tell()
+            file_stream.seek(0)
+        except (AttributeError, OSError, TypeError, ValueError):
+            try:
+                file_stream.seek(0)
+            except (AttributeError, OSError, TypeError, ValueError):
+                pass
+        if stream_size == 0:
             return jsonify({'success': False, 'error': '文件为空'})
         safe_key = _safe_version_key(version)
         key = f"{bucket_path}versions/{safe_key}.zip"
-        file_url = oss_upload_file(file_content, key)
+        file_url = oss_upload_fileobj(file_stream, key, VERSION_UPLOAD_MAX_BYTES,
+                                      'application/zip')
         conn = get_db()
         with conn.cursor() as cur:
             cur.execute(
@@ -2737,7 +2756,6 @@ def upload_version():
                 (version, file_url)
             )
         conn.commit()
-        conn.close()
         return jsonify({
             'success': True,
             'version': version,
@@ -2747,6 +2765,12 @@ def upload_version():
     except Exception as e:
         traceback.print_exc()
         return jsonify({'success': False, 'error': str(e)})
+    finally:
+        if conn is not None:
+            try:
+                conn.close()
+            except Exception:
+                pass
 
 
 # ---------- 店铺密钥管理 ----------
@@ -2972,6 +2996,7 @@ def export_dedupe_total_data():
             url,
             params=params,
             headers={'X-Internal-Token': _resolve_internal_token()},
+            stream=True,
             timeout=60,
         )
     except requests.RequestException:
@@ -2982,14 +3007,23 @@ def export_dedupe_total_data():
             error = data.get('message') or data.get('error') or '导出失败'
         except ValueError:
             error = '导出失败'
+        resp.close()
         return jsonify({'success': False, 'error': error}), resp.status_code
 
     headers = {}
     disposition = resp.headers.get('Content-Disposition')
     if disposition:
         headers['Content-Disposition'] = disposition
+    def generate():
+        try:
+            for chunk in resp.iter_content(chunk_size=1024 * 1024):
+                if chunk:
+                    yield chunk
+        finally:
+            resp.close()
+
     return Response(
-        resp.content,
+        stream_with_context(generate()),
         status=resp.status_code,
         headers=headers,
         content_type=resp.headers.get(

backend/tests/test_admin_dedupe_total_data.py

@@ -18,6 +18,12 @@ class _FakeExportResponse:
         'Content-Type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
     }
 
+    def iter_content(self, chunk_size=None):
+        yield self.content
+
+    def close(self):
+        return None
+
 
 class AdminDedupeTotalDataTest(unittest.TestCase):
     def setUp(self):
@@ -138,12 +144,14 @@ class AdminDedupeTotalDataTest(unittest.TestCase):
                 response = admin_api.export_dedupe_total_data.__wrapped__()
 
         self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.get_data(), b'xlsx')
         self.assertEqual(session.kwargs['params'], {
             'operatorId': 7,
             'username': 'operator',
             'groupId': 3,
         })
         self.assertEqual(session.kwargs['headers'], {'X-Internal-Token': 'token'})
+        self.assertTrue(session.kwargs['stream'])
 
     def test_import_requires_group(self):
         with self.app.test_request_context(