leaf/index.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...)
}