diff --git a/api/extensions/ext_storage.py b/api/extensions/ext_storage.py index 42422263c4dd03..588bdb2d2717e0 100644 --- a/api/extensions/ext_storage.py +++ b/api/extensions/ext_storage.py @@ -1,6 +1,6 @@ import logging from collections.abc import Callable, Generator -from typing import Union +from typing import Literal, Union, overload from flask import Flask @@ -79,6 +79,12 @@ def save(self, filename, data): logger.exception(f"Failed to save file {filename}") raise e + @overload + def load(self, filename: str, /, *, stream: Literal[False] = False) -> bytes: ... + + @overload + def load(self, filename: str, /, *, stream: Literal[True]) -> Generator: ... + def load(self, filename: str, /, *, stream: bool = False) -> Union[bytes, Generator]: try: if stream: diff --git a/api/tasks/ops_trace_task.py b/api/tasks/ops_trace_task.py index fa6b7afeab88ac..bb3b9e17ead6d2 100644 --- a/api/tasks/ops_trace_task.py +++ b/api/tasks/ops_trace_task.py @@ -1,6 +1,5 @@ import json import logging -from collections.abc import Generator from celery import shared_task # type: ignore from flask import current_app @@ -27,10 +26,7 @@ def process_trace_tasks(file_info): app_id = file_info.get("app_id") file_id = file_info.get("file_id") file_path = f"{OPS_FILE_PATH}{app_id}/{file_id}.json" - loaded_data = storage.load(file_path) - if isinstance(loaded_data, Generator): - loaded_data = b"".join(loaded_data) - file_data = json.loads(loaded_data) + file_data = json.loads(storage.load(file_path)) trace_info = file_data.get("trace_info") trace_info_type = file_data.get("trace_info_type") trace_instance = OpsTraceManager.get_ops_trace_instance(app_id)