Move save legacy file logic to S3Storage + add storage class support + add logging (#2624)

skyvern/forge/sdk/artifact/storage/base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import BinaryIO
 
 from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityType
 from skyvern.forge.sdk.models import Step
@@ -116,3 +117,9 @@ class BaseStorage(ABC):
         self, organization_id: str, task_id: str | None, workflow_run_id: str | None
     ) -> list[FileInfo]:
         pass
+
+    @abstractmethod
+    async def save_legacy_file(
+        self, *, organization_id: str, filename: str, fileObj: BinaryIO
+    ) -> tuple[str, str] | None:
+        pass

skyvern/forge/sdk/artifact/storage/local.py

@@ -2,6 +2,7 @@ import os
 import shutil
 from datetime import datetime
 from pathlib import Path
+from typing import BinaryIO
 
 import structlog
 
@@ -178,3 +179,10 @@ class LocalStorage(BaseStorage):
     def _create_directories_if_not_exists(path_including_file_name: Path) -> None:
         path = path_including_file_name.parent
         path.mkdir(parents=True, exist_ok=True)
+
+    async def save_legacy_file(
+        self, *, organization_id: str, filename: str, fileObj: BinaryIO
+    ) -> tuple[str, str] | None:
+        raise NotImplementedError(
+            "Legacy file storage is not implemented for LocalStorage. Please use a different storage backend."
+        )

skyvern/forge/sdk/artifact/storage/s3.py

@@ -1,6 +1,8 @@
 import os
 import shutil
-from datetime import datetime
+import uuid
+from datetime import datetime, timezone
+from typing import BinaryIO
 
 import structlog
 
@@ -210,3 +212,55 @@ class S3Storage(BaseStorage):
             file_infos.append(file_info)
 
         return file_infos
+
+    async def save_legacy_file(
+        self, *, organization_id: str, filename: str, fileObj: BinaryIO
+    ) -> 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)
+        # First try uploading with original filename
+        try:
+            sanitized_filename = os.path.basename(filename)  # Remove any path components
+            s3_uri = f"s3://{bucket}/{settings.ENV}/{organization_id}/{todays_date}/{sanitized_filename}"
+            uploaded_s3_uri = await self.async_client.upload_file_stream(s3_uri, fileObj, storage_class=sc)
+        except Exception:
+            LOG.error("Failed to upload file to S3", exc_info=True)
+            uploaded_s3_uri = None
+
+        # If upload fails, try again with UUID prefix
+        if not uploaded_s3_uri:
+            uuid_prefixed_filename = f"{str(uuid.uuid4())}_{filename}"
+            s3_uri = f"s3://{bucket}/{settings.ENV}/{organization_id}/{todays_date}/{uuid_prefixed_filename}"
+            fileObj.seek(0)  # Reset file pointer
+            uploaded_s3_uri = await self.async_client.upload_file_stream(s3_uri, fileObj, storage_class=sc)
+
+        if not uploaded_s3_uri:
+            LOG.error(
+                "Failed to upload file to S3 after retrying with UUID prefix",
+                organization_id=organization_id,
+                storage_class=sc,
+                filename=filename,
+                exc_info=True,
+            )
+            return None
+        LOG.debug(
+            "Legacy file upload",
+            organization_id=organization_id,
+            storage_class=sc,
+            filename=filename,
+            uploaded_s3_uri=uploaded_s3_uri,
+        )
+        # Generate a presigned URL for the uploaded file
+        presigned_urls = await self.async_client.create_presigned_urls([uploaded_s3_uri])
+        if not presigned_urls:
+            LOG.error(
+                "Failed to create presigned URL for uploaded file",
+                organization_id=organization_id,
+                storage_class=sc,
+                uploaded_s3_uri=uploaded_s3_uri,
+                filename=filename,
+                exc_info=True,
+            )
+            return None
+        return presigned_urls[0], uploaded_s3_uri