make lazy attachments save in a single couch request

dimagi · Aug 12, 2014 · 92b5bf1 · 92b5bf1
1 parent a04326a
commit 92b5bf1
Showing 1 changed file with 23 additions and 10 deletions.
diff --git a/dimagi/utils/couch/lazy_attachment_doc.py b/dimagi/utils/couch/lazy_attachment_doc.py
@@ -1,8 +1,14 @@
+from mimetypes import guess_type
 from couchdbkit.ext.django.schema import Document
+from couchdbkit.resource import encode_attachments
 
 
 class LazyAttachmentDoc(Document):
 
+    def __init__(self, *args, **kwargs):
+        super(LazyAttachmentDoc, self).__init__(*args, **kwargs)
+        self._LAZY_ATTACHMENTS = {}
+
     @classmethod
     def wrap(cls, data):
         self = super(LazyAttachmentDoc, cls).wrap(data)
@@ -15,8 +21,6 @@ def wrap(cls, data):
 
     def __store_lazy_attachment(self, content, name=None, content_type=None,
                                 content_length=None):
-        if not hasattr(self, '_LAZY_ATTACHMENTS'):
-            self._LAZY_ATTACHMENTS = {}
         info = {
             'content': content,
             'content_type': content_type,
@@ -41,14 +45,12 @@ def lazy_put_attachment(self, content, name=None, content_type=None,
         info = self.__store_lazy_attachment(content, name, content_type,
                                             content_length)
 
-        def put_attachment():
-            self.put_attachment(name=name, **info)
-
-        self.register_post_save(put_attachment)
+        # def put_attachment():
+        #     self.put_attachment(name=name, **info)
+        #
+        # self.register_post_save(put_attachment)
 
     def lazy_fetch_attachment(self, name):
-        if not hasattr(self, '_LAZY_ATTACHMENTS'):
-            self._LAZY_ATTACHMENTS = {}
         try:
             info = self._LAZY_ATTACHMENTS[name]
             return info['content']
@@ -80,8 +82,19 @@ def del_pre_save():
                 del self._PRE_SAVE
 
             self.register_post_save(del_pre_save)
-
-        super(LazyAttachmentDoc, self).save(**params)
+        _attachments = self._attachments.copy() if self._attachments else {}
+        for name, info in self._LAZY_ATTACHMENTS.items():
+            data = info['content']
+            content_type = (info['content_type']
+                            or ';'.join(filter(None, guess_type(name))))
+            if isinstance(data, unicode):
+                data = data.encode('utf8')
+            _attachments[name] = {
+                'content_type': content_type,
+                'data': data,
+            }
+        self._attachments = encode_attachments(_attachments)
+        super(LazyAttachmentDoc, self).save(encode_attachments=False, **params)
 
         if hasattr(self, '_POST_SAVE'):
             for post_save in self._POST_SAVE: