package storage import ( "context" "errors" "fmt" "github.com/spf13/viper" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" ) type MongoDBStorage struct { *mongo.Client DbName string Collections map[string]string } func NewMongoDBStorage(cfg *viper.Viper) (MongoDBStorage, error) { var ( mongodb_host = cfg.GetString("storage.db.mongodb.host") mongodb_port = cfg.GetString("storage.db.mongodb.port") mongodb_dbname = cfg.GetString("storage.db.mongodb.db_name") mongodb_groups = cfg.GetString("storage.db.mongodb.collections.groups") ) client, err := mongo.NewClient(options.Client().ApplyURI("mongodb://" + mongodb_host + ":" + mongodb_port)) if err != nil { return MongoDBStorage{}, err } err = client.Connect(context.TODO()) if err != nil { return MongoDBStorage{}, err } storage := MongoDBStorage{ Client: client, DbName: mongodb_dbname, Collections: map[string]string{ "groups": mongodb_groups, }, } //TODO Indexes return storage, err } func (s MongoDBStorage) CreateGroup(group Group) error { collection := s.Client.Database(s.DbName).Collection(s.Collections["groups"]) if _, err := collection.InsertOne(context.TODO(), group); err != nil { return err } return nil } func (s MongoDBStorage) GetGroup(id string) (*Group, error) { collection := s.Client.Database(s.DbName).Collection(s.Collections["groups"]) group := &Group{} if err := collection.FindOne(context.TODO(), bson.M{"_id": id}).Decode(group); err != nil { return nil, err } return group, nil } func (s MongoDBStorage) GetGroups(namespaces []string) (groups []Group, err error) { collection := s.Client.Database(s.DbName).Collection(s.Collections["groups"]) var cur *mongo.Cursor findOptions := options.Find() if len(namespaces) == 0 { cur, err = collection.Find(context.TODO(), bson.D{}, findOptions) if err != nil { return groups, err } } else { cur, err = collection.Find(context.TODO(), bson.M{"namespace": bson.M{"$in": namespaces}}, findOptions) if err != nil { return groups, err } } for cur.Next(context.TODO()) { var group Group var elem bson.M if err := cur.Decode(&elem); err != nil { return groups, err } bsonBytes, _ := bson.Marshal(elem) bson.Unmarshal(bsonBytes, &group) groups = append(groups, group) } return } func (s MongoDBStorage) GetGroupsByIds(groupids []string) (groups []Group, err error) { collection := s.Client.Database(s.DbName).Collection(s.Collections["groups"]) var cur *mongo.Cursor findOptions := options.Find() if len(groupids) == 0 { return groups, errors.New("no group id provided") } else { cur, err = collection.Find(context.TODO(), bson.M{"_id": bson.M{"$in": groupids}}, findOptions) if err != nil { return groups, err } } for cur.Next(context.TODO()) { var group Group var elem bson.M if err := cur.Decode(&elem); err != nil { return groups, err } bsonBytes, _ := bson.Marshal(elem) bson.Unmarshal(bsonBytes, &group) groups = append(groups, group) } return } func (s MongoDBStorage) UpdateGroup(group Group) error { collection := s.Client.Database(s.DbName).Collection(s.Collections["groups"]) if _, err := collection.ReplaceOne(context.TODO(), bson.M{"_id": group.ID}, group); err != nil { fmt.Println(err) return err } return nil }