fix task v2 download issue (#3220)

This commit is contained in:
LawyZheng
2025-08-18 14:24:18 +08:00
committed by GitHub
parent 35313508d0
commit 4617d2b5ca
10 changed files with 79 additions and 52 deletions

View File

@@ -124,15 +124,11 @@ class BaseStorage(ABC):
pass
@abstractmethod
async def save_downloaded_files(
self, organization_id: str, task_id: str | None, workflow_run_id: str | None
) -> None:
async def save_downloaded_files(self, organization_id: str, run_id: str | None) -> None:
pass
@abstractmethod
async def get_downloaded_files(
self, organization_id: str, task_id: str | None, workflow_run_id: str | None
) -> list[FileInfo]:
async def get_downloaded_files(self, organization_id: str, run_id: str | None) -> list[FileInfo]:
pass
@abstractmethod

View File

@@ -171,15 +171,11 @@ class LocalStorage(BaseStorage):
return None
return str(stored_folder_path)
async def save_downloaded_files(
self, organization_id: str, task_id: str | None, workflow_run_id: str | None
) -> None:
async def save_downloaded_files(self, organization_id: str, run_id: str | None) -> None:
pass
async def get_downloaded_files(
self, organization_id: str, task_id: str | None, workflow_run_id: str | None
) -> list[FileInfo]:
download_dir = get_download_dir(workflow_run_id=workflow_run_id, task_id=task_id)
async def get_downloaded_files(self, organization_id: str, run_id: str | None) -> list[FileInfo]:
download_dir = get_download_dir(run_id=run_id)
file_infos: list[FileInfo] = []
files_and_folders = os.listdir(download_dir)
for file_or_folder in files_and_folders:

View File

@@ -195,14 +195,14 @@ class S3Storage(BaseStorage):
temp_zip_file.close()
return temp_dir
async def save_downloaded_files(
self, organization_id: str, task_id: str | None, workflow_run_id: str | None
) -> None:
download_dir = get_download_dir(workflow_run_id=workflow_run_id, task_id=task_id)
async def save_downloaded_files(self, organization_id: str, run_id: str | None) -> None:
download_dir = get_download_dir(run_id=run_id)
files = os.listdir(download_dir)
sc = await self._get_storage_class_for_org(organization_id)
tags = await self._get_tags_for_org(organization_id)
base_uri = f"s3://{settings.AWS_S3_BUCKET_UPLOADS}/{DOWNLOAD_FILE_PREFIX}/{settings.ENV}/{organization_id}/{workflow_run_id or task_id}"
base_uri = (
f"s3://{settings.AWS_S3_BUCKET_UPLOADS}/{DOWNLOAD_FILE_PREFIX}/{settings.ENV}/{organization_id}/{run_id}"
)
for file in files:
fpath = os.path.join(download_dir, file)
if not os.path.isfile(fpath):
@@ -225,10 +225,8 @@ class S3Storage(BaseStorage):
tags=tags,
)
async def get_downloaded_files(
self, organization_id: str, task_id: str | None, workflow_run_id: str | None
) -> list[FileInfo]:
uri = f"s3://{settings.AWS_S3_BUCKET_UPLOADS}/{DOWNLOAD_FILE_PREFIX}/{settings.ENV}/{organization_id}/{workflow_run_id or task_id}"
async def get_downloaded_files(self, organization_id: str, run_id: str | None) -> list[FileInfo]:
uri = f"s3://{settings.AWS_S3_BUCKET_UPLOADS}/{DOWNLOAD_FILE_PREFIX}/{settings.ENV}/{organization_id}/{run_id}"
object_keys = await self.async_client.list_files(uri=uri)
if len(object_keys) == 0:
return []