jmap: fetch entire threads

Fetch an email's entire thread in the JMAP backend.

Changelog-added: Fetch entire threads in the JMAP backend.
Signed-off-by: Tristan Partin <tristan@partin.io>
Tested-by: Tim Culverhouse <tim@timculverhouse.com>
Acked-by: Robin Jarry <robin@jarry.cc>

4 files changed, 141 insertions, 36 deletions

diff --git a/worker/jmap/cache/state.go b/worker/jmap/cache/state.go
index 5fec5034..e7771075 100644
--- a/worker/jmap/cache/state.go
+++ b/worker/jmap/cache/state.go
@@ -24,7 +24,20 @@ func (c *JMAPCache) PutEmailState(state string) error {
 	return c.put(emailStateKey, []byte(state))
 }
 
+func (c *JMAPCache) GetThreadState() (string, error) {
+	buf, err := c.get(threadStateKey)
+	if err != nil {
+		return "", err
+	}
+	return string(buf), nil
+}
+
+func (c *JMAPCache) PutThreadState(state string) error {
+	return c.put(threadStateKey, []byte(state))
+}
+
 const (
 	mailboxStateKey = "state/mailbox"
 	emailStateKey   = "state/email"
+	threadStateKey  = "state/thread"
 )
diff --git a/worker/jmap/fetch.go b/worker/jmap/fetch.go
index 5d88e6a1..bbef1bb5 100644
--- a/worker/jmap/fetch.go
+++ b/worker/jmap/fetch.go
@@ -16,6 +16,7 @@ import (
 var headersProperties = []string{
 	"id",
 	"blobId",
+	"threadId",
 	"mailboxIds",
 	"keywords",
 	"size",
@@ -34,9 +35,8 @@ var headersProperties = []string{
 }
 
 func (w *JMAPWorker) handleFetchMessageHeaders(msg *types.FetchMessageHeaders) error {
-	var req jmap.Request
-
-	ids := make([]jmap.ID, 0, len(msg.Uids))
+	emailIdsToFetch := make([]jmap.ID, 0, len(msg.Uids))
+	currentEmails := make([]*email.Email, 0, len(msg.Uids))
 	for _, uid := range msg.Uids {
 		id, ok := w.uidStore.GetKey(uid)
 		if !ok {
@@ -45,47 +45,51 @@ func (w *JMAPWorker) handleFetchMessageHeaders(msg *types.FetchMessageHeaders) e
 		jid := jmap.ID(id)
 		m, err := w.cache.GetEmail(jid)
 		if err == nil {
-			w.w.PostMessage(&types.MessageInfo{
-				Message: types.RespondTo(msg),
-				Info:    w.translateMsgInfo(m),
-			}, nil)
-			continue
+			currentEmails = append(currentEmails, m)
+		} else {
+			emailIdsToFetch = append(emailIdsToFetch, jid)
 		}
-		ids = append(ids, jid)
 	}
 
-	if len(ids) == 0 {
-		return nil
-	}
+	if len(emailIdsToFetch) != 0 {
+		var req jmap.Request
 
-	req.Invoke(&email.Get{
-		Account:    w.AccountId(),
-		IDs:        ids,
-		Properties: headersProperties,
-	})
+		req.Invoke(&email.Get{
+			Account:    w.AccountId(),
+			IDs:        emailIdsToFetch,
+			Properties: []string{"threadId"},
+		})
 
-	resp, err := w.Do(&req)
+		resp, err := w.Do(&req)
+		if err != nil {
+			return err
+		}
+
+		for _, inv := range resp.Responses {
+			switch r := inv.Args.(type) {
+			case *email.GetResponse:
+				if err = w.cache.PutEmailState(r.State); err != nil {
+					w.w.Warnf("PutEmailState: %s", err)
+				}
+				currentEmails = append(currentEmails, r.List...)
+			case *jmap.MethodError:
+				return wrapMethodError(r)
+			}
+		}
+	}
+
+	allEmails, err := w.fetchEntireThreads(currentEmails)
 	if err != nil {
 		return err
 	}
 
-	for _, inv := range resp.Responses {
-		switch r := inv.Args.(type) {
-		case *email.GetResponse:
-			for _, m := range r.List {
-				w.w.PostMessage(&types.MessageInfo{
-					Message: types.RespondTo(msg),
-					Info:    w.translateMsgInfo(m),
-				}, nil)
-				if err := w.cache.PutEmail(m.ID, m); err != nil {
-					w.w.Warnf("PutEmail: %s", err)
-				}
-			}
-			if err = w.cache.PutEmailState(r.State); err != nil {
-				w.w.Warnf("PutEmailState: %s", err)
-			}
-		case *jmap.MethodError:
-			return wrapMethodError(r)
+	for _, m := range allEmails {
+		w.w.PostMessage(&types.MessageInfo{
+			Message: types.RespondTo(msg),
+			Info:    w.translateMsgInfo(m),
+		}, nil)
+		if err := w.cache.PutEmail(m.ID, m); err != nil {
+			w.w.Warnf("PutEmail: %s", err)
 		}
 	}
 
diff --git a/worker/jmap/push.go b/worker/jmap/push.go
index 54e6375e..0a85332e 100644
--- a/worker/jmap/push.go
+++ b/worker/jmap/push.go
@@ -253,7 +253,13 @@ func (w *JMAPWorker) refresh(newState jmap.TypeState) error {
 					selectedIds[id] = true
 				}
 			}
-			for _, m := range r.List {
+
+			emails, err := w.fetchEntireThreads(r.List)
+			if err != nil {
+				return err
+			}
+
+			for _, m := range emails {
 				err = w.cache.PutEmail(m.ID, m)
 				if err != nil {
 					w.w.Warnf("PutEmail: %s", err)
diff --git a/worker/jmap/threads.go b/worker/jmap/threads.go
new file mode 100644
index 00000000..b967e8cd
--- /dev/null
+++ b/worker/jmap/threads.go
@@ -0,0 +1,82 @@
+package jmap
+
+import (
+	"git.sr.ht/~rockorager/go-jmap"
+	"git.sr.ht/~rockorager/go-jmap/mail/email"
+	"git.sr.ht/~rockorager/go-jmap/mail/thread"
+)
+
+func (w *JMAPWorker) fetchEntireThreads(emailIds []*email.Email) ([]*email.Email, error) {
+	var req jmap.Request
+
+	if len(emailIds) == 0 {
+		return emailIds, nil
+	}
+
+	threadsToFetch := make([]jmap.ID, 0, len(emailIds))
+	for _, m := range emailIds {
+		threadsToFetch = append(threadsToFetch, m.ThreadID)
+	}
+
+	req.Invoke(&thread.Get{
+		Account: w.AccountId(),
+		IDs:     threadsToFetch,
+	})
+
+	resp, err := w.Do(&req)
+	if err != nil {
+		return nil, err
+	}
+
+	emailsToFetch := make([]jmap.ID, 0)
+	emailsToReturn := make([]*email.Email, 0)
+	for _, inv := range resp.Responses {
+		switch r := inv.Args.(type) {
+		case *thread.GetResponse:
+			for _, t := range r.List {
+				for _, emailId := range t.EmailIDs {
+					m, err := w.cache.GetEmail(emailId)
+					if err == nil || m == nil {
+						emailsToFetch = append(emailsToFetch, emailId)
+					} else {
+						emailsToReturn = append(emailsToReturn, m)
+					}
+				}
+			}
+			if err = w.cache.PutThreadState(r.State); err != nil {
+				w.w.Warnf("PutThreadState: %s", err)
+			}
+		case *jmap.MethodError:
+			return nil, wrapMethodError(r)
+		}
+	}
+
+	if len(emailsToFetch) == 0 {
+		return emailsToReturn, nil
+	}
+
+	req.Invoke(&email.Get{
+		Account:    w.AccountId(),
+		IDs:        emailsToFetch,
+		Properties: headersProperties,
+	})
+
+	resp, err = w.Do(&req)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, inv := range resp.Responses {
+		switch r := inv.Args.(type) {
+		case *email.GetResponse:
+			emailsToReturn = append(emailsToReturn, r.List...)
+			if err = w.cache.PutEmailState(r.State); err != nil {
+				w.w.Warnf("PutEmailState: %s", err)
+			}
+		case *jmap.MethodError:
+			return nil, wrapMethodError(r)
+		}
+	}
+
+	return emailsToReturn, nil
+}