super-graph/core/args.go

169 lines
3.3 KiB
Go

package core
import (
"bytes"
"encoding/json"
"fmt"
"io"
"github.com/dosco/super-graph/jsn"
)
// argMap function is used to string replace variables with values by
// the fasttemplate code
func (c *scontext) argMap() func(w io.Writer, tag string) (int, error) {
return func(w io.Writer, tag string) (int, error) {
switch tag {
case "user_id_provider":
if v := c.Value(UserIDProviderKey); v != nil {
return io.WriteString(w, v.(string))
}
return 0, argErr("user_id_provider")
case "user_id":
if v := c.Value(UserIDKey); v != nil {
return io.WriteString(w, v.(string))
}
return 0, argErr("user_id")
case "user_role":
if v := c.Value(UserRoleKey); v != nil {
return io.WriteString(w, v.(string))
}
return 0, argErr("user_role")
}
fields := jsn.Get(c.vars, [][]byte{[]byte(tag)})
if len(fields) == 0 {
return 0, argErr(tag)
}
v := fields[0].Value
// Open and close quotes
if len(v) >= 2 && v[0] == '"' && v[len(v)-1] == '"' {
fields[0].Value = v[1 : len(v)-1]
}
if tag == "cursor" {
if bytes.EqualFold(v, []byte("null")) {
return io.WriteString(w, ``)
}
v1, err := c.sg.decrypt(string(fields[0].Value))
if err != nil {
return 0, err
}
return w.Write(v1)
}
return w.Write(escSQuote(fields[0].Value))
}
}
// argList function is used to create a list of arguments to pass
// to a prepared statement. FYI no escaping of single quotes is
// needed here
func (c *scontext) argList(args [][]byte) ([]interface{}, error) {
vars := make([]interface{}, len(args))
var fields map[string]json.RawMessage
var err error
if len(c.vars) != 0 {
fields, _, err = jsn.Tree(c.vars)
if err != nil {
return nil, err
}
}
for i := range args {
av := args[i]
switch {
case bytes.Equal(av, []byte("user_id")):
if v := c.Value(UserIDKey); v != nil {
vars[i] = v.(string)
} else {
return nil, argErr("user_id")
}
case bytes.Equal(av, []byte("user_id_provider")):
if v := c.Value(UserIDProviderKey); v != nil {
vars[i] = v.(string)
} else {
return nil, argErr("user_id_provider")
}
case bytes.Equal(av, []byte("user_role")):
if v := c.Value(UserRoleKey); v != nil {
vars[i] = v.(string)
} else {
return nil, argErr("user_role")
}
case bytes.Equal(av, []byte("cursor")):
if v, ok := fields["cursor"]; ok && v[0] == '"' {
v1, err := c.sg.decrypt(string(v[1 : len(v)-1]))
if err != nil {
return nil, err
}
vars[i] = v1
} else {
return nil, argErr("cursor")
}
default:
if v, ok := fields[string(av)]; ok {
switch v[0] {
case '[', '{':
vars[i] = v
default:
var val interface{}
if err := json.Unmarshal(v, &val); err != nil {
return nil, err
}
vars[i] = val
}
} else {
return nil, argErr(string(av))
}
}
}
return vars, nil
}
//
func escSQuote(b []byte) []byte {
var buf *bytes.Buffer
s := 0
for i := range b {
if b[i] == '\'' {
if buf == nil {
buf = &bytes.Buffer{}
}
buf.Write(b[s:i])
buf.WriteString(`''`)
s = i + 1
}
}
if buf == nil {
return b
}
l := len(b)
if s < (l - 1) {
buf.Write(b[s:l])
}
return buf.Bytes()
}
func argErr(name string) error {
return fmt.Errorf("query requires variable '%s' to be set", name)
}