# -*- coding: utf-8 -*- """ 存储服务 提供文件上传/下载功能,支持 OSS 签名 URL 和 HTTP_REPLACE_MAP 环境变量。 """ import os import logging from typing import Optional import requests logger = logging.getLogger(__name__) def _apply_http_replace_map(url: str) -> str: """ 应用 HTTP_REPLACE_MAP 环境变量替换 URL Args: url: 原始 URL Returns: 替换后的 URL """ replace_map = os.getenv("HTTP_REPLACE_MAP", "") if not replace_map: return url new_url = url replace_list = [i.split("|", 1) for i in replace_map.split(",") if "|" in i] for src, dst in replace_list: new_url = new_url.replace(src, dst) if new_url != url: logger.debug(f"HTTP_REPLACE_MAP: {url} -> {new_url}") return new_url def upload_file(url: str, file_path: str, max_retries: int = 5, timeout: int = 60) -> bool: """ 使用签名 URL 上传文件到 OSS Args: url: 签名 URL file_path: 本地文件路径 max_retries: 最大重试次数 timeout: 超时时间(秒) Returns: 是否成功 """ if not os.path.exists(file_path): logger.error(f"File not found: {file_path}") return False file_size = os.path.getsize(file_path) logger.info(f"Uploading: {file_path} ({file_size} bytes)") # 检查是否使用 rclone 上传 if os.getenv("UPLOAD_METHOD") == "rclone": result = _upload_with_rclone(url, file_path) if result: return True # rclone 失败时回退到 HTTP # 应用 HTTP_REPLACE_MAP 替换 URL http_url = _apply_http_replace_map(url) retries = 0 while retries < max_retries: try: with open(file_path, 'rb') as f: response = requests.put( http_url, data=f, stream=True, timeout=timeout, headers={"Content-Type": "application/octet-stream"} ) response.raise_for_status() logger.info(f"Upload succeeded: {file_path}") return True except requests.exceptions.Timeout: retries += 1 logger.warning(f"Upload timed out. Retrying {retries}/{max_retries}...") except requests.exceptions.RequestException as e: retries += 1 logger.warning(f"Upload failed ({e}). Retrying {retries}/{max_retries}...") logger.error(f"Upload failed after {max_retries} retries: {file_path}") return False def _upload_with_rclone(url: str, file_path: str) -> bool: """ 使用 rclone 上传文件 Args: url: 目标 URL file_path: 本地文件路径 Returns: 是否成功 """ replace_map = os.getenv("RCLONE_REPLACE_MAP", "") if not replace_map: return False config_file = os.getenv("RCLONE_CONFIG_FILE", "") rclone_config = f"--config {config_file}" if config_file else "" # 替换 URL new_url = url replace_list = [i.split("|", 1) for i in replace_map.split(",") if "|" in i] for src, dst in replace_list: new_url = new_url.replace(src, dst) new_url = new_url.split("?", 1)[0] # 移除查询参数 if new_url == url: return False cmd = ( f"rclone copyto --no-check-dest --ignore-existing " f"--multi-thread-chunk-size 8M --multi-thread-streams 8 " f"{rclone_config} {file_path} {new_url}" ) logger.debug(f"rclone command: {cmd}") result = os.system(cmd) if result == 0: logger.info(f"rclone upload succeeded: {file_path}") return True logger.warning(f"rclone upload failed (code={result}): {file_path}") return False def download_file( url: str, file_path: str, max_retries: int = 5, timeout: int = 30, skip_if_exist: bool = False ) -> bool: """ 使用签名 URL 下载文件 Args: url: 签名 URL file_path: 本地文件路径 max_retries: 最大重试次数 timeout: 超时时间(秒) skip_if_exist: 如果文件存在则跳过 Returns: 是否成功 """ # 如果文件已存在且跳过 if skip_if_exist and os.path.exists(file_path): logger.debug(f"File exists, skipping download: {file_path}") return True logger.info(f"Downloading: {url}") # 确保目标目录存在 file_dir = os.path.dirname(file_path) if file_dir: os.makedirs(file_dir, exist_ok=True) # 应用 HTTP_REPLACE_MAP 替换 URL http_url = _apply_http_replace_map(url) retries = 0 while retries < max_retries: try: response = requests.get(http_url, timeout=timeout, stream=True) response.raise_for_status() with open(file_path, 'wb') as f: for chunk in response.iter_content(chunk_size=8192): if chunk: f.write(chunk) file_size = os.path.getsize(file_path) logger.info(f"Download succeeded: {file_path} ({file_size} bytes)") return True except requests.exceptions.Timeout: retries += 1 logger.warning(f"Download timed out. Retrying {retries}/{max_retries}...") except requests.exceptions.RequestException as e: retries += 1 logger.warning(f"Download failed ({e}). Retrying {retries}/{max_retries}...") logger.error(f"Download failed after {max_retries} retries: {url}") return False