From af08a9fae24cac37a17603d088105209f5ec9018 Mon Sep 17 00:00:00 2001 From: Sam Jewell <2903904+samjewell@users.noreply.github.com> Date: Tue, 1 Apr 2025 12:45:01 +0100 Subject: [PATCH] SQL Expressions: Add JSON support (#103157) - Support bi-directional mapping of frame JSON fields and GMS (go-mysql-server) columns - Permit GMS json functions Co-authored-by: Kyle Brandt --- pkg/expr/sql/db_test.go | 78 +++++++++++++++++++++++++++++++ pkg/expr/sql/frame_db_conv.go | 18 ++++++- pkg/expr/sql/frame_table.go | 19 +++++++- pkg/expr/sql/parser_allow.go | 6 +++ pkg/expr/sql/parser_allow_test.go | 17 +++++++ 5 files changed, 136 insertions(+), 2 deletions(-) diff --git a/pkg/expr/sql/db_test.go b/pkg/expr/sql/db_test.go index 77263d028de..96fa1d98c3c 100644 --- a/pkg/expr/sql/db_test.go +++ b/pkg/expr/sql/db_test.go @@ -4,6 +4,7 @@ package sql import ( "context" + "encoding/json" "testing" "time" @@ -205,6 +206,83 @@ func TestErrorsFromGoMySQLServerAreFlagged(t *testing.T) { require.Contains(t, err.Error(), "error in go-mysql-server") } +func TestFrameToSQLAndBack_JSONRoundtrip(t *testing.T) { + expectedFrame := &data.Frame{ + RefID: "json_test", + Name: "json_test", + Fields: []*data.Field{ + data.NewField("id", nil, []int64{1, 2}), + data.NewField("payload", nil, []json.RawMessage{ + json.RawMessage(`{"foo":1}`), + json.RawMessage(`{"bar":"baz"}`), + }), + }, + } + + db := DB{} + + query := `SELECT * FROM json_test` + + resultFrame, err := db.QueryFrames(context.Background(), "json_test", query, data.Frames{expectedFrame}) + require.NoError(t, err) + + // Use custom compare options that ignore Name and RefID + opts := append( + data.FrameTestCompareOptions(), + cmp.FilterPath(func(p cmp.Path) bool { + return p.String() == "Name" || p.String() == "RefID" + }, cmp.Ignore()), + ) + + if diff := cmp.Diff(expectedFrame, resultFrame, opts...); diff != "" { + require.FailNowf(t, "Frame mismatch (-want +got):\n%s", diff) + } +} + +func TestQueryFrames_JSONFilter(t *testing.T) { + input := &data.Frame{ + RefID: "A", + Name: "A", + Fields: []*data.Field{ + data.NewField("title", nil, []string{"Bug report", "Feature request"}), + data.NewField("labels", nil, []json.RawMessage{ + json.RawMessage(`["type/bug", "priority/high"]`), + json.RawMessage(`["type/feature", "priority/low"]`), + }), + }, + } + + expected := &data.Frame{ + RefID: "B", + Name: "B", + Fields: []*data.Field{ + data.NewField("title", nil, []string{"Bug report"}), + data.NewField("labels", nil, []json.RawMessage{ + json.RawMessage(`["type/bug", "priority/high"]`), + }), + }, + } + + db := DB{} + + query := `SELECT title, labels FROM A WHERE json_contains(labels, '"type/bug"')` + + result, err := db.QueryFrames(context.Background(), "B", query, data.Frames{input}) + require.NoError(t, err) + + // Use custom compare options that ignore Name and RefID + opts := append( + data.FrameTestCompareOptions(), + cmp.FilterPath(func(p cmp.Path) bool { + return p.String() == "Name" || p.String() == "RefID" + }, cmp.Ignore()), + ) + + if diff := cmp.Diff(expected, result, opts...); diff != "" { + require.FailNowf(t, "Result mismatch (-want +got):\n%s", diff) + } +} + // p is a utility for pointers from constants func p[T any](v T) *T { return &v diff --git a/pkg/expr/sql/frame_db_conv.go b/pkg/expr/sql/frame_db_conv.go index 302ff79d1bb..d9ce72506f0 100644 --- a/pkg/expr/sql/frame_db_conv.go +++ b/pkg/expr/sql/frame_db_conv.go @@ -3,6 +3,7 @@ package sql import ( + "encoding/json" "errors" "fmt" "io" @@ -92,6 +93,8 @@ func MySQLColToFieldType(col *mysql.Column) (data.FieldType, error) { fT = data.FieldTypeTime case types.Boolean: fT = data.FieldTypeBool + case types.JSON: + fT = data.FieldTypeJSON default: switch { case types.IsDecimal(col.Type): @@ -159,8 +162,21 @@ func fieldValFromRowVal(fieldType data.FieldType, val interface{}) (interface{}, case data.FieldTypeBool, data.FieldTypeNullableBool: return parseBoolFromInt8(val, nullable) + case data.FieldTypeJSON, data.FieldTypeNullableJSON: + switch v := val.(type) { + case types.JSONDocument: + raw := json.RawMessage(v.String()) + if nullable { + return &raw, nil + } + return raw, nil + + default: + return nil, fmt.Errorf("JSON field does not support val %v of type %T", val, val) + } + default: - return nil, fmt.Errorf("unsupported field type %s for val %v", fieldType, val) + return nil, fmt.Errorf("unsupported field type %s for val %v of type %T", fieldType, val, val) } } diff --git a/pkg/expr/sql/frame_table.go b/pkg/expr/sql/frame_table.go index 66b05cfaa60..b605ff68fd5 100644 --- a/pkg/expr/sql/frame_table.go +++ b/pkg/expr/sql/frame_table.go @@ -3,6 +3,7 @@ package sql import ( + "encoding/json" "fmt" "io" "strings" @@ -90,7 +91,21 @@ func (ri *rowIter) Next(_ *mysql.Context) (mysql.Row, error) { if field.NilAt(ri.row) { continue } - row[colIndex], _ = field.ConcreteAt(ri.row) + val, _ := field.ConcreteAt(ri.row) + + // If the field is JSON, convert json.RawMessage to types.JSONDocument + if raw, ok := val.(json.RawMessage); ok { + doc, inRange, err := types.JSON.Convert(raw) + if err != nil { + return nil, fmt.Errorf("failed to convert json.RawMessage to JSONDocument: %w", err) + } + if !inRange { + return nil, fmt.Errorf("invalid JSON value detected at row %d, column %s: value required type coercion", ri.row, ri.ft.Frame.Fields[colIndex].Name) + } + val = doc + } + + row[colIndex] = val } ri.row++ @@ -156,6 +171,8 @@ func convertDataType(fieldType data.FieldType) mysql.Type { return types.Boolean case data.FieldTypeTime, data.FieldTypeNullableTime: return types.Timestamp + case data.FieldTypeJSON, data.FieldTypeNullableJSON: + return types.JSON default: fmt.Printf("------- Unsupported field type: %v", fieldType) return types.JSON diff --git a/pkg/expr/sql/parser_allow.go b/pkg/expr/sql/parser_allow.go index 8d7b2d80e9e..e6ba193b834 100644 --- a/pkg/expr/sql/parser_allow.go +++ b/pkg/expr/sql/parser_allow.go @@ -210,6 +210,12 @@ func allowedFunction(f *sqlparser.FuncExpr) (b bool) { case "cast": return + // JSON functions + case "json_extract", "json_unquote", "json_contains", + "json_object", "json_array", "json_set", "json_remove", + "json_length", "json_search", "json_type": + return + default: return false } diff --git a/pkg/expr/sql/parser_allow_test.go b/pkg/expr/sql/parser_allow_test.go index 6144c30f9d5..5afa4769bbb 100644 --- a/pkg/expr/sql/parser_allow_test.go +++ b/pkg/expr/sql/parser_allow_test.go @@ -67,6 +67,11 @@ func TestAllowQuery(t *testing.T) { q: `SELECT __value__, SUBSTRING_INDEX(name, '.', -1) AS code FROM A`, err: nil, }, + { + name: "json functions", + q: example_json_functions, + err: nil, + }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { @@ -239,3 +244,15 @@ SELECT FROM sample_data GROUP BY name, value, created_at LIMIT 10` + +var example_json_functions = `SELECT + JSON_OBJECT('key1', 'value1', 'key2', 10) AS json_obj, + JSON_ARRAY(1, 'abc', NULL, TRUE) AS json_arr, + JSON_EXTRACT('{"id": 123, "name": "test"}', '$.id') AS json_ext, + JSON_UNQUOTE(JSON_EXTRACT('{"name": "test"}', '$.name')) AS json_unq, + JSON_CONTAINS('{"a": 1, "b": 2}', '{"a": 1}') AS json_contains, + JSON_SET('{"a": 1}', '$.b', 2) AS json_set, + JSON_REMOVE('{"a": 1, "b": 2}', '$.b') AS json_remove, + JSON_LENGTH('{"a": 1, "b": {"c": 3}}') AS json_len, + JSON_SEARCH('{"a": "xyz", "b": "abc"}', 'one', 'abc') AS json_search, + JSON_TYPE('{"a": 1}') AS json_type`