src/k8splugin/internal/db/mongo.go

   1 /*
   2  * Copyright 2018 Intel Corporation, Inc
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *     http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 package db
  18
  19 import (
  20         "github.com/mongodb/mongo-go-driver/bson"
  21         "github.com/mongodb/mongo-go-driver/bson/primitive"
  22         "github.com/mongodb/mongo-go-driver/mongo"
  23         "github.com/mongodb/mongo-go-driver/mongo/options"
  24         pkgerrors "github.com/pkg/errors"
  25         "golang.org/x/net/context"
  26         "log"
  27         "os"
  28 )
  29
  30 // MongoCollection defines the a subset of MongoDB operations
  31 // Note: This interface is defined mainly for mock testing
  32 type MongoCollection interface {
  33         InsertOne(ctx context.Context, document interface{},
  34                 opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error)
  35         FindOne(ctx context.Context, filter interface{},
  36                 opts ...*options.FindOneOptions) *mongo.SingleResult
  37         FindOneAndUpdate(ctx context.Context, filter interface{},
  38                 update interface{}, opts ...*options.FindOneAndUpdateOptions) *mongo.SingleResult
  39         DeleteOne(ctx context.Context, filter interface{},
  40                 opts ...*options.DeleteOptions) (*mongo.DeleteResult, error)
  41         Find(ctx context.Context, filter interface{},
  42                 opts ...*options.FindOptions) (mongo.Cursor, error)
  43 }
  44
  45 // MongoStore is an implementation of the db.Store interface
  46 type MongoStore struct {
  47         db *mongo.Database
  48 }
  49
  50 // This exists only for allowing us to mock the collection object
  51 // for testing purposes
  52 var getCollection = func(coll string, m *MongoStore) MongoCollection {
  53         return m.db.Collection(coll)
  54 }
  55
  56 // This exists only for allowing us to mock the DecodeBytes function
  57 // Mainly because we cannot construct a SingleResult struct from our
  58 // tests. All fields in that struct are private.
  59 var decodeBytes = func(sr *mongo.SingleResult) (bson.Raw, error) {
  60         return sr.DecodeBytes()
  61 }
  62
  63 // NewMongoStore initializes a Mongo Database with the name provided
  64 // If a database with that name exists, it will be returned
  65 func NewMongoStore(name string, store *mongo.Database) (Store, error) {
  66         if store == nil {
  67                 ip := "mongodb://" + os.Getenv("DATABASE_IP") + ":27017"
  68                 mongoClient, err := mongo.NewClient(ip)
  69                 if err != nil {
  70                         return nil, err
  71                 }
  72
  73                 err = mongoClient.Connect(context.Background())
  74                 if err != nil {
  75                         return nil, err
  76                 }
  77                 store = mongoClient.Database(name)
  78         }
  79
  80         return &MongoStore{
  81                 db: store,
  82         }, nil
  83 }
  84
  85 // HealthCheck verifies if the database is up and running
  86 func (m *MongoStore) HealthCheck() error {
  87
  88         _, err := decodeBytes(m.db.RunCommand(context.Background(), bson.D{{"serverStatus", 1}}))
  89         if err != nil {
  90                 return pkgerrors.Wrap(err, "Error getting server status")
  91         }
  92
  93         return nil
  94 }
  95
  96 // validateParams checks to see if any parameters are empty
  97 func (m *MongoStore) validateParams(args ...string) bool {
  98         for _, v := range args {
  99                 if v == "" {
 100                         return false
 101                 }
 102         }
 103
 104         return true
 105 }
 106
 107 // Create is used to create a DB entry
 108 func (m *MongoStore) Create(coll, key, tag string, data interface{}) error {
 109         if data == nil || !m.validateParams(coll, key, tag) {
 110                 return pkgerrors.New("No Data to store")
 111         }
 112
 113         c := getCollection(coll, m)
 114         ctx := context.Background()
 115
 116         //Insert the data and then add the objectID to the masterTable
 117         res, err := c.InsertOne(ctx, bson.D{
 118                 {tag, data},
 119         })
 120         if err != nil {
 121                 return pkgerrors.Errorf("Error inserting into database: %s", err.Error())
 122         }
 123
 124         //Add objectID of created data to masterKey document
 125         //Create masterkey document if it does not exist
 126         filter := bson.D{{"key", key}}
 127
 128         _, err = decodeBytes(
 129                 c.FindOneAndUpdate(
 130                         ctx,
 131                         filter,
 132                         bson.D{
 133                                 {"$set", bson.D{
 134                                         {tag, res.InsertedID},
 135                                 }},
 136                         },
 137                         options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)))
 138
 139         if err != nil {
 140                 return pkgerrors.Errorf("Error updating master table: %s", err.Error())
 141         }
 142
 143         return nil
 144 }
 145
 146 // Unmarshal implements an unmarshaler for bson data that
 147 // is produced from the mongo database
 148 func (m *MongoStore) Unmarshal(inp []byte, out interface{}) error {
 149         err := bson.Unmarshal(inp, out)
 150         if err != nil {
 151                 return pkgerrors.Wrap(err, "Unmarshaling bson")
 152         }
 153         return nil
 154 }
 155
 156 // Read method returns the data stored for this key and for this particular tag
 157 func (m *MongoStore) Read(coll, key, tag string) ([]byte, error) {
 158         if !m.validateParams(coll, key, tag) {
 159                 return nil, pkgerrors.New("Mandatory fields are missing")
 160         }
 161
 162         c := getCollection(coll, m)
 163         ctx := context.Background()
 164
 165         //Get the masterkey document based on given key
 166         filter := bson.D{{"key", key}}
 167         keydata, err := decodeBytes(c.FindOne(context.Background(), filter))
 168         if err != nil {
 169                 return nil, pkgerrors.Errorf("Error finding master table: %s", err.Error())
 170         }
 171
 172         //Read the tag objectID from document
 173         tagoid, ok := keydata.Lookup(tag).ObjectIDOK()
 174         if !ok {
 175                 return nil, pkgerrors.Errorf("Error finding objectID for tag %s", tag)
 176         }
 177
 178         //Use tag objectID to read the data from store
 179         filter = bson.D{{"_id", tagoid}}
 180         tagdata, err := decodeBytes(c.FindOne(ctx, filter))
 181         if err != nil {
 182                 return nil, pkgerrors.Errorf("Error reading found object: %s", err.Error())
 183         }
 184
 185         //Return the data as a byte array
 186         //Convert string data to byte array using the built-in functions
 187         switch tagdata.Lookup(tag).Type {
 188         case bson.TypeString:
 189                 return []byte(tagdata.Lookup(tag).StringValue()), nil
 190         default:
 191                 return tagdata.Lookup(tag).Value, nil
 192         }
 193 }
 194
 195 // Helper function that deletes an object by its ID
 196 func (m *MongoStore) deleteObjectByID(coll string, objID primitive.ObjectID) error {
 197
 198         c := getCollection(coll, m)
 199         ctx := context.Background()
 200
 201         _, err := c.DeleteOne(ctx, bson.D{{"_id", objID}})
 202         if err != nil {
 203                 return pkgerrors.Errorf("Error Deleting from database: %s", err.Error())
 204         }
 205
 206         log.Printf("Deleted Obj with ID %s", objID.String())
 207         return nil
 208 }
 209
 210 // Delete method removes a document from the Database that matches key
 211 // TODO: delete all referenced docs if tag is empty string
 212 func (m *MongoStore) Delete(coll, key, tag string) error {
 213         if !m.validateParams(coll, key, tag) {
 214                 return pkgerrors.New("Mandatory fields are missing")
 215         }
 216
 217         c := getCollection(coll, m)
 218         ctx := context.Background()
 219
 220         //Get the masterkey document based on given key
 221         filter := bson.D{{"key", key}}
 222         //Remove the tag ID entry from masterkey table
 223         update := bson.D{
 224                 {
 225                         "$unset", bson.D{
 226                                 {tag, ""},
 227                         },
 228                 },
 229         }
 230         keydata, err := decodeBytes(c.FindOneAndUpdate(ctx, filter, update,
 231                 options.FindOneAndUpdate().SetReturnDocument(options.Before)))
 232         if err != nil {
 233                 //No document was found. Return nil.
 234                 if err == mongo.ErrNoDocuments {
 235                         return nil
 236                 }
 237                 //Return any other error that was found.
 238                 return pkgerrors.Errorf("Error decoding master table after update: %s",
 239                         err.Error())
 240         }
 241
 242         //Read the tag objectID from document
 243         elems, err := keydata.Elements()
 244         if err != nil {
 245                 return pkgerrors.Errorf("Error reading elements from database: %s", err.Error())
 246         }
 247
 248         tagoid, ok := keydata.Lookup(tag).ObjectIDOK()
 249         if !ok {
 250                 return pkgerrors.Errorf("Error finding objectID for tag %s", tag)
 251         }
 252
 253         //Use tag objectID to read the data from store
 254         err = m.deleteObjectByID(coll, tagoid)
 255         if err != nil {
 256                 return pkgerrors.Errorf("Error deleting from database: %s", err.Error())
 257         }
 258
 259         //Delete master table if no more tags left
 260         //_id, key and tag should be elements in before doc
 261         //if master table needs to be removed too
 262         if len(elems) == 3 {
 263                 keyid, ok := keydata.Lookup("_id").ObjectIDOK()
 264                 if !ok {
 265                         return pkgerrors.Errorf("Error finding objectID for key %s", key)
 266                 }
 267                 err = m.deleteObjectByID(coll, keyid)
 268                 if err != nil {
 269                         return pkgerrors.Errorf("Error deleting master table from database: %s", err.Error())
 270                 }
 271         }
 272
 273         return nil
 274 }
 275
 276 // ReadAll is used to get all documents in db of a particular tag
 277 func (m *MongoStore) ReadAll(coll, tag string) (map[string][]byte, error) {
 278         if !m.validateParams(coll, tag) {
 279                 return nil, pkgerrors.New("Missing collection or tag name")
 280         }
 281
 282         c := getCollection(coll, m)
 283         ctx := context.Background()
 284
 285         //Get all master tables in this collection
 286         filter := bson.D{
 287                 {"key", bson.D{
 288                         {"$exists", true},
 289                 }},
 290         }
 291         cursor, err := c.Find(ctx, filter)
 292         if err != nil {
 293                 return nil, pkgerrors.Errorf("Error reading from database: %s", err.Error())
 294         }
 295         defer cursor.Close(ctx)
 296
 297         //Iterate over all the master tables
 298         result := make(map[string][]byte)
 299         for cursor.Next(ctx) {
 300                 d, err := cursor.DecodeBytes()
 301                 if err != nil {
 302                         log.Printf("Unable to decode data in Readall: %s", err.Error())
 303                         continue
 304                 }
 305
 306                 //Read key of each master table
 307                 key, ok := d.Lookup("key").StringValueOK()
 308                 if !ok {
 309                         log.Printf("Unable to read key string from mastertable %s", err.Error())
 310                         continue
 311                 }
 312
 313                 //Get objectID of tag document
 314                 tid, ok := d.Lookup(tag).ObjectIDOK()
 315                 if !ok {
 316                         log.Printf("Did not find tag: %s", tag)
 317                         continue
 318                 }
 319
 320                 //Find tag document and unmarshal it into []byte
 321                 tagData, err := decodeBytes(c.FindOne(ctx, bson.D{{"_id", tid}}))
 322                 if err != nil {
 323                         log.Printf("Unable to decode tag data %s", err.Error())
 324                         continue
 325                 }
 326                 result[key] = tagData.Lookup(tag).Value
 327         }
 328
 329         if len(result) == 0 {
 330                 return result, pkgerrors.Errorf("Did not find any objects with tag: %s", tag)
 331         }
 332
 333         return result, nil
 334 }