Skip to content

Commit

Permalink
fix: improve range type to support driver.Valuer and sql.Scanner
Browse files Browse the repository at this point in the history
  • Loading branch information
vmihailenco committed Nov 26, 2024
1 parent 5d55190 commit 856e12b
Show file tree
Hide file tree
Showing 5 changed files with 219 additions and 199 deletions.
4 changes: 2 additions & 2 deletions dialect/pgdialect/append.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,9 +67,9 @@ func appendMapStringString(b []byte, m map[string]string) []byte {
b = append(b, '\'')

for key, value := range m {
b = arrayAppendString(b, key)
b = appendStringElem(b, key)
b = append(b, '=', '>')
b = arrayAppendString(b, value)
b = appendStringElem(b, value)
b = append(b, ',')
}
if len(m) > 0 {
Expand Down
88 changes: 8 additions & 80 deletions dialect/pgdialect/array.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,11 @@ package pgdialect
import (
"database/sql"
"database/sql/driver"
"encoding/hex"
"fmt"
"math"
"reflect"
"strconv"
"time"
"unicode/utf8"

"github.com/uptrace/bun/dialect"
"github.com/uptrace/bun/internal"
Expand Down Expand Up @@ -146,52 +144,29 @@ func (d *Dialect) arrayElemAppender(typ reflect.Type) schema.AppenderFunc {
}
switch typ.Kind() {
case reflect.String:
return arrayAppendStringValue
return appendStringElemValue
case reflect.Slice:
if typ.Elem().Kind() == reflect.Uint8 {
return arrayAppendBytesValue
return appendBytesElemValue
}
}
return schema.Appender(d, typ)
}

func arrayAppend(fmter schema.Formatter, b []byte, v interface{}) []byte {
switch v := v.(type) {
case int64:
return strconv.AppendInt(b, v, 10)
case float64:
return arrayAppendFloat64(b, v)
case bool:
return dialect.AppendBool(b, v)
case []byte:
return arrayAppendBytes(b, v)
case string:
return arrayAppendString(b, v)
case time.Time:
b = append(b, '"')
b = appendTime(b, v)
b = append(b, '"')
return b
default:
err := fmt.Errorf("pgdialect: can't append %T", v)
return dialect.AppendError(b, err)
}
}

func arrayAppendStringValue(fmter schema.Formatter, b []byte, v reflect.Value) []byte {
return arrayAppendString(b, v.String())
func appendStringElemValue(fmter schema.Formatter, b []byte, v reflect.Value) []byte {
return appendStringElem(b, v.String())
}

func arrayAppendBytesValue(fmter schema.Formatter, b []byte, v reflect.Value) []byte {
return arrayAppendBytes(b, v.Bytes())
func appendBytesElemValue(fmter schema.Formatter, b []byte, v reflect.Value) []byte {
return appendBytesElem(b, v.Bytes())
}

func arrayAppendDriverValue(fmter schema.Formatter, b []byte, v reflect.Value) []byte {
iface, err := v.Interface().(driver.Valuer).Value()
if err != nil {
return dialect.AppendError(b, err)
}
return arrayAppend(fmter, b, iface)
return appendElem(b, iface)
}

func appendStringSliceValue(fmter schema.Formatter, b []byte, v reflect.Value) []byte {
Expand All @@ -208,7 +183,7 @@ func appendStringSlice(b []byte, ss []string) []byte {

b = append(b, '{')
for _, s := range ss {
b = arrayAppendString(b, s)
b = appendStringElem(b, s)
b = append(b, ',')
}
if len(ss) > 0 {
Expand Down Expand Up @@ -617,50 +592,3 @@ func toBytes(src interface{}) ([]byte, error) {
return nil, fmt.Errorf("pgdialect: got %T, wanted []byte or string", src)
}
}

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

func arrayAppendBytes(b []byte, bs []byte) []byte {
if bs == nil {
return dialect.AppendNull(b)
}

b = append(b, `"\\x`...)

s := len(b)
b = append(b, make([]byte, hex.EncodedLen(len(bs)))...)
hex.Encode(b[s:], bs)

b = append(b, '"')

return b
}

func arrayAppendString(b []byte, s string) []byte {
b = append(b, '"')
for _, r := range s {
switch r {
case 0:
// ignore
case '\'':
b = append(b, "''"...)
case '"':
b = append(b, '\\', '"')
case '\\':
b = append(b, '\\', '\\')
default:
if r < utf8.RuneSelf {
b = append(b, byte(r))
break
}
l := len(b)
if cap(b)-l < utf8.UTFMax {
b = append(b, make([]byte, utf8.UTFMax)...)
}
n := utf8.EncodeRune(b[l:l+utf8.UTFMax], r)
b = b[:l+n]
}
}
b = append(b, '"')
return b
}
87 changes: 87 additions & 0 deletions dialect/pgdialect/elem.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
package pgdialect

import (
"database/sql/driver"
"encoding/hex"
"fmt"
"strconv"
"time"
"unicode/utf8"

"github.com/uptrace/bun/dialect"
)

func appendElem(buf []byte, val interface{}) []byte {
switch val := val.(type) {
case int64:
return strconv.AppendInt(buf, val, 10)
case float64:
return arrayAppendFloat64(buf, val)
case bool:
return dialect.AppendBool(buf, val)
case []byte:
return appendBytesElem(buf, val)
case string:
return appendStringElem(buf, val)
case time.Time:
buf = append(buf, '"')
buf = appendTime(buf, val)
buf = append(buf, '"')
return buf
case driver.Valuer:
val2, err := val.Value()
if err != nil {
err := fmt.Errorf("pgdialect: can't append elem value: %w", err)
return dialect.AppendError(buf, err)
}
return appendElem(buf, val2)
default:
err := fmt.Errorf("pgdialect: can't append elem %T", val)
return dialect.AppendError(buf, err)
}
}

func appendBytesElem(b []byte, bs []byte) []byte {
if bs == nil {
return dialect.AppendNull(b)
}

b = append(b, `"\\x`...)

s := len(b)
b = append(b, make([]byte, hex.EncodedLen(len(bs)))...)
hex.Encode(b[s:], bs)

b = append(b, '"')

return b
}

func appendStringElem(b []byte, s string) []byte {
b = append(b, '"')
for _, r := range s {
switch r {
case 0:
// ignore
case '\'':
b = append(b, "''"...)
case '"':
b = append(b, '\\', '"')
case '\\':
b = append(b, '\\', '\\')
default:
if r < utf8.RuneSelf {
b = append(b, byte(r))
break
}
l := len(b)
if cap(b)-l < utf8.UTFMax {
b = append(b, make([]byte, utf8.UTFMax)...)
}
n := utf8.EncodeRune(b[l:l+utf8.UTFMax], r)
b = b[:l+n]
}
}
b = append(b, '"')
return b
}
107 changes: 107 additions & 0 deletions dialect/pgdialect/parser.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
package pgdialect

import (
"bytes"
"encoding/hex"

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

type pgparser struct {
parser.Parser
buf []byte
}

func newParser(b []byte) *pgparser {
p := new(pgparser)
p.Reset(b)
return p
}

func (p *pgparser) ReadLiteral(ch byte) []byte {
p.Unread()
lit, _ := p.ReadSep(',')
return lit
}

func (p *pgparser) ReadUnescapedSubstring(ch byte) ([]byte, error) {
return p.readSubstring(ch, false)
}

func (p *pgparser) ReadSubstring(ch byte) ([]byte, error) {
return p.readSubstring(ch, true)
}

func (p *pgparser) readSubstring(ch byte, escaped bool) ([]byte, error) {
ch, err := p.ReadByte()
if err != nil {
return nil, err
}

p.buf = p.buf[:0]
for {
if ch == '"' {
break
}

next, err := p.ReadByte()
if err != nil {
return nil, err
}

if ch == '\\' {
switch next {
case '\\', '"':
p.buf = append(p.buf, next)

ch, err = p.ReadByte()
if err != nil {
return nil, err
}
default:
p.buf = append(p.buf, '\\')
ch = next
}
continue
}

if escaped && ch == '\'' && next == '\'' {
p.buf = append(p.buf, next)
ch, err = p.ReadByte()
if err != nil {
return nil, err
}
continue
}

p.buf = append(p.buf, ch)
ch = next
}

if bytes.HasPrefix(p.buf, []byte("\\x")) && len(p.buf)%2 == 0 {
data := p.buf[2:]
buf := make([]byte, hex.DecodedLen(len(data)))
n, err := hex.Decode(buf, data)
if err != nil {
return nil, err
}
return buf[:n], nil
}

return p.buf, nil
}

func (p *pgparser) ReadRange(ch byte) ([]byte, error) {
p.buf = p.buf[:0]
p.buf = append(p.buf, ch)

for p.Valid() {
ch = p.Read()
p.buf = append(p.buf, ch)
if ch == ']' || ch == ')' {
break
}
}

return p.buf, nil
}
Loading

0 comments on commit 856e12b

Please sign in to comment.