package storage import ( "context" "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_vehicles = cfg.GetString("storage.db.mongodb.collections.vehicles") mongodb_bookings = cfg.GetString("storage.db.mongodb.collections.bookings") ) 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{ "vehicles": mongodb_vehicles, "bookings": mongodb_bookings, }, } //TODO Indexes return storage, err } func (s MongoDBStorage) CreateVehicle(vehicle Vehicle) error { collection := s.Client.Database(s.DbName).Collection(s.Collections["vehicles"]) if _, err := collection.InsertOne(context.TODO(), vehicle); err != nil { return err } return nil } func (s MongoDBStorage) GetVehicles(namespaces []string) (vehicles []Vehicle, err error) { collection := s.Client.Database(s.DbName).Collection(s.Collections["vehicles"]) var cur *mongo.Cursor findOptions := options.Find() if len(namespaces) == 0 { cur, err = collection.Find(context.TODO(), bson.D{}, findOptions) if err != nil { return vehicles, err } } else { cur, err = collection.Find(context.TODO(), bson.M{"namespace": bson.M{"$in": namespaces}}, findOptions) if err != nil { return vehicles, err } } for cur.Next(context.TODO()) { var vehicle Vehicle var elem bson.M if err := cur.Decode(&elem); err != nil { return vehicles, err } bsonBytes, _ := bson.Marshal(elem) bson.Unmarshal(bsonBytes, &vehicle) vehicles = append(vehicles, vehicle) } return } func (s MongoDBStorage) GetVehicle(id string) (*Vehicle, error) { collection := s.Client.Database(s.DbName).Collection(s.Collections["vehicles"]) vehicle := &Vehicle{} if err := collection.FindOne(context.TODO(), bson.M{"_id": id}).Decode(vehicle); err != nil { return nil, err } return vehicle, nil } func (s MongoDBStorage) CreateBooking(booking Booking) error { collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"]) if _, err := collection.InsertOne(context.TODO(), booking); err != nil { return err } return nil } func (s MongoDBStorage) GetBooking(id string) (*Booking, error) { collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"]) booking := &Booking{} if err := collection.FindOne(context.TODO(), bson.M{"_id": id}).Decode(booking); err != nil { return nil, err } return booking, nil } func (s MongoDBStorage) GetBookingsForVehicle(vehicleid string) (bookings []Booking, err error) { collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"]) var cur *mongo.Cursor findOptions := options.Find() cur, err = collection.Find(context.TODO(), bson.M{"vehicleid": vehicleid}, findOptions) if err != nil { return bookings, err } for cur.Next(context.TODO()) { var booking Booking var elem bson.M if err := cur.Decode(&elem); err != nil { return bookings, err } bsonBytes, _ := bson.Marshal(elem) bson.Unmarshal(bsonBytes, &booking) bookings = append(bookings, booking) } return } func (s MongoDBStorage) DeleteBooking(bookingid string) error { collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"]) _, err := collection.UpdateOne( context.TODO(), bson.M{"_id": bookingid}, bson.D{ {"$set", bson.D{{"deleted", true}}}, }, ) return err }