Form S3 URLs via storages APIs in TemporaryUpload

This fixes a bug where we would form URLs without the `private-media/` prefix that's set at the django-private-storages config.

Instead of directly appending the prefix, we now go through the storage to form the URL

Author: ipeterov

aiarena/core/models/temporary_upload.py

@@ -2,15 +2,16 @@
 
 from django.db import models
 
-from botocore.exceptions import ClientError
 from private_storage.fields import PrivateFileField
+from storages.backends.s3boto3 import S3Boto3Storage
+from storages.utils import clean_name
 
 from .user import User
 
 
 class TemporaryUpload(models.Model):
     """
-    Tracks temporary file uploads to S3 until they are consumed by result submission.
+    Tracks temporary file uploads until they are consumed by result submission.
     Records are deleted when the upload is used (copied to permanent location).
     """
 
@@ -22,74 +23,63 @@ class Meta:
         ordering = ["-created_at"]
 
     def __str__(self):
-        return f"TemporaryUpload({self.id}, {self.s3_key})"
-
-    @property
-    def _storage(self):
-        return self.file.storage
-
-    @property
-    def _client(self):
-        return self._storage.connection.meta.client
-
-    @property
-    def _bucket_name(self):
-        return self._storage.bucket_name
-
-    @property
-    def s3_key(self):
-        return self.file.name
+        return f"TemporaryUpload({self.id}, {self.file.name})"
 
     @classmethod
     def create_for_upload(cls, user: User) -> "TemporaryUpload":
-        """Create a TemporaryUpload record with a generated S3 key."""
+        """Create a TemporaryUpload record with a storage-generated path."""
         upload = cls(uploaded_by=user)
-        upload.file.name = f"temp-uploads/{uuid.uuid4().hex}"
+        upload.file.name = cls._meta.get_field("file").generate_filename(upload, f"{uuid.uuid4().hex}")
         upload.save()
         return upload
 
     def generate_presigned_put_url(self) -> str:
-        """Generate a presigned PUT URL for direct S3 upload."""
-        return self._client.generate_presigned_url(
+        """Generate a presigned PUT URL for direct S3 upload.
+
+        Only meaningful on S3 storage; dev/test environments don't presign.
+        """
+        storage = self.file.storage
+        if not isinstance(storage, S3Boto3Storage):
+            raise RuntimeError("Presigned URLs are only supported on S3 storage")
+        # Storage prepends `location` (e.g. "private-media/") to the field name when
+        # talking to S3; raw boto3 doesn't know that, so apply the same transform.
+        s3_key = storage._normalize_name(clean_name(self.file.name))
+        return storage.connection.meta.client.generate_presigned_url(
             "put_object",
-            Params={
-                "Bucket": self._bucket_name,
-                "Key": self.s3_key,
-            },
+            Params={"Bucket": storage.bucket.name, "Key": s3_key},
             ExpiresIn=3600,
         )
 
-    def copy_to(self, dest_key: str) -> None:
-        """Copy to permanent location. Does not delete the source."""
-        self._client.copy_object(
-            Bucket=self._bucket_name,
-            CopySource={"Bucket": self._bucket_name, "Key": self.s3_key},
-            Key=dest_key,
-        )
+    def exists_in_storage(self) -> bool:
+        return self.file.storage.exists(self.file.name)
 
-    def copy_to_file_field(self, file_field, filename: str) -> None:
-        """
-        Copy the uploaded file to a model's FileField.
+    def copy_to_file_field(self, dest_file_field, filename: str) -> None:
+        """Copy the uploaded file into a model's FileField via S3-to-S3 copy."""
+        if not self.exists_in_storage():
+            raise ValueError("Upload not found in storage")
 
-        Uses S3-to-S3 copy (no data through the web worker), then sets the field's
-        name to the destination key.
-        """
-        if not self.exists_in_s3():
-            raise ValueError("Upload not found in S3")
-        dest_key = file_field.field.upload_to(file_field.instance, filename)
-        self.copy_to(dest_key)
-        file_field.name = dest_key
-
-    def delete_s3_object(self) -> None:
-        """Delete the S3 object."""
-        self._client.delete_object(Bucket=self._bucket_name, Key=self.s3_key)
-
-    def exists_in_s3(self) -> bool:
-        """Check if the file was actually uploaded to S3."""
-        try:
-            self._client.head_object(Bucket=self._bucket_name, Key=self.s3_key)
-            return True
-        except ClientError as e:
-            if e.response["Error"]["Code"] == "404":
-                return False
-            raise
+        src_storage = self.file.storage
+        dst_storage = dest_file_field.storage
+
+        if not (isinstance(src_storage, S3Boto3Storage) and isinstance(dst_storage, S3Boto3Storage)):
+            raise RuntimeError("copy_to_file_field requires S3 storage on both source and destination")
+        if src_storage.bucket.name != dst_storage.bucket.name:
+            raise RuntimeError("Cross-bucket copy is not supported")
+
+        # Build the destination name through the field's upload_to and storage,
+        # so any path transforms (location prefix, name conflict resolution) apply.
+        dest_name = dest_file_field.field.generate_filename(dest_file_field.instance, filename)
+
+        client = src_storage.connection.meta.client
+        client.copy_object(
+            Bucket=dst_storage.bucket.name,
+            CopySource={
+                "Bucket": src_storage.bucket.name,
+                "Key": src_storage._normalize_name(clean_name(self.file.name)),
+            },
+            Key=dst_storage._normalize_name(clean_name(dest_name)),
+        )
+        dest_file_field.name = dest_name
+
+    def delete_from_storage(self) -> None:
+        self.file.storage.delete(self.file.name)

aiarena/graphql/mutations.py

@@ -535,7 +535,7 @@ def cleanup():
             # Delete temporary upload S3 objects and DB records
             for upload in temp_uploads:
                 if upload is not None:
-                    upload.delete_s3_object()
+                    upload.delete_from_storage()
                     upload.delete()
             # Delete old bot_data files that were replaced
             storage = Bot._meta.get_field("bot_data").storage