parse.go

package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/tidwall/gjson"
)

const fileTemplate = `// Code generated by '%s' DO NOT EDIT.
package %s

import (
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
)
`

const funcTemplate = `
func %s(%s) %s {
	return %s
}`

type paramSpec struct {
	name, typ string
	index     int
}

var paramsMap map[string]paramSpec

func buildFunctionFromFile(filename string) string {
	var inFile *os.File

	var err error
	if filename == "-" {
		inFile = os.Stdin
	} else {
		if fixSource {
			fixSourceFile(filename)
		}

		inFile, err = os.Open(filename)
		defer inFile.Close()
		// TODO err checking

	}

	paramsMap = make(map[string]paramSpec)

	bytes, err := ioutil.ReadAll(inFile)

	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: (%s) could not read an input file: %s\n", filename, err)
		os.Exit(1)
	}

	inPath := strings.Split(filename, string(os.PathSeparator))
	inFileBase := strings.Split(inPath[len(inPath)-1], ".")[0]
	fnName := "Get" + strings.Title(inFileBase) //strings.ToLower(fNameParts[0]))

	rval, err := buildFunction(bytes, fnName)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: (%s) failed to build aggregation function: %s\n", filename, err)
		os.Exit(1)
	}
	return rval
}

func buildFunction(bytes []byte, fnName string) (string, error) {
	result := gjson.ParseBytes(bytes)

	if !result.Exists() {
		return "", errors.New("Could not parse JSON - Malformed JSON or embedded JS, such as 'new Date(...)' will commonly cause this problem.")
	}

	body := recursiveParseAny(result)

	var returnType string = ""
	if result.IsArray() {
		// returnType = "bson.A"
		returnType = "mongo.Pipeline"
		var didCut bool
		if body, didCut = strings.CutPrefix(body, "bson.A"); !didCut {
			return "", errors.New("Internal Error: Top level array on input did not result in bson.A after initial parse.")
		}

		body = "mongo.Pipeline" + body

	} else if result.IsObject() {
		returnType = "bson.D"
	} else {
		return "", errors.New("Root of JSON was not an object or array")
	}

	// Convert the vars to a map, to eliminate duplicates
	fnParams := ""
	params := sortParams()

	for i, p := range params {
		fnParams += fmt.Sprint(p.name, " ", p.typ)
		if i < len(params)-1 {
			fnParams += ", "
		}
	}

	return fmt.Sprintf(funcTemplate, fnName, fnParams, returnType, body), nil
}

func recursiveParseArray(result []gjson.Result) string {
	rval := "bson.A{\n"

	for _, r := range result {
		content := recursiveParseAny(r)
		rval += content + ",\n"
	}
	rval += "}"

	return rval
}

func recursiveParseObject(result gjson.Result) string {
	rval := "bson.D{\n"

	result.ForEach(func(k, v gjson.Result) bool {
		content := recursiveParseAny(v)
		if structKeys {
			rval += fmt.Sprintf(`{Key: "%s", Value: %s},
			`, k.String(), content)
		} else {
			rval += fmt.Sprintf(`{"%s", %s},
			`, k.String(), content)
		}

		return true
	})

	rval += "}"
	return rval
}

func recursiveParseAny(result gjson.Result) string {
	switch result.Type {
	case gjson.JSON:
		if result.IsArray() {
			return recursiveParseArray(result.Array())
		} else if result.IsObject() {
			return recursiveParseObject(result)
		}
		panic("gjson type=JSON, but is not object or array!")

	case gjson.Number:
		return result.String()

	case gjson.String:
		rawVal := result.String()

		if strings.HasPrefix(rawVal, "%%") {

			return handleParam(rawVal)
		}
		return fmt.Sprintf(`"%s"`, result.String())

	case gjson.Null:
		return "nil"

	case gjson.True:
		return "true"

	case gjson.False:
		return "false"

	default:
		panic("unhandled json type! " + result.String())
	}

}

func handleParam(paramString string) string {
	var spec paramSpec
	var err error
	parts := strings.Split(paramString[2:], "%")

	if len(parts) == 1 {
		if parts[0] == "" { // "%%"
			parts = []string{fmt.Sprint("p", len(paramsMap)), "string"}
		} else { // "%%paramName"
			parts = append(parts, "string")
		}
	} else if parts[1] == "" {
		parts[1] = "string"
	}

	if len(parts) == 2 { // %%paramName%string
		parts = append(parts, fmt.Sprint(len(paramsMap)+1024)) // Add 1024 to avoid collisions with user-indexed params
	} else if parts[2] == "" {
		parts[2] = fmt.Sprint(len(paramsMap) + 1024)
	}

	spec.name = parts[0]
	spec.typ = parts[1]
	spec.index, err = strconv.Atoi(parts[2])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Could not parse %s to an index (param name %s); %s", parts[2], parts[0], err.Error())
		os.Exit(1)
	}

	// Before checking/setting the index, find out if this is a repeated parameter
	if existingSpec, ok := paramsMap[spec.name]; ok {
		if existingSpec.typ != spec.typ {
			fmt.Fprintf(os.Stderr,
				"ERROR: Parameter type conflict. Parameter %s (type %s) was previously declared as type %s\n",
				spec.name, spec.typ, existingSpec.typ,
			)
			os.Exit(2)
		}

	} else {
		paramsMap[spec.name] = spec
	}

	return spec.name
}

func sortParams() []paramSpec {
	rval := make(paramSlice, 0)
	for _, v := range paramsMap {
		rval = append(rval, v)
	}
	sort.Sort(rval)

	return rval
}

type paramSlice []paramSpec

func (ps paramSlice) Len() int {
	return len(ps)
}
func (ps paramSlice) Swap(i, j int) {
	ps[i], ps[j] = ps[j], ps[i]
}
func (ps paramSlice) Less(i, j int) bool {
	return ps[i].index < ps[j].index
}