implement JSON_QUOTE()

enginetest/queries/queries.go

@@ -5554,6 +5554,23 @@ Select * from (
 		Query:    `SELECT JSON_UNQUOTE(JSON_EXTRACT('{"xid":"hello"}', '$.xid')) = "hello"`,
 		Expected: []sql.Row{{true}},
 	},
+
+	{
+		Query:    `SELECT JSON_QUOTE('"foo"')`,
+		Expected: []sql.Row{{`"\"foo\""`}},
+	},
+	{
+		Query:    `SELECT JSON_QUOTE('[1, 2, 3]')`,
+		Expected: []sql.Row{{`"[1, 2, 3]"`}},
+	},
+	{
+		Query:    `SELECT JSON_QUOTE('"\t\u0032"')`,
+		Expected: []sql.Row{{`"\"\tu0032\""`}},
+	},
+	{
+		Query:    `SELECT JSON_QUOTE('"\t\\u0032"')`,
+		Expected: []sql.Row{{`"\"\t\\u0032\""`}},
+	},
 	{
 		Query:    `SELECT JSON_EXTRACT('{"xid":"hello"}', '$.xid') = "hello"`,
 		Expected: []sql.Row{{true}},

internal/strings/unquote.go

@@ -8,6 +8,7 @@ import (
 	"unicode/utf8"
 )
 
+// Unquote returns a json unquoted string.
 // The implementation is taken from TiDB
 // https://github.com/pingcap/tidb/blob/a594287e9f402037b06930026906547000006bb6/types/json/binary_functions.go#L89
 func Unquote(s string) (string, error) {
@@ -16,7 +17,8 @@ func Unquote(s string) (string, error) {
 		if s[i] == '\\' {
 			i++
 			if i == len(s) {
-				return "", fmt.Errorf("Missing a closing quotation mark in string")
+				ret.WriteByte('\\')
+				break
 			}
 			switch s[i] {
 			case '"':
@@ -84,3 +86,57 @@ func decodeEscapedUnicode(s []byte) (char [4]byte, size int, err error) {
 	utf8.EncodeRune(char[0:size], rune(unicode))
 	return
 }
+
+// Quote returns a json quoted string with escape characters.
+// The implementation is taken from TiDB:
+// https://github.com/pingcap/tidb/blob/a594287e9f402037b06930026906547000006bb6/types/json/binary_functions.go#L155
+func Quote(s string) string {
+	var escapeByteMap = map[byte]string{
+		'\\': "\\\\",
+		'"':  "\\\"",
+		'\b': "\\b",
+		'\f': "\\f",
+		'\n': "\\n",
+		'\r': "\\r",
+		'\t': "\\t",
+	}
+
+	ret := new(bytes.Buffer)
+	ret.WriteByte('"')
+
+	start := 0
+	for i := 0; i < len(s); {
+		if b := s[i]; b < utf8.RuneSelf {
+			escaped, ok := escapeByteMap[b]
+			if ok {
+				if start < i {
+					ret.WriteString(s[start:i])
+				}
+				ret.WriteString(escaped)
+				i++
+				start = i
+			} else {
+				i++
+			}
+		} else {
+			c, size := utf8.DecodeRune([]byte(s[i:]))
+			if c == utf8.RuneError && size == 1 { // refer to codes of `binary.marshalStringTo`
+				if start < i {
+					ret.WriteString(s[start:i])
+				}
+				ret.WriteString(`\ufffd`)
+				i += size
+				start = i
+				continue
+			}
+			i += size
+		}
+	}
+
+	if start < len(s) {
+		ret.WriteString(s[start:])
+	}
+
+	ret.WriteByte('"')
+	return ret.String()
+}

sql/expression/function/json/json_quote.go

@@ -0,0 +1,107 @@
+// Copyright 2024 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package json
+
+import (
+	"fmt"
+	"reflect"
+
+	"github.com/dolthub/go-mysql-server/internal/strings"
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/expression"
+	"github.com/dolthub/go-mysql-server/sql/types"
+)
+
+// JSONQuote (string)
+//
+// JSONQuote Quotes a string as a JSON value by wrapping it with double quote characters and escaping interior quote and
+// other characters, then returning the result as a utf8mb4 string. Returns NULL if the argument is NULL. This function
+// is typically used to produce a valid JSON string literal for inclusion within a JSON document. Certain special
+// characters are escaped with backslashes per the escape sequences shown in Table 12.23, “JSON_UNQUOTE() Special
+// Character Escape Sequences”:
+// https://dev.mysql.com/doc/refman/8.0/en/json-modification-functions.html#json-unquote-character-escape-sequences
+//
+// https://dev.mysql.com/doc/refman/8.0/en/json-creation-functions.html#function_json-quote
+type JSONQuote struct {
+	expression.UnaryExpression
+}
+
+var _ sql.FunctionExpression = (*JSONQuote)(nil)
+var _ sql.CollationCoercible = (*JSONQuote)(nil)
+
+// NewJSONQuote creates a new JSONQuote UDF.
+func NewJSONQuote(json sql.Expression) sql.Expression {
+	return &JSONQuote{expression.UnaryExpression{Child: json}}
+}
+
+// FunctionName implements sql.FunctionExpression
+func (js *JSONQuote) FunctionName() string {
+	return "json_quote"
+}
+
+// Description implements sql.FunctionExpression
+func (js *JSONQuote) Description() string {
+	return "quotes a string as a JSON value and returns the result as a utf8mb4 string."
+}
+
+// String implements the fmt.Stringer interface.
+func (js *JSONQuote) String() string {
+	return fmt.Sprintf("%s(%s)", js.FunctionName(), js.Child)
+}
+
+// Type implements the Expression interface.
+func (*JSONQuote) Type() sql.Type {
+	return types.LongText
+}
+
+// CollationCoercibility implements the interface sql.CollationCoercible.
+func (*JSONQuote) CollationCoercibility(ctx *sql.Context) (collation sql.CollationID, coercibility byte) {
+	return ctx.GetCharacterSet().BinaryCollation(), 4
+}
+
+// WithChildren implements the Expression interface.
+func (js *JSONQuote) WithChildren(children ...sql.Expression) (sql.Expression, error) {
+	if len(children) != 1 {
+		return nil, sql.ErrInvalidChildrenNumber.New(js, len(children), 1)
+	}
+	return NewJSONQuote(children[0]), nil
+}
+
+// Eval implements the Expression interface.
+func (js *JSONQuote) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	typ := js.Child.Type()
+	if typ != types.Null && !types.IsText(typ) {
+		return nil, sql.ErrInvalidType.New(typ)
+	}
+
+	val, err := js.Child.Eval(ctx, row)
+	if val == nil || err != nil {
+		return val, err
+	}
+
+	ex, _, err := types.LongText.Convert(val)
+	if err != nil {
+		return nil, err
+	}
+	if ex == nil {
+		return nil, nil
+	}
+	str, ok := ex.(string)
+	if !ok {
+		return nil, sql.ErrInvalidType.New(reflect.TypeOf(ex).String())
+	}
+
+	return strings.Quote(str), nil
+}

sql/expression/function/json/json_quote_test.go

@@ -0,0 +1,89 @@
+// Copyright 2024 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package json
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/expression"
+	"github.com/dolthub/go-mysql-server/sql/types"
+)
+
+func TestJSONQuote(t *testing.T) {
+	testCases := []struct {
+		arg sql.Expression
+		exp interface{}
+		err bool
+	}{
+		{
+			arg: expression.NewLiteral(true, types.Boolean),
+			err: true,
+		},
+		{
+			arg: expression.NewLiteral(123, types.Int64),
+			err: true,
+		},
+		{
+			arg: expression.NewLiteral(123.0, types.Float64),
+			err: true,
+		},
+		{
+			arg: expression.NewLiteral(types.MustJSON(`{"a": 1}`), types.JSON),
+			err: true,
+		},
+		{
+			arg: expression.NewLiteral(nil, types.Null),
+			exp: nil,
+		},
+		{
+			arg: expression.NewLiteral("abc", types.Text),
+			exp: `"abc"`,
+		},
+		{
+			arg: expression.NewLiteral(`[1, 2, 3]`, types.Text),
+			exp: `"[1, 2, 3]"`,
+		},
+		{
+			arg: expression.NewLiteral(`"\t\u0032"`, types.Text),
+			exp: `"\"\\t\\u0032\""`,
+		},
+		{
+			arg: expression.NewLiteral(`\`, types.Text),
+			exp: `"\\"`,
+		},
+		{
+			arg: expression.NewLiteral(`\b\f\n\r\t\"`, types.Text),
+			exp: `"\\b\\f\\n\\r\\t\\\""`,
+		},
+	}
+
+	for _, tt := range testCases {
+		t.Run(fmt.Sprintf("%v", tt.arg), func(t *testing.T) {
+			require := require.New(t)
+			js := NewJSONQuote(tt.arg)
+			res, err := js.Eval(sql.NewEmptyContext(), nil)
+			if tt.err {
+				require.Error(err)
+				return
+			}
+			require.NoError(err)
+			require.Equal(tt.exp, res)
+		})
+	}
+}

sql/expression/function/json/json_unquote.go

@@ -78,6 +78,11 @@ func (js *JSONUnquote) WithChildren(children ...sql.Expression) (sql.Expression,
 
 // Eval implements the Expression interface.
 func (js *JSONUnquote) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	typ := js.Child.Type()
+	if typ != types.Null && !types.IsText(typ) && !types.IsJSON(typ) {
+		return nil, sql.ErrInvalidType.New(typ)
+	}
+
 	json, err := js.Child.Eval(ctx, row)
 	if json == nil || err != nil {
 		return json, err