Configure indices in struct tags

4 years ago · eeb83daf4b
parent 6e6a042a16
commit eeb83daf4b
3 changed files with 291 additions and 0 deletions
--- a/mongox/database/index.go
+++ b/mongox/database/index.go
@ -0,0 +1,134 @@
+package database
+
+import (
+	"bytes"
+	"fmt"
+	"reflect"
+	"strconv"
+	"strings"
+	"text/template"
+
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+// IndexEnsure function ensures index in mongo collection of document
+//   `index:""` -- https://docs.mongodb.com/manual/indexes/#create-an-index
+//   `index:"-"` -- (descending)
+//   `index:"-,+foo,+-bar"` -- https://docs.mongodb.com/manual/core/index-compound
+//   `index:"-,+foo,+-bar,unique"` -- https://docs.mongodb.com/manual/core/index-unique
+//   `index:"-,+foo,+-bar,unique,allowNull"` -- https://docs.mongodb.com/manual/core/index-partial
+//   `index:"-,unique,allowNull,expireAfter=86400"` -- https://docs.mongodb.com/manual/core/index-ttl
+//   `index:"-,unique,allowNull,expireAfter={{.Expire}}"` -- evaluate index as a golang template with `cfg` arguments
+func (d *Database) IndexEnsure(cfg interface{}, document interface{}) error {
+
+	el := reflect.ValueOf(document).Elem().Type()
+	numField := el.NumField()
+	documents := d.GetCollectionOf(document)
+
+	for i := 0; i < numField; i++ {
+
+		field := el.Field(i)
+		tag := field.Tag
+
+		indexTag, ok := tag.Lookup("index")
+		if !ok {
+			continue
+		}
+		bsonTag, ok := tag.Lookup("bson")
+		if !ok {
+			return fmt.Errorf("bson tag is not defined for field:%v document:%v", field, document)
+		}
+
+		tmpBuffer := &bytes.Buffer{}
+		tpl, err := template.New("").Parse(indexTag)
+		if err != nil {
+			panic(fmt.Errorf("invalid prop template, %v", indexTag))
+		}
+		err = tpl.Execute(tmpBuffer, cfg)
+		if err != nil {
+			panic(fmt.Errorf("failed to evaluate prop template, %v", indexTag))
+		}
+
+		indexString := tmpBuffer.String()
+		indexValues := strings.Split(indexString, ",")
+		bsonValues := strings.Split(bsonTag, ",")
+
+		var f = false
+		var t = true
+		var key = bsonValues[0]
+		var name = fmt.Sprintf("%s_%s_", indexString, key)
+
+		if len(key) == 0 {
+			panic(fmt.Errorf("cannot evaluate index key"))
+		}
+
+		index := primitive.M{key: 1}
+		opts := &options.IndexOptions{
+			Background: &f,
+			Unique:     &f,
+			Name:       &name,
+		}
+
+		if indexValues[0] == "-" {
+			index[key] = -1
+		}
+
+		for _, prop := range indexValues[1:] {
+			var left string
+			var right string
+
+			pair := strings.SplitN(prop, "=", 2)
+			left = pair[0]
+			if len(pair) > 1 {
+				right = pair[1]
+			}
+
+			switch {
+			case left == "unique":
+				opts.Unique = &t
+
+			case left == "allowNull":
+				expression, isMap := opts.PartialFilterExpression.(primitive.M)
+				if !isMap || expression == nil {
+					expression = primitive.M{}
+				}
+
+				expression[key] = primitive.M{"$exists": true}
+				opts.PartialFilterExpression = expression
+
+			case left == "expireAfter":
+				expireAfter, err := strconv.Atoi(right)
+				if err != nil || expireAfter < 1 {
+					panic(fmt.Errorf("invalid expireAfter value, err: %w", err))
+				}
+
+				expireAfterInt32 := int32(expireAfter)
+				opts.ExpireAfterSeconds = &expireAfterInt32
+
+			case len(left) > 0 && left[0] == '+':
+				compoundValue := left[1:]
+				if len(compoundValue) == 0 {
+					panic(fmt.Errorf("invalid compound value"))
+				}
+
+				if compoundValue[0] == '-' {
+					index[compoundValue[1:]] = -1
+				} else {
+					index[compoundValue] = 1
+				}
+
+			default:
+				panic(fmt.Errorf("unsupported flag:%q in tag:%q of type:%s", prop, tag, el))
+			}
+		}
+
+		_, err = documents.Indexes().CreateOne(d.Context(), mongo.IndexModel{Keys: index, Options: opts})
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
--- a/mongox/database/index_test.go
+++ b/mongox/database/index_test.go
@ -0,0 +1,156 @@
+package database_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/mainnika/mongox-go-driver/v2/mongox-testing/database"
+	"github.com/mainnika/mongox-go-driver/v2/mongox/base/oidbased"
+)
+
+func TestDatabase_Ensure(t *testing.T) {
+
+	db, err := database.NewEphemeral("mongodb://localhost")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	testvalues := []struct {
+		doc      interface{}
+		settings map[string]interface{}
+		index    map[string]interface{}
+	}{
+		{
+			doc: &struct {
+				oidbased.Primary `bson:",inline" json:",inline" collection:"1"`
+
+				Foobar int `bson:"foobar" json:"foobar" index:"-,unique,allowNull,expireAfter=86400"`
+				Foo    int `bson:"foo" json:"foo"`
+				Bar    int `bson:"bar" json:"bar"`
+			}{},
+			index: map[string]interface{}{
+				"background":         false,
+				"expireAfterSeconds": int32(86400),
+				"key": map[string]interface{}{
+					"foobar": int32(-1),
+				},
+				"name": "-,unique,allowNull,expireAfter=86400_foobar_",
+				"partialFilterExpression": map[string]interface{}{
+					"foobar": map[string]interface{}{"$exists": true},
+				},
+				"unique": true,
+			},
+		},
+		{
+			doc: &struct {
+				oidbased.Primary `bson:",inline" json:",inline" collection:"2"`
+
+				Foobar int `bson:"foobar" json:"foobar" index:",unique"`
+			}{},
+			index: map[string]interface{}{
+				"background": false,
+				"key": map[string]interface{}{
+					"foobar": int32(1),
+				},
+				"name":   ",unique_foobar_",
+				"unique": true,
+			},
+		},
+		{
+			doc: &struct {
+				oidbased.Primary `bson:",inline" json:",inline" collection:"3"`
+
+				Foobar int `bson:"foobar" json:"foobar" index:"-,+foo,+-bar,unique,allowNull"`
+				Foo    int `bson:"foo" json:"foo"`
+				Bar    int `bson:"bar" json:"bar"`
+			}{},
+			index: map[string]interface{}{
+				"background": false,
+				"key": map[string]interface{}{
+					"foobar": int32(-1),
+					"foo":    int32(1),
+					"bar":    int32(-1),
+				},
+				"name": "-,+foo,+-bar,unique,allowNull_foobar_",
+				"partialFilterExpression": map[string]interface{}{
+					"foobar": map[string]interface{}{"$exists": true},
+				},
+				"unique": true,
+			},
+		},
+		{
+			doc: &struct {
+				oidbased.Primary `bson:",inline" json:",inline" collection:"4"`
+
+				Foobar int `bson:"foobar" json:"foobar" index:""`
+				Foo    int `bson:"foo" json:"foo"`
+				Bar    int `bson:"bar" json:"bar"`
+			}{},
+			index: map[string]interface{}{
+				"background": false,
+				"key": map[string]interface{}{
+					"foobar": int32(1),
+				},
+				"name": "_foobar_",
+			},
+		},
+		{
+			doc: &struct {
+				oidbased.Primary `bson:",inline" json:",inline" collection:"5"`
+
+				Foobar int `bson:"foobar" json:"foobar" index:"-"`
+				Foo    int `bson:"foo" json:"foo"`
+				Bar    int `bson:"bar" json:"bar"`
+			}{},
+			index: map[string]interface{}{
+				"background": false,
+				"key": map[string]interface{}{
+					"foobar": int32(-1),
+				},
+				"name": "-_foobar_",
+			},
+		},
+		{
+			doc: &struct {
+				oidbased.Primary `bson:",inline" json:",inline" collection:"1"`
+
+				Foobar int `bson:"foobar" json:"foobar" index:"-,unique,allowNull,expireAfter={{.Expire}}"`
+				Foo    int `bson:"foo" json:"foo"`
+				Bar    int `bson:"bar" json:"bar"`
+			}{},
+			settings: map[string]interface{}{
+				"Expire": 86400,
+			},
+			index: map[string]interface{}{
+				"background":         false,
+				"expireAfterSeconds": int32(86400),
+				"key": map[string]interface{}{
+					"foobar": int32(-1),
+				},
+				"name": "-,unique,allowNull,expireAfter=86400_foobar_",
+				"partialFilterExpression": map[string]interface{}{
+					"foobar": map[string]interface{}{"$exists": true},
+				},
+				"unique": true,
+			},
+		},
+	}
+
+	for _, tt := range testvalues {
+		err = db.IndexEnsure(tt.settings, tt.doc)
+		assert.NoError(t, err)
+
+		indexes, _ := db.GetCollectionOf(tt.doc).Indexes().List(db.Context())
+		index := new(map[string]interface{})
+
+		indexes.Next(db.Context()) // skip _id_
+		indexes.Next(db.Context())
+		indexes.Decode(&index)
+
+		for k, v := range tt.index {
+			assert.Equal(t, v, (*index)[k])
+		}
+	}
+}
--- a/mongox/mongox.go
+++ b/mongox/mongox.go
@ -28,6 +28,7 @@ type Database interface {
 	LoadOne(target interface{}, filters ...interface{}) error
 	LoadStream(target interface{}, filters ...interface{}) (StreamLoader, error)
 	SaveOne(source interface{}) error
+	IndexEnsure(cfg interface{}, document interface{}) error
 }

 // StreamLoader is a interface to control database cursor