From 75dc1f103b40a97b652a5aaaaca9797de6c83870 Mon Sep 17 00:00:00 2001
From: Benjamin Dauvergne <bdauvergne@entrouvert.com>
Date: Mon, 6 Apr 2020 17:09:47 +0200
Subject: [PATCH] api_views: use atomicwrites for writing petale content
 (#41379)

---
 debian/control      |  1 +
 petale/api_views.py | 36 ++++++++++++++++++++----------------
 setup.py            |  1 +
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/debian/control b/debian/control
index ba6dd92..a5224bb 100644
--- a/debian/control
+++ b/debian/control
@@ -12,6 +12,7 @@ Depends: ${misc:Depends}, ${python3:Depends},
     python3-django-jsonfield,
     python3-djangorestframework,
     python3-requests,
+    python3-atomicwrites
 Description: Simple key value datastore
 
 Package: petale
diff --git a/petale/api_views.py b/petale/api_views.py
index b1b4240..82c9244 100644
--- a/petale/api_views.py
+++ b/petale/api_views.py
@@ -38,6 +38,8 @@ from rest_framework import status
 from rest_framework.views import APIView
 from rest_framework.response import Response
 
+from atomicwrites import atomic_write
+
 from .models import CUT, Petal, Partner, AccessControlList
 from .utils import logit, StreamingHash
 from .exceptions import (PartnerNotFound, CutNotFound, KeyNotFound, NotFound, MissingContentType,
@@ -259,26 +261,28 @@ class PetalAPIView(APIView):
 
         petal.check_limits(content_length)
 
+        streaming_digest = StreamingHash(request)
+
+        def update_meta():
+            # update metadata
+            petal.content_type = content_type
+            petal.etag = streaming_digest.etag()
+            petal.size = content_length
+
+            # update partner size
+            Partner.objects.filter(id=petal.partner.id).update(size=F('size') + size_delta)
+            petal.save()
+
         if created:
             status_code = status.HTTP_201_CREATED
-            old_name = None
+            petal.data.save(petal_name, streaming_digest, save=False)
+            update_meta()
         else:
-            old_name = petal.data.name
+            with atomic_write(petal.data.path, overwrite=True, mode='wb') as fd:
+                for block in iter(lambda: streaming_digest.read(512), b''):
+                    fd.write(block)
+                update_meta()
 
-        streaming_digest = StreamingHash(request)
-        petal.data.save(petal_name, streaming_digest, save=False)
-
-        # update metadata
-        petal.content_type = content_type
-        petal.etag = streaming_digest.etag()
-        petal.size = content_length
-
-        # update partner size
-        Partner.objects.filter(id=petal.partner.id).update(size=F('size') + size_delta)
-
-        petal.save()
-        if old_name:
-            petal.data.storage.delete(old_name)
         return Response(
             {},
             status=status_code,
diff --git a/setup.py b/setup.py
index 1ca3a67..5c3ebaa 100644
--- a/setup.py
+++ b/setup.py
@@ -126,6 +126,7 @@ setup(
         'django-jsonfield',
         'djangorestframework>=3.3',
         'requests',
+        'atomicwrites',
     ],
     zip_safe=False,
     cmdclass={