carpool-service/storage/mongodb.go

package storage

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"

	"git.coopgo.io/coopgo-platform/carpool-service/internal"
	"github.com/paulmach/orb/geojson"
	"github.com/rs/zerolog/log"
	"github.com/spf13/viper"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"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_uri                          = cfg.GetString("storage.db.mongodb.uri")
		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_regular_routes               = cfg.GetString("storage.db.mongodb.collections.regular_routes")
		mongodb_punctual_routes              = cfg.GetString("storage.db.mongodb.collections.punctual_routes")
		mongodb_bookings                     = cfg.GetString("storage.db.mongodb.collections.bookings")
		mongodb_driver_candidate_journeys    = cfg.GetString("storage.db.mongodb.collections.driver_candidate_journeys")
		mongodb_passenger_candidate_journeys = cfg.GetString("storage.db.mongodb.collections.passenger_candidate_journeys")
		mongodb_persisted_kv                 = cfg.GetString("storage.db.mongodb.collections.persisted_kv")
	)

	if mongodb_uri == "" {
		mongodb_uri = fmt.Sprintf("mongodb://%s:%s/%s", mongodb_host, mongodb_port, mongodb_dbname)
	}

	client, err := mongo.NewClient(options.Client().ApplyURI(mongodb_uri))
	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{
			"regular_routes":               mongodb_regular_routes,
			"punctual_routes":              mongodb_punctual_routes,
			"bookings":                     mongodb_bookings,
			"driver_candidate_journeys":    mongodb_driver_candidate_journeys,
			"passenger_candidate_journeys": mongodb_passenger_candidate_journeys,
			"persisted_kv":                 mongodb_persisted_kv,
		},
	}

	return storage, nil
}

func (s MongoDBStorage) CreateRegularRoutes(routes []*geojson.FeatureCollection) error {

	log.Debug().Msg("Storage - CreateRegularRoutes")

	documents := []any{}

	for _, fc := range routes {
		if fc != nil {
			fc.ExtraMembers["_id"] = fc.ExtraMembers.MustString("id")
			delete(fc.ExtraMembers, "id")
			documents = append(documents, fc)
		}
	}

	collection := s.Client.Database(s.DbName).Collection(s.Collections["regular_routes"])
	if _, err := collection.InsertMany(context.TODO(), documents); err != nil {
		return err
	}

	return nil
}

func (s MongoDBStorage) GetUserRegularRoutes(userid string) ([]*geojson.FeatureCollection, error) {
	findOptions := options.Find()
	collection := s.Client.Database(s.DbName).Collection(s.Collections["regular_routes"])
	cur, err := collection.Find(context.TODO(), bson.M{"properties.user.id": userid}, findOptions)
	if err != nil {
		return nil, err
	}

	results := []*geojson.FeatureCollection{}

	for cur.Next(context.TODO()) {
		var elem bson.M

		if err := cur.Decode(&elem); err != nil {
			return nil, err
		}
		bsonBytes, _ := bson.Marshal(elem)
		fc := geojson.NewFeatureCollection()
		err := fc.UnmarshalBSON(bsonBytes)
		if err != nil {
			return nil, err
		}
		fc.ExtraMembers["id"] = fc.ExtraMembers.MustString("_id")
		delete(fc.ExtraMembers, "_id")

		for k, v := range fc.ExtraMembers {
			if val, ok := v.(primitive.D); ok {
				em := map[string]any{}
				jsonbytes, _ := bson.MarshalExtJSON(val, true, true)
				json.Unmarshal(jsonbytes, &em)
				fc.ExtraMembers[k] = em
			}
		}

		results = append(results, fc)
	}

	return results, nil
}

func (s MongoDBStorage) GetDriverRegularRoutesForTile(day string, gridId int64) ([]*geojson.FeatureCollection, error) {
	log.Debug().
		Str("day", day).
		Int64("grid id", gridId).
		Msg("GetDriverRegularRoutesForTile")

	findOptions := options.Find()
	collection := s.Client.Database(s.DbName).Collection(s.Collections["regular_routes"])
	cur, err := collection.Find(context.TODO(), bson.M{
		"properties.schedules.day": day,
		"grid_ids":                 gridId,
		"properties.is_driver":     true,
	}, findOptions)
	if err != nil {
		log.Error().Err(err).Msg("GetDriverRegularRoutesForTile - mongodb issue")
		return nil, err
	}

	results := []*geojson.FeatureCollection{}
	for cur.Next(context.TODO()) {
		var elem bson.M

		if err := cur.Decode(&elem); err != nil {
			log.Error().Err(err).Msg("GetDriverRegularRoutesForTile - decode error")
			return nil, err
		}
		bsonBytes, _ := bson.Marshal(elem)
		fc := geojson.NewFeatureCollection()
		err := fc.UnmarshalBSON(bsonBytes)
		if err != nil {
			log.Error().Err(err).Msg("GetDriverRegularRoutesForTile - unmarshalling error")
			return nil, err
		}
		fc.ExtraMembers["id"] = fc.ExtraMembers.MustString("_id")
		delete(fc.ExtraMembers, "_id")

		for k, v := range fc.ExtraMembers {
			if val, ok := v.(primitive.D); ok {
				em := map[string]any{}
				jsonbytes, _ := bson.MarshalExtJSON(val, true, true)
				json.Unmarshal(jsonbytes, &em)
				fc.ExtraMembers[k] = em
			}
		}

		results = append(results, fc)
	}

	return results, nil
}

func (s MongoDBStorage) GetPassengerRegularRoutesForTile(day string, gridId int64) ([]*geojson.FeatureCollection, error) {
	findOptions := options.Find()
	collection := s.Client.Database(s.DbName).Collection(s.Collections["regular_routes"])
	cur, err := collection.Find(context.TODO(), bson.M{
		"properties.schedules.day": day,
		"grid_ids":                 gridId,
		"properties.is_passenger":  true,
	}, findOptions)
	if err != nil {
		return nil, err
	}

	results := []*geojson.FeatureCollection{}
	for cur.Next(context.TODO()) {
		var elem bson.M

		if err := cur.Decode(&elem); err != nil {
			return nil, err
		}
		bsonBytes, _ := bson.Marshal(elem)
		fc := geojson.NewFeatureCollection()
		err := fc.UnmarshalBSON(bsonBytes)
		if err != nil {
			return nil, err
		}
		fc.ExtraMembers["id"] = fc.ExtraMembers.MustString("_id")
		delete(fc.ExtraMembers, "_id")

		for k, v := range fc.ExtraMembers {
			if val, ok := v.(primitive.D); ok {
				em := map[string]any{}
				jsonbytes, _ := bson.MarshalExtJSON(val, true, true)
				json.Unmarshal(jsonbytes, &em)
				fc.ExtraMembers[k] = em
			}
		}

		results = append(results, fc)
	}

	return results, nil
}

func (s MongoDBStorage) CreateBooking(booking internal.Booking) error {

	collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"])
	_, err := collection.InsertOne(context.TODO(), booking)
	if err != nil {
		return err
	}
	return nil
}
func (s MongoDBStorage) GetBooking(id string) (booking *internal.Booking, err error) {
	var b internal.Booking
	log.Debug().Str("booking id", id).Msg("get booking in DB")
	collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"])
	err = collection.FindOne(context.TODO(), bson.M{"_id": id}).Decode(&b)
	if err != nil {
		log.Error().Err(err).Msg("error")
		return nil, err
	}

	return &b, nil
}

func (s MongoDBStorage) GetUserBookings(userid string) ([]internal.Booking, error) {
	findOptions := options.Find()
	collection := s.Client.Database(s.DbName).Collection(s.Collections["bookings"])
	cur, err := collection.Find(context.TODO(), bson.M{
		"$or": []bson.M{
			{"driver.id": userid},
			{"passenger.id": userid},
		},
	}, findOptions)
	if err != nil {
		return nil, err
	}

	var bookings []internal.Booking
	for cur.Next(context.TODO()) {
		var elem internal.Booking

		if err := cur.Decode(&elem); err != nil {
			log.Error().Err(err).Msg("error reading bookings response")
			return nil, err
		}

		bookings = append(bookings, elem)
	}

	return bookings, nil
}

func (s MongoDBStorage) UpdateBookingStatus(bookingid string, status string) error {
	//TODO implement UpdateBookingStatus
	return errors.New("MongoDB Storage - UpdateBookingStatus not implemented")
}

func (s MongoDBStorage) PersistedKVPut(documents []any) error {
	collection := s.Client.Database(s.DbName).Collection(s.Collections["persisted_kv"])
	if _, err := collection.InsertMany(context.TODO(), documents); err != nil {
		return err
	}

	return nil
}
func (s MongoDBStorage) PersistedKVGet(id string, document any) error {
	collection := s.Client.Database(s.DbName).Collection(s.Collections["persisted_kv"])
	err := collection.FindOne(context.TODO(), bson.M{"_id": id}).Decode(document)
	if err != nil {
		return err
	}
	return nil
}

func (s MongoDBStorage) StoreRouteSchedules(js []internal.PlannedRouteSchedule) error {

	log.Debug().Msg("Storage - StoreRouteSchedules")

	documents := []any{}
	for _, sr := range js {
		documents = append(documents, sr)
	}

	return s.PersistedKVPut(documents)
}

func (s MongoDBStorage) GetRouteSchedule(id string) (*internal.PlannedRouteSchedule, error) {

	var result internal.PlannedRouteSchedule
	err := s.PersistedKVGet(id, &result)
	if err != nil {
		return nil, err
	}

	for k, v := range result.Route.ExtraMembers {
		if val, ok := v.(primitive.D); ok {
			em := map[string]any{}
			jsonbytes, _ := bson.MarshalExtJSON(val, true, true)
			json.Unmarshal(jsonbytes, &em)
			result.Route.ExtraMembers[k] = em
		}
	}

	return &result, nil
}

// func (s MongoDBStorage) CreatePassengerRegularTrips(trips []*geojson.FeatureCollection) error {

// 	log.Debug().Msg("Storage - CreatePassengerRegularTrips")

// 	documents := []any{}

// 	for _, fc := range trips {
// 		if fc != nil {
// 			fc.ExtraMembers["_id"] = fc.ExtraMembers.MustString("id")
// 			delete(fc.ExtraMembers, "id")
// 			documents = append(documents, fc)
// 		}
// 	}

// 	collection := s.Client.Database(s.DbName).Collection(s.Collections["passenger_regular_trips"])
// 	if _, err := collection.InsertMany(context.TODO(), documents); err != nil {
// 		return err
// 	}

// 	return nil
// }

func (s MongoDBStorage) Migrate() error {
	return nil
}