123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266 |
- 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
- }
|