-
Notifications
You must be signed in to change notification settings - Fork 786
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
11811e7
commit ffc4b35
Showing
4 changed files
with
60 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,42 @@ | ||
from datetime import datetime | ||
|
||
from skyvern.config import settings | ||
from skyvern.forge.sdk.api.aws import AsyncAWSClient | ||
from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType | ||
from skyvern.forge.sdk.artifact.storage.base import FILE_EXTENTSION_MAP, BaseStorage | ||
from skyvern.forge.sdk.models import Step | ||
|
||
|
||
class S3Storage(BaseStorage): | ||
def __init__(self, bucket: str | None = None) -> None: | ||
self.async_client = AsyncAWSClient() | ||
self.bucket = bucket or settings.AWS_S3_BUCKET_ARTIFACTS | ||
|
||
def build_uri(self, artifact_id: str, step: Step, artifact_type: ArtifactType) -> str: | ||
file_ext = FILE_EXTENTSION_MAP[artifact_type] | ||
return f"s3://{self.bucket}/{settings.ENV}/{step.task_id}/{step.order:02d}_{step.retry_index}_{step.step_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}" | ||
|
||
async def store_artifact(self, artifact: Artifact, data: bytes) -> None: | ||
await self.async_client.upload_file(artifact.uri, data) | ||
|
||
async def retrieve_artifact(self, artifact: Artifact) -> bytes | None: | ||
return await self.async_client.download_file(artifact.uri) | ||
|
||
async def get_share_link(self, artifact: Artifact) -> str | None: | ||
share_urls = await self.async_client.create_presigned_urls([artifact.uri]) | ||
return share_urls[0] if share_urls else None | ||
|
||
async def get_share_links(self, artifacts: list[Artifact]) -> list[str] | None: | ||
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: | ||
await self.async_client.upload_file_from_path(artifact.uri, path) | ||
|
||
async def save_streaming_file(self, organization_id: str, file_name: str) -> None: | ||
from_path = f"{settings.STREAMING_FILE_BASE_PATH}/{organization_id}/{file_name}" | ||
to_path = f"s3://{settings.AWS_S3_BUCKET_SCREENSHOTS}/{settings.ENV}/{organization_id}/{file_name}" | ||
await self.async_client.upload_file_from_path(to_path, from_path) | ||
|
||
async def get_streaming_file(self, organization_id: str, file_name: str, use_default: bool = True) -> bytes | None: | ||
path = f"s3://{settings.AWS_S3_BUCKET_SCREENSHOTS}/{settings.ENV}/{organization_id}/{file_name}" | ||
return await self.async_client.download_file(path, log_exception=False) |