4 "golang.org/x/net/context"
7 "bpa-restapi-agent/internal/config"
9 pkgerrors "github.com/pkg/errors"
10 "go.mongodb.org/mongo-driver/bson"
11 "go.mongodb.org/mongo-driver/bson/primitive"
12 "go.mongodb.org/mongo-driver/mongo"
13 "go.mongodb.org/mongo-driver/mongo/options"
16 // MongoCollection defines the a subset of MongoDB operations
17 // Note: This interface is defined mainly for mock testing
18 type MongoCollection interface {
19 InsertOne(ctx context.Context, document interface{},
20 opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error)
21 FindOne(ctx context.Context, filter interface{},
22 opts ...*options.FindOneOptions) *mongo.SingleResult
23 FindOneAndUpdate(ctx context.Context, filter interface{},
24 update interface{}, opts ...*options.FindOneAndUpdateOptions) *mongo.SingleResult
25 DeleteOne(ctx context.Context, filter interface{},
26 opts ...*options.DeleteOptions) (*mongo.DeleteResult, error)
27 Find(ctx context.Context, filter interface{},
28 opts ...*options.FindOptions) (*mongo.Cursor, error)
31 // MongoStore is an implementation of the db.Store interface
32 type MongoStore struct {
36 // This exists only for allowing us to mock the collection object
37 // for testing purposes
38 var getCollection = func(coll string, m *MongoStore) MongoCollection {
39 return m.db.Collection(coll)
42 // This exists only for allowing us to mock the DecodeBytes function
43 // Mainly because we cannot construct a SingleResult struct from our
44 // tests. All fields in that struct are private.
45 var decodeBytes = func(sr *mongo.SingleResult) (bson.Raw, error) {
46 return sr.DecodeBytes()
49 // These exists only for allowing us to mock the cursor.Next function
50 // Mainly because we cannot construct a mongo.Cursor struct from our
51 // tests. All fields in that struct are private and there is no public
52 // constructor method.
53 var cursorNext = func(ctx context.Context, cursor *mongo.Cursor) bool {
54 return cursor.Next(ctx)
56 var cursorClose = func(ctx context.Context, cursor *mongo.Cursor) error {
57 return cursor.Close(ctx)
60 // NewMongoStore initializes a Mongo Database with the name provided
61 // If a database with that name exists, it will be returned
62 func NewMongoStore(name string, store *mongo.Database) (Store, error) {
64 ip := "mongodb://" + config.GetConfiguration().DatabaseAddress + ":27017"
65 clientOptions := options.Client()
66 clientOptions.ApplyURI(ip)
67 mongoClient, err := mongo.NewClient(clientOptions)
72 err = mongoClient.Connect(context.Background())
76 store = mongoClient.Database(name)
84 // HealthCheck verifies if the database is up and running
85 func (m *MongoStore) HealthCheck() error {
87 _, err := decodeBytes(m.db.RunCommand(context.Background(), bson.D{{"serverStatus", 1}}))
89 return pkgerrors.Wrap(err, "Error getting server status")
95 // validateParams checks to see if any parameters are empty
96 func (m *MongoStore) validateParams(args ...interface{}) bool {
97 for _, v := range args {
113 // Create is used to create a DB entry
114 func (m *MongoStore) Create(coll string, key Key, tag string, data interface{}) error {
115 if data == nil || !m.validateParams(coll, key, tag) {
116 return pkgerrors.New("No Data to store")
119 c := getCollection(coll, m)
120 ctx := context.Background()
122 //Insert the data and then add the objectID to the masterTable
123 res, err := c.InsertOne(ctx, bson.D{
127 return pkgerrors.Errorf("Error inserting into database: %s", err.Error())
130 //Add objectID of created data to masterKey document
131 //Create masterkey document if it does not exist
132 filter := bson.D{{"key", key}}
134 _, err = decodeBytes(
140 {tag, res.InsertedID},
143 options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)))
146 return pkgerrors.Errorf("Error updating master table: %s", err.Error())
152 // Update is used to update a DB entry
153 func (m *MongoStore) Update(coll string, key Key, tag string, data interface{}) error {
154 if data == nil || !m.validateParams(coll, key, tag) {
155 return pkgerrors.New("No Data to update")
158 c := getCollection(coll, m)
159 ctx := context.Background()
161 //Get the masterkey document based on given key
162 filter := bson.D{{"key", key}}
163 keydata, err := decodeBytes(c.FindOne(context.Background(), filter))
165 return pkgerrors.Errorf("Error finding master table: %s", err.Error())
168 //Read the tag objectID from document
169 tagoid, ok := keydata.Lookup(tag).ObjectIDOK()
171 return pkgerrors.Errorf("Error finding objectID for tag %s", tag)
174 //Update the document with new data
175 filter = bson.D{{"_id", tagoid}}
177 _, err = decodeBytes(
186 options.FindOneAndUpdate().SetReturnDocument(options.After)))
189 return pkgerrors.Errorf("Error updating record: %s", err.Error())
195 // Unmarshal implements an unmarshaler for bson data that
196 // is produced from the mongo database
197 func (m *MongoStore) Unmarshal(inp []byte, out interface{}) error {
198 err := bson.Unmarshal(inp, out)
200 return pkgerrors.Wrap(err, "Unmarshaling bson")
205 // Read method returns the data stored for this key and for this particular tag
206 func (m *MongoStore) Read(coll string, key Key, tag string) ([]byte, error) {
207 if !m.validateParams(coll, key, tag) {
208 return nil, pkgerrors.New("Mandatory fields are missing")
211 c := getCollection(coll, m)
212 ctx := context.Background()
214 //Get the masterkey document based on given key
215 filter := bson.D{{"key", key}}
216 keydata, err := decodeBytes(c.FindOne(context.Background(), filter))
218 return nil, pkgerrors.Errorf("Error finding master table: %s", err.Error())
221 //Read the tag objectID from document
222 tagoid, ok := keydata.Lookup(tag).ObjectIDOK()
224 return nil, pkgerrors.Errorf("Error finding objectID for tag %s", tag)
227 //Use tag objectID to read the data from store
228 filter = bson.D{{"_id", tagoid}}
229 tagdata, err := decodeBytes(c.FindOne(ctx, filter))
231 return nil, pkgerrors.Errorf("Error reading found object: %s", err.Error())
234 //Return the data as a byte array
235 //Convert string data to byte array using the built-in functions
236 switch tagdata.Lookup(tag).Type {
237 case bson.TypeString:
238 return []byte(tagdata.Lookup(tag).StringValue()), nil
240 return tagdata.Lookup(tag).Value, nil
244 // Helper function that deletes an object by its ID
245 func (m *MongoStore) deleteObjectByID(coll string, objID primitive.ObjectID) error {
247 c := getCollection(coll, m)
248 ctx := context.Background()
250 _, err := c.DeleteOne(ctx, bson.D{{"_id", objID}})
252 return pkgerrors.Errorf("Error Deleting from database: %s", err.Error())
255 log.Printf("Deleted Obj with ID %s", objID.String())
259 // Delete method removes a document from the Database that matches key
260 // TODO: delete all referenced docs if tag is empty string
261 func (m *MongoStore) Delete(coll string, key Key, tag string) error {
262 if !m.validateParams(coll, key, tag) {
263 return pkgerrors.New("Mandatory fields are missing")
266 c := getCollection(coll, m)
267 ctx := context.Background()
269 //Get the masterkey document based on given key
270 filter := bson.D{{"key", key}}
271 //Remove the tag ID entry from masterkey table
279 keydata, err := decodeBytes(c.FindOneAndUpdate(ctx, filter, update,
280 options.FindOneAndUpdate().SetReturnDocument(options.Before)))
282 //No document was found. Return nil.
283 if err == mongo.ErrNoDocuments {
286 //Return any other error that was found.
287 return pkgerrors.Errorf("Error decoding master table after update: %s",
291 //Read the tag objectID from document
292 elems, err := keydata.Elements()
294 return pkgerrors.Errorf("Error reading elements from database: %s", err.Error())
297 tagoid, ok := keydata.Lookup(tag).ObjectIDOK()
299 return pkgerrors.Errorf("Error finding objectID for tag %s", tag)
302 //Use tag objectID to read the data from store
303 err = m.deleteObjectByID(coll, tagoid)
305 return pkgerrors.Errorf("Error deleting from database: %s", err.Error())
308 //Delete master table if no more tags left
309 //_id, key and tag should be elements in before doc
310 //if master table needs to be removed too
312 keyid, ok := keydata.Lookup("_id").ObjectIDOK()
314 return pkgerrors.Errorf("Error finding objectID for key %s", key)
316 err = m.deleteObjectByID(coll, keyid)
318 return pkgerrors.Errorf("Error deleting master table from database: %s", err.Error())
325 // ReadAll is used to get all documents in db of a particular tag
326 func (m *MongoStore) ReadAll(coll, tag string) (map[string][]byte, error) {
327 if !m.validateParams(coll, tag) {
328 return nil, pkgerrors.New("Missing collection or tag name")
331 c := getCollection(coll, m)
332 ctx := context.Background()
334 //Get all master tables in this collection
340 cursor, err := c.Find(ctx, filter)
342 return nil, pkgerrors.Errorf("Error reading from database: %s", err.Error())
344 defer cursorClose(ctx, cursor)
346 //Iterate over all the master tables
347 result := make(map[string][]byte)
348 for cursorNext(ctx, cursor) {
351 //Read key of each master table
352 key, ok := d.Lookup("key").DocumentOK()
354 //Throw error if key is not found
355 pkgerrors.New("Unable to read key from mastertable")
358 //Get objectID of tag document
359 tid, ok := d.Lookup(tag).ObjectIDOK()
361 log.Printf("Did not find tag: %s", tag)
365 //Find tag document and unmarshal it into []byte
366 tagData, err := decodeBytes(c.FindOne(ctx, bson.D{{"_id", tid}}))
368 log.Printf("Unable to decode tag data %s", err.Error())
371 result[key.String()] = tagData.Lookup(tag).Value
374 if len(result) == 0 {
375 return result, pkgerrors.Errorf("Did not find any objects with tag: %s", tag)