carpool-service/storage/mongodb.go

package storage

import (
	"context"
	"encoding/json"

	"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_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")
	)

	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{
			"regular_routes":  mongodb_regular_routes,
			"punctual_routes": mongodb_punctual_routes,
			"bookings":        mongodb_bookings,
		},
	}

	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) {
	findOptions := options.Find()
	collection := s.Client.Database(s.DbName).Collection(s.Collections["regular_routes"])
	cur, err := collection.Find(context.TODO(), bson.M{
		"properties.schedule.day": day,
		"grid_ids":                gridId,
		"properties.is_driver":    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) 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.schedule.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) 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
// }