sql: enable update and upsert in udf bodies

This PR enables UPDATE in UDFs and adds some logic test coverage. Since
both INSERT and UPDATE are supported, this also enabled UPSERT.

Epic: CRDB-19255
Informs: cockroachdb#87289

Release note(sql change): Enables UPDATE and UPSERT commands in UDF
statement bodies.

pkg/sql/logictest/testdata/logic_test/udf

@@ -2672,10 +2672,7 @@ CREATE FUNCTION err() RETURNS VOID LANGUAGE SQL AS 'DROP TABLE t'
 
 subtest mutation
 
-# Mutations are not currently supported in UDF bodies.
-statement error pgcode 0A000 unimplemented: UPDATE usage inside a function definition
-CREATE FUNCTION err() RETURNS VOID LANGUAGE SQL AS 'UPDATE t SET a = 1'
-
+# Not all mutations are currently supported in UDF bodies.
 statement error pgcode 0A000 unimplemented: DELETE usage inside a function definition
 CREATE FUNCTION err() RETURNS VOID LANGUAGE SQL AS 'DELETE FROM t WHERE a = 1'
 

pkg/sql/logictest/testdata/logic_test/udf_update

@@ -0,0 +1,176 @@
+statement ok
+CREATE TABLE t (a INT PRIMARY KEY, b INT);
+
+statement ok
+CREATE FUNCTION f1(i INT, j INT) RETURNS RECORD AS
+$$
+  UPDATE t SET a = i, b = j RETURNING *;
+$$ LANGUAGE SQL;
+
+query T
+SELECT f1(1,1);
+----
+NULL
+
+query T
+SELECT f1(1,1);
+----
+NULL
+
+statement ok
+INSERT INTO t VALUES (1, 2);
+
+query T
+SELECT f1(1,1);
+----
+(1,1)
+
+query T
+SELECT f1(3,4);
+----
+(3,4)
+
+query II
+SELECT * FROM t;
+----
+3 4
+
+statement ok
+INSERT INTO t VALUES (1, 2),(5,6),(7,8);
+
+statement error duplicate key value violates unique constraint "t_pkey"
+SELECT f1(1, 11);
+
+statement error duplicate key value violates unique constraint "t_pkey"
+SELECT f1(14, 14);
+
+query II
+SELECT * FROM t;
+----
+1 2
+3 4
+5 6
+7 8
+
+statement ok
+CREATE FUNCTION f2(i INT, j INT) RETURNS RECORD AS
+$$
+  UPDATE t SET b = j WHERE a = i RETURNING *;
+$$ LANGUAGE SQL;
+
+query T
+SELECT f2(11,2);
+----
+NULL
+
+query T
+SELECT f2(5,32);
+----
+(5,32)
+
+statement error pq: multiple modification subqueries of the same table \"t\" are not supported
+SELECT f2(5,9), f2(7,11);
+
+statement ok
+CREATE TABLE t2 (a INT, b INT, c INT);
+INSERT INTO t2 VALUES (1,2,3),(4,5,6),(7,8,9);
+
+statement ok
+CREATE FUNCTION f3(i INT, j INT, k INT) RETURNS SETOF RECORD AS
+$$
+  UPDATE t2 SET (a, b) = (i, j) WHERE b < k RETURNING *;
+$$ LANGUAGE SQL;
+
+query III
+SELECT * FROM f3(3, 3, 7) AS foo(a INT, b INT, c INT);
+----
+3 3 3
+3 3 6
+
+statement ok
+CREATE TABLE t3(a) AS SELECT 1::INT;
+
+statement error aggregate functions are not allowed in UPDATE SET
+CREATE FUNCTION f4() RETURNS RECORD AS
+$$
+  UPDATE t3 SET a = count(a);
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f5(i INT) RETURNS INT AS
+$$
+  UPDATE t3 SET a = i RETURNING *;
+$$ LANGUAGE SQL;
+
+# Aggregate functions are ok as input parameters to a UDF performing an update.
+query I
+SELECT f5(count(a)) FROM t3;
+----
+1
+
+subtest constraints
+
+statement ok
+CREATE TABLE t_check(x DECIMAL(1,0) CHECK (x >= 1))
+
+statement ok
+INSERT INTO t_check VALUES (2)
+
+statement ok
+CREATE FUNCTION f_check1(d DECIMAL) RETURNS RECORD AS
+$$
+  UPDATE t_check SET x = d RETURNING *;
+$$ LANGUAGE SQL;
+
+query T
+SELECT f_check1(3.5);
+----
+(4)
+
+query T
+SELECT f_check1(0.5);
+----
+(1)
+
+statement error pq: failed to satisfy CHECK constraint \(x >= 1:::DECIMAL\)
+SELECT f_check1(0);
+
+# subtest generated_as_identity
+statement ok
+CREATE TABLE generated_as_id_t (
+  a INT UNIQUE,
+  b INT GENERATED ALWAYS AS IDENTITY,
+  c INT GENERATED BY DEFAULT AS IDENTITY
+);
+
+statement ok
+INSERT INTO generated_as_id_t (a) VALUES (7), (8), (9);
+
+statement error column "b" can only be updated to DEFAULT
+CREATE FUNCTION f_err(i INT, j INT) RETURNS RECORD AS
+$$
+  UPDATE generated_as_id_t SET b=i, c=j WHERE a > 6 RETURNING *;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_generated_as_id(j INT) RETURNS RECORD AS
+$$
+  UPDATE generated_as_id_t SET b=DEFAULT, c=j WHERE a > 6 RETURNING *;
+$$ LANGUAGE SQL;
+
+query III
+SELECT * FROM generated_as_id_t ORDER BY a;
+----
+7  1  1
+8  2  2
+9  3  3
+
+statement ok
+SELECT f_generated_as_id(1+1);
+
+query III
+SELECT * FROM generated_as_id_t ORDER BY a;
+----
+7  4  2
+8  5  2
+9  6  2

pkg/sql/opt/optbuilder/builder.go

@@ -320,7 +320,6 @@ func (b *Builder) buildStmt(
 			panic(unimplemented.NewWithIssuef(87289, "%s usage inside a function definition", stmt.StatementTag()))
 		case *tree.Insert:
 		case *tree.Update:
-			panic(unimplemented.NewWithIssuef(87289, "%s usage inside a function definition", stmt.StatementTag()))
 		default:
 			panic(unimplemented.Newf("user-defined functions", "%s usage inside a function definition", stmt.StatementTag()))
 		}