vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix s3 _get_storage_class_for_org (#4424)

Browse Source
This commit is contained in:
Shuchang Zheng
2026-01-09 11:12:06 -08:00
committed by GitHub
parent 9410bfe6f3
commit 2e35276431
3 changed files with 11 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
skyvern/forge/sdk/api/aws.py
View File

@@ -28,6 +28,7 @@ class S3StorageClass(StrEnum):
# INTELLIGENT_TIERING = "INTELLIGENT_TIERING"
ONEZONE_IA = "ONEZONE_IA"
GLACIER = "GLACIER"
GLACIER_IR = "GLACIER_IR" # Glacier Instant Retrieval
# DEEP_ARCHIVE = "DEEP_ARCHIVE"
# OUTPOSTS = "OUTPOSTS"
# STANDARD_IA = "STANDARD_IA"

18
skyvern/forge/sdk/artifact/storage/s3.py
View File

@@ -113,7 +113,7 @@ class S3Storage(BaseStorage):
uri = uri.replace(f".{file_ext}", f".{file_ext}.zst")
artifact.uri = uri
sc = await self._get_storage_class_for_org(artifact.organization_id)
sc = await self._get_storage_class_for_org(artifact.organization_id, self.bucket)
tags = await self._get_tags_for_org(artifact.organization_id)
LOG.debug(
"Storing artifact",
@@ -125,7 +125,7 @@ class S3Storage(BaseStorage):
)
await self.async_client.upload_file(uri, data, storage_class=sc, tags=tags)
async def _get_storage_class_for_org(self, organization_id: str) -> S3StorageClass:
async def _get_storage_class_for_org(self, organization_id: str, bucket: str) -> S3StorageClass:
return S3StorageClass.STANDARD
async def _get_tags_for_org(self, organization_id: str) -> dict[str, str]:
@@ -147,7 +147,7 @@ class S3Storage(BaseStorage):
return await self.async_client.create_presigned_urls([artifact.uri for artifact in artifacts])
async def store_artifact_from_path(self, artifact: Artifact, path: str) -> None:
sc = await self._get_storage_class_for_org(artifact.organization_id)
sc = await self._get_storage_class_for_org(artifact.organization_id, self.bucket)
tags = await self._get_tags_for_org(artifact.organization_id)
LOG.debug(
"Storing artifact from path",
@@ -163,7 +163,7 @@ class S3Storage(BaseStorage):
async def save_streaming_file(self, organization_id: str, file_name: str) -> None:
from_path = f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}"
to_path = f"s3://{settings.AWS_S3_BUCKET_SCREENSHOTS}/{settings.ENV}/{organization_id}/{file_name}"
sc = await self._get_storage_class_for_org(organization_id)
sc = await self._get_storage_class_for_org(organization_id, settings.AWS_S3_BUCKET_SCREENSHOTS)
tags = await self._get_tags_for_org(organization_id)
LOG.debug(
"Saving streaming file",
@@ -185,7 +185,7 @@ class S3Storage(BaseStorage):
temp_zip_file = create_named_temporary_file()
zip_file_path = shutil.make_archive(temp_zip_file.name, "zip", directory)
browser_session_uri = f"s3://{settings.AWS_S3_BUCKET_BROWSER_SESSIONS}/{settings.ENV}/{organization_id}/{workflow_permanent_id}.zip"
sc = await self._get_storage_class_for_org(organization_id)
sc = await self._get_storage_class_for_org(organization_id, settings.AWS_S3_BUCKET_BROWSER_SESSIONS)
tags = await self._get_tags_for_org(organization_id)
LOG.debug(
"Storing browser session",
@@ -219,7 +219,7 @@ class S3Storage(BaseStorage):
profile_uri = (
f"s3://{settings.AWS_S3_BUCKET_BROWSER_SESSIONS}/{settings.ENV}/{organization_id}/profiles/{profile_id}.zip"
)
sc = await self._get_storage_class_for_org(organization_id)
sc = await self._get_storage_class_for_org(organization_id, settings.AWS_S3_BUCKET_BROWSER_SESSIONS)
tags = await self._get_tags_for_org(organization_id)
LOG.debug(
"Storing browser profile",
@@ -371,7 +371,7 @@ class S3Storage(BaseStorage):
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)
sc = await self._get_storage_class_for_org(organization_id, settings.AWS_S3_BUCKET_UPLOADS)
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}/{run_id}"
@@ -435,7 +435,7 @@ class S3Storage(BaseStorage):
) -> tuple[str, str] | None:
todays_date = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d")
bucket = settings.AWS_S3_BUCKET_UPLOADS
sc = await self._get_storage_class_for_org(organization_id)
sc = await self._get_storage_class_for_org(organization_id, bucket)
tags = await self._get_tags_for_org(organization_id)
# First try uploading with original filename
try:
@@ -508,7 +508,7 @@ class S3Storage(BaseStorage):
) -> str:
"""Sync a file from local browser session to S3."""
uri = self._build_browser_session_uri(organization_id, browser_session_id, artifact_type, remote_path, date)
sc = await self._get_storage_class_for_org(organization_id)
sc = await self._get_storage_class_for_org(organization_id, self.bucket)
tags = await self._get_tags_for_org(organization_id)
await self.async_client.upload_file_from_path(uri, local_file_path, storage_class=sc, tags=tags)
return uri

2
skyvern/forge/sdk/artifact/storage/test_s3_storage.py
View File

@@ -36,7 +36,7 @@ class S3StorageForTests(S3Storage):
async def _get_tags_for_org(self, organization_id: str) -> dict[str, str]:
return {"dummy": f"org-{organization_id}", "test": "jerry"}
async def _get_storage_class_for_org(self, organization_id: str) -> S3StorageClass:
async def _get_storage_class_for_org(self, organization_id: str, bucket: str) -> S3StorageClass:
return S3StorageClass.ONEZONE_IA