Fix the issue #5673

If the expression is simple, such as, SELECT function() or PEFORM function()
in PL/PgSQL code, PL engine does a simple expression evaluation which can't
interpret the Citus CustomScan Node. Code checks for simple expressions when
executing an UDF but missed the DO-Block scenario, this commit fixes it.

src/backend/distributed/planner/function_call_delegation.c

@@ -363,14 +363,16 @@ TryToDelegateFunctionCall(DistributedPlanningContext *planContext)
 			}
 
 			/*
-			 * If the expression is simple, such as, SELECT fn() in
-			 * PL/PgSQL code, PL engine is doing simple expression
-			 * evaluation, which can't interpret the CustomScan Node.
-			 * Function from FROM clause is not simple, so it's ok.
+			 * If the expression is simple, such as, SELECT function() or PEFORM function()
+			 * in PL/PgSQL code, PL engine does a simple expression evaluation which can't
+			 * interpret the Citus CustomScan Node.
+			 * Note: Function from FROM clause is not simple, so it's ok to pushdown.
 			 */
-			if (MaybeExecutingUDF() && IsQuerySimple(planContext->query) && !fromFuncExpr)
+			if ((MaybeExecutingUDF() || DoBlockLevel > 0) &&
+				IsQuerySimple(planContext->query) &&
+				!fromFuncExpr)
 			{
-				ereport(DEBUG1, (errmsg("Skipping delegation of function "
+				ereport(DEBUG1, (errmsg("Skipping pushdown of function "
 										"from a PL/PgSQL simple expression")));
 				return NULL;
 			}

src/test/regress/expected/forcedelegation_functions.out

@@ -627,7 +627,7 @@ DETAIL:  A distributed function is created. To make sure subsequent commands see
 (1 row)
 
 SELECT outer_emp();
-DEBUG:  Skipping delegation of function from a PL/PgSQL simple expression
+DEBUG:  Skipping pushdown of function from a PL/PgSQL simple expression
 CONTEXT:  SQL statement "SELECT inner_emp('hello')"
 PL/pgSQL function outer_emp() line XX at PERFORM
  outer_emp
@@ -1388,10 +1388,47 @@ SELECT COUNT(*) FROM table_test_prepare;
     28
 (1 row)
 
+CREATE TABLE test_perform(i int);
+SELECT create_distributed_table('test_perform', 'i', colocate_with := 'none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE OR REPLACE FUNCTION test(x int)
+RETURNS int
+AS $$
+DECLARE
+BEGIN
+    RAISE NOTICE 'INPUT %', x;
+    RETURN x;
+END;
+$$ LANGUAGE plpgsql;
+SELECT create_distributed_function('test(int)', 'x',
+                colocate_with := 'test_perform', force_delegation := true);
+DEBUG:  switching to sequential query execution mode
+DETAIL:  A distributed function is created. To make sure subsequent commands see the type correctly we need to make sure to use only one connection for all future commands
+ create_distributed_function
+---------------------------------------------------------------------
+
+(1 row)
+
+DO $$
+BEGIN
+    PERFORM test(3);
+END;
+$$ LANGUAGE plpgsql;
+DEBUG:  Skipping pushdown of function from a PL/PgSQL simple expression
+CONTEXT:  SQL statement "SELECT test(3)"
+PL/pgSQL function inline_code_block line XX at PERFORM
+NOTICE:  INPUT 3
+CONTEXT:  PL/pgSQL function test(integer) line XX at RAISE
+SQL statement "SELECT test(3)"
+PL/pgSQL function inline_code_block line XX at PERFORM
 RESET client_min_messages;
 SET citus.log_remote_commands TO off;
 DROP SCHEMA forcepushdown_schema CASCADE;
-NOTICE:  drop cascades to 36 other objects
+NOTICE:  drop cascades to 38 other objects
 DETAIL:  drop cascades to table test_forcepushdown
 drop cascades to table test_forcepushdown_noncolocate
 drop cascades to function insert_data(integer)
@@ -1428,3 +1465,5 @@ drop cascades to function insert_data_cte_nondist(integer)
 drop cascades to table table_test_prepare
 drop cascades to function test_prepare(integer,integer)
 drop cascades to function outer_test_prepare(integer,integer)
+drop cascades to table test_perform
+drop cascades to function test(integer)

src/test/regress/sql/forcedelegation_functions.sql

@@ -661,6 +661,27 @@ SELECT outer_test_prepare(1,2);
 
 SELECT COUNT(*) FROM table_test_prepare;
 
+CREATE TABLE test_perform(i int);
+SELECT create_distributed_table('test_perform', 'i', colocate_with := 'none');
+
+CREATE OR REPLACE FUNCTION test(x int)
+RETURNS int
+AS $$
+DECLARE
+BEGIN
+    RAISE NOTICE 'INPUT %', x;
+    RETURN x;
+END;
+$$ LANGUAGE plpgsql;
+
+SELECT create_distributed_function('test(int)', 'x',
+                colocate_with := 'test_perform', force_delegation := true);
+DO $$
+BEGIN
+    PERFORM test(3);
+END;
+$$ LANGUAGE plpgsql;
+
 RESET client_min_messages;
 SET citus.log_remote_commands TO off;
 DROP SCHEMA forcepushdown_schema CASCADE;