Merge "Adding Remove function to Mongo"

4 files changed, 50 insertions, 4 deletions

diff --git a/src/orchestrator/pkg/infra/db/README.md b/src/orchestrator/pkg/infra/db/README.md
index ff482b98..71da1e0a 100644
--- a/src/orchestrator/pkg/infra/db/README.md
+++ b/src/orchestrator/pkg/infra/db/README.md
@@ -130,7 +130,7 @@ key := CompositeAppKey{
 
 NOTE: Key structure can be different from the original key and can include Query fields also. ANY operation is not supported for Query fields.
 
-### Remove
+### RemoveAll
 
 Arguments:
 ```go
@@ -139,6 +139,15 @@ key interface
 ```
 Similar to find. This will remove one or more documents based on the key structure.
 
+### Remove
+
+Arguments:
+```go
+collection string
+key interface
+```
+This will remove one document based on the key structure. If child refrences exist for the key then the document will not be removed.
+
 ### Unmarshal
 
 Data in mongo is stored as `bson` which is a compressed form of `json`. We need mongo to convert the stored `bson` data to regular `json`
@@ -147,3 +156,4 @@ that we can use in our code when returned.
 `bson.Unmarshal` API is used to achieve this.
 
 
+
diff --git a/src/orchestrator/pkg/infra/db/mongo.go b/src/orchestrator/pkg/infra/db/mongo.go
index 30eb899f..a344aa1c 100644
--- a/src/orchestrator/pkg/infra/db/mongo.go
+++ b/src/orchestrator/pkg/infra/db/mongo.go
@@ -49,6 +49,8 @@ type MongoCollection interface {
 		opts ...*options.FindOptions) (*mongo.Cursor, error)
 	UpdateOne(ctx context.Context, filter interface{}, update interface{},
 		opts ...*options.UpdateOptions) (*mongo.UpdateResult, error)
+        CountDocuments(ctx context.Context, filter interface{},
+		opts ...*options.CountOptions) (int64, error)
 }
 
 // MongoStore is an implementation of the db.Store interface
@@ -543,8 +545,8 @@ func (m *MongoStore) Find(coll string, key Key, tag string) ([][]byte, error) {
 	return result, nil
 }
 
-// Remove method to remove the documet by key
-func (m *MongoStore) Remove(coll string, key Key) error {
+// RemoveAll method to removes all the documet matching key
+func (m *MongoStore) RemoveAll(coll string, key Key) error {
 	if !m.validateParams(coll, key) {
 		return pkgerrors.New("Mandatory fields are missing")
 	}
@@ -560,3 +562,29 @@ func (m *MongoStore) Remove(coll string, key Key) error {
 	}
 	return nil
 }
+
+// Remove method to remove the documet by key if no child references
+func (m *MongoStore) Remove(coll string, key Key) error {
+	if !m.validateParams(coll, key) {
+		return pkgerrors.New("Mandatory fields are missing")
+	}
+	c := getCollection(coll, m)
+	ctx := context.Background()
+	filter, err := m.findFilter(key)
+	if err != nil {
+		return err
+	}
+	count, err := c.CountDocuments(context.Background(), filter)
+	if err != nil {
+		return pkgerrors.Errorf("Error finding: %s", err.Error())
+	}
+	if count > 1 {
+		return pkgerrors.Errorf("Can't delete parent without deleting child references first")
+	}
+	_, err = c.DeleteOne(ctx, filter)
+	if err != nil {
+		return pkgerrors.Errorf("Error Deleting from database: %s", err.Error())
+	}
+	return nil
+}
+
diff --git a/src/orchestrator/pkg/infra/db/mongo_test.go b/src/orchestrator/pkg/infra/db/mongo_test.go
index 9f813ca4..d57c19dd 100644
--- a/src/orchestrator/pkg/infra/db/mongo_test.go
+++ b/src/orchestrator/pkg/infra/db/mongo_test.go
@@ -79,6 +79,11 @@ func (c *mockCollection) UpdateOne(ctx context.Context, filter interface{}, upda
 	return nil, c.Err
 }
 
+func (c *mockCollection) CountDocuments(ctx context.Context, filter interface{},
+	opts ...*options.CountOptions) (int64, error) {
+	return 1, c.Err
+}
+
 func TestCreate(t *testing.T) {
 	testCases := []struct {
 		label         string
diff --git a/src/orchestrator/pkg/infra/db/store.go b/src/orchestrator/pkg/infra/db/store.go
index f8cb4466..e87722cd 100644
--- a/src/orchestrator/pkg/infra/db/store.go
+++ b/src/orchestrator/pkg/infra/db/store.go
@@ -59,8 +59,11 @@ type Store interface {
 	// Find the document(s) with key and get the tag values from the document(s)
 	Find(coll string, key Key, tag string) ([][]byte, error)
 
-	// Removes the document(s) matching the key
+        // Removes the document(s) matching the key if no child reference in collection
 	Remove(coll string, key Key) error
+
+	// Remove all the document(s) matching the key
+	RemoveAll(coll string, key Key) error
 }
 
 // CreateDBClient creates the DB client