sql: handle single clause UPDATE ... SET for JSONB descriptors

docs/generated/sql/bnf/on_conflict.bnf

@@ -1,6 +1,6 @@
 on_conflict ::=
 	'ON' 'CONFLICT' 'DO' 'NOTHING'
 	| 'ON' 'CONFLICT' '(' ( ( name ) ( ( ',' name ) )* ) ')'  'DO' 'NOTHING'
-	| 'ON' 'CONFLICT' '(' ( ( name ) ( ( ',' name ) )* ) ')'  'DO' 'UPDATE' 'SET' ( ( ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) 
+	| 'ON' 'CONFLICT' '(' ( ( name ) ( ( ',' name ) )* ) ')'  'DO' 'UPDATE' 'SET' ( ( ( ( column_ref '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_ref '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) 
 	| 'ON' 'CONFLICT' 'ON' 'CONSTRAINT' constraint_name 'DO' 'NOTHING'
-	| 'ON' 'CONFLICT' 'ON' 'CONSTRAINT' constraint_name 'DO' 'UPDATE' 'SET' ( ( ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) 
+	| 'ON' 'CONFLICT' 'ON' 'CONSTRAINT' constraint_name 'DO' 'UPDATE' 'SET' ( ( ( ( column_ref '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_ref '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) 

docs/generated/sql/bnf/stmt_block.bnf

@@ -2394,7 +2394,7 @@ partition ::=
 	'PARTITION' partition_name
 
 single_set_clause ::=
-	column_name '=' a_expr
+	column_ref '=' a_expr
 
 multiple_set_clause ::=
 	'(' insert_column_list ')' '=' in_expr
@@ -2862,6 +2862,10 @@ var_list ::=
 schema_wildcard ::=
 	wildcard_pattern
 
+column_ref ::=
+	column_name
+	| column_name array_subscripts
+
 opt_ordinality ::=
 	'WITH' 'ORDINALITY'
 	| 

docs/generated/sql/bnf/update_stmt.bnf

@@ -1,2 +1,2 @@
 update_stmt ::=
-	( ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) |  ) 'UPDATE' ( ( ( 'ONLY' |  ) table_name opt_index_flags ( '*' |  ) ) | ( ( 'ONLY' |  ) table_name opt_index_flags ( '*' |  ) ) table_alias_name | ( ( 'ONLY' |  ) table_name opt_index_flags ( '*' |  ) ) 'AS' table_alias_name ) 'SET' ( ( ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) ( 'FROM' ( ( table_ref ) ( ( ',' table_ref ) )* ) |  ) ( ( 'WHERE' a_expr ) |  ) ( sort_clause |  ) ( limit_clause |  ) ( 'RETURNING' target_list | 'RETURNING' 'NOTHING' |  )
+	( ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) |  ) 'UPDATE' ( ( ( 'ONLY' |  ) table_name opt_index_flags ( '*' |  ) ) | ( ( 'ONLY' |  ) table_name opt_index_flags ( '*' |  ) ) table_alias_name | ( ( 'ONLY' |  ) table_name opt_index_flags ( '*' |  ) ) 'AS' table_alias_name ) 'SET' ( ( ( ( column_ref '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_ref '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) ( 'FROM' ( ( table_ref ) ( ( ',' table_ref ) )* ) |  ) ( ( 'WHERE' a_expr ) |  ) ( sort_clause |  ) ( limit_clause |  ) ( 'RETURNING' target_list | 'RETURNING' 'NOTHING' |  )

pkg/internal/sqlsmith/relational.go

@@ -937,8 +937,8 @@ func (s *Smither) makeUpdate(refs colRefs) (*tree.Update, *tableRef, bool) {
 			}
 		}
 		update.Exprs = append(update.Exprs, &tree.UpdateExpr{
-			Names: tree.NameList{ref.item.ColumnName},
-			Expr:  expr,
+			ColumnRefs: tree.ColumnRefList{{Name: ref.item.ColumnName}},
+			Expr:       expr,
 		})
 	}
 	if len(update.Exprs) == 0 {

pkg/sql/logictest/testdata/logic_test/json

@@ -918,15 +918,15 @@ SELECT ('{"a": 1}'::jsonb)['a':'b']
 # Check it works from a JSON table.
 statement ok
 CREATE TABLE json_subscript_test (
-  id SERIAL PRIMARY KEY,
+  id INT 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)
+INSERT INTO json_subscript_test (id, j, extract_field, extract_int_field) VALUES
+  (1, '{"other_field": 2}', 'other_field', 1),
+  (2, '{"field": {"field": 2}}', 'field', 0),
+  (3, '[1, 2, 3]', 'nothing_to_fetch', 1)
 
 # Test subscripts with fields using other columns.
 query TTITTTT
@@ -942,3 +942,30 @@ query T
 SELECT j FROM json_subscript_test WHERE j['other_field'] = '2' ORDER BY id
 ----
 {"other_field": 2}
+
+# Test array update.
+statement ok
+UPDATE json_subscript_test SET j[extract_int_field] = '3' WHERE id = 3
+
+query T
+SELECT j FROM json_subscript_test WHERE id = 3
+----
+[1, 3, 3]
+
+# Test updates using other field names.
+statement ok
+UPDATE json_subscript_test SET j[extract_field]['field'] = '[1]' WHERE id = 2
+
+query T
+SELECT j FROM json_subscript_test WHERE id = 2
+----
+{"field": {"field": [1]}}
+
+# Test mixed string and int fields.
+statement ok
+UPDATE json_subscript_test SET j[extract_field]['field'][0] = '2' WHERE id = 2
+
+query T
+SELECT j FROM json_subscript_test WHERE id = 2
+----
+{"field": {"field": [2]}}

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -452,6 +452,25 @@ func (mb *mutationBuilder) addTargetColsByName(names tree.NameList) {
 	}
 }
 
+// addTargetColsByRefs adds one target column for each of the names in the given
+// list.
+func (mb *mutationBuilder) addTargetColsByRefs(refs tree.ColumnRefList) {
+	for _, ref := range refs {
+		name := ref.Name
+		// Determine the ordinal position of the named column in the table and
+		// add it as a target column.
+		if ord := findPublicTableColumnByName(mb.tab, name); ord != -1 {
+			// System columns are invalid target columns.
+			if mb.tab.Column(ord).Kind() == cat.System {
+				panic(pgerror.Newf(pgcode.InvalidColumnReference, "cannot modify system column %q", name))
+			}
+			mb.addTargetCol(ord)
+			continue
+		}
+		panic(colinfo.NewUndefinedColumnError(string(name)))
+	}
+}
+
 // addTargetCol adds a target column by its ordinal position in the target
 // table. It raises an error if a mutation or computed column is targeted, or if
 // the same column is targeted multiple times.

pkg/sql/opt/optbuilder/testdata/update

@@ -97,6 +97,14 @@ CREATE TABLE generated_as_identity (
 )
 ----
 
+exec-ddl
+CREATE TABLE jsonb_table (
+  j JSONB,
+  k JSONB,
+  ref TEXT
+)
+----
+
 # ------------------------------------------------------------------------------
 # Basic tests.
 # ------------------------------------------------------------------------------
@@ -613,6 +621,24 @@ UPDATE abcde SET b=1 ORDER BY c
 ----
 error (42601): UPDATE statement requires LIMIT when ORDER BY is used
 
+# JSONB single clause update.
+build
+UPDATE jsonb_table SET j['a'] = '1', k[ref]['b'][0] = '2'
+----
+update jsonb_table
+ ├── columns: <none>
+ ├── fetch columns: j:7 k:8 ref:9 rowid:10
+ ├── update-mapping:
+ │    ├── j_new:13 => j:1
+ │    └── k_new:14 => k:2
+ └── project
+      ├── columns: j_new:13 k_new:14 j:7 k:8 ref:9 rowid:10!null crdb_internal_mvcc_timestamp:11 tableoid:12
+      ├── scan jsonb_table
+      │    └── columns: j:7 k:8 ref:9 rowid:10!null crdb_internal_mvcc_timestamp:11 tableoid:12
+      └── projections
+           ├── json_set(j:7, ARRAY['a'], '1', true) [as=j_new:13]
+           └── json_set(k:8, ARRAY[ref:9, 'b', 0::STRING], '2', true) [as=k_new:14]
+
 # ------------------------------------------------------------------------------
 # Test RETURNING.
 # ------------------------------------------------------------------------------

pkg/sql/opt/optbuilder/update.go

@@ -126,7 +126,7 @@ func (mb *mutationBuilder) addTargetColsForUpdate(exprs tree.UpdateExprs) {
 	}
 
 	for _, expr := range exprs {
-		mb.addTargetColsByName(expr.Names)
+		mb.addTargetColsByRefs(expr.ColumnRefs)
 
 		if expr.Tuple {
 			n := -1
@@ -137,8 +137,8 @@ func (mb *mutationBuilder) addTargetColsForUpdate(exprs tree.UpdateExprs) {
 				// Use the data types of the target columns to resolve expressions
 				// with ambiguous types (e.g. should 1 be interpreted as an INT or
 				// as a FLOAT).
-				desiredTypes := make([]*types.T, len(expr.Names))
-				targetIdx := len(mb.targetColList) - len(expr.Names)
+				desiredTypes := make([]*types.T, len(expr.ColumnRefs))
+				targetIdx := len(mb.targetColList) - len(expr.ColumnRefs)
 				for i := range desiredTypes {
 					desiredTypes[i] = mb.md.ColumnMeta(mb.targetColList[targetIdx+i]).Type
 				}
@@ -153,10 +153,10 @@ func (mb *mutationBuilder) addTargetColsForUpdate(exprs tree.UpdateExprs) {
 				panic(unimplementedWithIssueDetailf(35713, fmt.Sprintf("%T", expr.Expr),
 					"source for a multiple-column UPDATE item must be a sub-SELECT or ROW() expression; not supported: %T", expr.Expr))
 			}
-			if len(expr.Names) != n {
+			if len(expr.ColumnRefs) != n {
 				panic(pgerror.Newf(pgcode.Syntax,
 					"number of columns (%d) does not match number of values (%d)",
-					len(expr.Names), n))
+					len(expr.ColumnRefs), n))
 			}
 		}
 	}
@@ -272,7 +272,34 @@ func (mb *mutationBuilder) addUpdateCols(exprs tree.UpdateExprs) {
 				}
 			}
 		} else {
-			addCol(set.Expr, mb.targetColList[n])
+			expr := set.Expr
+			if len(set.ColumnRefs) > 1 {
+				panic(errors.AssertionFailedf("expected <= 1 column ref, found %d", len(set.ColumnRefs)))
+			}
+			for _, ref := range set.ColumnRefs {
+				// For JSONB subscripts, replace with json_set.
+				if len(ref.Subscripts) > 0 {
+					arr := &tree.Array{}
+					for _, t := range ref.Subscripts {
+						if t.Slice {
+							panic(pgerror.Newf(pgcode.DatatypeMismatch, "cannot reference a slice in UPDATE"))
+						}
+						// Cast all expressions to strings to support mixing strings and ints.
+						// JSONB automatically knows how to evaluate strings to int indexes and vice versa.
+						arr.Exprs = append(arr.Exprs, &tree.CastExpr{Expr: t.Begin, Type: types.String})
+					}
+					expr = &tree.FuncExpr{
+						Func: tree.WrapFunction("json_set"),
+						Exprs: tree.Exprs{
+							&tree.UnresolvedName{NumParts: 1, Parts: tree.NameParts{string(ref.Name)}},
+							arr,
+							expr,
+							tree.DBoolTrue,
+						},
+					}
+				}
+			}
+			addCol(expr, mb.targetColList[n])
 			n++
 		}
 	}

pkg/sql/parser/sql.y

@@ -326,6 +326,9 @@ func (u *sqlSymUnion) tableNames() tree.TableNames {
 func (u *sqlSymUnion) indexFlags() *tree.IndexFlags {
     return u.val.(*tree.IndexFlags)
 }
+func (u *sqlSymUnion) columnRef() tree.ColumnRef {
+    return u.val.(tree.ColumnRef)
+}
 func (u *sqlSymUnion) arraySubscript() *tree.ArraySubscript {
     return u.val.(*tree.ArraySubscript)
 }
@@ -1302,6 +1305,7 @@ func (u *sqlSymUnion) asTenantClause() tree.TenantID {
 %type <tree.UpdateExprs> set_clause_list
 %type <*tree.UpdateExpr> set_clause multiple_set_clause
 %type <tree.ArraySubscripts> array_subscripts
+%type <tree.ColumnRef> column_ref
 %type <tree.GroupBy> group_clause
 %type <tree.Exprs> group_by_list
 %type <tree.Expr> group_by_item
@@ -9944,17 +9948,33 @@ set_clause:
   single_set_clause
 | multiple_set_clause
 
-single_set_clause:
-  column_name '=' a_expr
+column_ref:
+  column_name
+  {
+    $$.val = tree.ColumnRef{Name: tree.Name($1)}
+  }
+| column_name array_subscripts
   {
-    $$.val = &tree.UpdateExpr{Names: tree.NameList{tree.Name($1)}, Expr: $3.expr()}
+    $$.val = tree.ColumnRef{Name: tree.Name($1), Subscripts: $2.arraySubscripts()}
   }
 | column_name '.' error { return unimplementedWithIssue(sqllex, 27792) }
 
+single_set_clause:
+  column_ref '=' a_expr
+  {
+    $$.val = &tree.UpdateExpr{ColumnRefs: tree.ColumnRefList{$1.columnRef()}, Expr: $3.expr()}
+  }
+
 multiple_set_clause:
   '(' insert_column_list ')' '=' in_expr
   {
-    $$.val = &tree.UpdateExpr{Tuple: true, Names: $2.nameList(), Expr: $5.expr()}
+    // TODO(#77434): make this work for the multiple_set_clause case.
+    names := $2.nameList()
+    refs := make([]tree.ColumnRef, len(names))
+    for idx, name := range names {
+      refs[idx] = tree.ColumnRef{Name: name}
+    }
+    $$.val = &tree.UpdateExpr{Tuple: true, ColumnRefs: refs, Expr: $5.expr()}
   }
 
 // %Help: REASSIGN OWNED BY - change ownership of all objects

pkg/sql/parser/testdata/update

@@ -200,11 +200,35 @@ UPDATE a SET b = (3) -- fully parenthesized
 UPDATE a SET b = _ -- literals removed
 UPDATE _ SET _ = 3 -- identifiers removed
 
-error
+parse
 UPDATE kv SET k[0] = 9
 ----
+UPDATE kv SET k[0] = 9
+UPDATE kv SET k[(0)] = (9) -- fully parenthesized
+UPDATE kv SET k[_] = _ -- literals removed
+UPDATE _ SET _[0] = 9 -- identifiers removed
+
+parse
+UPDATE a SET k['a'] = 1
+----
+UPDATE a SET k['a'] = 1
+UPDATE a SET k[('a')] = (1) -- fully parenthesized
+UPDATE a SET k['_'] = _ -- literals removed
+UPDATE _ SET _['a'] = 1 -- identifiers removed
+
+parse
+UPDATE a SET k['a'][0][column_ref] = 1
+----
+UPDATE a SET k['a'][0][column_ref] = 1
+UPDATE a SET k[('a')][(0)][(column_ref)] = (1) -- fully parenthesized
+UPDATE a SET k['_'][_][column_ref] = _ -- literals removed
+UPDATE _ SET _['a'][0][_] = 1 -- identifiers removed
+
+error
+UPDATE a SET (k['a'], k['b']) = ('1', '2')
+----
 at or near "[": syntax error
 DETAIL: source SQL:
-UPDATE kv SET k[0] = 9
+UPDATE a SET (k['a'], k['b']) = ('1', '2')
                ^
 HINT: try \h UPDATE

pkg/sql/sem/tree/pretty.go

@@ -1217,7 +1217,7 @@ func (node *Order) doc(p *PrettyCfg) pretty.Doc {
 }
 
 func (node *UpdateExpr) doc(p *PrettyCfg) pretty.Doc {
-	d := p.Doc(&node.Names)
+	d := p.Doc(&node.ColumnRefs)
 	if node.Tuple {
 		d = p.bracket("(", d, ")")
 	}

pkg/sql/sem/tree/update.go

@@ -73,11 +73,36 @@ func (node *UpdateExprs) Format(ctx *FmtCtx) {
 	}
 }
 
+// ColumnRef is a reference to a column with a possible subscript.
+type ColumnRef struct {
+	Name       Name
+	Subscripts ArraySubscripts
+}
+
+// Format implements the NodeFormatter interface.
+func (node *ColumnRef) Format(ctx *FmtCtx) {
+	ctx.FormatNode(&node.Name)
+	ctx.FormatNode(&node.Subscripts)
+}
+
+// A ColumnRefList is a list of column references.
+type ColumnRefList []ColumnRef
+
+// Format implements the NodeFormatter interface.
+func (l *ColumnRefList) Format(ctx *FmtCtx) {
+	for i := range *l {
+		if i > 0 {
+			ctx.WriteString(", ")
+		}
+		ctx.FormatNode(&(*l)[i])
+	}
+}
+
 // UpdateExpr represents an update expression.
 type UpdateExpr struct {
-	Tuple bool
-	Names NameList
-	Expr  Expr
+	Tuple      bool
+	ColumnRefs ColumnRefList
+	Expr       Expr
 }
 
 // Format implements the NodeFormatter interface.
@@ -87,7 +112,7 @@ func (node *UpdateExpr) Format(ctx *FmtCtx) {
 		open, close = "(", ")"
 	}
 	ctx.WriteString(open)
-	ctx.FormatNode(&node.Names)
+	ctx.FormatNode(&node.ColumnRefs)
 	ctx.WriteString(close)
 	ctx.WriteString(" = ")
 	ctx.FormatNode(node.Expr)