memdb: testing compound index

nomad/memdb/index.go

@@ -174,3 +174,55 @@ func (u *UUIDFieldIndex) parseString(s string) ([]byte, error) {
 	copy(buf[10:16], part5)
 	return buf, nil
 }
+
+// CompoundIndex is used to build an index using multiple sub-indexes
+// Prefix based iteration is supported as long as the appropriate prefix
+// of indexers support it. All sub-indexers are only assumed to expect
+// a single argument.
+type CompoundIndex struct {
+	Indexes []Indexer
+
+	// AllowMissing results in an index based on only the indexers
+	// that return data. If true, you may end up with 2/3 columns
+	// indexed which might be useful for an index scan. Otherwise,
+	// the CompoundIndex requires all indexers to be satisfied.
+	AllowMissing bool
+}
+
+func (c *CompoundIndex) FromObject(raw interface{}) (bool, []byte, error) {
+	var out []byte
+	for i, idx := range c.Indexes {
+		ok, val, err := idx.FromObject(raw)
+		if err != nil {
+			return false, nil, fmt.Errorf("sub-index %d error: %v", i, err)
+		}
+		if !ok {
+			if c.AllowMissing {
+				break
+			} else {
+				return false, nil, nil
+			}
+		}
+		out = append(out, val...)
+	}
+	return true, out, nil
+}
+
+func (c *CompoundIndex) FromArgs(args ...interface{}) ([]byte, error) {
+	if len(args) != len(c.Indexes) {
+		return nil, fmt.Errorf("less arguments than index fields")
+	}
+	return c.PrefixFromArgs(args...)
+}
+
+func (c *CompoundIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
+	var out []byte
+	for i, arg := range args {
+		val, err := c.Indexes[i].FromArgs(arg)
+		if err != nil {
+			return nil, fmt.Errorf("sub-index %d error: %v", i, err)
+		}
+		out = append(out, val...)
+	}
+	return out, nil
+}

nomad/memdb/index_test.go

@@ -224,3 +224,113 @@ func generateUUID() ([]byte, string) {
 		buf[10:16])
 	return buf, uuid
 }
+
+func TestCompoundIndex_FromObject(t *testing.T) {
+	obj := testObj()
+	indexer := &CompoundIndex{
+		Indexes: []Indexer{
+			&StringFieldIndex{"ID", false},
+			&StringFieldIndex{"Foo", false},
+			&StringFieldIndex{"Baz", false},
+		},
+		AllowMissing: false,
+	}
+
+	ok, val, err := indexer.FromObject(obj)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if string(val) != "my-cool-obj\x00Testing\x00yep\x00" {
+		t.Fatalf("bad: %s", val)
+	}
+	if !ok {
+		t.Fatalf("should be ok")
+	}
+
+	missing := &CompoundIndex{
+		Indexes: []Indexer{
+			&StringFieldIndex{"ID", false},
+			&StringFieldIndex{"Foo", true},
+			&StringFieldIndex{"Empty", false},
+		},
+		AllowMissing: true,
+	}
+	ok, val, err = missing.FromObject(obj)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if string(val) != "my-cool-obj\x00testing\x00" {
+		t.Fatalf("bad: %s", val)
+	}
+	if !ok {
+		t.Fatalf("should be ok")
+	}
+
+	// Test when missing not allowed
+	missing.AllowMissing = false
+	ok, _, err = missing.FromObject(obj)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if ok {
+		t.Fatalf("should not be okay")
+	}
+}
+
+func TestCompoundIndex_FromArgs(t *testing.T) {
+	indexer := &CompoundIndex{
+		Indexes: []Indexer{
+			&StringFieldIndex{"ID", false},
+			&StringFieldIndex{"Foo", false},
+			&StringFieldIndex{"Baz", false},
+		},
+		AllowMissing: false,
+	}
+	_, err := indexer.FromArgs()
+	if err == nil {
+		t.Fatalf("should get err")
+	}
+
+	_, err = indexer.FromArgs(42, 42, 42)
+	if err == nil {
+		t.Fatalf("should get err")
+	}
+
+	val, err := indexer.FromArgs("foo", "bar", "baz")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if string(val) != "foo\x00bar\x00baz\x00" {
+		t.Fatalf("bad: %s", val)
+	}
+}
+
+func TestCompoundIndex_PrefixFromArgs(t *testing.T) {
+	indexer := &CompoundIndex{
+		Indexes: []Indexer{
+			&UUIDFieldIndex{"ID"},
+			&StringFieldIndex{"Foo", false},
+			&StringFieldIndex{"Baz", false},
+		},
+		AllowMissing: false,
+	}
+	val, err := indexer.PrefixFromArgs()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if len(val) != 0 {
+		t.Fatalf("bad: %s", val)
+	}
+
+	uuidBuf, uuid := generateUUID()
+	val, err = indexer.PrefixFromArgs(uuid, "foo")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if !bytes.Equal(val[:16], uuidBuf) {
+		t.Fatalf("bad prefix")
+	}
+	if string(val[16:]) != "foo\x00" {
+		t.Fatalf("bad: %s", val)
+	}
+}