defer fetch to thread

also, dispatch query for all headers to its own method.

4 files changed, 70 insertions, 17 deletions

diff --git a/src/leap/mail/imap/mailbox.py b/src/leap/mail/imap/mailbox.py
index fa97512..21f0554 100644
--- a/src/leap/mail/imap/mailbox.py
+++ b/src/leap/mail/imap/mailbox.py
@@ -211,6 +211,9 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
                 fields.TYPE_MBOX_VAL, self.mbox)
             if query:
                 return query.pop()
+            else:
+                logger.error("Could not find mbox document for %r" %
+                             (self.mbox,))
         except Exception as exc:
             logger.exception("Unhandled error %r" % exc)
 
@@ -576,10 +579,30 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
                     otherwise.
         :type uid: bool
 
+        :rtype: deferred
+        """
+        d = defer.Deferred()
+        self.reactor.callInThread(self._do_fetch, messages_asked, uid, d)
+        if PROFILE_CMD:
+            do_profile_cmd(d, "FETCH")
+        return d
+
+    # called in thread
+    def _do_fetch(self, messages_asked, uid, d):
+        """
+        :param messages_asked: IDs of the messages to retrieve information
+                               about
+        :type messages_asked: MessageSet
+
+        :param uid: If true, the IDs are UIDs. They are message sequence IDs
+                    otherwise.
+        :type uid: bool
+        :param d: deferred whose callback will be called with result.
+        :type d: Deferred
+
         :rtype: A tuple of two-tuples of message sequence numbers and
                 LeapMessage
         """
-        from twisted.internet import reactor
         # For the moment our UID is sequential, so we
         # can treat them all the same.
         # Change this to the flag that twisted expects when we
@@ -597,9 +620,11 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
             logger.debug("Getting msg by index: INEFFICIENT call!")
             raise NotImplementedError
         else:
-            result = ((msgid, getmsg(msgid)) for msgid in seq_messg)
-            reactor.callLater(0, self.unset_recent_flags, seq_messg)
-        return result
+            got_msg = [(msgid, getmsg(msgid)) for msgid in seq_messg]
+            result = ((msgid, msg) for msgid, msg in got_msg
+                      if msg is not None)
+            self.reactor.callLater(0, self.unset_recent_flags, seq_messg)
+            self.reactor.callFromThread(d.callback, result)
 
     def fetch_flags(self, messages_asked, uid):
         """
@@ -668,6 +693,8 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
                 MessagePart.
         :rtype: tuple
         """
+        # TODO how often is thunderbird doing this?
+
         class headersPart(object):
             def __init__(self, uid, headers):
                 self.uid = uid
@@ -685,10 +712,9 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         messages_asked = self._bound_seq(messages_asked)
         seq_messg = self._filter_msg_seq(messages_asked)
 
-        all_chash = self.messages.all_flags_chash()
         all_headers = self.messages.all_headers()
         result = ((msgid, headersPart(
-            msgid, all_headers.get(all_chash.get(msgid, 'nil'), {})))
+            msgid, all_headers.get(msgid, {})))
             for msgid in seq_messg)
         return result
 
diff --git a/src/leap/mail/imap/memorystore.py b/src/leap/mail/imap/memorystore.py
index 2835826..e8e8152 100644
--- a/src/leap/mail/imap/memorystore.py
+++ b/src/leap/mail/imap/memorystore.py
@@ -434,6 +434,8 @@ class MemoryStore(object):
             hdoc = self._hdoc_store[chash]
             if empty(hdoc):
                 hdoc = self._permanent_store.get_headers_doc(chash)
+                if empty(hdoc):
+                    return None
                 if not empty(hdoc.content):
                     self._hdoc_store[chash] = hdoc.content
                     hdoc = hdoc.content
@@ -699,6 +701,31 @@ class MemoryStore(object):
                 continue
         return flags_dict
 
+    def all_headers(self, mbox):
+        """
+        Return a dictionary with all the header docs for a given mbox.
+
+        :param mbox: the mailbox
+        :type mbox: str or unicode
+        :rtype: dict
+        """
+        headers_dict = {}
+        uids = self.get_uids(mbox)
+        fdoc_store = self._fdoc_store[mbox]
+        hdoc_store = self._hdoc_store
+
+        for uid in uids:
+            try:
+                chash = fdoc_store[uid][fields.CONTENT_HASH_KEY]
+                hdoc = hdoc_store[chash]
+                if not empty(hdoc):
+                    headers_dict[uid] = hdoc
+            except KeyError:
+                continue
+
+        import pprint; pprint.pprint(headers_dict)
+        return headers_dict
+
     # Counting sheeps...
 
     def count_new_mbox(self, mbox):
diff --git a/src/leap/mail/imap/messages.py b/src/leap/mail/imap/messages.py
index bbc9deb..7884fb0 100644
--- a/src/leap/mail/imap/messages.py
+++ b/src/leap/mail/imap/messages.py
@@ -28,7 +28,6 @@ from functools import partial
 
 from twisted.mail import imap4
 from twisted.internet import defer
-from twisted.python import log
 from zope.interface import implements
 from zope.proxy import sameProxiedObjects
 
@@ -1248,12 +1247,14 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
         all_flags = dict(all_docs)
         return all_flags
 
-    # TODO get from memstore
     def all_headers(self):
         """
-        Return a dict with all the headers documents for this
+        Return a dict with all the header documents for this
         mailbox.
+
+        :rtype: dict
         """
+        return self.memstore.all_headers(self.mbox)
 
     def count(self):
         """
diff --git a/src/leap/mail/imap/server.py b/src/leap/mail/imap/server.py
index 3497a8b..7c09784 100644
--- a/src/leap/mail/imap/server.py
+++ b/src/leap/mail/imap/server.py
@@ -119,15 +119,14 @@ class LeapIMAPServer(imap4.IMAP4Server):
                 cbFetch, tag, query, uid
             ).addErrback(ebFetch, tag)
 
-        # XXX not implemented yet --- should hit memstore
-        #elif len(query) == 1 and str(query[0]) == "rfc822.header":
-            #self._oldTimeout = self.setTimeout(None)
+        elif len(query) == 1 and str(query[0]) == "rfc822.header":
+            self._oldTimeout = self.setTimeout(None)
             # no need to call iter, we get a generator
-            #maybeDeferred(
-                #self.mbox.fetch_headers, messages, uid=uid
-            #).addCallback(
-                #cbFetch, tag, query, uid
-            #).addErrback(ebFetch, tag)
+            maybeDeferred(
+                self.mbox.fetch_headers, messages, uid=uid
+            ).addCallback(
+                cbFetch, tag, query, uid
+            ).addErrback(ebFetch, tag)
         else:
             self._oldTimeout = self.setTimeout(None)
             # no need to call iter, we get a generator