CREATE SCHEMA multi_update_select;
SET search_path TO multi_update_select;
SET citus.next_shard_id TO 751000;
-- specific tests related to get_update_query_targetlist_def
-- we test only queries with sublinks, like:
-- ( ... SET (...) = (SELECT ...))
-- Reference tables
CREATE TABLE test_ref_indirection (
    id bigint primary key
  , col_bool bool , col_date date , col_int integer , col_text text
  );
SELECT create_reference_table('test_ref_indirection');
 create_reference_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE test_ref_indirection_new (
    id bigint primary key
  , col_bool bool , col_date date , col_int integer , col_text text
  );
SELECT create_reference_table('test_ref_indirection_new');
 create_reference_table
---------------------------------------------------------------------

(1 row)

-- Distributed tables
CREATE TABLE test_dist_indirection (
    id bigint primary key
  , col_bool bool , col_date date , col_int integer , col_text text
  );
SELECT create_distributed_table('test_dist_indirection', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

CREATE TABLE test_dist_indirection_new (
    id bigint primary key
  , col_bool bool , col_date date , col_int integer , col_text text
  );
SELECT create_distributed_table('test_dist_indirection_new', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- those should work:
INSERT INTO test_ref_indirection (id, col_bool, col_date, col_int, col_text)
    SELECT 1, true, '1970-01-01'::date, 1, 'one';
INSERT INTO test_dist_indirection (id, col_bool, col_date, col_int, col_text)
    SELECT 1, true, '1970-01-01'::date, 1, 'one';
INSERT INTO test_ref_indirection (id, col_text, col_bool, col_date, col_int)
    SELECT 2, 'two', false, '1970-01-01'::date, 2;
INSERT INTO test_dist_indirection (id, col_text, col_bool, col_date, col_int)
    SELECT 2, 'two', false, '1970-01-01'::date, 2;
INSERT INTO test_ref_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
INSERT INTO test_dist_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
INSERT INTO test_ref_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
INSERT INTO test_dist_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
INSERT INTO test_ref_indirection_new SELECT * FROM test_ref_indirection;
INSERT INTO test_dist_indirection_new SELECT * FROM test_dist_indirection;
SELECT * FROM test_ref_indirection ORDER BY id;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1970 |       1 | one
  2 | f        | 01-01-1970 |       2 | two
  3 | f        | 01-01-1970 |       0 | empty
  4 | f        | 01-01-1970 |       0 | empty
(4 rows)

SELECT * FROM test_dist_indirection ORDER BY id;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1970 |       1 | one
  2 | f        | 01-01-1970 |       2 | two
  3 | f        | 01-01-1970 |       0 | empty
  4 | f        | 01-01-1970 |       0 | empty
(4 rows)

SELECT * FROM test_ref_indirection_new ORDER BY id;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1970 |       1 | one
  2 | f        | 01-01-1970 |       2 | two
  3 | f        | 01-01-1970 |       0 | empty
  4 | f        | 01-01-1970 |       0 | empty
(4 rows)

SELECT * FROM test_dist_indirection_new ORDER BY id;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1970 |       1 | one
  2 | f        | 01-01-1970 |       2 | two
  3 | f        | 01-01-1970 |       0 | empty
  4 | f        | 01-01-1970 |       0 | empty
(4 rows)

-- now UPDATEs
UPDATE test_ref_indirection
    SET (col_bool, col_date, col_int, col_text)
        = (SELECT true, '1970-01-01'::date, 1, 'ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1970 |       1 | ok
  2 | t        | 01-01-1970 |       1 | ok
  3 | t        | 01-01-1970 |       1 | ok
  4 | t        | 01-01-1970 |       1 | ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_bool, col_date, col_int, col_text)
        = (SELECT true, '1970-01-01'::date, 1, 'ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1970 |       1 | ok
  2 | t        | 01-01-1970 |       1 | ok
  3 | t        | 01-01-1970 |       1 | ok
  4 | t        | 01-01-1970 |       1 | ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
      , (col_int, col_text) = (select 2, '2 ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1971 |       2 | 2 ok
  2 | f        | 01-01-1971 |       2 | 2 ok
  3 | f        | 01-01-1971 |       2 | 2 ok
  4 | f        | 01-01-1971 |       2 | 2 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
      , (col_int, col_text) = (select 2, '2 ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1971 |       2 | 2 ok
  2 | f        | 01-01-1971 |       2 | 2 ok
  3 | f        | 01-01-1971 |       2 | 2 ok
  4 | f        | 01-01-1971 |       2 | 2 ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_bool, col_int)  = (select true, 3)
      , (col_text) = (select '3 ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1971 |       3 | 3 ok
  2 | t        | 01-01-1971 |       3 | 3 ok
  3 | t        | 01-01-1971 |       3 | 3 ok
  4 | t        | 01-01-1971 |       3 | 3 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_bool, col_int)  = (select true, 3)
      , (col_text) = (select '3 ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1971 |       3 | 3 ok
  2 | t        | 01-01-1971 |       3 | 3 ok
  3 | t        | 01-01-1971 |       3 | 3 ok
  4 | t        | 01-01-1971 |       3 | 3 ok
(4 rows)

-- but those should work since 13.X
UPDATE test_ref_indirection
    SET (col_date, col_text, col_int, col_bool)
        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1972 |       4 | 4 ok
  2 | f        | 01-01-1972 |       4 | 4 ok
  3 | f        | 01-01-1972 |       4 | 4 ok
  4 | f        | 01-01-1972 |       4 | 4 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_date, col_text, col_int, col_bool)
        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1972 |       4 | 4 ok
  2 | f        | 01-01-1972 |       4 | 4 ok
  3 | f        | 01-01-1972 |       4 | 4 ok
  4 | f        | 01-01-1972 |       4 | 4 ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_int, col_text)  = (select 5, '5 ok')
      , (col_bool) = (select true)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1972 |       5 | 5 ok
  2 | t        | 01-01-1972 |       5 | 5 ok
  3 | t        | 01-01-1972 |       5 | 5 ok
  4 | t        | 01-01-1972 |       5 | 5 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_int, col_text)  = (select 5, '5 ok')
      , (col_bool) = (select true)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1972 |       5 | 5 ok
  2 | t        | 01-01-1972 |       5 | 5 ok
  3 | t        | 01-01-1972 |       5 | 5 ok
  4 | t        | 01-01-1972 |       5 | 5 ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
      , (col_text, col_bool) = (select '6 ok', false)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1973 |       6 | 6 ok
  2 | f        | 01-01-1973 |       6 | 6 ok
  3 | f        | 01-01-1973 |       6 | 6 ok
  4 | f        | 01-01-1973 |       6 | 6 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
      , (col_text, col_bool) = (select '6 ok', false)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1973 |       6 | 6 ok
  2 | f        | 01-01-1973 |       6 | 6 ok
  3 | f        | 01-01-1973 |       6 | 6 ok
  4 | f        | 01-01-1973 |       6 | 6 ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
      , (col_bool) = (select true)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1974 |       7 | 7 ok
  2 | t        | 01-01-1974 |       7 | 7 ok
  3 | t        | 01-01-1974 |       7 | 7 ok
  4 | t        | 01-01-1974 |       7 | 7 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
      , (col_bool) = (select true)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1974 |       7 | 7 ok
  2 | t        | 01-01-1974 |       7 | 7 ok
  3 | t        | 01-01-1974 |       7 | 7 ok
  4 | t        | 01-01-1974 |       7 | 7 ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
      , (col_int) = (select 8)
      , (col_bool) = (select false)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1975 |       8 | 8 ok
  2 | f        | 01-01-1975 |       8 | 8 ok
  3 | f        | 01-01-1975 |       8 | 8 ok
  4 | f        | 01-01-1975 |       8 | 8 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
      , (col_int) = (select 8)
      , (col_bool) = (select false)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1975 |       8 | 8 ok
  2 | f        | 01-01-1975 |       8 | 8 ok
  3 | f        | 01-01-1975 |       8 | 8 ok
  4 | f        | 01-01-1975 |       8 | 8 ok
(4 rows)

--
-- more restrictive ones, just in case we miss a wrong value
--
-- those should work
UPDATE test_ref_indirection
    SET (col_bool, col_text) = (SELECT true, '9 ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1975 |       8 | 9 ok
  2 | t        | 01-01-1975 |       8 | 9 ok
  3 | t        | 01-01-1975 |       8 | 9 ok
  4 | t        | 01-01-1975 |       8 | 9 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_bool, col_text) = (SELECT true, '9 ok')
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1975 |       8 | 9 ok
  2 | t        | 01-01-1975 |       8 | 9 ok
  3 | t        | 01-01-1975 |       8 | 9 ok
  4 | t        | 01-01-1975 |       8 | 9 ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_bool, col_text) = (SELECT false, '10 ok')
WHERE id = 1
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1975 |       8 | 10 ok
(1 row)

UPDATE test_dist_indirection
    SET (col_bool, col_text) = (SELECT false, '10 ok')
WHERE id = 1
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | f        | 01-01-1975 |       8 | 10 ok
(1 row)

UPDATE test_ref_indirection
    SET (col_text, col_bool) = (SELECT '11 ok', true)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1975 |       8 | 11 ok
  2 | t        | 01-01-1975 |       8 | 11 ok
  3 | t        | 01-01-1975 |       8 | 11 ok
  4 | t        | 01-01-1975 |       8 | 11 ok
(4 rows)

UPDATE test_dist_indirection
    SET (col_text, col_bool) = (SELECT '11 ok', true)
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  1 | t        | 01-01-1975 |       8 | 11 ok
  2 | t        | 01-01-1975 |       8 | 11 ok
  3 | t        | 01-01-1975 |       8 | 11 ok
  4 | t        | 01-01-1975 |       8 | 11 ok
(4 rows)

UPDATE test_ref_indirection
    SET (col_text, col_bool) = (SELECT '12 ok', false)
WHERE id = 2
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  2 | f        | 01-01-1975 |       8 | 12 ok
(1 row)

UPDATE test_dist_indirection
    SET (col_text, col_bool) = (SELECT '12 ok', false)
WHERE id = 2
RETURNING *;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  2 | f        | 01-01-1975 |       8 | 12 ok
(1 row)

-- several updates in CTE shoult not work
with qq3 as (
    update test_ref_indirection
    SET (col_text, col_bool)
      = (SELECT '13', true)
    where id = 3
    returning *
),
qq4 as (
    update test_ref_indirection
    SET (col_text, col_bool)
      = (SELECT '14', false)
    where id = 4
    returning *
)
select * from qq3 union all select * from qq4;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  3 | t        | 01-01-1975 |       8 | 13
  4 | f        | 01-01-1975 |       8 | 14
(2 rows)

with qq3 as (
    update test_dist_indirection
    SET (col_text, col_bool)
      = (SELECT '13', true)
    where id = 3
    returning *
),
qq4 as (
    update test_dist_indirection
    SET (col_text, col_bool)
      = (SELECT '14', false)
    where id = 4
    returning *
)
select * from qq3 union all select * from qq4;
 id | col_bool |  col_date  | col_int | col_text
---------------------------------------------------------------------
  3 | t        | 01-01-1975 |       8 | 13
  4 | f        | 01-01-1975 |       8 | 14
(2 rows)

DROP TABLE test_dist_indirection;
DROP TABLE test_dist_indirection_new;
DROP TABLE test_ref_indirection;
DROP TABLE test_ref_indirection_new;
-- https://github.com/citusdata/citus/issues/4092
CREATE TABLE update_test (
    a   INT DEFAULT 10,
    b   INT,
    c   TEXT
);
SELECT create_reference_table('update_test');
 create_reference_table
---------------------------------------------------------------------

(1 row)

INSERT INTO update_test VALUES (11, 41, 'car');
INSERT INTO update_test VALUES (100, 20, 'bike');
INSERT INTO update_test VALUES (100, 20, 'tractor');
SELECT * FROM update_test;
  a  | b  |    c
---------------------------------------------------------------------
  11 | 41 | car
 100 | 20 | bike
 100 | 20 | tractor
(3 rows)

UPDATE update_test
SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
WHERE a = 100 AND b = 20
RETURNING *;
 a  | b  |    c
---------------------------------------------------------------------
 41 | 11 | bike
 41 | 11 | tractor
(2 rows)

-- Test that multiple out of order columns and multiple sublinks are handled correctly.
CREATE TABLE upd2_test (a1 int, b1 int, c1 int, d1 int, e1 int, f1 int, g1 int);
SELECT create_reference_table('upd2_test');
 create_reference_table
---------------------------------------------------------------------

(1 row)

INSERT INTO upd2_test SELECT 1, 1, 1, 1, 1, 1, 1 FROM generate_series(1,5) c(i);
UPDATE upd2_test set (b1, a1) = (SELECT 200, 100), (g1, f1, e1) = (SELECT 700, 600, 500), (d1, c1) = (SELECT 400, 300);
SELECT * FROM upd2_test;
 a1  | b1  | c1  | d1  | e1  | f1  | g1
---------------------------------------------------------------------
 100 | 200 | 300 | 400 | 500 | 600 | 700
 100 | 200 | 300 | 400 | 500 | 600 | 700
 100 | 200 | 300 | 400 | 500 | 600 | 700
 100 | 200 | 300 | 400 | 500 | 600 | 700
 100 | 200 | 300 | 400 | 500 | 600 | 700
(5 rows)

UPDATE upd2_test set (g1, a1) = (SELECT 77, 11), (f1, b1) = (SELECT 66, 22), (e1, c1) = (SELECT 55, 33), (d1) = (SELECT 44);
SELECT * FROM upd2_test;
 a1 | b1 | c1 | d1 | e1 | f1 | g1
---------------------------------------------------------------------
 11 | 22 | 33 | 44 | 55 | 66 | 77
 11 | 22 | 33 | 44 | 55 | 66 | 77
 11 | 22 | 33 | 44 | 55 | 66 | 77
 11 | 22 | 33 | 44 | 55 | 66 | 77
 11 | 22 | 33 | 44 | 55 | 66 | 77
(5 rows)

UPDATE upd2_test set (g1, a1) = (SELECT 7, 1), (f1) = (SELECT 6), (c1, e1) = (SELECT 3, 5), (b1) = (SELECT 2), (d1) = (SELECT 4);
SELECT * FROM upd2_test;
 a1 | b1 | c1 | d1 | e1 | f1 | g1
---------------------------------------------------------------------
  1 |  2 |  3 |  4 |  5 |  6 |  7
  1 |  2 |  3 |  4 |  5 |  6 |  7
  1 |  2 |  3 |  4 |  5 |  6 |  7
  1 |  2 |  3 |  4 |  5 |  6 |  7
  1 |  2 |  3 |  4 |  5 |  6 |  7
(5 rows)

-- Test out of order updates on distributed table
CREATE TABLE dist_test(a1 int, b1 numeric, c1 text, d1 int);
SELECT create_distributed_table('dist_test', 'a1');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO dist_test VALUES (1, 10.0, 'xxx', 4);
INSERT INTO dist_test VALUES (1, 10.0, 'xxx', 4);
INSERT INTO dist_test VALUES (2, 10.0, 'xxx', 4);
INSERT INTO dist_test VALUES (2, 10.0, 'xxx', 4);
INSERT INTO dist_test VALUES (3, 10.0, 'xxx', 4);
INSERT INTO dist_test VALUES (3, 10.0, 'xxx', 4);
INSERT INTO dist_test VALUES (3, 10.0, 'xxx', 4);
-- Router plan:
UPDATE dist_test
SET (d1, c1, b1) = (SELECT 7, 'tractor', 4.2)
WHERE a1=1
RETURNING *;
 a1 | b1  |   c1    | d1
---------------------------------------------------------------------
  1 | 4.2 | tractor |  7
  1 | 4.2 | tractor |  7
(2 rows)

-- Pushdown plan:
UPDATE dist_test
SET (d1, c1, b1) = (SELECT X, 'car', Y)
FROM (SELECT r.a * d1 as X, r.b * b1 as Y FROM update_test r, dist_test WHERE r.c=c1) upd
WHERE dist_test.a1 > 2
RETURNING *;
 a1 |  b1  | c1  | d1  |  x  |  y
---------------------------------------------------------------------
  3 | 46.2 | car | 287 | 287 | 46.2
  3 | 46.2 | car | 287 | 287 | 46.2
  3 | 46.2 | car | 287 | 287 | 46.2
(3 rows)

-- Test subscripting updates
CREATE TABLE jsonb_subscript_update (id INT, data JSONB);
SELECT create_distributed_table('jsonb_subscript_update', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}');
UPDATE jsonb_subscript_update
SET data['b'] = updated_vals.b::TEXT::jsonb,
    data['c'] = updated_vals.c::TEXT::jsonb,
    data['d'] = updated_vals.d::TEXT::jsonb
FROM (
  SELECT id,
         data['a'] AS a,
         data['a']::NUMERIC + 1 AS b,
         data['a']::NUMERIC + 2 AS c,
         data['a']::NUMERIC + 3 AS d
  FROM jsonb_subscript_update
) updated_vals
WHERE jsonb_subscript_update.id = updated_vals.id;
SELECT * FROM jsonb_subscript_update ORDER BY 1,2;
 id |               data
---------------------------------------------------------------------
  1 | {"a": 1, "b": 2, "c": 3, "d": 4}
  2 | {"a": 2, "b": 3, "c": 4, "d": 5}
(2 rows)

TRUNCATE jsonb_subscript_update;
INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}');
-- test router update with jsonb subscript
UPDATE jsonb_subscript_update
SET data['b'] = updated_vals.b::TEXT::jsonb,
    data['c'] = updated_vals.c::TEXT::jsonb,
    data['d'] = updated_vals.d::TEXT::jsonb
FROM (
  SELECT id,
         data['a'] AS a,
         data['a']::NUMERIC + 1 AS b,
         data['a']::NUMERIC + 2 AS c,
         data['a']::NUMERIC + 3 AS d
  FROM jsonb_subscript_update
) updated_vals
WHERE jsonb_subscript_update.id = updated_vals.id
    AND jsonb_subscript_update.id = 1;
SELECT * FROM jsonb_subscript_update WHERE id = 1 ORDER BY 1,2;
 id |               data
---------------------------------------------------------------------
  1 | {"a": 1, "b": 2, "c": 3, "d": 4}
(1 row)

TRUNCATE jsonb_subscript_update;
-- Test updates on nested json objects
INSERT INTO jsonb_subscript_update VALUES (1, '{"a": {"c":20, "d" : 200}}'), (2, '{"a": {"d":10, "c" : 100}}');
BEGIN;
UPDATE jsonb_subscript_update
SET DATA['a']['c'] = (updated_vals.d + updated_vals.a::NUMERIC)::TEXT::JSONB
FROM
  (SELECT id,
          DATA['a']['c'] AS a,
                   DATA['a']['c']::NUMERIC + 1 AS b,
                            DATA['a']['c']::NUMERIC + 2 AS c,
                                     DATA['a']['d']::NUMERIC + 3 AS d
   FROM jsonb_subscript_update) updated_vals
WHERE jsonb_subscript_update.id = updated_vals.id;
SELECT * FROM jsonb_subscript_update ORDER BY 1,2;
 id |            data
---------------------------------------------------------------------
  1 | {"a": {"c": 223, "d": 200}}
  2 | {"a": {"c": 113, "d": 10}}
(2 rows)

ROLLBACK;
BEGIN;
-- Router plan
UPDATE jsonb_subscript_update
SET DATA['a']['c'] = (updated_vals.d + updated_vals.a::NUMERIC)::TEXT::JSONB
FROM
  (SELECT id,
          DATA['a']['c'] AS a,
                   DATA['a']['c']::NUMERIC + 1 AS b,
                            DATA['a']['c']::NUMERIC + 2 AS c,
                                     DATA['a']['d']::NUMERIC + 3 AS d
   FROM jsonb_subscript_update) updated_vals
WHERE jsonb_subscript_update.id = updated_vals.id
    AND jsonb_subscript_update.id = 1;
SELECT * FROM jsonb_subscript_update WHERE id = 1 ORDER BY 1,2;
 id |            data
---------------------------------------------------------------------
  1 | {"a": {"c": 223, "d": 200}}
(1 row)

ROLLBACK;
TRUNCATE jsonb_subscript_update;
INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}'), (4, '{"a": 4, "b": 10}');
ALTER TABLE jsonb_subscript_update ADD CONSTRAINT pkey PRIMARY KEY (id, data);
INSERT INTO jsonb_subscript_update VALUES (1, '{"a": 1}'), (2, '{"a": 2}')
ON CONFLICT (id, data)
DO UPDATE SET data['d']=(jsonb_subscript_update.data['a']::INT*100)::TEXT::JSONB,
              data['b']=(jsonb_subscript_update.data['a']::INT*-100)::TEXT::JSONB;
SELECT * FROM jsonb_subscript_update ORDER BY 1,2;
 id |             data
---------------------------------------------------------------------
  1 | {"a": 1, "b": -100, "d": 100}
  2 | {"a": 2, "b": -200, "d": 200}
  4 | {"a": 4, "b": 10}
(3 rows)

CREATE TABLE nested_obj_update(id INT, data JSONB, text_col TEXT);
SELECT create_distributed_table('nested_obj_update', 'id');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO nested_obj_update VALUES
  (1, '{"a": [1,2,3], "b": [4,5,6], "c": [7,8,9], "d": [1,2,1,2]}', '4'),
  (2, '{"a": [10,20,30], "b": [41,51,61], "c": [72,82,92], "d": [11,21,11,21]}', '6');
BEGIN;
-- Pushdown plan
UPDATE nested_obj_update
SET data['a'][0] = (updated_vals.b * 1)::TEXT::JSONB,
    data['b'][2] = (updated_vals.c * 2)::TEXT::JSONB,
    data['c'][0] = (updated_vals.d * 3)::TEXT::JSONB,
    text_col = (nested_obj_update.id*1000)::TEXT,
    data['a'][0] = (text_col::INT * data['a'][0]::INT)::TEXT::JSONB,
    data['d'][6] = (nested_obj_update.id*1)::TEXT::JSONB,
    data['d'][4] = (nested_obj_update.id*2)::TEXT::JSONB
FROM (
  SELECT id,
         data['a'][0] AS a,
         data['b'][0]::NUMERIC + 1 AS b,
         data['c'][0]::NUMERIC + 2 AS c,
         data['c'][1]::NUMERIC + 3 AS d
  FROM nested_obj_update
) updated_vals
WHERE nested_obj_update.id = updated_vals.id;
SELECT * FROM nested_obj_update ORDER BY 1,2,3;
 id |                                              data                                              | text_col
---------------------------------------------------------------------
  1 | {"a": [4, 2, 3], "b": [4, 5, 18], "c": [33, 8, 9], "d": [1, 2, 1, 2, 2, null, 1]}              | 1000
  2 | {"a": [60, 20, 30], "b": [41, 51, 148], "c": [255, 82, 92], "d": [11, 21, 11, 21, 4, null, 2]} | 2000
(2 rows)

ROLLBACK;
BEGIN;
-- Router plan
UPDATE nested_obj_update
SET data['a'][0] = (updated_vals.b * 1)::TEXT::JSONB,
    data['b'][2] = (updated_vals.c * 2)::TEXT::JSONB,
    data['c'][0] = (updated_vals.d * 3)::TEXT::JSONB,
    text_col = (nested_obj_update.id*1000)::TEXT,
    data['a'][0] = (text_col::INT * data['a'][0]::INT)::TEXT::JSONB,
    data['d'][6] = (nested_obj_update.id*1)::TEXT::JSONB,
    data['d'][4] = (nested_obj_update.id*2)::TEXT::JSONB
FROM (
  SELECT id,
         data['a'][0] AS a,
         data['b'][0]::NUMERIC + 1 AS b,
         data['c'][0]::NUMERIC + 2 AS c,
         data['c'][1]::NUMERIC + 3 AS d
  FROM nested_obj_update
) updated_vals
WHERE nested_obj_update.id = updated_vals.id
    AND nested_obj_update.id = 2;
SELECT * FROM nested_obj_update WHERE id = 2 ORDER BY 1,2,3;
 id |                                              data                                              | text_col
---------------------------------------------------------------------
  2 | {"a": [60, 20, 30], "b": [41, 51, 148], "c": [255, 82, 92], "d": [11, 21, 11, 21, 4, null, 2]} | 2000
(1 row)

ROLLBACK;
-- suppress cascade messages
SET client_min_messages to ERROR;
DROP SCHEMA multi_update_select CASCADE;
RESET client_min_messages;
