package leaf import ( "context" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" ) type IndexMessage struct { Name string Keys interface{} Message string } func ensureIndexes(ctx context.Context, collection *mongo.Collection, models []mongo.IndexModel, createIndexes bool) ([]IndexMessage, error) { existing, err := indexNames(ctx, collection) if err != nil { return []IndexMessage{}, err } indexes, messages := filterIndexes(models, existing, []IndexMessage{}) if len(indexes) < 1 { return messages, nil } if !createIndexes { for _, im := range indexes { messages = append(messages, IndexMessage{Name: *im.Options.Name, Message: "index declared but not present"}) } return messages, nil } _, err = collection.Indexes().CreateMany(ctx, indexes) if err != nil { return messages, err } for _, im := range indexes { messages = append(messages, IndexMessage{Name: *im.Options.Name, Keys: im.Keys, Message: "index created"}) } return messages, nil } func indexNames(ctx context.Context, collection *mongo.Collection) (map[string]struct{}, error) { cursor, err := collection.Indexes().List(ctx) if err != nil { return nil, err } defer cursor.Close(ctx) names := make(map[string]struct{}) for cursor.Next(ctx) { var idx bson.M if err := cursor.Decode(&idx); err != nil { continue } name, ok := idx["name"].(string) if !ok { continue } names[name] = struct{}{} } return names, nil } func filterIndexes(models []mongo.IndexModel, names map[string]struct{}, messages []IndexMessage) ([]mongo.IndexModel, []IndexMessage) { var filtered []mongo.IndexModel for _, im := range models { if im.Options.Name == nil { messages = append(messages, IndexMessage{Keys: im.Keys, Message: "must specify index name explicitly"}) continue } if _, ok := names[*im.Options.Name]; !ok { filtered = append(filtered, im) } } return filtered, messages } func withTimestampIndexes(models []mongo.IndexModel) []mongo.IndexModel { tsModels := []mongo.IndexModel{ { Keys: bson.D{bson.E{Key: "created_at", Value: 1}}, Options: options.Index().SetName("created_at_1"), }, { Keys: bson.D{bson.E{Key: "updated_at", Value: 1}}, Options: options.Index().SetName("updated_at_1"), }, } return append(models, tsModels...) }