From 63319b84feab0ca23b946bf613fffb6b8cca0af6 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Tue, 7 Jan 2014 17:28:15 -0800
Subject: [PATCH] Support blank fields

---
 consul/mdb_table.go      |  22 ++++++---
 consul/mdb_table_test.go | 103 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 119 insertions(+), 6 deletions(-)

diff --git a/consul/mdb_table.go b/consul/mdb_table.go
index f7e27b4c65..0968d1aa41 100644
--- a/consul/mdb_table.go
+++ b/consul/mdb_table.go
@@ -32,9 +32,10 @@ type MDBTable struct {
 // An Index is named, and uses a series of column values to
 // map to the row-id containing the table
 type MDBIndex struct {
-	Unique  bool      // Controls if values are unique
-	Fields  []string  // Fields are used to build the index
-	IdxFunc IndexFunc // Can be used to provide custom indexing
+	AllowBlank bool      // Can fields be blank
+	Unique     bool      // Controls if values are unique
+	Fields     []string  // Fields are used to build the index
+	IdxFunc    IndexFunc // Can be used to provide custom indexing
 
 	table   *MDBTable
 	name    string
@@ -62,9 +63,11 @@ type RowID uint64
 type IndexFunc func(...string) string
 
 // DefaultIndexFunc is used if no IdxFunc is provided. It joins
-// the columns using '||' which is reasonably unlikely to occur
+// the columns using '||' which is reasonably unlikely to occur.
+// We also prefix with a byte to ensure we never have a zero length
+// key
 func DefaultIndexFunc(parts ...string) string {
-	return strings.Join(parts, "||")
+	return "_" + strings.Join(parts, "||")
 }
 
 // Init is used to initialize the MDBTable and ensure it's ready
@@ -87,6 +90,9 @@ func (t *MDBTable) Init() error {
 	if !id.Unique {
 		return fmt.Errorf("id index must be unique")
 	}
+	if id.AllowBlank {
+		return fmt.Errorf("id index must not allow blanks")
+	}
 
 	// Create the table
 	if err := t.createTable(); err != nil {
@@ -332,7 +338,11 @@ func (i *MDBIndex) keyFromObject(obj interface{}) ([]byte, error) {
 		if !fv.IsValid() {
 			return nil, fmt.Errorf("Field '%s' for %#v is invalid", field, obj)
 		}
-		parts = append(parts, fv.String())
+		val := fv.String()
+		if !i.AllowBlank && val == "" {
+			return nil, fmt.Errorf("Field '%s' must be set: %#v", field, obj)
+		}
+		parts = append(parts, val)
 	}
 	key := i.IdxFunc(parts...)
 	return []byte(key), nil
diff --git a/consul/mdb_table_test.go b/consul/mdb_table_test.go
index 5918b7eb6a..95b270d34d 100644
--- a/consul/mdb_table_test.go
+++ b/consul/mdb_table_test.go
@@ -162,3 +162,106 @@ func TestMDBTableInsert(t *testing.T) {
 		t.Fatalf("bad: %#v", res[2])
 	}
 }
+
+func TestMDBTableInsert_MissingFields(t *testing.T) {
+	dir, env := testMDBEnv(t)
+	defer os.RemoveAll(dir)
+	defer env.Close()
+
+	table := &MDBTable{
+		Env:  env,
+		Name: "test",
+		Indexes: map[string]*MDBIndex{
+			"id": &MDBIndex{
+				Unique: true,
+				Fields: []string{"Key"},
+			},
+			"name": &MDBIndex{
+				Fields: []string{"First", "Last"},
+			},
+			"country": &MDBIndex{
+				Fields: []string{"Country"},
+			},
+		},
+		Encoder: MockEncoder,
+		Decoder: MockDecoder,
+	}
+	if err := table.Init(); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	objs := []*MockData{
+		&MockData{
+			First:   "Kevin",
+			Last:    "Smith",
+			Country: "USA",
+		},
+		&MockData{
+			Key:     "2",
+			Last:    "Wang",
+			Country: "USA",
+		},
+		&MockData{
+			Key:     "2",
+			First:   "Kevin",
+			Country: "USA",
+		},
+		&MockData{
+			Key:   "3",
+			First: "Bernardo",
+			Last:  "Torres",
+		},
+	}
+
+	// Insert some mock objects
+	for _, obj := range objs {
+		if err := table.Insert(obj); err == nil {
+			t.Fatalf("expected err")
+		}
+	}
+}
+
+func TestMDBTableInsert_AllowBlank(t *testing.T) {
+	dir, env := testMDBEnv(t)
+	defer os.RemoveAll(dir)
+	defer env.Close()
+
+	table := &MDBTable{
+		Env:  env,
+		Name: "test",
+		Indexes: map[string]*MDBIndex{
+			"id": &MDBIndex{
+				Unique: true,
+				Fields: []string{"Key"},
+			},
+			"name": &MDBIndex{
+				Fields: []string{"First", "Last"},
+			},
+			"country": &MDBIndex{
+				Fields:     []string{"Country"},
+				AllowBlank: true,
+			},
+		},
+		Encoder: MockEncoder,
+		Decoder: MockDecoder,
+	}
+	if err := table.Init(); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	objs := []*MockData{
+		&MockData{
+			Key:     "1",
+			First:   "Kevin",
+			Last:    "Smith",
+			Country: "",
+		},
+	}
+
+	// Insert some mock objects
+	for _, obj := range objs {
+		if err := table.Insert(obj); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+	}
+}