~rogpeppe/juju-core/256-more-status

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
package state

import (
	"fmt"
	"labix.org/v2/mgo/txn"
	"strings"
)

// annotator stores annotations and information required to query MongoDB.
type annotator struct {
	annotations *map[string]string
	st          *State
	coll        string
	id          string
}

// SetAnnotation adds a key/value pair to annotations in MongoDB and the annotator.
func (a *annotator) SetAnnotation(key, value string) error {
	if strings.Contains(key, ".") {
		return fmt.Errorf("invalid key %q", key)
	}
	if *a.annotations == nil {
		*a.annotations = make(map[string]string)
	}
	if value == "" {
		// Delete a key/value pair in MongoDB.
		ops := []txn.Op{{
			C:      a.coll,
			Id:     a.id,
			Assert: isAliveDoc,
			Update: D{{"$unset", D{{"annotations." + key, true}}}},
		}}
		if err := a.st.runner.Run(ops, "", nil); err != nil {
			return fmt.Errorf("cannot delete annotation %q: %v", key, onAbort(err, errNotAlive))
		}
		delete(*a.annotations, key)
	} else {
		// Set a key/value pair in MongoDB.
		ops := []txn.Op{{
			C:      a.coll,
			Id:     a.id,
			Assert: isAliveDoc,
			Update: D{{"$set", D{{"annotations." + key, value}}}},
		}}
		if err := a.st.runner.Run(ops, "", nil); err != nil {
			return fmt.Errorf("cannot set annotation %q = %q: %v", key, value, onAbort(err, errNotAlive))
		}
		(*a.annotations)[key] = value
	}
	return nil
}

// Annotation returns the annotation value corresponding to the given key.
func (a annotator) Annotation(key string) string {
	return (*a.annotations)[key]
}