fix: enhance load method with overloads for better type hinting and s…

…implify data loading in process_trace_tasks

Signed-off-by: -LAN- <[email protected]>

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:

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)