99 lines
2.2 KiB
Go
99 lines
2.2 KiB
Go
package leaf
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/rs/zerolog"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/bson/primitive"
|
|
driver "go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
func ensureIndexes(ctx context.Context, collection *driver.Collection, models []driver.IndexModel, createIndexes bool, logger *zerolog.Logger) error {
|
|
existing, err := indexNames(ctx, collection)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
indexes := filterIndexes(models, existing, logger)
|
|
if len(indexes) < 1 {
|
|
return nil
|
|
}
|
|
|
|
if !createIndexes {
|
|
for _, im := range indexes {
|
|
logger.Warn().Str("name", *im.Options.Name).Msg("index declared but not present")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
_, err = collection.Indexes().CreateMany(ctx, indexes)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, im := range indexes {
|
|
logger.Info().Str("name", *im.Options.Name).Msg("index created")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func indexNames(ctx context.Context, collection *driver.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 []driver.IndexModel, names map[string]struct{}, logger *zerolog.Logger) []driver.IndexModel {
|
|
var filtered []driver.IndexModel
|
|
for _, im := range models {
|
|
if im.Options.Name == nil {
|
|
logger.Warn().Interface("keys", im.Keys).Msg("must specify index name explicitly")
|
|
continue
|
|
}
|
|
|
|
if _, ok := names[*im.Options.Name]; !ok {
|
|
filtered = append(filtered, im)
|
|
}
|
|
}
|
|
|
|
return filtered
|
|
}
|
|
|
|
func withTimestampIndexes(models []driver.IndexModel) []driver.IndexModel {
|
|
tsModels := []driver.IndexModel{
|
|
{
|
|
Keys: bson.D{primitive.E{Key: "created_at", Value: 1}},
|
|
Options: options.Index().SetName("created_at_1"),
|
|
},
|
|
{
|
|
Keys: bson.D{primitive.E{Key: "updated_at", Value: 1}},
|
|
Options: options.Index().SetName("updated_at_1"),
|
|
},
|
|
}
|
|
|
|
return append(models, tsModels...)
|
|
}
|