Mongo json filter

query.go

// Helper for dealing with mongo quries as strings
// Example:
// 
// f := mongo.MustCompileQuery(`{
// 	"id" : "$1",
// 	"start": { "$lte": "$2" },
// 	"$or": [
// 		{ "end": { "$exists": false } },
// 		{ "end": null },
// 		{ "end": { "$gte": "$2" } }
// 	]}`,
// 	"$1", vppID,
// 	"$2", date,
// )
// var result model.Object
// err := collection.FindOne(ctx, f).Decode(&result)
package mongo

import (
	"errors"
	"fmt"
	"strings"
	"sync"

	"go.mongodb.org/mongo-driver/bson"
)

var (
	queryCache = make(map[string]interface{})
	queryLock  = sync.RWMutex{}
)

// MustCompileQuery creates bson request from specified JSON and parameters. Same as CompileQuery
// but panics if there was an error during parsing JSON..
func MustCompileQuery(query string, keyValues ...interface{}) interface{} {
	i, err := CompileQuery(query, keyValues...)
	if err != nil {
		panic(fmt.Sprintf("can not compile query: \"%s\", error: %v", query, err))
	}

	return i
}

// CompileQuery creates bson request from specified JSON and parameters.
func CompileQuery(query string, keyValues ...interface{}) (interface{}, error) {
	queryLock.RLock()
	cq, ok := queryCache[query]
	queryLock.RUnlock()

	if ok {
		cq = clone(cq)
		err := setParams(cq, keyValues...)
		return cq, err
	}

	queryLock.Lock()
	defer queryLock.Unlock()

	cq, ok = queryCache[query]
	if ok {
		cq = clone(cq)
		err := setParams(cq, keyValues...)
		return cq, err
	}

	var i interface{}

	err := bson.UnmarshalExtJSON([]byte(query), true, &i)
	if err != nil {
		return nil, err
	}

	queryCache[query] = i

	i = clone(i)
	err = setParams(i, keyValues...)
	return i, err
}

func setParams(query interface{}, keyValues ...interface{}) error {
	prmMap := map[string]interface{}{}

	if len(keyValues) == 0 {
		return nil
	}

	if len(keyValues)%2 != 0 {
		return errors.New("keyValues should be pairs of string key and any value")
	}

	for i := 0; i < len(keyValues); i += 2 {
		s, ok := keyValues[i].(string)
		if !ok {
			return fmt.Errorf("parameter key %v must be string", keyValues[i])
		}

		prmMap[s] = keyValues[i+1]
	}

	traverse(query, prmMap)
	return nil
}

func traverse(node interface{}, prmMap map[string]interface{}) {
	switch v := node.(type) {
	case bson.D:
		for i, e := range v {
			s, ok := e.Value.(string)
			if ok && strings.HasPrefix(s, "$") {
				pv, ok := prmMap[s]
				if !ok {
					continue
				}

				v[i] = bson.E{Key: e.Key, Value: pv}

				continue
			}

			traverse(e.Value, prmMap)
		}
	case bson.A:
		for _, e := range v {
			traverse(e, prmMap)
		}
	}
}

func clone(node interface{}) interface{} {
	switch v := node.(type) {
	case bson.D:
		c := bson.D{}
		for _, e := range v {
			c = append(c, bson.E{Key: e.Key, Value: clone(e.Value)})
		}
		return c
	case bson.A:
		a := bson.A{}
		for _, e := range v {
			a = append(a, clone(e))
		}
		return a
	}

	return node
}