fleets/storage/mongodb.go

package storage

import (
	"context"
	"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_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) UpdateBooking(booking Booking) error {
	collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"])
	if _, err := collection.ReplaceOne(context.TODO(), bson.M{"_id": booking.ID}, booking); err != nil {
		fmt.Println(err)
		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) GetBookings() (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.D{}, 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) 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) GetBookingsForDriver(driver 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{"driver": driver}, 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.DeleteOne(
		context.TODO(),
		bson.M{"_id": bookingid},
	)

	return err

}

/////////////////////update///////////////
func (s MongoDBStorage) UpdateVehicle(vehicle Vehicle) error {
	collection := s.Client.Database(s.DbName).Collection(s.Collections["vehicles"])
	if _, err := collection.ReplaceOne(context.TODO(), bson.M{"_id": vehicle.ID}, vehicle); err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}