1 files changed, 79 insertions, 12 deletions

diff --git a/worker/notmuch/message.go b/worker/notmuch/message.go
index 077fb92..aa16cee 100644
--- a/worker/notmuch/message.go
+++ b/worker/notmuch/message.go
@@ -17,13 +17,15 @@ import (
 )
 
 type Message struct {
-	uid uint32
-	key string
-	msg *notmuch.Message
+	uid     uint32
+	key     string
+	msg     *notmuch.Message
+	rwDB    func() (*notmuch.DB, error) // used to open a db for writing
+	refresh func(*Message) error        // called after msg modification
 }
 
 // NewReader reads a message into memory and returns an io.Reader for it.
-func (m Message) NewReader() (io.Reader, error) {
+func (m *Message) NewReader() (io.Reader, error) {
 	f, err := os.Open(m.msg.Filename())
 	if err != nil {
 		return nil, err
@@ -37,13 +39,13 @@ func (m Message) NewReader() (io.Reader, error) {
 }
 
 // MessageInfo populates a models.MessageInfo struct for the message.
-func (m Message) MessageInfo() (*models.MessageInfo, error) {
+func (m *Message) MessageInfo() (*models.MessageInfo, error) {
 	return lib.MessageInfo(m)
 }
 
 // NewBodyPartReader creates a new io.Reader for the requested body part(s) of
 // the message.
-func (m Message) NewBodyPartReader(requestedParts []int) (io.Reader, error) {
+func (m *Message) NewBodyPartReader(requestedParts []int) (io.Reader, error) {
 	f, err := os.Open(m.msg.Filename())
 	if err != nil {
 		return nil, err
@@ -57,7 +59,7 @@ func (m Message) NewBodyPartReader(requestedParts []int) (io.Reader, error) {
 }
 
 // MarkRead either adds or removes the maildir.FlagSeen flag from the message.
-func (m Message) MarkRead(seen bool) error {
+func (m *Message) MarkRead(seen bool) error {
 	haveUnread := false
 	for _, t := range m.tags() {
 		if t == "unread" {
@@ -71,14 +73,14 @@ func (m Message) MarkRead(seen bool) error {
 	}
 
 	if haveUnread {
-		err := m.msg.RemoveTag("unread")
+		err := m.RemoveTag("unread")
 		if err != nil {
 			return err
 		}
 		return nil
 	}
 
-	err := m.msg.AddTag("unread")
+	err := m.AddTag("unread")
 	if err != nil {
 		return err
 	}
@@ -86,7 +88,7 @@ func (m Message) MarkRead(seen bool) error {
 }
 
 // tags returns the notmuch tags of a message
-func (m Message) tags() []string {
+func (m *Message) tags() []string {
 	ts := m.msg.Tags()
 	var tags []string
 	var tag *notmuch.Tag
@@ -96,7 +98,72 @@ func (m Message) tags() []string {
 	return tags
 }
 
-func (m Message) ModelFlags() ([]models.Flag, error) {
+func (m *Message) modify(cb func(*notmuch.Message) error) error {
+	db, err := m.rwDB()
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+	msg, err := db.FindMessage(m.key)
+	if err != nil {
+		return err
+	}
+	err = cb(msg)
+	if err != nil {
+		return err
+	}
+	// we need to explicitly close here, else we don't commit
+	dcerr := db.Close()
+	if dcerr != nil && err == nil {
+		err = dcerr
+	}
+	// next we need to refresh the notmuch msg, else we serve stale tags
+	rerr := m.refresh(m)
+	if rerr != nil && err == nil {
+		err = rerr
+	}
+	return err
+}
+
+func (m *Message) AddTag(tag string) error {
+	err := m.modify(func(msg *notmuch.Message) error {
+		return msg.AddTag(tag)
+	})
+	return err
+}
+
+func (m *Message) AddTags(tags []string) error {
+	err := m.modify(func(msg *notmuch.Message) error {
+		ierr := msg.Atomic(func(msg *notmuch.Message) {
+			for _, t := range tags {
+				msg.AddTag(t)
+			}
+		})
+		return ierr
+	})
+	return err
+}
+
+func (m *Message) RemoveTag(tag string) error {
+	err := m.modify(func(msg *notmuch.Message) error {
+		return msg.RemoveTag(tag)
+	})
+	return err
+}
+
+func (m *Message) RemoveTags(tags []string) error {
+	err := m.modify(func(msg *notmuch.Message) error {
+		ierr := msg.Atomic(func(msg *notmuch.Message) {
+			for _, t := range tags {
+				msg.RemoveTag(t)
+			}
+		})
+		return ierr
+	})
+	return err
+}
+
+func (m *Message) ModelFlags() ([]models.Flag, error) {
 	var flags []models.Flag
 	seen := true
 
@@ -118,6 +185,6 @@ func (m Message) ModelFlags() ([]models.Flag, error) {
 	return flags, nil
 }
 
-func (m Message) UID() uint32 {
+func (m *Message) UID() uint32 {
 	return m.uid
 }