Card #30 - migrates attachments controller to twisted

1 files changed, 16 insertions, 8 deletions

diff --git a/service/pixelated/controllers/attachments_controller.py b/service/pixelated/controllers/attachments_controller.py
index 1d5360f7..4965f166 100644
--- a/service/pixelated/controllers/attachments_controller.py
+++ b/service/pixelated/controllers/attachments_controller.py
@@ -14,12 +14,11 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
 
-from flask import send_file
-from flask import request
-
 import io
 import re
-from twisted.web.server import NOT_DONE_YET
+
+from twisted.protocols.basic import FileSender
+from twisted.python.log import err
 
 
 class AttachmentsController:
@@ -28,12 +27,21 @@ class AttachmentsController:
         self.querier = querier
 
     def attachment(self, request, attachment_id):
-        encoding = request.args.get('encoding', [''])[0]
+        encoding = request.args.get('encoding')
         attachment = self.querier.attachment(attachment_id, encoding)
-        request.setRawHeader('Content-Type', self._extract_mimetype(attachment['content-type']))
-        request.write(io.BytesIO(attachment['content']))
 
-        return NOT_DONE_YET
+        content_type = self._extract_mimetype(attachment['content-type'])
+        request.setHeader('Content-Type', content_type)
+        bytes_io = io.BytesIO(attachment['content'])
+        d = FileSender().beginFileTransfer(bytes_io, request)
+
+        def cbFinished(ignored):
+            bytes_io.close()
+            request.finish()
+
+        d.addErrback(err).addCallback(cbFinished)
+
+        return d
 
     def _extract_mimetype(self, content_type):
         match = re.compile('([A-Za-z-]+\/[A-Za-z-]+)').search(content_type)