-- when using distributed functions in mx for function call delegation you could define
-- functions that have references to distributed tables in their signature or declare
-- blocks.
-- This has caused issues with adding nodes to the cluster as functions will be
-- distributed to the new node before shards or metadata gets synced to the new now,
-- causing referenced tables to not be available during dependency distribution time,
SET citus.next_shard_id TO 20060000;
CREATE SCHEMA function_table_reference;
SET search_path TO function_table_reference;
SET citus.shard_count TO 4;
SET citus.shard_replication_factor TO 1;
SET citus.replication_model TO streaming;
SELECT start_metadata_sync_to_node('localhost', :worker_1_port);
 start_metadata_sync_to_node
---------------------------------------------------------------------

(1 row)

SELECT start_metadata_sync_to_node('localhost', :worker_2_port);
 start_metadata_sync_to_node
---------------------------------------------------------------------

(1 row)

-- SET citus.log_remote_commands TO on;
-- SET client_min_messages TO log;
-- remove worker 2, so we can add it after we have created some functions that caused
-- problems
SELECT master_remove_node('localhost', :worker_2_port);
 master_remove_node
---------------------------------------------------------------------

(1 row)

-- reproduction case as described in #3378
CREATE TABLE zoop_table (x int, y decimal(4, 4));
SELECT create_distributed_table('zoop_table','x');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

-- Create a function that refers to the distributed table
CREATE OR REPLACE FUNCTION zoop(a int)
    RETURNS int
    LANGUAGE plpgsql
    -- setting the search path makes the table name resolve on the worker during initial
    -- distribution
    SET search_path FROM CURRENT
AS $$
DECLARE
    b zoop_table.x%TYPE := 3;
BEGIN
    return a + b;
END;
$$;
SELECT create_distributed_function('zoop(int)', '$1');
 create_distributed_function
---------------------------------------------------------------------

(1 row)

-- now add the worker back, this triggers function distribution which should not fail.
SELECT 1 FROM master_add_node('localhost', :worker_2_port);
 ?column?
---------------------------------------------------------------------
        1
(1 row)

SELECT public.wait_until_metadata_sync(30000);
 wait_until_metadata_sync
---------------------------------------------------------------------

(1 row)

-- verify typmod of zoop_table.b was propagated
-- see numerictypmodin in postgres for how typmod is derived
SELECT run_command_on_workers($$SELECT atttypmod FROM pg_attribute WHERE attnum = 2 AND attrelid = (SELECT typrelid FROM pg_type WHERE typname = 'zoop_table');$$);
   run_command_on_workers
---------------------------------------------------------------------
 (localhost,57637,t,262152)
 (localhost,57638,t,262152)
(2 rows)

-- test that a distributed function can be colocated with a reference table
CREATE TABLE ref(groupid int);
SELECT create_reference_table('ref');
 create_reference_table
---------------------------------------------------------------------

(1 row)

CREATE OR REPLACE FUNCTION my_group_id()
RETURNS void
LANGUAGE plpgsql
SET search_path FROM CURRENT
AS $$
DECLARE
    gid int;
BEGIN
    SELECT groupid INTO gid
    FROM pg_dist_local_group;

    INSERT INTO ref(groupid) VALUES (gid);
END;
$$;
SELECT create_distributed_function('my_group_id()', colocate_with := 'ref');
 create_distributed_function
---------------------------------------------------------------------

(1 row)

SELECT my_group_id();
 my_group_id
---------------------------------------------------------------------

(1 row)

SELECT my_group_id();
 my_group_id
---------------------------------------------------------------------

(1 row)

SELECT DISTINCT(groupid) FROM ref ORDER BY 1;
 groupid
---------------------------------------------------------------------
      14
(1 row)

TRUNCATE TABLE ref;
-- test round robin task assignment policy uses different workers on consecutive function calls.
SET citus.task_assignment_policy TO 'round-robin';
SELECT my_group_id();
 my_group_id
---------------------------------------------------------------------

(1 row)

SELECT my_group_id();
 my_group_id
---------------------------------------------------------------------

(1 row)

SELECT my_group_id();
 my_group_id
---------------------------------------------------------------------

(1 row)

SELECT DISTINCT(groupid) FROM ref ORDER BY 1;
 groupid
---------------------------------------------------------------------
      14
      18
(2 rows)

TRUNCATE TABLE ref;
-- clean up after testing
RESET citus.task_assignment_policy;
DROP SCHEMA function_table_reference CASCADE;
NOTICE:  drop cascades to 4 other objects
DETAIL:  drop cascades to table zoop_table
drop cascades to function zoop(integer)
drop cascades to table ref
drop cascades to function my_group_id()
-- make sure the worker is added at the end irregardless of anything failing to not make
-- subsequent tests fail as well. All artifacts created during this test should have been
-- dropped by the drop cascade above.
SELECT 1 FROM master_add_node('localhost', :worker_2_port);
 ?column?
---------------------------------------------------------------------
        1
(1 row)