Instrumentation for master_append_table_to_shard

deinstrument.sql

BEGIN;

DROP FUNCTION public.master_append_table_to_shard(bigint,text,text,int);

ALTER FUNCTION pg_catalog.real_master_append_table_to_shard(bigint,text,text,int)
RENAME TO master_append_table_to_shard;

END;

instrument.sql

BEGIN;

DROP TABLE IF EXISTS pg_dist_append_log;

CREATE TABLE pg_dist_append_log(
  logicalrelid regclass not null,
  shardid bigint not null,
  stagetable text not null,
  bytes_added bigint,
  shardlength bigint,
  time timestamptz default now()
);

ALTER FUNCTION pg_catalog.master_append_table_to_shard(bigint,text,text,int)
RENAME TO real_master_append_table_to_shard;

CREATE OR REPLACE FUNCTION public.master_append_table_to_shard(shard_id bigint, table_name text, node_name text, node_port integer)
RETURNS real
LANGUAGE plpgsql
AS $function$
DECLARE
        pre_shardlength bigint;
        result real;
BEGIN
        SELECT shardlength INTO pre_shardlength
        FROM pg_dist_shard_placement
        WHERE shardid = shard_id;

        SELECT pg_catalog.real_master_append_table_to_shard(shard_id, table_name, node_name, node_port)
        INTO result;

        INSERT INTO pg_dist_append_log
        SELECT logicalrelid, shardid, table_name, shardlength - pre_shardlength, shardlength
        FROM pg_dist_shard JOIN pg_dist_shard_placement USING (shardid)
        WHERE shardid = shard_id
        LIMIT 1;

        RETURN result;
END;
$function$;

END;