diff --git a/nomad/memdb/schema_test.go b/nomad/memdb/schema_test.go index 139a00e00..07c4dfc36 100644 --- a/nomad/memdb/schema_test.go +++ b/nomad/memdb/schema_test.go @@ -13,6 +13,10 @@ func testValidSchema() *DBSchema { Unique: true, Indexer: &StringFieldIndex{Field: "ID"}, }, + "foo": &IndexSchema{ + Name: "foo", + Indexer: &StringFieldIndex{Field: "Foo"}, + }, }, }, }, diff --git a/nomad/memdb/txn.go b/nomad/memdb/txn.go index ce8d56919..942121934 100644 --- a/nomad/memdb/txn.go +++ b/nomad/memdb/txn.go @@ -130,7 +130,7 @@ func (txn *Txn) Insert(table string, obj interface{}) error { // Get the primary ID of the object idSchema := tableSchema.Indexes["id"] - ok, val, err := idSchema.Indexer.FromObject(obj) + ok, idVal, err := idSchema.Indexer.FromObject(obj) if err != nil { return fmt.Errorf("failed to build primary index: %v", err) } @@ -140,7 +140,7 @@ func (txn *Txn) Insert(table string, obj interface{}) error { // Lookup the object by ID first, to see if this is an update idTxn := txn.writableIndex(table, "id") - existing, update := idTxn.Get(val) + existing, update := idTxn.Get(idVal) // On an update, there is an existing object with the given // primary ID. We do the update by deleting the current object @@ -155,7 +155,12 @@ func (txn *Txn) Insert(table string, obj interface{}) error { return fmt.Errorf("failed to build index '%s': %v", name, err) } if ok { - // TODO: Handle non-unique index + // Handle non-unique index by computing a unique index. + // This is done by appending the primary key which must + // be unique anyways. + if !indexSchema.Unique { + val = append(val, idVal...) + } indexTxn.Delete(val) } } @@ -173,7 +178,12 @@ func (txn *Txn) Insert(table string, obj interface{}) error { } } - // TODO: Handle non-unique index + // Handle non-unique index by computing a unique index. + // This is done by appending the primary key which must + // be unique anyways. + if !indexSchema.Unique { + val = append(val, idVal...) + } indexTxn.Insert(val, obj) } return nil @@ -209,13 +219,24 @@ func (txn *Txn) First(table, index string, args ...interface{}) (interface{}, er indexTxn := txn.readableIndex(table, index) // Do an exact lookup - obj, ok := indexTxn.Get(val) - if !ok { - return nil, nil + if indexSchema.Unique { + obj, ok := indexTxn.Get(val) + if !ok { + return nil, nil + } + return obj, nil } - // TODO: handle non-unique index - return obj, nil + // Handle non-unique index by doing a prefix walk + // and getting the first value + // TODO: Optimize this + var firstVal interface{} + indexRoot := indexTxn.Root() + indexRoot.WalkPrefix(val, func(key []byte, val interface{}) bool { + firstVal = val + return true + }) + return firstVal, nil } type ResultIterator interface { diff --git a/nomad/memdb/txn_test.go b/nomad/memdb/txn_test.go index 759b952cc..b08482fb8 100644 --- a/nomad/memdb/txn_test.go +++ b/nomad/memdb/txn_test.go @@ -98,3 +98,105 @@ func TestTxn_InsertUpdate_First(t *testing.T) { t.Fatalf("bad: %#v %#v", raw, obj) } } + +func TestTxn_InsertUpdate_First_NonUnique(t *testing.T) { + db := testDB(t) + txn := db.Txn(true) + + obj := &TestObject{ + ID: "my-object", + Foo: "abc", + } + err := txn.Insert("main", obj) + if err != nil { + t.Fatalf("err: %v", err) + } + + raw, err := txn.First("main", "foo", obj.Foo) + if err != nil { + t.Fatalf("err: %v", err) + } + + if raw != obj { + t.Fatalf("bad: %#v %#v", raw, obj) + } + + // Update the object + obj2 := &TestObject{ + ID: "my-object", + Foo: "xyz", + } + err = txn.Insert("main", obj2) + if err != nil { + t.Fatalf("err: %v", err) + } + + raw, err = txn.First("main", "foo", obj2.Foo) + if err != nil { + t.Fatalf("err: %v", err) + } + + if raw != obj2 { + t.Fatalf("bad: %#v %#v", raw, obj2) + } + + // Lookup of the old value should fail + raw, err = txn.First("main", "foo", obj.Foo) + if err != nil { + t.Fatalf("err: %v", err) + } + + if raw != nil { + t.Fatalf("bad: %#v", raw) + } +} + +func TestTxn_First_NonUnique_Multiple(t *testing.T) { + db := testDB(t) + txn := db.Txn(true) + + obj := &TestObject{ + ID: "my-object", + Foo: "abc", + } + obj2 := &TestObject{ + ID: "my-cool-thing", + Foo: "xyz", + } + obj3 := &TestObject{ + ID: "my-other-cool-thing", + Foo: "xyz", + } + + err := txn.Insert("main", obj) + 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) + } + + // The first object has a unique secondary value + raw, err := txn.First("main", "foo", obj.Foo) + if err != nil { + t.Fatalf("err: %v", err) + } + if raw != obj { + t.Fatalf("bad: %#v %#v", raw, obj) + } + + // Second and third object share secondary value, + // but the primary ID of obj2 should be first + raw, err = txn.First("main", "foo", obj2.Foo) + if err != nil { + t.Fatalf("err: %v", err) + } + if raw != obj2 { + t.Fatalf("bad: %#v %#v", raw, obj2) + } +}