COPY with upsert on Citus via triggers

upsert-trigger.sql

DROP TABLE IF EXISTS data, data_writes;

CREATE TABLE data (
    key int primary key,
    value int
);
SELECT create_distributed_table('data','key');

CREATE TABLE data_writes (
    LIKE data
);
SELECT create_distributed_table('data_writes','key', colocate_with := 'data');

CREATE OR REPLACE FUNCTION data_trigger() RETURNS trigger AS $fn$
BEGIN
    INSERT INTO data (key, value) VALUES (NEW.key, NEW.value) ON CONFLICT (key) DO UPDATE SET value = NEW.value;
    RETURN NULL;
END;
$fn$ LANGUAGE plpgsql;

SET citus.enable_unsafe_triggers TO on;
CREATE TRIGGER data_trigger BEFORE INSERT ON data_writes
    FOR EACH ROW EXECUTE FUNCTION data_trigger();

\COPY (SELECT s, s*2 FROM generate_series(1,10000) s) TO 'data.csv'
\COPY data_writes FROM 'data.csv'