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 // }