CREATE SERVER lw_update_admin FOREIGN DATA WRAPPER clickhouse_fdw OPTIONS(driver 'http');
CREATE USER MAPPING FOR CURRENT_USER SERVER lw_update_admin;
\set ECHO errors
\if :old_ch
SET client_min_messages = warning;
DROP SERVER lw_update_admin CASCADE;
\echo 'SKIP: lightweight UPDATE requires ClickHouse 25.7'
\quit
\endif
CREATE SERVER lw_update_svr FOREIGN DATA WRAPPER clickhouse_fdw
    OPTIONS(dbname 'lw_update_test', driver 'http');
CREATE USER MAPPING FOR CURRENT_USER SERVER lw_update_svr;
CREATE SERVER lw_update_bin FOREIGN DATA WRAPPER clickhouse_fdw
    OPTIONS(dbname 'lw_update_test', driver 'binary');
CREATE USER MAPPING FOR CURRENT_USER SERVER lw_update_bin;
CALL clickhouse_perform('lw_update_admin', 'DROP DATABASE IF EXISTS lw_update_test');
CALL clickhouse_perform('lw_update_admin', 'CREATE DATABASE lw_update_test');
CALL clickhouse_perform('lw_update_admin', $$
    CREATE TABLE lw_update_test.items (
        id Int32, name String, qty Int32, note Nullable(String)
    ) ENGINE = MergeTree ORDER BY id
    SETTINGS enable_block_number_column = 1, enable_block_offset_column = 1
$$);
CALL clickhouse_perform('lw_update_admin', $$
    INSERT INTO lw_update_test.items
    SELECT number, 'item' || toString(number), number * 10, NULL FROM numbers(1, 6)
$$);
CREATE SCHEMA lw_update;
SET search_path = lw_update, public;
CREATE FOREIGN TABLE items (id int, name text, qty int, note text)
    SERVER lw_update_svr;
CREATE FOREIGN TABLE items_bin (id int, name text, qty int, note text)
    SERVER lw_update_bin OPTIONS (table_name 'items');
CREATE FOREIGN TABLE items_renamed (key int OPTIONS (column_name 'id'), label text OPTIONS (column_name 'name'))
    SERVER lw_update_svr OPTIONS (table_name 'items');
-- Run UPDATE as one ClickHouse statement
EXPLAIN (VERBOSE, COSTS OFF) UPDATE items SET qty = qty + 1 WHERE id = 1;
                                      QUERY PLAN                                      
--------------------------------------------------------------------------------------
 Update on lw_update.items
   ->  Foreign Update on lw_update.items
         Remote SQL: UPDATE lw_update_test.items SET qty = (qty + 1) WHERE ((id = 1))
(3 rows)

UPDATE items SET qty = qty + 1 WHERE id = 1;
SELECT * FROM items ORDER BY id;
 id | name  | qty | note 
----+-------+-----+------
  1 | item1 |  11 | 
  2 | item2 |  20 | 
  3 | item3 |  30 | 
  4 | item4 |  40 | 
  5 | item5 |  50 | 
  6 | item6 |  60 | 
(6 rows)

EXPLAIN (VERBOSE, COSTS OFF)
UPDATE items SET name = upper(name), note = 'it''s' WHERE qty > 30 AND note IS NULL;
                                                             QUERY PLAN                                                             
------------------------------------------------------------------------------------------------------------------------------------
 Update on lw_update.items
   ->  Foreign Update on lw_update.items
         Remote SQL: UPDATE lw_update_test.items SET name = upperUTF8(name), note = 'it''s' WHERE ((note IS NULL)) AND ((qty > 30))
(3 rows)

UPDATE items SET name = upper(name), note = 'it''s' WHERE qty > 30 AND note IS NULL;
SELECT * FROM items ORDER BY id;
 id | name  | qty | note 
----+-------+-----+------
  1 | item1 |  11 | 
  2 | item2 |  20 | 
  3 | item3 |  30 | 
  4 | ITEM4 |  40 | it's
  5 | ITEM5 |  50 | it's
  6 | ITEM6 |  60 | it's
(6 rows)

UPDATE items SET note = NULL WHERE id = 4;
SELECT * FROM items WHERE id = 4;
 id | name  | qty | note 
----+-------+-----+------
  4 | ITEM4 |  40 | 
(1 row)

-- Binary driver
UPDATE items_bin SET qty = -qty WHERE name = 'item2';
SELECT * FROM items_bin ORDER BY id;
 id | name  | qty | note 
----+-------+-----+------
  1 | item1 |  11 | 
  2 | item2 | -20 | 
  3 | item3 |  30 | 
  4 | ITEM4 |  40 | 
  5 | ITEM5 |  50 | it's
  6 | ITEM6 |  60 | it's
(6 rows)

-- column_name option
EXPLAIN (VERBOSE, COSTS OFF) UPDATE items_renamed SET label = 'third' WHERE key = 3;
                                     QUERY PLAN                                      
-------------------------------------------------------------------------------------
 Update on lw_update.items_renamed
   ->  Foreign Update on lw_update.items_renamed
         Remote SQL: UPDATE lw_update_test.items SET name = 'third' WHERE ((id = 3))
(3 rows)

UPDATE items_renamed SET label = 'third' WHERE key = 3;
SELECT * FROM items_renamed ORDER BY key;
 key | label 
-----+-------
   1 | item1
   2 | item2
   3 | third
   4 | ITEM4
   5 | ITEM5
   6 | ITEM6
(6 rows)

-- Send prepared statement arguments as ClickHouse query parameters
SET plan_cache_mode = force_generic_plan;
PREPARE upd(int, text) AS UPDATE items SET note = $2 WHERE id = $1;
EXPLAIN (VERBOSE, COSTS OFF) EXECUTE upd(5, 'five');
                                            QUERY PLAN                                            
--------------------------------------------------------------------------------------------------
 Update on lw_update.items
   ->  Foreign Update on lw_update.items
         Remote SQL: UPDATE lw_update_test.items SET note = {p1:String} WHERE ((id = {p2:Int32}))
(3 rows)

EXECUTE upd(5, 'five');
PREPARE upd_bin(int, text) AS UPDATE items_bin SET note = $2 WHERE id = $1;
EXECUTE upd_bin(6, 'six');
RESET plan_cache_mode;
SELECT * FROM items ORDER BY id;
 id | name  | qty | note 
----+-------+-----+------
  1 | item1 |  11 | 
  2 | item2 | -20 | 
  3 | third |  30 | 
  4 | ITEM4 |  40 | 
  5 | ITEM5 |  50 | five
  6 | ITEM6 |  60 | six
(6 rows)

-- Run subquery in Postgres and send result as query parameter
CREATE TABLE local_ids (id int, qty int);
INSERT INTO local_ids VALUES (1, 100);
UPDATE items SET qty = (SELECT max(qty) FROM local_ids) WHERE id = 1;
SELECT * FROM items WHERE id = 1;
 id | name  | qty | note 
----+-------+-----+------
  1 | item1 | 100 | 
(1 row)

-- Update local and ClickHouse partitions in same statement
SET datestyle = 'ISO';
CALL clickhouse_perform('lw_update_admin', $$
    CREATE TABLE lw_update_test.events (id Int32, ts Date, qty Int32)
    ENGINE = MergeTree ORDER BY id
    SETTINGS enable_block_number_column = 1, enable_block_offset_column = 1
$$);
CALL clickhouse_perform('lw_update_admin', $$
    INSERT INTO lw_update_test.events VALUES (1, '2023-01-15', 10), (2, '2023-06-01', 20)
$$);
CREATE TABLE events (id int, ts date, qty int) PARTITION BY RANGE (ts);
CREATE FOREIGN TABLE events_cold PARTITION OF events
    FOR VALUES FROM ('2023-01-01') TO ('2024-01-01')
    SERVER lw_update_svr OPTIONS (table_name 'events');
CREATE TABLE events_hot PARTITION OF events
    FOR VALUES FROM ('2024-01-01') TO ('2025-01-01');
INSERT INTO events_hot VALUES (3, '2024-01-10', 30), (4, '2024-02-15', 40);
EXPLAIN (VERBOSE, COSTS OFF) UPDATE events SET qty = qty + 1 WHERE id IN (2, 4);
                                            QUERY PLAN                                            
--------------------------------------------------------------------------------------------------
 Update on lw_update.events
   Foreign Update on lw_update.events_cold events_1
   Update on lw_update.events_hot events_2
   ->  Append
         ->  Foreign Update on lw_update.events_cold events_1
               Remote SQL: UPDATE lw_update_test.events SET qty = (qty + 1) WHERE ((id IN (2,4)))
         ->  Seq Scan on lw_update.events_hot events_2
               Output: (events_2.qty + 1), events_2.tableoid, NULL::record, events_2.ctid
               Filter: (events_2.id = ANY ('{2,4}'::integer[]))
(9 rows)

UPDATE events SET qty = qty + 1 WHERE id IN (2, 4);
SELECT id, ts, qty FROM events ORDER BY id;
 id |     ts     | qty 
----+------------+-----
  1 | 2023-01-15 |  10
  2 | 2023-06-01 |  21
  3 | 2024-01-10 |  30
  4 | 2024-02-15 |  41
(4 rows)

-- Reject partition key updates through parent and leaf
UPDATE events SET ts = '2024-06-01' WHERE ts = '2023-06-01' AND id = 2;
ERROR:  pg_clickhouse: cannot update partition key of foreign table "events_cold"
DETAIL:  Moving rows out of foreign partitions is not supported.
UPDATE events_cold SET ts = '2024-06-01' WHERE id = 2;
ERROR:  pg_clickhouse: cannot update partition key of foreign table "events_cold"
DETAIL:  Moving rows out of foreign partitions is not supported.
SELECT tableoid::regclass, id, ts, qty FROM events WHERE ts = '2023-06-01';
  tableoid   | id |     ts     | qty 
-------------+----+------------+-----
 events_cold |  2 | 2023-06-01 |  21
(1 row)

SELECT id FROM events WHERE ts = '2024-06-01';
 id 
----
(0 rows)

-- Include ancestor expressions and default partition constraints
CREATE TABLE events_expr (id int, ts date, qty int)
    PARTITION BY RANGE ((extract(year FROM ts)));
CREATE TABLE events_expr_2023 PARTITION OF events_expr
    FOR VALUES FROM (2023) TO (2024) PARTITION BY LIST (id);
CREATE TABLE events_expr_other PARTITION OF events_expr_2023 FOR VALUES IN (99);
CREATE FOREIGN TABLE events_expr_default PARTITION OF events_expr_2023 DEFAULT
    SERVER lw_update_bin OPTIONS (table_name 'events');
UPDATE events_expr_default SET ts = '2024-06-01' WHERE id = 2;
ERROR:  pg_clickhouse: cannot update partition key of foreign table "events_expr_default"
DETAIL:  Moving rows out of foreign partitions is not supported.
UPDATE events_expr_default SET id = 99 WHERE id = 2;
ERROR:  pg_clickhouse: cannot update partition key of foreign table "events_expr_default"
DETAIL:  Moving rows out of foreign partitions is not supported.
UPDATE events_expr SET qty = qty + 1 WHERE id = 1;
SELECT id, ts, qty FROM events_expr WHERE id = 1;
 id |     ts     | qty 
----+------------+-----
  1 | 2023-01-15 |  11
(1 row)

-- Reject statements that cannot run entirely in ClickHouse
UPDATE items SET qty = 0 WHERE id = 1 RETURNING id;
ERROR:  pg_clickhouse: cannot run UPDATE on foreign table "items" entirely in ClickHouse
HINT:  Avoid RETURNING, joins, row triggers, and conditions or SET expressions that require Postgres.
CREATE FUNCTION local_double(int) RETURNS int LANGUAGE plpgsql AS 'BEGIN RETURN $1 * 2; END';
UPDATE items SET qty = local_double(qty) WHERE id = 1;
ERROR:  pg_clickhouse: cannot run UPDATE on foreign table "items" entirely in ClickHouse
HINT:  Avoid RETURNING, joins, row triggers, and conditions or SET expressions that require Postgres.
UPDATE items SET qty = 0 WHERE local_double(id) = 2;
ERROR:  pg_clickhouse: cannot run UPDATE on foreign table "items" entirely in ClickHouse
HINT:  Avoid RETURNING, joins, row triggers, and conditions or SET expressions that require Postgres.
UPDATE items SET qty = local_ids.qty FROM local_ids WHERE items.id = local_ids.id;
ERROR:  pg_clickhouse: cannot run UPDATE on foreign table "items" entirely in ClickHouse
HINT:  Avoid RETURNING, joins, row triggers, and conditions or SET expressions that require Postgres.
CREATE FUNCTION no_op() RETURNS trigger LANGUAGE plpgsql AS 'BEGIN RETURN NEW; END';
CREATE TRIGGER items_upd BEFORE UPDATE ON items FOR EACH ROW EXECUTE FUNCTION no_op();
UPDATE items SET qty = 0 WHERE id = 1;
ERROR:  pg_clickhouse: cannot run UPDATE on foreign table "items" entirely in ClickHouse
HINT:  Avoid RETURNING, joins, row triggers, and conditions or SET expressions that require Postgres.
DROP TRIGGER items_upd ON items;
SELECT * FROM items ORDER BY id;
 id | name  | qty | note 
----+-------+-----+------
  1 | item1 | 100 | 
  2 | item2 | -20 | 
  3 | third |  30 | 
  4 | ITEM4 |  40 | 
  5 | ITEM5 |  50 | five
  6 | ITEM6 |  60 | six
(6 rows)

-- ClickHouse requires WHERE
EXPLAIN (VERBOSE, COSTS OFF) UPDATE items SET qty = 0;
                               QUERY PLAN                               
------------------------------------------------------------------------
 Update on lw_update.items
   ->  Foreign Update on lw_update.items
         Remote SQL: UPDATE lw_update_test.items SET qty = 0 WHERE true
(3 rows)

UPDATE items SET qty = 0;
SELECT id, qty FROM items ORDER BY id;
 id | qty 
----+-----
  1 |   0
  2 |   0
  3 |   0
  4 |   0
  5 |   0
  6 |   0
(6 rows)

-- Support ReplacingMergeTree and CollapsingMergeTree.
CALL clickhouse_perform('lw_update_admin', $$
    CREATE TABLE lw_update_test.replacing (id Int32, name String)
    ENGINE = ReplacingMergeTree ORDER BY id
    SETTINGS enable_block_number_column = 1, enable_block_offset_column = 1
$$);
CALL clickhouse_perform('lw_update_admin', $$
    INSERT INTO lw_update_test.replacing VALUES (1, 'Sushi'), (2, 'Bagel')
$$);
CALL clickhouse_perform('lw_update_admin', $$
    CREATE TABLE lw_update_test.collapsing (id Int32, name String, sign Int8)
    ENGINE = CollapsingMergeTree(sign) ORDER BY id
    SETTINGS enable_block_number_column = 1, enable_block_offset_column = 1
$$);
CALL clickhouse_perform('lw_update_admin', $$
    INSERT INTO lw_update_test.collapsing VALUES (1, 'Little Bit', 1), (2, 'Biscuit', 1)
$$);
CREATE FOREIGN TABLE replacing (id int, name text)
    SERVER lw_update_bin OPTIONS (table_name 'replacing');
CREATE FOREIGN TABLE collapsing (id int, name text)
    SERVER lw_update_bin OPTIONS (table_name 'collapsing');
UPDATE replacing SET name = 'Sushi May' WHERE id = 1;
SELECT * FROM replacing ORDER BY id;
 id |   name    
----+-----------
  1 | Sushi May
  2 | Bagel
(2 rows)

UPDATE collapsing SET name = 'Sweet Pea' WHERE id = 1;
SELECT * FROM collapsing ORDER BY id;
 id |   name    
----+-----------
  1 | Sweet Pea
  2 | Biscuit
(2 rows)

CALL clickhouse_perform('lw_update_admin', 'DROP DATABASE lw_update_test');
SET client_min_messages = warning;
DROP SCHEMA lw_update CASCADE;
DROP SERVER lw_update_svr CASCADE;
DROP SERVER lw_update_bin CASCADE;
DROP SERVER lw_update_admin CASCADE;