diff --git a/util/oss.py b/util/oss.py index 05b3866..cc0e293 100644 --- a/util/oss.py +++ b/util/oss.py @@ -10,6 +10,24 @@ from telemetry import get_tracer logger = logging.getLogger(__name__) +def _apply_http_replace_map(url): + """ + 应用 HTTP_REPLACE_MAP 环境变量替换 URL + :param str url: 原始 URL + :return str: 替换后的 URL + """ + replace_map = os.getenv("HTTP_REPLACE_MAP", "") + if not replace_map: + return url + replace_list = [i.split("|", 1) for i in replace_map.split(",")] + new_url = url + 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_to_oss(url, file_path): """ 使用签名URL上传文件到OSS @@ -47,14 +65,17 @@ def upload_to_oss(url, file_path): return True else: span.set_status(Status(StatusCode.ERROR)) + # 应用 HTTP_REPLACE_MAP 替换 URL + http_url = _apply_http_replace_map(url) + span.set_attribute("file.http_url", http_url) while retries < max_retries: with tracer.start_as_current_span("upload_to_oss.request") as req_span: req_span.set_attribute("http.retry_count", retries) try: req_span.set_attribute("http.method", "PUT") - req_span.set_attribute("http.url", url) + req_span.set_attribute("http.url", http_url) with open(file_path, 'rb') as f: - response = requests.put(url, data=f, stream=True, timeout=60, headers={"Content-Type": "video/mp4"}) + response = requests.put(http_url, data=f, stream=True, timeout=60, headers={"Content-Type": "video/mp4"}) req_span.set_attribute("http.status_code", response.status_code) req_span.set_attribute("http.response", response.text) response.raise_for_status() @@ -101,6 +122,9 @@ def download_from_oss(url, file_path, skip_if_exist=None): if file_dir: if not os.path.exists(file_dir): os.makedirs(file_dir) + # 应用 HTTP_REPLACE_MAP 替换 URL + http_url = _apply_http_replace_map(url) + span.set_attribute("file.http_url", http_url) max_retries = 5 retries = 0 while retries < max_retries: @@ -108,8 +132,8 @@ def download_from_oss(url, file_path, skip_if_exist=None): req_span.set_attribute("http.retry_count", retries) try: req_span.set_attribute("http.method", "GET") - req_span.set_attribute("http.url", url) - response = requests.get(url, timeout=15) # 设置超时时间 + req_span.set_attribute("http.url", http_url) + response = requests.get(http_url, timeout=15) # 设置超时时间 req_span.set_attribute("http.status_code", response.status_code) with open(file_path, 'wb') as f: f.write(response.content)