266 lines
6.1 KiB
Go
266 lines
6.1 KiB
Go
package sqlx
|
|
|
|
import (
|
|
"bytes"
|
|
"database/sql/driver"
|
|
"errors"
|
|
"reflect"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/jmoiron/sqlx/reflectx"
|
|
)
|
|
|
|
// Bindvar types supported by Rebind, BindMap and BindStruct.
|
|
const (
|
|
UNKNOWN = iota
|
|
QUESTION
|
|
DOLLAR
|
|
NAMED
|
|
AT
|
|
)
|
|
|
|
var defaultBinds = map[int][]string{
|
|
DOLLAR: []string{"postgres", "pgx", "pq-timeouts", "cloudsqlpostgres", "ql", "nrpostgres", "cockroach"},
|
|
QUESTION: []string{"mysql", "sqlite3", "nrmysql", "nrsqlite3"},
|
|
NAMED: []string{"oci8", "ora", "goracle", "godror"},
|
|
AT: []string{"sqlserver"},
|
|
}
|
|
|
|
var binds sync.Map
|
|
|
|
func init() {
|
|
for bind, drivers := range defaultBinds {
|
|
for _, driver := range drivers {
|
|
BindDriver(driver, bind)
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
// BindType returns the bindtype for a given database given a drivername.
|
|
func BindType(driverName string) int {
|
|
itype, ok := binds.Load(driverName)
|
|
if !ok {
|
|
return UNKNOWN
|
|
}
|
|
return itype.(int)
|
|
}
|
|
|
|
// BindDriver sets the BindType for driverName to bindType.
|
|
func BindDriver(driverName string, bindType int) {
|
|
binds.Store(driverName, bindType)
|
|
}
|
|
|
|
// FIXME: this should be able to be tolerant of escaped ?'s in queries without
|
|
// losing much speed, and should be to avoid confusion.
|
|
|
|
// Rebind a query from the default bindtype (QUESTION) to the target bindtype.
|
|
func Rebind(bindType int, query string) string {
|
|
switch bindType {
|
|
case QUESTION, UNKNOWN:
|
|
return query
|
|
}
|
|
|
|
// Add space enough for 10 params before we have to allocate
|
|
rqb := make([]byte, 0, len(query)+10)
|
|
|
|
var i, j int
|
|
|
|
for i = strings.Index(query, "?"); i != -1; i = strings.Index(query, "?") {
|
|
rqb = append(rqb, query[:i]...)
|
|
|
|
switch bindType {
|
|
case DOLLAR:
|
|
rqb = append(rqb, '$')
|
|
case NAMED:
|
|
rqb = append(rqb, ':', 'a', 'r', 'g')
|
|
case AT:
|
|
rqb = append(rqb, '@', 'p')
|
|
}
|
|
|
|
j++
|
|
rqb = strconv.AppendInt(rqb, int64(j), 10)
|
|
|
|
query = query[i+1:]
|
|
}
|
|
|
|
return string(append(rqb, query...))
|
|
}
|
|
|
|
// Experimental implementation of Rebind which uses a bytes.Buffer. The code is
|
|
// much simpler and should be more resistant to odd unicode, but it is twice as
|
|
// slow. Kept here for benchmarking purposes and to possibly replace Rebind if
|
|
// problems arise with its somewhat naive handling of unicode.
|
|
func rebindBuff(bindType int, query string) string {
|
|
if bindType != DOLLAR {
|
|
return query
|
|
}
|
|
|
|
b := make([]byte, 0, len(query))
|
|
rqb := bytes.NewBuffer(b)
|
|
j := 1
|
|
for _, r := range query {
|
|
if r == '?' {
|
|
rqb.WriteRune('$')
|
|
rqb.WriteString(strconv.Itoa(j))
|
|
j++
|
|
} else {
|
|
rqb.WriteRune(r)
|
|
}
|
|
}
|
|
|
|
return rqb.String()
|
|
}
|
|
|
|
func asSliceForIn(i interface{}) (v reflect.Value, ok bool) {
|
|
if i == nil {
|
|
return reflect.Value{}, false
|
|
}
|
|
|
|
v = reflect.ValueOf(i)
|
|
t := reflectx.Deref(v.Type())
|
|
|
|
// Only expand slices
|
|
if t.Kind() != reflect.Slice {
|
|
return reflect.Value{}, false
|
|
}
|
|
|
|
// []byte is a driver.Value type so it should not be expanded
|
|
if t == reflect.TypeOf([]byte{}) {
|
|
return reflect.Value{}, false
|
|
|
|
}
|
|
|
|
return v, true
|
|
}
|
|
|
|
// In expands slice values in args, returning the modified query string
|
|
// and a new arg list that can be executed by a database. The `query` should
|
|
// use the `?` bindVar. The return value uses the `?` bindVar.
|
|
func In(query string, args ...interface{}) (string, []interface{}, error) {
|
|
// argMeta stores reflect.Value and length for slices and
|
|
// the value itself for non-slice arguments
|
|
type argMeta struct {
|
|
v reflect.Value
|
|
i interface{}
|
|
length int
|
|
}
|
|
|
|
var flatArgsCount int
|
|
var anySlices bool
|
|
|
|
var stackMeta [32]argMeta
|
|
|
|
var meta []argMeta
|
|
if len(args) <= len(stackMeta) {
|
|
meta = stackMeta[:len(args)]
|
|
} else {
|
|
meta = make([]argMeta, len(args))
|
|
}
|
|
|
|
for i, arg := range args {
|
|
if a, ok := arg.(driver.Valuer); ok {
|
|
var err error
|
|
arg, err = a.Value()
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
}
|
|
|
|
if v, ok := asSliceForIn(arg); ok {
|
|
meta[i].length = v.Len()
|
|
meta[i].v = v
|
|
|
|
anySlices = true
|
|
flatArgsCount += meta[i].length
|
|
|
|
if meta[i].length == 0 {
|
|
return "", nil, errors.New("empty slice passed to 'in' query")
|
|
}
|
|
} else {
|
|
meta[i].i = arg
|
|
flatArgsCount++
|
|
}
|
|
}
|
|
|
|
// don't do any parsing if there aren't any slices; note that this means
|
|
// some errors that we might have caught below will not be returned.
|
|
if !anySlices {
|
|
return query, args, nil
|
|
}
|
|
|
|
newArgs := make([]interface{}, 0, flatArgsCount)
|
|
|
|
var buf strings.Builder
|
|
buf.Grow(len(query) + len(", ?")*flatArgsCount)
|
|
|
|
var arg, offset int
|
|
|
|
for i := strings.IndexByte(query[offset:], '?'); i != -1; i = strings.IndexByte(query[offset:], '?') {
|
|
if arg >= len(meta) {
|
|
// if an argument wasn't passed, lets return an error; this is
|
|
// not actually how database/sql Exec/Query works, but since we are
|
|
// creating an argument list programmatically, we want to be able
|
|
// to catch these programmer errors earlier.
|
|
return "", nil, errors.New("number of bindVars exceeds arguments")
|
|
}
|
|
|
|
argMeta := meta[arg]
|
|
arg++
|
|
|
|
// not a slice, continue.
|
|
// our questionmark will either be written before the next expansion
|
|
// of a slice or after the loop when writing the rest of the query
|
|
if argMeta.length == 0 {
|
|
offset = offset + i + 1
|
|
newArgs = append(newArgs, argMeta.i)
|
|
continue
|
|
}
|
|
|
|
// write everything up to and including our ? character
|
|
buf.WriteString(query[:offset+i+1])
|
|
|
|
for si := 1; si < argMeta.length; si++ {
|
|
buf.WriteString(", ?")
|
|
}
|
|
|
|
newArgs = appendReflectSlice(newArgs, argMeta.v, argMeta.length)
|
|
|
|
// slice the query and reset the offset. this avoids some bookkeeping for
|
|
// the write after the loop
|
|
query = query[offset+i+1:]
|
|
offset = 0
|
|
}
|
|
|
|
buf.WriteString(query)
|
|
|
|
if arg < len(meta) {
|
|
return "", nil, errors.New("number of bindVars less than number arguments")
|
|
}
|
|
|
|
return buf.String(), newArgs, nil
|
|
}
|
|
|
|
func appendReflectSlice(args []interface{}, v reflect.Value, vlen int) []interface{} {
|
|
switch val := v.Interface().(type) {
|
|
case []interface{}:
|
|
args = append(args, val...)
|
|
case []int:
|
|
for i := range val {
|
|
args = append(args, val[i])
|
|
}
|
|
case []string:
|
|
for i := range val {
|
|
args = append(args, val[i])
|
|
}
|
|
default:
|
|
for si := 0; si < vlen; si++ {
|
|
args = append(args, v.Index(si).Interface())
|
|
}
|
|
}
|
|
|
|
return args
|
|
}
|