Merge "Add update method to db interface"

5 files changed, 132 insertions, 1 deletions

diff --git a/src/k8splugin/internal/db/consul.go b/src/k8splugin/internal/db/consul.go
index 23d2ae88..0977cfba 100644
--- a/src/k8splugin/internal/db/consul.go
+++ b/src/k8splugin/internal/db/consul.go
@@ -88,6 +88,11 @@ func (c *ConsulStore) Create(root string, key Key, tag string, data interface{})
 	return err
 }
 
+// Update is used to update a DB entry
+func (c *ConsulStore) Update(root string, key Key, tag string, data interface{}) error {
+	return c.Create(root, key, tag, data)
+}
+
 // Read method returns the internalID for a particular externalID
 func (c *ConsulStore) Read(root string, key Key, tag string) ([]byte, error) {
 
diff --git a/src/k8splugin/internal/db/mongo.go b/src/k8splugin/internal/db/mongo.go
index 8c422380..a9e9d98a 100644
--- a/src/k8splugin/internal/db/mongo.go
+++ b/src/k8splugin/internal/db/mongo.go
@@ -164,6 +164,49 @@ func (m *MongoStore) Create(coll string, key Key, tag string, data interface{})
 	return nil
 }
 
+// Update is used to update a DB entry
+func (m *MongoStore) Update(coll string, key Key, tag string, data interface{}) error {
+	if data == nil || !m.validateParams(coll, key, tag) {
+		return pkgerrors.New("No Data to update")
+	}
+
+	c := getCollection(coll, m)
+	ctx := context.Background()
+
+	//Get the masterkey document based on given key
+	filter := bson.D{{"key", key}}
+	keydata, err := decodeBytes(c.FindOne(context.Background(), filter))
+	if err != nil {
+		return pkgerrors.Errorf("Error finding master table: %s", err.Error())
+	}
+
+	//Read the tag objectID from document
+	tagoid, ok := keydata.Lookup(tag).ObjectIDOK()
+	if !ok {
+		return pkgerrors.Errorf("Error finding objectID for tag %s", tag)
+	}
+
+	//Update the document with new data
+	filter = bson.D{{"_id", tagoid}}
+
+	_, err = decodeBytes(
+		c.FindOneAndUpdate(
+			ctx,
+			filter,
+			bson.D{
+				{"$set", bson.D{
+					{tag, data},
+				}},
+			},
+			options.FindOneAndUpdate().SetReturnDocument(options.After)))
+
+	if err != nil {
+		return pkgerrors.Errorf("Error updating record: %s", err.Error())
+	}
+
+	return nil
+}
+
 // Unmarshal implements an unmarshaler for bson data that
 // is produced from the mongo database
 func (m *MongoStore) Unmarshal(inp []byte, out interface{}) error {
diff --git a/src/k8splugin/internal/db/mongo_test.go b/src/k8splugin/internal/db/mongo_test.go
index baa442ce..3a7e4251 100644
--- a/src/k8splugin/internal/db/mongo_test.go
+++ b/src/k8splugin/internal/db/mongo_test.go
@@ -143,6 +143,84 @@ func TestCreate(t *testing.T) {
 	}
 }
 
+func TestUpdate(t *testing.T) {
+	testCases := []struct {
+		label         string
+		input         map[string]interface{}
+		mockColl      *mockCollection
+		bson          bson.Raw
+		expectedError string
+	}{
+		{
+			label: "Successfull update of entry",
+			input: map[string]interface{}{
+				"coll": "collname",
+				"key":  mockKey{Key: "keyvalue"},
+				"tag":  "metadata",
+				"data": "Data In String Format",
+			},
+			// Binary form of
+			// {
+			//	"_id" : ObjectId("5c115156777ff85654248ae1"),
+			//  "key" : bson.D{{"name","testdef"},{"version","v1"}},
+			//  "metadata" : ObjectId("5c115156c9755047e318bbfd")
+			// }
+			bson: bson.Raw{
+				'\x58', '\x00', '\x00', '\x00', '\x03', '\x6b', '\x65', '\x79',
+				'\x00', '\x27', '\x00', '\x00', '\x00', '\x02', '\x6e', '\x61',
+				'\x6d', '\x65', '\x00', '\x08', '\x00', '\x00', '\x00', '\x74',
+				'\x65', '\x73', '\x74', '\x64', '\x65', '\x66', '\x00', '\x02',
+				'\x76', '\x65', '\x72', '\x73', '\x69', '\x6f', '\x6e', '\x00',
+				'\x03', '\x00', '\x00', '\x00', '\x76', '\x31', '\x00', '\x00',
+				'\x07', '\x6d', '\x65', '\x74', '\x61', '\x64', '\x61', '\x74',
+				'\x61', '\x00', '\x5c', '\x11', '\x51', '\x56', '\x77', '\x7f',
+				'\xf8', '\x56', '\x54', '\x24', '\x8a', '\xe1', '\x07', '\x5f',
+				'\x69', '\x64', '\x00', '\x5c', '\x11', '\x51', '\x56', '\x77',
+				'\x7f', '\xf8', '\x56', '\x54', '\x24', '\x8a', '\xe1', '\x00',
+			},
+			mockColl: &mockCollection{},
+		},
+		{
+			label: "Entry does not exist",
+			input: map[string]interface{}{
+				"coll": "collname",
+				"key":  mockKey{Key: "keyvalue"},
+				"tag":  "tagName",
+				"data": "Data In String Format",
+			},
+			mockColl: &mockCollection{
+				Err: pkgerrors.New("DB Error"),
+			},
+			expectedError: "DB Error",
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.label, func(t *testing.T) {
+			m, _ := NewMongoStore("name", &mongo.Database{})
+			// Override the getCollection function with our mocked version
+			getCollection = func(coll string, m *MongoStore) MongoCollection {
+				return testCase.mockColl
+			}
+
+			decodeBytes = func(sr *mongo.SingleResult) (bson.Raw, error) {
+				return testCase.bson, testCase.mockColl.Err
+			}
+
+			err := m.Update(testCase.input["coll"].(string), testCase.input["key"].(Key),
+				testCase.input["tag"].(string), testCase.input["data"])
+			if err != nil {
+				if testCase.expectedError == "" {
+					t.Fatalf("Create method returned an un-expected (%s)", err)
+				}
+				if !strings.Contains(string(err.Error()), testCase.expectedError) {
+					t.Fatalf("Create method returned an error (%s)", err)
+				}
+			}
+		})
+	}
+}
+
 func TestRead(t *testing.T) {
 	testCases := []struct {
 		label         string
diff --git a/src/k8splugin/internal/db/store.go b/src/k8splugin/internal/db/store.go
index 148e078e..0b869c0c 100644
--- a/src/k8splugin/internal/db/store.go
+++ b/src/k8splugin/internal/db/store.go
@@ -45,7 +45,8 @@ type Store interface {
 	// Reads data for a particular key with specific tag.
 	Read(table string, key Key, tag string) ([]byte, error)
 
-	//TODO: Update(context.Context, string, interface{}) error
+	// Update data for particular key with specific tag
+	Update(table string, key Key, tag string, data interface{}) error
 
 	// Deletes a specific tag data for key.
 	// TODO: If tag is empty, it will delete all tags under key.
diff --git a/src/k8splugin/internal/db/testing.go b/src/k8splugin/internal/db/testing.go
index e7e7436a..ded5067c 100644
--- a/src/k8splugin/internal/db/testing.go
+++ b/src/k8splugin/internal/db/testing.go
@@ -41,6 +41,10 @@ func (m *MockDB) Create(table string, key Key, tag string, data interface{}) err
 	return m.Err
 }
 
+func (m *MockDB) Update(table string, key Key, tag string, data interface{}) error {
+	return m.Err
+}
+
 // MockDB uses simple JSON and not BSON
 func (m *MockDB) Unmarshal(inp []byte, out interface{}) error {
 	err := json.Unmarshal(inp, out)