From cd0b2ad5b50b35ef78e72ca0efa01869d6de654f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philip=20Dub=C3=A9?= <gunkmute@gmail.com>
Date: Tue, 23 Jun 2020 17:07:30 +0000
Subject: [PATCH] citus_evaluate_expression: call expand_function_arguments
 beforehand to avoid segfaulting on implicit parameters

---
 src/backend/distributed/utils/citus_clauses.c | 46 +++++++++++++++++++
 .../expected/multi_function_evaluation.out    | 29 +++++++++---
 .../regress/sql/multi_function_evaluation.sql | 14 ++++++
 3 files changed, 83 insertions(+), 6 deletions(-)

diff --git a/src/backend/distributed/utils/citus_clauses.c b/src/backend/distributed/utils/citus_clauses.c
index b93839212..9e7e7c517 100644
--- a/src/backend/distributed/utils/citus_clauses.c
+++ b/src/backend/distributed/utils/citus_clauses.c
@@ -21,9 +21,13 @@
 #include "nodes/nodes.h"
 #include "nodes/primnodes.h"
 #include "optimizer/clauses.h"
+#if PG_VERSION_NUM >= PG_VERSION_12
+#include "optimizer/optimizer.h"
+#endif
 #include "optimizer/planmain.h"
 #include "utils/datum.h"
 #include "utils/lsyscache.h"
+#include "utils/syscache.h"
 
 
 /* private function declarations */
@@ -36,6 +40,8 @@ static bool CitusIsMutableFunctionIdChecker(Oid func_id, void *context);
 static bool ShouldEvaluateExpression(Expr *expression);
 static bool ShouldEvaluateFunctionWithMasterContext(MasterEvaluationContext *
 													evaluationContext);
+static void FixFunctionArguments(Node *expr);
+static bool FixFunctionArgumentsWalker(Node *expr, void *context);
 
 /*
  * RequiresMasterEvaluation returns the executor needs to reparse and
@@ -304,6 +310,9 @@ citus_evaluate_expr(Expr *expr, Oid result_type, int32 result_typmod,
 	/* We can use the estate's working context to avoid memory leaks. */
 	MemoryContext oldcontext = MemoryContextSwitchTo(estate->es_query_cxt);
 
+	/* handles default values */
+	FixFunctionArguments((Node *) expr);
+
 	/* Make sure any opfuncids are filled in. */
 	fix_opfuncids((Node *) expr);
 
@@ -452,3 +461,40 @@ CitusIsMutableFunction(Node *node)
 
 	return false;
 }
+
+
+/* FixFunctionArguments applies expand_function_arguments to all function calls. */
+static void
+FixFunctionArguments(Node *expr)
+{
+	FixFunctionArgumentsWalker(expr, NULL);
+}
+
+
+/* FixFunctionArgumentsWalker is the helper function for fix_funcargs. */
+static bool
+FixFunctionArgumentsWalker(Node *expr, void *context)
+{
+	if (expr == NULL)
+	{
+		return false;
+	}
+
+	if (IsA(expr, FuncExpr))
+	{
+		FuncExpr *funcExpr = castNode(FuncExpr, expr);
+		HeapTuple func_tuple =
+			SearchSysCache1(PROCOID, ObjectIdGetDatum(funcExpr->funcid));
+		if (!HeapTupleIsValid(func_tuple))
+		{
+			elog(ERROR, "cache lookup failed for function %u", funcExpr->funcid);
+		}
+
+		funcExpr->args = expand_function_arguments(funcExpr->args,
+												   funcExpr->funcresulttype, func_tuple);
+
+		ReleaseSysCache(func_tuple);
+	}
+
+	return expression_tree_walker(expr, FixFunctionArgumentsWalker, NULL);
+}
diff --git a/src/test/regress/expected/multi_function_evaluation.out b/src/test/regress/expected/multi_function_evaluation.out
index deeeb2049..7d546c15f 100644
--- a/src/test/regress/expected/multi_function_evaluation.out
+++ b/src/test/regress/expected/multi_function_evaluation.out
@@ -220,10 +220,27 @@ NOTICE:  stable_fn called
 DETAIL:  from localhost:xxxxx
 NOTICE:  stable_fn called
 DETAIL:  from localhost:xxxxx
+-- https://github.com/citusdata/citus/issues/3939
+CREATE TABLE test_table(id int);
+SELECT create_distributed_table('test_table', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE OR REPLACE FUNCTION f(val text DEFAULT 'default')
+RETURNS int AS $$ BEGIN RETURN length(val); END; $$ LANGUAGE 'plpgsql' IMMUTABLE;
+INSERT INTO test_table VALUES (f('test'));
+INSERT INTO test_table VALUES (f());
+INSERT INTO test_table VALUES (f(f()::text));
+SELECT * FROM test_table ORDER BY 1;
+ id
+---------------------------------------------------------------------
+  1
+  4
+  7
+(3 rows)
+
+\set VERBOSITY terse
 DROP SCHEMA multi_function_evaluation CASCADE;
-NOTICE:  drop cascades to 5 other objects
-DETAIL:  drop cascades to table example
-drop cascades to sequence example_value_seq
-drop cascades to function stable_fn()
-drop cascades to table table_1
-drop cascades to function stable_squared(integer)
+NOTICE:  drop cascades to 7 other objects
diff --git a/src/test/regress/sql/multi_function_evaluation.sql b/src/test/regress/sql/multi_function_evaluation.sql
index 7e5480bab..cad63fe9d 100644
--- a/src/test/regress/sql/multi_function_evaluation.sql
+++ b/src/test/regress/sql/multi_function_evaluation.sql
@@ -165,4 +165,18 @@ FROM (SELECT key, stable_squared((count(*) OVER ())::int) y FROM example GROUP B
 UPDATE example SET value = timestamp '10-10-2000 00:00'
 FROM (SELECT key, stable_squared(grouping(key)) y FROM example GROUP BY key) a WHERE example.key = a.key;
 
+-- https://github.com/citusdata/citus/issues/3939
+CREATE TABLE test_table(id int);
+SELECT create_distributed_table('test_table', 'id');
+
+CREATE OR REPLACE FUNCTION f(val text DEFAULT 'default')
+RETURNS int AS $$ BEGIN RETURN length(val); END; $$ LANGUAGE 'plpgsql' IMMUTABLE;
+
+INSERT INTO test_table VALUES (f('test'));
+INSERT INTO test_table VALUES (f());
+INSERT INTO test_table VALUES (f(f()::text));
+
+SELECT * FROM test_table ORDER BY 1;
+
+\set VERBOSITY terse
 DROP SCHEMA multi_function_evaluation CASCADE;