From b251b85e2b0a6b157f5fa806327a8ffe4dcd668c Mon Sep 17 00:00:00 2001
From: Antoine Nguyen <tonio@ngyn.org>
Date: Sun, 25 Jan 2015 16:18:57 +0100
Subject: [PATCH] Use a BinaryField and build the content.

see #677
---
 modoboa/extensions/amavis/models.py        |  2 +-
 modoboa/extensions/amavis/sql_connector.py | 27 +++++++++-------------
 modoboa/extensions/amavis/sql_email.py     |  5 +---
 modoboa/extensions/amavis/views.py         | 10 ++------
 4 files changed, 15 insertions(+), 29 deletions(-)

diff --git a/modoboa/extensions/amavis/models.py b/modoboa/extensions/amavis/models.py
index 7a711fdd8..1688d6fed 100644
--- a/modoboa/extensions/amavis/models.py
+++ b/modoboa/extensions/amavis/models.py
@@ -162,7 +162,7 @@ class Quarantine(models.Model):
     partition_tag = models.IntegerField(null=True, blank=True)
     mail = models.ForeignKey(Msgs, primary_key=True)
     chunk_ind = models.IntegerField()
-    mail_text = models.TextField()
+    mail_text = models.BinaryField()
 
     class Meta:
         db_table = u'quarantine'
diff --git a/modoboa/extensions/amavis/sql_connector.py b/modoboa/extensions/amavis/sql_connector.py
index 46f0075ae..1be908b03 100644
--- a/modoboa/extensions/amavis/sql_connector.py
+++ b/modoboa/extensions/amavis/sql_connector.py
@@ -1,6 +1,7 @@
 """
 SQL connector module.
 """
+
 import datetime
 
 from django.db.models import Q
@@ -209,9 +210,14 @@ class SQLconnector(object):
         return Msgrcpt.objects.filter(rq).count()
 
     def get_mail_content(self, mailid):
-        """Retrieve the content of a message.
-        """
-        return Quarantine.objects.filter(mail=mailid)
+        """Retrieve the content of a message."""
+        content = "".join([
+            str(qmail.mail_text)
+            for qmail in Quarantine.objects.filter(mail=mailid)
+        ])
+        if isinstance(content, unicode):
+            content = content.encode("utf-8")
+        return content
 
 
 class PgSQLconnector(SQLconnector):
@@ -229,9 +235,7 @@ class PgSQLconnector(SQLconnector):
     """
 
     def _apply_msgrcpt_filters(self, flt):
-        """Return filters based on user's role.
-
-        """
+        """Return filters based on user's role. """
         self._where = []
         if self.user.group == 'SimpleUsers':
             rcpts = [self.user.email]
@@ -248,10 +252,7 @@ class PgSQLconnector(SQLconnector):
         return flt
 
     def _apply_extra_search_filter(self, crit, pattern):
-        """Apply search filters using additional criterias.
-
-
-        """
+        """Apply search filters using additional criterias."""
         if crit == "to":
             self._where.append(
                 "convert_from(maddr.email, 'UTF8') LIKE '%%{0}%%'".format(
@@ -300,12 +301,6 @@ class PgSQLconnector(SQLconnector):
             ).count()
         return Msgrcpt.objects.filter(rq).count()
 
-    def get_mail_content(self, mailid):
-        """Retrieve the content of a message."""
-        return Quarantine.objects.filter(mail=mailid).extra(
-            select={'mail_text': "convert_from(mail_text, 'UTF8')"}
-        )
-
 
 def get_connector(**kwargs):
     """Return the appropriate *connector* class.
diff --git a/modoboa/extensions/amavis/sql_email.py b/modoboa/extensions/amavis/sql_email.py
index 2722f09a8..067004d7d 100644
--- a/modoboa/extensions/amavis/sql_email.py
+++ b/modoboa/extensions/amavis/sql_email.py
@@ -40,10 +40,7 @@ class SQLemail(Email):
         import email
 
         if self._msg is None:
-            qmails = get_connector().get_mail_content(self.mailid)
-            mail_text = "".join([qm.mail_text for qm in qmails])
-            if type(mail_text) is unicode:
-                mail_text = mail_text.encode("utf-8")
+            mail_text = get_connector().get_mail_content(self.mailid)
             self._msg = email.message_from_string(mail_text)
             self._parse(self._msg)
         return self._msg
diff --git a/modoboa/extensions/amavis/views.py b/modoboa/extensions/amavis/views.py
index 8a05008a9..a15165090 100644
--- a/modoboa/extensions/amavis/views.py
+++ b/modoboa/extensions/amavis/views.py
@@ -176,11 +176,7 @@ def viewmail(request, mail_id):
 @login_required
 def viewheaders(request, mail_id):
     """Display message headers."""
-    content = ""
-    for qm in get_connector().get_mail_content(mail_id):
-        content += qm.mail_text
-    if type(content) is unicode:
-        content = content.encode("utf-8")
+    content = get_connector().get_mail_content(mail_id)
     msg = email.message_from_string(content)
     headers = []
     for name, value in msg.items():
@@ -348,9 +344,7 @@ def mark_messages(request, selection, mtype, recipient_db=None):
     saclient = SpamassassinClient(request.user, recipient_db)
     for item in selection:
         rcpt, mail_id = item.split()
-        content = "".join(
-            [msg.mail_text for msg in connector.get_mail_content(mail_id)]
-        )
+        content = connector.get_mail_content(mail_id)
         result = saclient.learn_spam(rcpt, content) if mtype == "spam" \
             else saclient.learn_ham(rcpt, content)
         if not result: