tamar-file-hub-client 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

file_hub_client/utils/upload_helper.py

@@ -9,11 +9,61 @@ import asyncio
 import aiohttp
 import requests
 from pathlib import Path
-from typing import Union, BinaryIO, Optional, Callable, Dict, Any, AsyncGenerator
+from typing import Union, BinaryIO, Optional, Callable, Dict, Any, AsyncGenerator, Tuple
 from dataclasses import dataclass
 import hashlib
 
 
+def detect_storage_type(url: str) -> str:
+    """
+    根据URL检测存储类型
+
+    Args:
+        url: 上传URL
+
+    Returns:
+        存储类型: 'gcs'、'oss' 或 'unknown'
+    """
+    if 'storage.googleapis.com' in url or 'storage.cloud.google.com' in url:
+        return 'gcs'
+    elif 'aliyuncs.com' in url:
+        return 'oss'
+    else:
+        return 'gcs'
+
+
+def get_forbid_overwrite_headers(url: str, forbid_overwrite: bool = False) -> Dict[str, str]:
+    """
+    获取防止覆盖的headers
+
+    Args:
+        url: 上传URL
+        forbid_overwrite: 是否防止覆盖
+
+    Returns:
+        包含防止覆盖header的字典
+
+    Note:
+        - GCS: 需要在HTTP header中添加 x-goog-if-generation-match: 0
+        - OSS: 需要在HTTP header中添加 x-oss-forbid-overwrite: true
+    """
+    if not forbid_overwrite:
+        return {}
+
+    storage_type = detect_storage_type(url)
+
+    if storage_type == 'gcs':
+        return {
+            'x-goog-if-generation-match': '0'
+        }
+    elif storage_type == 'oss':
+        return {
+            'x-oss-forbid-overwrite': 'true'
+        }
+    else:
+        return {}
+
+
 @dataclass
 class UploadProgress:
     """上传进度信息"""
@@ -127,10 +177,11 @@ class HttpUploader:
             progress_callback: Optional[Callable[[UploadProgress], None]] = None,
             total_size: Optional[int] = None,
             is_resume: bool = False,
+            forbid_overwrite: bool = False,
     ) -> requests.Response:
         """
         上传文件到指定URL
-        
+
         Args:
             url: 上传URL
             content: 文件内容
@@ -138,9 +189,15 @@ class HttpUploader:
             headers: 请求头
             progress_callback: 进度回调函数
             is_resume: 是否断点续传
+            forbid_overwrite: 是否防止覆盖（添加相应的header）
         """
         headers = headers or {}
 
+        # 添加防止覆盖的headers
+        if forbid_overwrite:
+            forbid_headers = get_forbid_overwrite_headers(url, forbid_overwrite)
+            headers.update(forbid_headers)
+
         # 获取文件大小（不生成 chunk，避免提前读取）
         final_total_size = self._calculate_total_size(content) if total_size is None else total_size
 
@@ -360,19 +417,26 @@ class AsyncHttpUploader:
             progress_callback: Optional[Callable[[UploadProgress], None]] = None,
             total_size: Optional[int] = None,
             is_resume: bool = False,
+            forbid_overwrite: bool = False,
     ) -> aiohttp.ClientResponse:
         """
         异步上传文件到指定URL
-        
+
         Args:
             url: 上传URL
             content: 文件内容
             headers: 请求头
             progress_callback: 进度回调函数
             is_resume: 是否断点续传
+            forbid_overwrite: 是否防止覆盖（添加相应的header）
         """
         headers = headers or {}
 
+        # 添加防止覆盖的headers
+        if forbid_overwrite:
+            forbid_headers = get_forbid_overwrite_headers(url, forbid_overwrite)
+            headers.update(forbid_headers)
+
         # 获取文件大小（避免读取内容）
         final_total_size = total_size or await self._calculate_total_size(content)