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 pgcode 23505 duplicate key value violates unique constraint "t_pkey"
SELECT f1(1, 11);

statement error pgcode 23505 duplicate key value violates unique constraint "t_pkey"
SELECT f1(14, 14);

query II rowsort
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)

query TT
SELECT f2(5,9), f2(7,11);
----
(5,9)  (7,11)

query T nosort
SELECT f2(x,y) FROM (VALUES (1,16),(1,17)) v(x,y);
----
(1,16)
(1,17)

query II rowsort
SELECT * FROM t;
----
1  17
3  4
5  9
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 rowsort
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 pgcode 42803 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 pgcode 23514 pq: failed to satisfy CHECK constraint \(x >= 1:::DECIMAL\)
SELECT f_check1(0);


subtest end

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 pgcode 428C9 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


subtest end

subtest repro_103693

statement ok
CREATE TABLE t1_103693(k1 PRIMARY KEY) AS SELECT 1::INT;
CREATE TABLE t2_103693(k2 PRIMARY KEY, v2) AS SELECT 1::INT, 0::INT;

statement ok
CREATE FUNCTION f_103693(k INT, v INT) RETURNS INT AS
$$
  UPDATE t2_103693 SET v2 = v + 1 WHERE k2 = k2 RETURNING v2;
$$ LANGUAGE SQL;

statement ok
SET streamer_enabled = true

query I
SELECT f_103693(k2, v2) FROM t1_103693 INNER LOOKUP JOIN t2_103693 ON k1 = k2;
----
1

subtest end

subtest select_for_update

statement ok
CREATE TABLE kv (k INT PRIMARY KEY, v INT);
INSERT INTO kv VALUES (1,2), (3,4), (5,6);

statement ok
GRANT SELECT, UPDATE ON kv TO testuser;

statement ok
CREATE FUNCTION f_select_for_update(i INT) RETURNS kv AS
$$
  SELECT * FROM kv WHERE k = i FOR UPDATE;
$$ LANGUAGE SQL;

statement ok
CREATE FUNCTION f_update(i INT, j INT) RETURNS VOID AS
$$
  UPDATE kv SET v = j WHERE k = i;
$$ LANGUAGE SQL;

statement ok
CREATE FUNCTION f_select_and_update(i INT, j INT) RETURNS VOID AS
$$
  SELECT * FROM kv WHERE k = i FOR UPDATE;
  UPDATE kv SET v = j WHERE k = i;
$$ LANGUAGE SQL;

statement ok
SELECT f_select_and_update(1,1);

query II rowsort
SELECT * FROM kv;
----
1 1
3 4
5 6

statement ok
BEGIN; SELECT * FROM kv WHERE k = 3 FOR UPDATE;

statement ok
SELECT f_update(3,3);

statement ok
END;

query II rowsort
SELECT * FROM kv;
----
1 1
3 3
5 6

statement ok
BEGIN; SELECT * FROM kv WHERE k = 5 FOR UPDATE;

statement ok
SELECT f_update(5,5);

statement ok
ROLLBACK;

query II rowsort
SELECT * FROM kv;
----
1 1
3 3
5 6

statement ok
BEGIN;

statement ok
SELECT f_select_for_update(5);

user testuser

statement ok
SET statement_timeout = '10ms'

# Make sure that the UDF did acquire locks by attempting to acquire them in a
# different session.
query error pgcode 57014 query execution canceled due to statement timeout
SELECT * FROM kv WHERE k = 5 FOR UPDATE;

statement ok
SET statement_timeout = 0

user root

statement ok
SELECT f_update(5,5);

statement ok
END;

query II rowsort
SELECT * FROM kv;
----
1 1
3 3
5 5

subtest end
