From 2ea248c02e5ad0fb2ed43a3f40a66ffa89a8f431 Mon Sep 17 00:00:00 2001
From: Jerry Yan <792602257@qq.com>
Date: Sun, 16 Feb 2025 18:15:35 +0800
Subject: [PATCH] =?UTF-8?q?=E4=B8=8A=E4=BC=A0=E6=96=87=E4=BB=B6=E4=B9=9F?=
 =?UTF-8?q?=E5=BC=84=E4=B8=AA=E8=B6=85=E6=97=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 util/oss.py | 42 +++++++++++++++++++++++++++++-------------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/util/oss.py b/util/oss.py
index 140dba1..e41f6d1 100644
--- a/util/oss.py
+++ b/util/oss.py
@@ -13,13 +13,22 @@ def upload_to_oss(url, file_path):
     :param str file_path: 文件路径
     :return bool: 是否成功
     """
-    with open(file_path, 'rb') as f:
+    max_retries = 5
+    retries = 0
+    while retries < max_retries:
         try:
-            response = requests.put(url, data=f)
-            return response.status_code == 200
+            with open(file_path, 'rb') as f:
+                response = requests.put(url, data=f, timeout=60)  # 设置超时时间为1分钟
+                if response.status_code == 200:
+                    return True
+        except requests.exceptions.Timeout:
+            retries += 1
+            logger.warning(f"Upload timed out. Retrying {retries}/{max_retries}...")
         except Exception as e:
-            print(e)
-            return False
+            logger.warning(f"Upload failed. Retrying {retries}/{max_retries}...")
+            retries += 1
+    return False
+
 
 def download_from_oss(url, file_path):
     """
@@ -33,11 +42,18 @@ def download_from_oss(url, file_path):
     if file_dir:
         if not os.path.exists(file_dir):
             os.makedirs(file_dir)
-    try:
-        response = requests.get(url)
-        with open(file_path, 'wb') as f:
-            f.write(response.content)
-        return True
-    except Exception as e:
-        print(e)
-        return False
+    max_retries = 5
+    retries = 0
+    while retries < max_retries:
+        try:
+            response = requests.get(url, timeout=15)  # 设置超时时间
+            with open(file_path, 'wb') as f:
+                f.write(response.content)
+            return True
+        except requests.exceptions.Timeout:
+            retries += 1
+            logger.warning(f"Download timed out. Retrying {retries}/{max_retries}...")
+        except Exception as e:
+            logger.warning(f"Download failed. Retrying {retries}/{max_retries}...")
+            retries += 1
+    return False