mgokv: add Update method

This allows atomic updates to a value.

mgokv/persist.go

@@ -121,6 +121,34 @@ func (s *Session) PutInitial(key string, val interface{}) (bool, error) {
 	return s.putInitialAtTime(key, val, time.Now())
 }
 
+// Update updates the value using the MongoDB update operation
+// specified in update. The value is stored in the "value" field
+// in the document.
+//
+// For example, if a value of type struct { N int } is associated
+// with a key, then:
+//
+//	s.Update(key, bson.M{"$inc": bson.M{"value.n": 1}})
+//
+// will atomically increment the N value.
+//
+// If there is no value associated with the key, Update
+// returns ErrNotFound.
+func (s *Session) Update(key string, update interface{}) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if err := s.coll.UpdateId(key, update); err != nil {
+		if err == mgo.ErrNotFound {
+			return ErrNotFound
+		}
+		return errgo.Mask(err)
+	}
+	// We can't easily find the new value so just delete the
+	// item from the cache so it will be fetched next time.
+	delete(s.entries, key)
+	return nil
+}
+
 // putInitialAtTime is the internal version of PutInitial - it takes the current time
 // as an argument for testing.
 func (s *Session) putInitialAtTime(key string, val interface{}, now time.Time) (bool, error) {

mgokv/persist_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/juju/testing"
 	gc "gopkg.in/check.v1"
 	"gopkg.in/errgo.v1"
+	"gopkg.in/mgo.v2/bson"
 
 	"github.com/juju/utils/mgokv"
 )
@@ -196,3 +197,25 @@ func (s *suite) TestRefresh(c *gc.C) {
 	c.Check(err, gc.Equals, nil)
 	c.Assert(v, gc.Equals, val{88, 99})
 }
+
+func (s *suite) TestUpdate(c *gc.C) {
+	type val struct {
+		N int
+	}
+	store := mgokv.NewStore(time.Minute, s.Session.DB("foo").C("x")).Session(s.Session)
+	_, err := store.PutInitial("somekey", val{2})
+	c.Assert(err, gc.Equals, nil)
+	err = store.Update("somekey", bson.M{"$inc": bson.M{"value.n": 1}})
+	c.Assert(err, gc.Equals, nil)
+
+	var v val
+	err = store.Get("somekey", &v)
+	c.Assert(err, gc.IsNil)
+	c.Assert(v, gc.Equals, val{N: 3})
+}
+
+func (s *suite) TestUpdateNotFound(c *gc.C) {
+	store := mgokv.NewStore(time.Minute, s.Session.DB("foo").C("x")).Session(s.Session)
+	err := store.Update("somekey", bson.M{"$inc": bson.M{"value.n": 1}})
+	c.Assert(err, gc.Equals, mgokv.ErrNotFound)
+}