relspecgo/vendor/github.com/uptrace/bun/schema/querygen.go

package schema

import (
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/uptrace/bun/dialect"
	"github.com/uptrace/bun/dialect/feature"
	"github.com/uptrace/bun/internal"
	"github.com/uptrace/bun/internal/parser"
)

var nopQueryGen = QueryGen{
	dialect: newNopDialect(),
}

type QueryGen struct {
	dialect Dialect
	args    *namedArgList
}

func NewQueryGen(dialect Dialect) QueryGen {
	return QueryGen{
		dialect: dialect,
	}
}

func NewNopQueryGen() QueryGen {
	return nopQueryGen
}

func (f QueryGen) IsNop() bool {
	return f.dialect.Name() == dialect.Invalid
}

func (f QueryGen) Dialect() Dialect {
	return f.dialect
}

func (f QueryGen) IdentQuote() byte {
	return f.dialect.IdentQuote()
}

func (gen QueryGen) Append(b []byte, v any) []byte {
	switch v := v.(type) {
	case nil:
		return dialect.AppendNull(b)
	case bool:
		return dialect.AppendBool(b, v)
	case int:
		return strconv.AppendInt(b, int64(v), 10)
	case int32:
		return strconv.AppendInt(b, int64(v), 10)
	case int64:
		return strconv.AppendInt(b, v, 10)
	case uint:
		return strconv.AppendInt(b, int64(v), 10)
	case uint32:
		return gen.Dialect().AppendUint32(b, v)
	case uint64:
		return gen.Dialect().AppendUint64(b, v)
	case float32:
		return dialect.AppendFloat32(b, v)
	case float64:
		return dialect.AppendFloat64(b, v)
	case string:
		return gen.Dialect().AppendString(b, v)
	case time.Time:
		return gen.Dialect().AppendTime(b, v)
	case []byte:
		return gen.Dialect().AppendBytes(b, v)
	case QueryAppender:
		return AppendQueryAppender(gen, b, v)
	default:
		vv := reflect.ValueOf(v)
		if vv.Kind() == reflect.Ptr && vv.IsNil() {
			return dialect.AppendNull(b)
		}
		appender := Appender(gen.Dialect(), vv.Type())
		return appender(gen, b, vv)
	}
}

func (f QueryGen) AppendName(b []byte, name string) []byte {
	return dialect.AppendName(b, name, f.IdentQuote())
}

func (f QueryGen) AppendIdent(b []byte, ident string) []byte {
	return dialect.AppendIdent(b, ident, f.IdentQuote())
}

func (f QueryGen) AppendValue(b []byte, v reflect.Value) []byte {
	if v.Kind() == reflect.Ptr && v.IsNil() {
		return dialect.AppendNull(b)
	}
	appender := Appender(f.dialect, v.Type())
	return appender(f, b, v)
}

func (f QueryGen) HasFeature(feature feature.Feature) bool {
	return f.dialect.Features().Has(feature)
}

func (f QueryGen) WithArg(arg NamedArgAppender) QueryGen {
	return QueryGen{
		dialect: f.dialect,
		args:    f.args.WithArg(arg),
	}
}

func (f QueryGen) WithNamedArg(name string, value any) QueryGen {
	return QueryGen{
		dialect: f.dialect,
		args:    f.args.WithArg(&namedArg{name: name, value: value}),
	}
}

func (f QueryGen) FormatQuery(query string, args ...any) string {
	if f.IsNop() || (args == nil && f.args == nil) || strings.IndexByte(query, '?') == -1 {
		return query
	}
	return internal.String(f.AppendQuery(nil, query, args...))
}

func (f QueryGen) AppendQuery(dst []byte, query string, args ...any) []byte {
	if f.IsNop() || (args == nil && f.args == nil) || strings.IndexByte(query, '?') == -1 {
		return append(dst, query...)
	}
	return f.append(dst, parser.NewString(query), args)
}

func (f QueryGen) append(dst []byte, p *parser.Parser, args []any) []byte {
	var namedArgs NamedArgAppender
	if len(args) == 1 {
		if v, ok := args[0].(NamedArgAppender); ok {
			namedArgs = v
		} else if v, ok := newStructArgs(f, args[0]); ok {
			namedArgs = v
		}
	}

	var argIndex int
	for p.Valid() {
		b, ok := p.ReadSep('?')
		if !ok {
			dst = append(dst, b...)
			continue
		}
		if len(b) > 0 && b[len(b)-1] == '\\' {
			dst = append(dst, b[:len(b)-1]...)
			dst = append(dst, '?')
			continue
		}
		dst = append(dst, b...)

		name, numeric := p.ReadIdentifier()
		if name != "" {
			if numeric {
				idx, err := strconv.Atoi(name)
				if err != nil {
					goto restore_arg
				}

				if idx >= len(args) {
					goto restore_arg
				}

				dst = f.appendArg(dst, args[idx])
				continue
			}

			if namedArgs != nil {
				dst, ok = namedArgs.AppendNamedArg(f, dst, name)
				if ok {
					continue
				}
			}

			dst, ok = f.args.AppendNamedArg(f, dst, name)
			if ok {
				continue
			}

		restore_arg:
			dst = append(dst, '?')
			dst = append(dst, name...)
			continue
		}

		if argIndex >= len(args) {
			dst = append(dst, '?')
			continue
		}

		arg := args[argIndex]
		argIndex++

		dst = f.appendArg(dst, arg)
	}

	return dst
}

func (gen QueryGen) appendArg(b []byte, arg any) []byte {
	switch arg := arg.(type) {
	case QueryAppender:
		bb, err := arg.AppendQuery(gen, b)
		if err != nil {
			return dialect.AppendError(b, err)
		}
		return bb
	default:
		return gen.Append(b, arg)
	}
}

//------------------------------------------------------------------------------

type NamedArgAppender interface {
	AppendNamedArg(gen QueryGen, b []byte, name string) ([]byte, bool)
}

type namedArgList struct {
	arg  NamedArgAppender
	next *namedArgList
}

func (l *namedArgList) WithArg(arg NamedArgAppender) *namedArgList {
	return &namedArgList{
		arg:  arg,
		next: l,
	}
}

func (l *namedArgList) AppendNamedArg(gen QueryGen, b []byte, name string) ([]byte, bool) {
	for l != nil && l.arg != nil {
		if b, ok := l.arg.AppendNamedArg(gen, b, name); ok {
			return b, true
		}
		l = l.next
	}
	return b, false
}

//------------------------------------------------------------------------------

type namedArg struct {
	name  string
	value any
}

var _ NamedArgAppender = (*namedArg)(nil)

func (a *namedArg) AppendNamedArg(gen QueryGen, b []byte, name string) ([]byte, bool) {
	if a.name == name {
		return gen.appendArg(b, a.value), true
	}
	return b, false
}

//------------------------------------------------------------------------------

type structArgs struct {
	table *Table
	strct reflect.Value
}

var _ NamedArgAppender = (*structArgs)(nil)

func newStructArgs(gen QueryGen, strct any) (*structArgs, bool) {
	v := reflect.ValueOf(strct)
	if !v.IsValid() {
		return nil, false
	}

	v = reflect.Indirect(v)
	if v.Kind() != reflect.Struct {
		return nil, false
	}

	return &structArgs{
		table: gen.Dialect().Tables().Get(v.Type()),
		strct: v,
	}, true
}

func (m *structArgs) AppendNamedArg(gen QueryGen, b []byte, name string) ([]byte, bool) {
	return m.table.AppendNamedArg(gen, b, name, m.strct)
}