Merge pull request #20 from hashicorp/f-longest-prefix

Adds longest prefix matching for custom indexes.

txn.go

@@ -334,6 +334,39 @@ func (txn *Txn) First(table, index string, args ...interface{}) (interface{}, er
 	return value, nil
 }
 
+// LongestPrefix is used to fetch the longest prefix match for the given
+// constraints on the index. Note that this will not work with the memdb
+// StringFieldIndex because it adds null terminators which prevent the
+// algorithm from correctly finding a match (it will get to right before the
+// null and fail to find a leaf node). This should only be used where the prefix
+// given is capable of matching indexed entries directly, which typically only
+// applies to a custom indexer. See the unit test for an example.
+func (txn *Txn) LongestPrefix(table, index string, args ...interface{}) (interface{}, error) {
+	// Enforce that this only works on prefix indexes.
+	if !strings.HasSuffix(index, "_prefix") {
+		return nil, fmt.Errorf("must use '%s_prefix' on index", index)
+	}
+
+	// Get the index value.
+	indexSchema, val, err := txn.getIndexValue(table, index, args...)
+	if err != nil {
+		return nil, err
+	}
+
+	// This algorithm only makes sense against a unique index, otherwise the
+	// index keys will have the IDs appended to them.
+	if !indexSchema.Unique {
+		return nil, fmt.Errorf("index '%s' is not unique", index)
+	}
+
+	// Find the longest prefix match with the given index.
+	indexTxn := txn.readableIndex(table, indexSchema.Name)
+	if _, value, ok := indexTxn.Root().LongestPrefix(val); ok {
+		return value, nil
+	}
+	return nil, nil
+}
+
 // getIndexValue is used to get the IndexSchema and the value
 // used to scan the index given the parameters. This handles prefix based
 // scans when the index has the "_prefix" suffix. The index must support

txn_test.go

@@ -1,6 +1,10 @@
 package memdb
 
-import "testing"
+import (
+	"fmt"
+	"strings"
+	"testing"
+)
 
 func testDB(t *testing.T) *MemDB {
 	db, err := NewMemDB(testValidSchema())
@@ -585,6 +589,187 @@ func TestTxn_InsertGet_Prefix(t *testing.T) {
 	checkResult(txn)
 }
 
+// CustomIndex is a simple custom indexer that doesn't add any suffixes to its
+// object keys; this is compatible with the LongestPrefixMatch algorithm.
+type CustomIndex struct{}
+
+// FromObject takes the Foo field of a TestObject and prepends a null.
+func (*CustomIndex) FromObject(obj interface{}) (bool, []byte, error) {
+	t, ok := obj.(*TestObject)
+	if !ok {
+		return false, nil, fmt.Errorf("not a test object")
+	}
+
+	// Prepend a null so we can address an empty Foo field.
+	out := "\x00" + t.Foo
+	return true, []byte(out), nil
+}
+
+// FromArgs always returns an error.
+func (*CustomIndex) FromArgs(args ...interface{}) ([]byte, error) {
+	return nil, fmt.Errorf("only prefix lookups are supported")
+}
+
+// Prefix from args takes the argument as a string and prepends a null.
+func (*CustomIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
+	if len(args) != 1 {
+		return nil, fmt.Errorf("must provide only a single argument")
+	}
+	arg, ok := args[0].(string)
+	if !ok {
+		return nil, fmt.Errorf("argument must be a string: %#v", args[0])
+	}
+	arg = "\x00" + arg
+	return []byte(arg), nil
+}
+
+func TestTxn_InsertGet_LongestPrefix(t *testing.T) {
+	schema := &DBSchema{
+		Tables: map[string]*TableSchema{
+			"main": &TableSchema{
+				Name: "main",
+				Indexes: map[string]*IndexSchema{
+					"id": &IndexSchema{
+						Name:   "id",
+						Unique: true,
+						Indexer: &StringFieldIndex{
+							Field: "ID",
+						},
+					},
+					"foo": &IndexSchema{
+						Name:    "foo",
+						Unique:  true,
+						Indexer: &CustomIndex{},
+					},
+					"nope": &IndexSchema{
+						Name:    "nope",
+						Indexer: &CustomIndex{},
+					},
+				},
+			},
+		},
+	}
+
+	db, err := NewMemDB(schema)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	txn := db.Txn(true)
+
+	obj1 := &TestObject{
+		ID:  "object1",
+		Foo: "foo",
+	}
+	obj2 := &TestObject{
+		ID:  "object2",
+		Foo: "foozipzap",
+	}
+	obj3 := &TestObject{
+		ID:  "object3",
+		Foo: "",
+	}
+
+	err = txn.Insert("main", obj1)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	err = txn.Insert("main", obj2)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	err = txn.Insert("main", obj3)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	checkResult := func(txn *Txn) {
+		raw, err := txn.LongestPrefix("main", "foo_prefix", "foo")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj1 {
+			t.Fatalf("bad: %#v", raw)
+		}
+
+		raw, err = txn.LongestPrefix("main", "foo_prefix", "foobar")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj1 {
+			t.Fatalf("bad: %#v", raw)
+		}
+
+		raw, err = txn.LongestPrefix("main", "foo_prefix", "foozip")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj1 {
+			t.Fatalf("bad: %#v", raw)
+		}
+
+		raw, err = txn.LongestPrefix("main", "foo_prefix", "foozipza")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj1 {
+			t.Fatalf("bad: %#v", raw)
+		}
+
+		raw, err = txn.LongestPrefix("main", "foo_prefix", "foozipzap")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj2 {
+			t.Fatalf("bad: %#v", raw)
+		}
+
+		raw, err = txn.LongestPrefix("main", "foo_prefix", "foozipzapzone")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj2 {
+			t.Fatalf("bad: %#v", raw)
+		}
+
+		raw, err = txn.LongestPrefix("main", "foo_prefix", "funky")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj3 {
+			t.Fatalf("bad: %#v", raw)
+		}
+
+		raw, err = txn.LongestPrefix("main", "foo_prefix", "")
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if raw != obj3 {
+			t.Fatalf("bad: %#v", raw)
+		}
+	}
+
+	// Check the results within the txn
+	checkResult(txn)
+
+	// Commit and start a new read transaction
+	txn.Commit()
+	txn = db.Txn(false)
+
+	// Check the results in a new txn
+	checkResult(txn)
+
+	// Try some disallowed index types.
+	_, err = txn.LongestPrefix("main", "foo", "")
+	if err == nil || !strings.Contains(err.Error(), "must use 'foo_prefix' on index") {
+		t.Fatalf("bad: %v", err)
+	}
+	_, err = txn.LongestPrefix("main", "nope_prefix", "")
+	if err == nil || !strings.Contains(err.Error(), "index 'nope_prefix' is not unique") {
+		t.Fatalf("bad: %v", err)
+	}
+}
+
 func TestTxn_Defer(t *testing.T) {
 	db := testDB(t)
 	txn := db.Txn(true)