leaf/index.go

103 lines
2.3 KiB
Go
Raw Normal View History

2024-01-07 05:40:41 +00:00
package leaf
import (
"context"
"go.mongodb.org/mongo-driver/bson"
2024-01-12 04:13:30 +00:00
"go.mongodb.org/mongo-driver/mongo"
2024-01-07 05:40:41 +00:00
"go.mongodb.org/mongo-driver/mongo/options"
)
2024-01-12 04:13:30 +00:00
type IndexMessage struct {
Name string
Keys interface{}
Message string
}
func ensureIndexes(ctx context.Context, collection *mongo.Collection, models []mongo.IndexModel, createIndexes bool) ([]IndexMessage, error) {
2024-01-07 05:40:41 +00:00
existing, err := indexNames(ctx, collection)
if err != nil {
2024-01-12 04:13:30 +00:00
return []IndexMessage{}, err
2024-01-07 05:40:41 +00:00
}
2024-01-12 04:13:30 +00:00
indexes, messages := filterIndexes(models, existing, []IndexMessage{})
2024-01-07 05:40:41 +00:00
if len(indexes) < 1 {
2024-01-12 04:13:30 +00:00
return messages, nil
2024-01-07 05:40:41 +00:00
}
if !createIndexes {
for _, im := range indexes {
2024-01-12 04:13:30 +00:00
messages = append(messages, IndexMessage{Name: *im.Options.Name, Message: "index declared but not present"})
2024-01-07 05:40:41 +00:00
}
2024-01-12 04:13:30 +00:00
return messages, nil
2024-01-07 05:40:41 +00:00
}
_, err = collection.Indexes().CreateMany(ctx, indexes)
if err != nil {
2024-01-12 04:13:30 +00:00
return messages, err
2024-01-07 05:40:41 +00:00
}
for _, im := range indexes {
2024-01-12 04:13:30 +00:00
messages = append(messages, IndexMessage{Name: *im.Options.Name, Keys: im.Keys, Message: "index created"})
2024-01-07 05:40:41 +00:00
}
2024-01-12 04:13:30 +00:00
return messages, nil
2024-01-07 05:40:41 +00:00
}
2024-01-12 04:13:30 +00:00
func indexNames(ctx context.Context, collection *mongo.Collection) (map[string]struct{}, error) {
2024-01-07 05:40:41 +00:00
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
}
2024-01-12 04:13:30 +00:00
func filterIndexes(models []mongo.IndexModel, names map[string]struct{}, messages []IndexMessage) ([]mongo.IndexModel, []IndexMessage) {
var filtered []mongo.IndexModel
2024-01-07 05:40:41 +00:00
for _, im := range models {
if im.Options.Name == nil {
2024-01-12 04:13:30 +00:00
messages = append(messages, IndexMessage{Keys: im.Keys, Message: "must specify index name explicitly"})
2024-01-07 05:40:41 +00:00
continue
}
if _, ok := names[*im.Options.Name]; !ok {
filtered = append(filtered, im)
}
}
2024-01-12 04:13:30 +00:00
return filtered, messages
2024-01-07 05:40:41 +00:00
}
2024-01-12 04:13:30 +00:00
func withTimestampIndexes(models []mongo.IndexModel) []mongo.IndexModel {
tsModels := []mongo.IndexModel{
2024-01-07 05:40:41 +00:00
{
2024-01-12 04:13:30 +00:00
Keys: bson.D{bson.E{Key: "created_at", Value: 1}},
2024-01-07 05:40:41 +00:00
Options: options.Index().SetName("created_at_1"),
},
{
2024-01-12 04:13:30 +00:00
Keys: bson.D{bson.E{Key: "updated_at", Value: 1}},
2024-01-07 05:40:41 +00:00
Options: options.Index().SetName("updated_at_1"),
},
}
return append(models, tsModels...)
}