103 lines
2.3 KiB
Go
103 lines
2.3 KiB
Go
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...)
|
|
}
|