CREATE SCHEMA upgrade_distributed_function_before;
SET search_path TO upgrade_distributed_function_before, public;
SET citus.shard_replication_factor TO 1;
CREATE TABLE t1 (a int PRIMARY KEY, b int);
SELECT create_distributed_table('t1','a');
 create_distributed_table
---------------------------------------------------------------------

(1 row)

INSERT INTO t1 VALUES (11), (12);
-- create a very simple distributed function colocated with the table
CREATE FUNCTION count_values(input int) RETURNS int AS
$$
    DECLARE
       cnt int := 0;
    BEGIN
        SELECT count(*) INTO cnt FROM upgrade_distributed_function_before.t1 WHERE a = $1;
        RETURN cnt;
    END;
$$ LANGUAGE plpgsql;
SELECT create_distributed_function('count_values(int)', '$1', colocate_with:='t1');
 create_distributed_function
---------------------------------------------------------------------

(1 row)

SET client_min_messages TO DEBUG1;
SELECT count_values(11);
DEBUG:  pushing down the function call
 count_values
---------------------------------------------------------------------
            1
(1 row)

SELECT count_values(12);
DEBUG:  pushing down the function call
 count_values
---------------------------------------------------------------------
            1
(1 row)

