From 706caa44bfda2261681d81ed5fcc01394a95db00 Mon Sep 17 00:00:00 2001 From: Dan Sosedoff Date: Wed, 29 Dec 2021 11:03:50 -0600 Subject: [PATCH] Serialize binary bytea cols into hex/base64 (#537) - Adds binary serialization into hex/base64 - Default codec is base64 - Codec can be changed via `--binary-codec` CLI option --- pkg/cli/cli.go | 6 ++++ pkg/client/client.go | 2 +- pkg/client/codec.go | 40 ++++++++++++++++++++++++ pkg/client/codec_test.go | 49 +++++++++++++++++++++++++++++ pkg/client/result.go | 16 +++++----- pkg/client/result_test.go | 65 +++++++++++++++++++++++++-------------- pkg/client/util.go | 12 ++++++++ pkg/command/options.go | 1 + 8 files changed, 159 insertions(+), 32 deletions(-) create mode 100644 pkg/client/codec.go create mode 100644 pkg/client/codec_test.go diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go index e91e60c..a8637f4 100644 --- a/pkg/cli/cli.go +++ b/pkg/cli/cli.go @@ -150,6 +150,12 @@ func initOptions() { fmt.Println(readonlyWarning) } + if options.BinaryCodec != "" { + if err := client.SetBinaryCodec(options.BinaryCodec); err != nil { + exitWithMessage(err.Error()) + } + } + printVersion() } diff --git a/pkg/client/client.go b/pkg/client/client.go index 837c3ed..865a2e3 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -416,7 +416,7 @@ func (client *Client) query(query string, args ...interface{}) (*Result, error) } } - result.PrepareBigints() + result.PostProcess() return &result, nil } diff --git a/pkg/client/codec.go b/pkg/client/codec.go new file mode 100644 index 0000000..5d11bce --- /dev/null +++ b/pkg/client/codec.go @@ -0,0 +1,40 @@ +package client + +import ( + "encoding/base64" + "encoding/hex" + "fmt" +) + +const ( + CodecNone = "none" + CodecHex = "hex" + CodecBase64 = "base64" +) + +var ( + // BinaryEncodingFormat specifies the default serialization format of binary data + BinaryCodec = CodecBase64 +) + +func SetBinaryCodec(codec string) error { + switch codec { + case CodecNone, CodecHex, CodecBase64: + BinaryCodec = codec + default: + return fmt.Errorf("invalid binary codec: %v", codec) + } + + return nil +} + +func encodeBinaryData(data []byte, codec string) string { + switch codec { + case CodecHex: + return hex.EncodeToString(data) + case CodecBase64: + return base64.StdEncoding.EncodeToString(data) + default: + return string(data) + } +} diff --git a/pkg/client/codec_test.go b/pkg/client/codec_test.go new file mode 100644 index 0000000..1c25f1d --- /dev/null +++ b/pkg/client/codec_test.go @@ -0,0 +1,49 @@ +package client + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestSetBinaryCodec(t *testing.T) { + examples := []struct { + input string + err error + }{ + {input: CodecNone, err: nil}, + {input: CodecBase64, err: nil}, + {input: CodecHex, err: nil}, + {input: "foobar", err: errors.New("invalid binary codec: foobar")}, + } + + for _, ex := range examples { + t.Run(ex.input, func(t *testing.T) { + val := BinaryCodec + defer func() { + BinaryCodec = val + }() + + assert.Equal(t, ex.err, SetBinaryCodec(ex.input)) + }) + } +} + +func Test_encodeBinaryData(t *testing.T) { + examples := []struct { + input string + expected string + encoding string + }{ + {input: "hello world", expected: "hello world", encoding: CodecNone}, + {input: "hello world", expected: "aGVsbG8gd29ybGQ=", encoding: CodecBase64}, + {input: "hello world", expected: "68656c6c6f20776f726c64", encoding: CodecHex}, + } + + for _, ex := range examples { + t.Run(ex.input, func(t *testing.T) { + assert.Equal(t, ex.expected, encodeBinaryData([]byte(ex.input), ex.encoding)) + }) + } +} diff --git a/pkg/client/result.go b/pkg/client/result.go index 80cfb5f..4cef669 100644 --- a/pkg/client/result.go +++ b/pkg/client/result.go @@ -6,7 +6,6 @@ import ( "encoding/json" "fmt" "math" - "reflect" "strconv" "time" @@ -37,22 +36,19 @@ type Objects struct { // Due to big int number limitations in javascript, numbers should be encoded // as strings so they could be properly loaded on the frontend. -func (res *Result) PrepareBigints() { +func (res *Result) PostProcess() { for i, row := range res.Rows { for j, col := range row { if col == nil { continue } - switch reflect.TypeOf(col).Kind() { - case reflect.Int64: - val := col.(int64) + switch val := col.(type) { + case int64: if val < -9007199254740991 || val > 9007199254740991 { res.Rows[i][j] = strconv.FormatInt(col.(int64), 10) } - case reflect.Float64: - val := col.(float64) - + case float64: // json.Marshal panics when dealing with NaN/Inf values // issue: https://github.com/golang/go/issues/25721 if math.IsNaN(val) { @@ -63,6 +59,10 @@ func (res *Result) PrepareBigints() { if val < -999999999999999 || val > 999999999999999 { res.Rows[i][j] = strconv.FormatFloat(val, 'e', -1, 64) } + case string: + if hasBinary(val, 8) && BinaryCodec != CodecNone { + res.Rows[i][j] = encodeBinaryData([]byte(val), BinaryCodec) + } } } } diff --git a/pkg/client/result_test.go b/pkg/client/result_test.go index 56b5d2d..b1bdd29 100644 --- a/pkg/client/result_test.go +++ b/pkg/client/result_test.go @@ -7,33 +7,52 @@ import ( "github.com/stretchr/testify/assert" ) -func Test_PrepareBigints(t *testing.T) { - result := Result{ - Columns: []string{"value"}, - Rows: []Row{ - Row{int(1234)}, - Row{int64(9223372036854775807)}, - Row{int64(-9223372036854775808)}, - Row{float64(9223372036854775808.9223372036854775808)}, - Row{float64(999999999999999.9)}, - }, - } +func TestPostProcess(t *testing.T) { + t.Run("large numbers", func(t *testing.T) { + result := Result{ + Columns: []string{"value"}, + Rows: []Row{ + {int(1234)}, + {int64(9223372036854775807)}, + {int64(-9223372036854775808)}, + {float64(9223372036854775808.9223372036854775808)}, + {float64(999999999999999.9)}, + }, + } - result.PrepareBigints() + result.PostProcess() - assert.Equal(t, 1234, result.Rows[0][0]) - assert.Equal(t, "9223372036854775807", result.Rows[1][0]) - assert.Equal(t, "-9223372036854775808", result.Rows[2][0]) - assert.Equal(t, "9.223372036854776e+18", result.Rows[3][0]) - assert.Equal(t, "9.999999999999999e+14", result.Rows[4][0]) + assert.Equal(t, 1234, result.Rows[0][0]) + assert.Equal(t, "9223372036854775807", result.Rows[1][0]) + assert.Equal(t, "-9223372036854775808", result.Rows[2][0]) + assert.Equal(t, "9.223372036854776e+18", result.Rows[3][0]) + assert.Equal(t, "9.999999999999999e+14", result.Rows[4][0]) + }) + + t.Run("binary encoding", func(t *testing.T) { + result := Result{ + Columns: []string{"data"}, + Rows: []Row{ + {"text value"}, + {"text with symbols !@#$%"}, + {string([]byte{10, 11, 12, 13})}, + }, + } + + result.PostProcess() + + assert.Equal(t, "text value", result.Rows[0][0]) + assert.Equal(t, "text with symbols !@#$%", result.Rows[1][0]) + assert.Equal(t, "CgsMDQ==", result.Rows[2][0]) + }) } -func Test_CSV(t *testing.T) { +func TestCSV(t *testing.T) { result := Result{ Columns: []string{"id", "name", "email"}, Rows: []Row{ - Row{1, "John", "john@example.com"}, - Row{2, "Bob", "bob@example.com"}, + {1, "John", "john@example.com"}, + {2, "Bob", "bob@example.com"}, }, } @@ -43,12 +62,12 @@ func Test_CSV(t *testing.T) { assert.Equal(t, expected, output) } -func Test_JSON(t *testing.T) { +func TestJSON(t *testing.T) { result := Result{ Columns: []string{"id", "name", "email"}, Rows: []Row{ - Row{1, "John", "john@example.com"}, - Row{2, "Bob", "bob@example.com"}, + {1, "John", "john@example.com"}, + {2, "Bob", "bob@example.com"}, }, } diff --git a/pkg/client/util.go b/pkg/client/util.go index ced4c4d..d83bf1d 100644 --- a/pkg/client/util.go +++ b/pkg/client/util.go @@ -31,3 +31,15 @@ func containsRestrictedKeywords(str string) bool { return reRestrictedKeywords.MatchString(str) } + +func hasBinary(data string, checkLen int) bool { + for idx, chr := range data { + if int(chr) < 32 || int(chr) > 126 { + return true + } + if idx >= checkLen { + break + } + } + return false +} diff --git a/pkg/command/options.go b/pkg/command/options.go index 6ea3872..85f6490 100644 --- a/pkg/command/options.go +++ b/pkg/command/options.go @@ -42,6 +42,7 @@ type Options struct { ConnectionIdleTimeout int `long:"idle-timeout" description:"Set connection idle timeout in minutes" default:"180"` Cors bool `long:"cors" description:"Enable Cross-Origin Resource Sharing (CORS)"` CorsOrigin string `long:"cors-origin" description:"Allowed CORS origins" default:"*"` + BinaryCodec string `long:"binary-codec" description:"Codec for binary data serialization"` } var Opts Options