diff --git a/pkg/ccl/changefeedccl/changefeed_test.go b/pkg/ccl/changefeedccl/changefeed_test.go index 38ab35952456..8b4dc0cf765f 100644 --- a/pkg/ccl/changefeedccl/changefeed_test.go +++ b/pkg/ccl/changefeedccl/changefeed_test.go @@ -3212,7 +3212,12 @@ func TestChangefeedJobUpdateFailsIfNotClaimed(t *testing.T) { sqlDB.Exec(t, `INSERT INTO foo (a, b) VALUES (1, 1)`) cf := feed(t, f, "CREATE CHANGEFEED FOR TABLE foo") - defer closeFeed(t, cf) + jobID := cf.(cdctest.EnterpriseTestFeed).JobID() + defer func() { + // Manually update job status to avoid closeFeed waitng for the registry to cancel it + sqlDB.Exec(t, `UPDATE system.jobs SET status = $1 WHERE id = $2`, jobs.StatusFailed, jobID) + closeFeed(t, cf) + }() assertPayloads(t, cf, []string{ `foo: [1]->{"after": {"a": 1, "b": 1}}`, @@ -3220,7 +3225,6 @@ func TestChangefeedJobUpdateFailsIfNotClaimed(t *testing.T) { // Mimic the claim dying and being cleaned up by // another node. - jobID := cf.(cdctest.EnterpriseTestFeed).JobID() sqlDB.Exec(t, `UPDATE system.jobs SET claim_session_id = NULL WHERE id = $1`, jobID) // Expect that the distflow fails since it can't diff --git a/pkg/ccl/changefeedccl/testfeed_test.go b/pkg/ccl/changefeedccl/testfeed_test.go index 940de37470ba..00ffdd37ba5c 100644 --- a/pkg/ccl/changefeedccl/testfeed_test.go +++ b/pkg/ccl/changefeedccl/testfeed_test.go @@ -418,8 +418,17 @@ func (f *jobFeed) Close() error { close(f.shutdown) return nil } + if status == string(jobs.StatusFailed) { + f.mu.Lock() + defer f.mu.Unlock() + f.mu.terminalErr = errors.New("changefeed failed") + close(f.shutdown) + return nil + } if _, err := f.db.Exec(`CANCEL JOB $1`, f.jobID); err != nil { log.Infof(context.Background(), `could not cancel feed %d: %v`, f.jobID, err) + } else { + return f.WaitForStatus(func(s jobs.Status) bool { return s == jobs.StatusCanceled }) } } diff --git a/pkg/sql/logictest/testdata/logic_test/json b/pkg/sql/logictest/testdata/logic_test/json index f1229ffdc74d..7ac2cad8dd27 100644 --- a/pkg/sql/logictest/testdata/logic_test/json +++ b/pkg/sql/logictest/testdata/logic_test/json @@ -884,3 +884,61 @@ query TBTB SELECT j, j ? 'a', j-1, (j-1) ? 'a' FROM t81647 ---- ["a", "b"] true ["a"] true + +# +# Test JSONB subscripting. +# + +# Constant folding. +query TTT +SELECT + ('{"a": {"b": {"c": 1}}}'::jsonb)['a'], + ('{"a": {"b": {"c": 1}}}'::jsonb)['a']['b']['c'], + ('[1, "2", null]'::jsonb)[1] +---- +{"b": {"c": 1}} 1 "2" + +# Referencing subscript which does not exist. +query TTTT +SELECT + ('{"a": 1}'::jsonb)['b'], + ('{"a": {"b": {"c": 1}}}'::jsonb)['c']['b']['c'], + ('[1, "2", null]'::jsonb)[4], + ('{"a": 1}'::jsonb)[NULL] +---- +NULL NULL NULL NULL + +# Error cases. +statement error unexpected JSON subscript type: TIMESTAMPTZ +SELECT ('{"a": 1}'::jsonb)[now()] + +statement error jsonb subscript does not support slices +SELECT ('{"a": 1}'::jsonb)['a':'b'] + +# Check it works from a JSON table. +statement ok +CREATE TABLE json_subscript_test ( + id SERIAL PRIMARY KEY, + j JSONB, + extract_field TEXT, + extract_int_field INT +); +INSERT INTO json_subscript_test (j, extract_field, extract_int_field) VALUES + ('{"other_field": 2}', 'other_field', 1), + ('{"field": {"field": 2}}', 'field', 0), + ('[1, 2, 3]', 'nothing_to_fetch', 1) + +# Test subscripts with fields using other columns. +query TTITTTT +SELECT j, extract_field, extract_int_field, j['field'], j[extract_field], j[extract_field][extract_field], j[extract_int_field] +FROM json_subscript_test ORDER BY id +---- +{"other_field": 2} other_field 1 NULL 2 NULL NULL +{"field": {"field": 2}} field 0 {"field": 2} {"field": 2} 2 NULL +[1, 2, 3] nothing_to_fetch 1 NULL NULL NULL 2 + +# Test use in a WHERE clause. +query T +SELECT j FROM json_subscript_test WHERE j['other_field'] = '2' ORDER BY id +---- +{"other_field": 2} diff --git a/pkg/sql/opt/memo/typing.go b/pkg/sql/opt/memo/typing.go index 931e1665852a..3c787f7ed3c1 100644 --- a/pkg/sql/opt/memo/typing.go +++ b/pkg/sql/opt/memo/typing.go @@ -247,9 +247,18 @@ func typeArrayAgg(e opt.ScalarExpr) *types.T { return types.MakeArray(typ) } -// typeIndirection returns the type of the element of the array. +// typeIndirection returns the type of the element after the indirection +// is applied. func typeIndirection(e opt.ScalarExpr) *types.T { - return e.Child(0).(opt.ScalarExpr).DataType().ArrayContents() + t := e.Child(0).(opt.ScalarExpr).DataType() + switch t.Family() { + case types.JsonFamily: + return t + case types.ArrayFamily: + return t.ArrayContents() + default: + panic(errors.AssertionFailedf("unknown type indirection type %s", t.SQLString())) + } } // typeCollate returns the collated string typed with the given locale. diff --git a/pkg/sql/opt/norm/fold_constants_funcs.go b/pkg/sql/opt/norm/fold_constants_funcs.go index b516c0865572..e08e5893dc56 100644 --- a/pkg/sql/opt/norm/fold_constants_funcs.go +++ b/pkg/sql/opt/norm/fold_constants_funcs.go @@ -498,10 +498,19 @@ func (c *CustomFuncs) FoldIndirection(input, index opt.ScalarExpr) (_ opt.Scalar return nil, false } - // Case 2: The input is a constant DArray. + // Case 2: The input is a constant DArray or DJSON. if memo.CanExtractConstDatum(input) { + var resolvedType *types.T + switch input.DataType().Family() { + case types.JsonFamily: + resolvedType = input.DataType() + case types.ArrayFamily: + resolvedType = input.DataType().ArrayContents() + default: + panic(errors.AssertionFailedf("expected array or json; found %s", input.DataType().SQLString())) + } inputD := memo.ExtractConstDatum(input) - texpr := tree.NewTypedIndirectionExpr(inputD, indexD, input.DataType().ArrayContents()) + texpr := tree.NewTypedIndirectionExpr(inputD, indexD, resolvedType) result, err := eval.Expr(c.f.evalCtx, texpr) if err == nil { return c.f.ConstructConstVal(result, texpr.ResolvedType()), true diff --git a/pkg/sql/opt/norm/testdata/rules/fold_constants b/pkg/sql/opt/norm/testdata/rules/fold_constants index 5cc07396eec9..be1b49840066 100644 --- a/pkg/sql/opt/norm/testdata/rules/fold_constants +++ b/pkg/sql/opt/norm/testdata/rules/fold_constants @@ -1197,6 +1197,28 @@ project └── projections └── a.arr:6[0] [as=arr:9, outer=(6)] +# Fold JSONB constants. +norm expect=FoldIndirection +SELECT ('{"a": 1}'::jsonb)['a'] AS other_col FROM a +---- +project + ├── columns: other_col:9!null + ├── fd: ()-->(9) + ├── scan a + └── projections + └── '1' [as=other_col:9] + +# JSONB is dynamically constructured. +norm expect-not=FoldIndirection +SELECT j['field'] FROM a +---- +project + ├── columns: j:9 + ├── scan a + │ └── columns: a.j:5 + └── projections + └── a.j:5['field'] [as=j:9, outer=(5)] + # Regression test for #40404. norm expect=FoldIndirection SELECT (SELECT x[1]) FROM (VALUES(null::oid[])) v(x) diff --git a/pkg/sql/opt/optbuilder/scalar.go b/pkg/sql/opt/optbuilder/scalar.go index 82fa834b63d6..1d82a33112ef 100644 --- a/pkg/sql/opt/optbuilder/scalar.go +++ b/pkg/sql/opt/optbuilder/scalar.go @@ -170,21 +170,22 @@ func (b *Builder) buildScalar( out = b.factory.ConstructArrayFlatten(s.node, &subqueryPrivate) case *tree.IndirectionExpr: - expr := b.buildScalar(t.Expr.(tree.TypedExpr), inScope, nil, nil, colRefs) - - if len(t.Indirection) != 1 { + if len(t.Indirection) != 1 && t.Expr.(tree.TypedExpr).ResolvedType().Family() == types.ArrayFamily { panic(unimplementedWithIssueDetailf(32552, "ind", "multidimensional indexing is not supported")) } - subscript := t.Indirection[0] - if subscript.Slice { - panic(unimplementedWithIssueDetailf(32551, "", "array slicing is not supported")) - } + out = b.buildScalar(t.Expr.(tree.TypedExpr), inScope, nil, nil, colRefs) - out = b.factory.ConstructIndirection( - expr, - b.buildScalar(subscript.Begin.(tree.TypedExpr), inScope, nil, nil, colRefs), - ) + for _, subscript := range t.Indirection { + if subscript.Slice { + panic(unimplementedWithIssueDetailf(32551, "", "array slicing is not supported")) + } + + out = b.factory.ConstructIndirection( + out, + b.buildScalar(subscript.Begin.(tree.TypedExpr), inScope, nil, nil, colRefs), + ) + } case *tree.IfErrExpr: cond := b.buildScalar(t.Cond.(tree.TypedExpr), inScope, nil, nil, colRefs) diff --git a/pkg/sql/opt/optbuilder/testdata/select b/pkg/sql/opt/optbuilder/testdata/select index 057e43cf66da..45786e812e9a 100644 --- a/pkg/sql/opt/optbuilder/testdata/select +++ b/pkg/sql/opt/optbuilder/testdata/select @@ -312,7 +312,7 @@ error (42P01): no data source matches pattern: bar.kv.* build SELECT kv.*[1] FROM kv ---- -error (42804): cannot subscript type tuple{char AS k, char AS v} because it is not an array +error (42804): cannot subscript type tuple{char AS k, char AS v} because it is not an array or json object build SELECT ARRAY[] diff --git a/pkg/sql/sem/eval/expr.go b/pkg/sql/sem/eval/expr.go index 200b37ea95fc..c5e63dc5f269 100644 --- a/pkg/sql/sem/eval/expr.go +++ b/pkg/sql/sem/eval/expr.go @@ -15,6 +15,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree/treecmp" + "github.com/cockroachdb/cockroach/pkg/sql/types" "github.com/cockroachdb/errors" ) @@ -275,20 +276,6 @@ func (e *evaluator) EvalIndexedVar(iv *tree.IndexedVar) (tree.Datum, error) { func (e *evaluator) EvalIndirectionExpr(expr *tree.IndirectionExpr) (tree.Datum, error) { var subscriptIdx int - for i, t := range expr.Indirection { - if t.Slice || i > 0 { - return nil, errors.AssertionFailedf("unsupported feature should have been rejected during planning") - } - - d, err := t.Begin.(tree.TypedExpr).Eval(e) - if err != nil { - return nil, err - } - if d == tree.DNull { - return d, nil - } - subscriptIdx = int(tree.MustBeDInt(d)) - } d, err := expr.Expr.(tree.TypedExpr).Eval(e) if err != nil { @@ -298,17 +285,68 @@ func (e *evaluator) EvalIndirectionExpr(expr *tree.IndirectionExpr) (tree.Datum, return d, nil } - // Index into the DArray, using 1-indexing. - arr := tree.MustBeDArray(d) + switch d.ResolvedType().Family() { + case types.ArrayFamily: + for i, t := range expr.Indirection { + if t.Slice || i > 0 { + return nil, errors.AssertionFailedf("unsupported feature should have been rejected during planning") + } - // VECTOR types use 0-indexing. - if arr.FirstIndex() == 0 { - subscriptIdx++ - } - if subscriptIdx < 1 || subscriptIdx > arr.Len() { - return tree.DNull, nil + beginDatum, err := t.Begin.(tree.TypedExpr).Eval(e) + if err != nil { + return nil, err + } + if beginDatum == tree.DNull { + return tree.DNull, nil + } + subscriptIdx = int(tree.MustBeDInt(beginDatum)) + } + + // Index into the DArray, using 1-indexing. + arr := tree.MustBeDArray(d) + + // VECTOR types use 0-indexing. + if arr.FirstIndex() == 0 { + subscriptIdx++ + } + if subscriptIdx < 1 || subscriptIdx > arr.Len() { + return tree.DNull, nil + } + return arr.Array[subscriptIdx-1], nil + case types.JsonFamily: + j := tree.MustBeDJSON(d) + curr := j.JSON + for _, t := range expr.Indirection { + if t.Slice { + return nil, errors.AssertionFailedf("unsupported feature should have been rejected during planning") + } + + field, err := t.Begin.(tree.TypedExpr).Eval(e) + if err != nil { + return nil, err + } + if field == tree.DNull { + return tree.DNull, nil + } + switch field.ResolvedType().Family() { + case types.StringFamily: + if curr, err = curr.FetchValKeyOrIdx(string(tree.MustBeDString(field))); err != nil { + return nil, err + } + case types.IntFamily: + if curr, err = curr.FetchValIdx(int(tree.MustBeDInt(field))); err != nil { + return nil, err + } + default: + return nil, errors.AssertionFailedf("unsupported feature should have been rejected during planning") + } + if curr == nil { + return tree.DNull, nil + } + } + return tree.NewDJSON(curr), nil } - return arr.Array[subscriptIdx-1], nil + return nil, errors.AssertionFailedf("unsupported feature should have been rejected during planning") } func (e *evaluator) EvalDefaultVal(expr *tree.DefaultVal) (tree.Datum, error) { diff --git a/pkg/sql/sem/tree/type_check.go b/pkg/sql/sem/tree/type_check.go index 5e01859075c6..e1a1f08f87a3 100644 --- a/pkg/sql/sem/tree/type_check.go +++ b/pkg/sql/sem/tree/type_check.go @@ -33,6 +33,7 @@ import ( // type checking of the relevant function is made. var ( OnTypeCheckArraySubscript func() + OnTypeCheckJSONBSubscript func() OnTypeCheckIfErr func() OnTypeCheckArrayConstructor func() OnTypeCheckArrayFlatten func() @@ -600,34 +601,64 @@ func (expr *CastExpr) TypeCheck( func (expr *IndirectionExpr) TypeCheck( ctx context.Context, semaCtx *SemaContext, desired *types.T, ) (TypedExpr, error) { - for i, t := range expr.Indirection { - if t.Slice { - return nil, unimplemented.NewWithIssuef(32551, "ARRAY slicing in %s", expr) - } - if i > 0 { - return nil, unimplemented.NewWithIssueDetailf(32552, "ind", "multidimensional indexing: %s", expr) - } - - beginExpr, err := typeCheckAndRequire(ctx, semaCtx, t.Begin, types.Int, "ARRAY subscript") - if err != nil { - return nil, err - } - t.Begin = beginExpr - } - subExpr, err := expr.Expr.TypeCheck(ctx, semaCtx, types.MakeArray(desired)) if err != nil { return nil, err } typ := subExpr.ResolvedType() - if typ.Family() != types.ArrayFamily { - return nil, pgerror.Newf(pgcode.DatatypeMismatch, "cannot subscript type %s because it is not an array", typ) - } expr.Expr = subExpr - expr.typ = typ.ArrayContents() - if OnTypeCheckArraySubscript != nil { - OnTypeCheckArraySubscript() + switch typ.Family() { + case types.ArrayFamily: + expr.typ = typ.ArrayContents() + for i, t := range expr.Indirection { + if t.Slice { + return nil, unimplemented.NewWithIssuef(32551, "ARRAY slicing in %s", expr) + } + if i > 0 { + return nil, unimplemented.NewWithIssueDetailf(32552, "ind", "multidimensional indexing: %s", expr) + } + + beginExpr, err := typeCheckAndRequire(ctx, semaCtx, t.Begin, types.Int, "ARRAY subscript") + if err != nil { + return nil, err + } + t.Begin = beginExpr + } + + if OnTypeCheckArraySubscript != nil { + OnTypeCheckArraySubscript() + } + case types.JsonFamily: + expr.typ = typ + for _, t := range expr.Indirection { + if t.Slice { + return nil, pgerror.Newf(pgcode.DatatypeMismatch, "jsonb subscript does not support slices") + } + beginExpr, err := t.Begin.TypeCheck(ctx, semaCtx, types.Any) + if err != nil { + return nil, err + } + switch beginExpr.ResolvedType().Family() { + case types.IntFamily, types.StringFamily, types.UnknownFamily: + default: + return nil, errors.WithHint( + pgerror.Newf( + pgcode.DatatypeMismatch, + "unexpected JSON subscript type: %s", + beginExpr.ResolvedType().SQLString(), + ), + "subscript type must be integer or text", + ) + } + t.Begin = beginExpr + } + + if OnTypeCheckJSONBSubscript != nil { + OnTypeCheckJSONBSubscript() + } + default: + return nil, pgerror.Newf(pgcode.DatatypeMismatch, "cannot subscript type %s because it is not an array or json object", typ) } return expr, nil } diff --git a/pkg/sql/sqltelemetry/scalar.go b/pkg/sql/sqltelemetry/scalar.go index b7030d706e7d..f7144cd5187c 100644 --- a/pkg/sql/sqltelemetry/scalar.go +++ b/pkg/sql/sqltelemetry/scalar.go @@ -72,6 +72,10 @@ var ArrayFlattenCounter = telemetry.GetCounterOnce("sql.plan.ops.array.flatten") // array subscript expression x[...]. var ArraySubscriptCounter = telemetry.GetCounterOnce("sql.plan.ops.array.ind") +// JSONBSubscriptCounter is to be incremented upon type checking an +// JSONB subscript expression x[...]. +var JSONBSubscriptCounter = telemetry.GetCounterOnce("sql.plan.ops.jsonb.subscript") + // IfErrCounter is to be incremented upon type checking an // IFERROR(...) expression or analogous. var IfErrCounter = telemetry.GetCounterOnce("sql.plan.ops.iferr") diff --git a/pkg/sql/telemetry.go b/pkg/sql/telemetry.go index 2f17533b07f9..706c932f26e5 100644 --- a/pkg/sql/telemetry.go +++ b/pkg/sql/telemetry.go @@ -74,6 +74,9 @@ func init() { tree.OnTypeCheckArraySubscript = func() { telemetry.Inc(sqltelemetry.ArraySubscriptCounter) } + tree.OnTypeCheckJSONBSubscript = func() { + telemetry.Inc(sqltelemetry.JSONBSubscriptCounter) + } tree.OnTypeCheckArrayFlatten = func() { telemetry.Inc(sqltelemetry.ArrayFlattenCounter) } diff --git a/pkg/sql/testdata/telemetry/json b/pkg/sql/testdata/telemetry/json new file mode 100644 index 000000000000..11924ccf01c6 --- /dev/null +++ b/pkg/sql/testdata/telemetry/json @@ -0,0 +1,8 @@ +feature-allowlist +sql.plan.ops.jsonb.subscript +---- + +feature-usage +SELECT ('{"a": {"b": {"c": 1}}}'::jsonb)['a']['b'] +---- +sql.plan.ops.jsonb.subscript