review fixes

Author: sylvinus

src/backend/core/management/commands/verify_tiered_storage.py

@@ -13,7 +13,9 @@
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
-from core.enums import BlobStorageLocationChoices
+import pyzstd
+
+from core.enums import BlobStorageLocationChoices, CompressionTypeChoices
 from core.models import Blob
 from core.services.tiered_storage import TieredStorageService
 
@@ -277,10 +279,6 @@ def _verify_blob_hash(self, obj_name: str, expected_sha_bytes: bytes) -> bool:
 
             # The decrypted content is still compressed
             # We need to decompress to verify the original hash
-            import pyzstd
-
-            from core.enums import CompressionTypeChoices
-
             if blob.compression == CompressionTypeChoices.ZSTD:
                 original = pyzstd.decompress(decrypted)
             else:
@@ -451,7 +449,13 @@ def _re_encrypt_single_blob(self, blob: Blob, target_key_id: int) -> str:
 
                 blob.raw_content = encrypted
                 blob.encryption_key_id = new_key_id
-                blob.save(update_fields=["raw_content", "encryption_key_id"])
+                blob.save(
+                    update_fields=[
+                        "raw_content",
+                        "encryption_key_id",
+                        "size_compressed",
+                    ]
+                )
 
             self.stdout.write(
                 f"  Re-encrypted blob {blob.id} (POSTGRES): "
@@ -470,52 +474,74 @@ def _re_encrypt_single_blob(self, blob: Blob, target_key_id: int) -> str:
 
             storage_key = self.service.compute_storage_key(bytes(blob.sha256))
 
-            # Keep old content in memory so we can restore S3 if the
-            # transaction fails after the S3 overwrite.
-            old_encrypted = None
-            s3_updated = False
-
+            # Read & re-encrypt outside the DB transaction.
             try:
-                with transaction.atomic():
-                    blob = Blob.objects.select_for_update().get(id=blob.id)
-                    if blob.encryption_key_id == target_key_id:
-                        return "skipped"
+                with self.service.storage.open(storage_key, "rb") as f:
+                    old_encrypted = f.read()
+            except FileNotFoundError:
+                self.stderr.write(
+                    self.style.ERROR(
+                        f"  ERROR blob {blob.id}: not found in storage at {storage_key}"
+                    )
+                )
+                return "error"
 
-                    old_key_id = blob.encryption_key_id
+            decrypted = self.service.decrypt(old_encrypted, blob.encryption_key_id)
+            encrypted, new_key_id = self.service.encrypt(decrypted)
 
+            # Stage the new ciphertext at a temp key so the canonical object
+            # is not touched until the DB transaction commits. If the DB
+            # update fails, we drop the temp object and storage stays
+            # consistent with the DB.
+            temp_key = f"{storage_key}.tmp.{new_key_id}"
+            self.service.storage.save(temp_key, ContentFile(encrypted))
+
+            promote_done = {"value": False}
+
+            def _promote_temp():
+                try:
+                    self.service.storage.save(storage_key, ContentFile(encrypted))
+                finally:
                     try:
-                        with self.service.storage.open(storage_key, "rb") as f:
-                            old_encrypted = f.read()
-                    except FileNotFoundError:
+                        self.service.storage.delete(temp_key)
+                    except Exception as cleanup_err:  # pylint: disable=broad-except
                         self.stderr.write(
                             self.style.ERROR(
-                                f"  ERROR blob {blob.id}: not found in storage at {storage_key}"
+                                f"  WARN blob {blob.id}: failed to delete temp "
+                                f"{temp_key}: {cleanup_err}"
                             )
                         )
-                        return "error"
-
-                    decrypted = self.service.decrypt(old_encrypted, old_key_id)
-                    encrypted, new_key_id = self.service.encrypt(decrypted)
+                promote_done["value"] = True
 
-                    self.service.storage.save(storage_key, ContentFile(encrypted))
-                    s3_updated = True
+            try:
+                with transaction.atomic():
+                    blob = Blob.objects.select_for_update().get(id=blob.id)
+                    if blob.encryption_key_id == target_key_id:
+                        # Another worker beat us to it; clean up our temp.
+                        try:
+                            self.service.storage.delete(temp_key)
+                        except Exception as cleanup_err:  # pylint: disable=broad-except
+                            self.stderr.write(
+                                self.style.ERROR(
+                                    f"  WARN blob {blob.id}: failed to delete temp "
+                                    f"{temp_key}: {cleanup_err}"
+                                )
+                            )
+                        return "skipped"
 
+                    old_key_id = blob.encryption_key_id
                     blob.encryption_key_id = new_key_id
                     blob.save(update_fields=["encryption_key_id"])
+                    transaction.on_commit(_promote_temp)
             except Exception:
-                # If S3 was overwritten but the transaction failed (DB error
-                # or commit failure), restore the old S3 content to prevent
-                # leaving the blob in a corrupted state.
-                if s3_updated and old_encrypted is not None:
+                if not promote_done["value"]:
                     try:
-                        self.service.storage.save(
-                            storage_key, ContentFile(old_encrypted)
-                        )
-                    except Exception as restore_err:  # pylint: disable=broad-except
+                        self.service.storage.delete(temp_key)
+                    except Exception as cleanup_err:  # pylint: disable=broad-except
                         self.stderr.write(
                             self.style.ERROR(
-                                f"  CRITICAL: failed to restore S3 content for "
-                                f"blob {blob.id}: {restore_err}"
+                                f"  WARN blob {blob.id}: failed to delete temp "
+                                f"{temp_key}: {cleanup_err}"
                             )
                         )
                 raise

src/backend/core/services/tiered_storage.py

@@ -31,8 +31,7 @@ class TieredStorageService:
     def __init__(self):
         """Initialize the service, checking if object storage is configured."""
         self._storage = None
-        opts = settings.STORAGES.get("message-blobs", {}).get("OPTIONS", {})
-        self.enabled = bool(opts.get("endpoint_url") or opts.get("access_key"))
+        self.enabled = bool(settings.STORAGES.get("message-blobs"))
         # encryption_keys is a dict: {"1": "key1", "2": "key2"}
         self.encryption_keys = settings.MESSAGES_BLOB_ENCRYPTION_KEYS or {}
         self.active_key_id = settings.MESSAGES_BLOB_ENCRYPTION_ACTIVE_KEY_ID

src/backend/core/signals.py

@@ -144,11 +144,20 @@ def cleanup_blob_storage(sender, instance, **kwargs):
 
     Uses a post_delete signal instead of a model delete() override to ensure
     cleanup also runs during CASCADE and QuerySet bulk deletes.
+
+    Defers the actual storage deletion to transaction.on_commit so that an
+    enclosing transaction rollback does not leave us with an object deleted
+    from S3 while the blob row is still in the DB.
     """
-    if instance.storage_location == BlobStorageLocationChoices.OBJECT_STORAGE:
-        service = TieredStorageService()
-        if service.enabled:
-            service.delete_if_orphaned(bytes(instance.sha256))
+    if instance.storage_location != BlobStorageLocationChoices.OBJECT_STORAGE:
+        return
+
+    service = TieredStorageService()
+    if not service.enabled:
+        return
+
+    sha256 = bytes(instance.sha256)
+    transaction.on_commit(lambda: service.delete_if_orphaned(sha256))
 
 
 @receiver(pre_delete, sender=models.Message)

src/backend/core/tests/commands/test_verify_tiered_storage.py

@@ -34,9 +34,7 @@ def test_command_disabled_when_no_storage(self):
         stderr = StringIO()
 
         with patch("core.services.tiered_storage.settings") as mock_settings:
-            mock_settings.STORAGES = {
-                "message-blobs": {"OPTIONS": {"endpoint_url": ""}}
-            }
+            mock_settings.STORAGES = {}
             mock_settings.MESSAGES_BLOB_ENCRYPTION_KEYS = {}
             mock_settings.MESSAGES_BLOB_ENCRYPTION_ACTIVE_KEY_ID = 0
 
@@ -465,6 +463,7 @@ def test_re_encrypt_postgres_blob(self):
         decrypted = service.decrypt(bytes(blob.raw_content), blob.encryption_key_id)
         assert pyzstd.decompress(decrypted) == original_content
 
+    @pytest.mark.django_db(transaction=True)
     def test_re_encrypt_object_storage_blob(self):
         """Test re-encrypting an object storage blob with real encryption."""
         import pyzstd

src/backend/core/tests/services/test_tiered_storage.py

@@ -315,6 +315,7 @@ def test_upload_download_with_encryption(self):
         finally:
             service.storage.delete(storage_key)
 
+    @pytest.mark.django_db(transaction=True)
     def test_deduplication_single_upload(self):
         """Test that two blobs with same content result in single storage object."""
         service = TieredStorageService()
@@ -357,9 +358,12 @@ def test_deduplication_single_upload(self):
             if service.storage.exists(storage_key):
                 service.storage.delete(storage_key)
 
+    @pytest.mark.django_db(transaction=True)
     def test_full_offload_workflow(self):
         """Test the complete offload workflow: create blob, offload, read content."""
-        from core.services.tiered_storage_tasks import offload_single_blob_task
+        from core.services.tiered_storage_tasks import (
+            offload_single_blob_task,
+        )
 
         mailbox = factories.MailboxFactory()
         content = b"Content for full offload workflow test" * 20
@@ -386,14 +390,17 @@ def test_full_offload_workflow(self):
         MESSAGES_BLOB_ENCRYPTION_KEYS={"1": _TEST_ENCRYPTION_KEY},
         MESSAGES_BLOB_ENCRYPTION_ACTIVE_KEY_ID=1,
     )
+    @pytest.mark.django_db(transaction=True)
     def test_offload_with_encryption_roundtrip(self):
         """
         Test full offload workflow with encryption: create encrypted blob,
         offload to storage, read back content.
 
         This is a critical regression test for the double-encryption bug.
         """
-        from core.services.tiered_storage_tasks import offload_single_blob_task
+        from core.services.tiered_storage_tasks import (
+            offload_single_blob_task,
+        )
 
         mailbox = factories.MailboxFactory()
         original_content = b"Test content for encryption offload roundtrip" * 50
@@ -419,6 +426,7 @@ def test_offload_with_encryption_roundtrip(self):
         finally:
             blob.delete()
 
+    @pytest.mark.django_db(transaction=True)
     def test_delete_if_orphaned(self):
         """Test that delete_if_orphaned correctly handles references."""
         service = TieredStorageService()
@@ -562,7 +570,7 @@ def test_deduplication_with_different_encryption_keys(self):
                     service.storage.delete(storage_key)
 
 
-@pytest.mark.django_db
+@pytest.mark.django_db(transaction=True)
 class TestTieredStorageCascadeDelete:
     """Tests that S3 cleanup works during cascade and bulk deletes."""
 

src/backend/core/tests/tasks/test_tiered_storage_tasks.py

@@ -44,9 +44,7 @@ def test_task_disabled_by_setting(self):
     def test_task_disabled_when_no_storage(self):
         """Test that task returns disabled status when storage not configured."""
         with patch("core.services.tiered_storage.settings") as mock_settings:
-            mock_settings.STORAGES = {
-                "message-blobs": {"OPTIONS": {"endpoint_url": ""}}
-            }
+            mock_settings.STORAGES = {}
             mock_settings.MESSAGES_BLOB_ENCRYPTION_KEYS = {}
             mock_settings.MESSAGES_BLOB_ENCRYPTION_ACTIVE_KEY_ID = 0
 
@@ -207,9 +205,7 @@ def test_task_disabled_by_setting(self):
     def test_task_disabled_when_no_storage(self):
         """Test that task returns disabled status when storage not configured."""
         with patch("core.services.tiered_storage.settings") as mock_settings:
-            mock_settings.STORAGES = {
-                "message-blobs": {"OPTIONS": {"endpoint_url": ""}}
-            }
+            mock_settings.STORAGES = {}
             mock_settings.MESSAGES_BLOB_ENCRYPTION_KEYS = {}
             mock_settings.MESSAGES_BLOB_ENCRYPTION_ACTIVE_KEY_ID = 0
 
@@ -310,6 +306,7 @@ def test_handles_upload_error(self):
         assert blob.storage_location == BlobStorageLocationChoices.POSTGRES
         assert blob.raw_content is not None
 
+    @pytest.mark.django_db(transaction=True)
     def test_deduplication_during_offload(self):
         """Test that offload uses deduplication for identical content."""
         service = TieredStorageService()