From 6ddd862e231960710b9a5a03f6544c5e1ff33cb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Thu, 2 Nov 2023 13:15:24 +0100
Subject: [PATCH] Fix #7242, CALL(@0) crash backend (#7288)

When executing a prepared CALL, which is not pure SQL but available with
some drivers like npgsql and jpgdbc, Citus entered a code path where a
plan is not defined, while trying to increase its cost. Thus SIG11 when
plan is a NULL pointer.

Fix by only increasing plan cost when plan is not null.

However, it is a bit suspicious to get here with a NULL plan and maybe a
better change will be to not call
ShardPlacementForFunctionColocatedWithDistTable() with a NULL plan at
all (in call.c:134)

bug hit with for example:
```
CallableStatement proc = con.prepareCall("{CALL p(?)}");
proc.registerOutParameter(1, java.sql.Types.BIGINT);
proc.setInt(1, -100);
proc.execute();
```

where `p(bigint)` is a distributed "function" and the param the
distribution key (also in a distributed table), see #7242 for details

Fixes #7242

(cherry picked from commit 0678a2fd895670e0dd0a3b594915144d468f20e4)
---
 .../distributed/planner/distributed_planner.c |  1 +
 .../planner/function_call_delegation.c        | 12 ++++++--
 src/test/regress/citus_tests/common.py        |  8 +++++
 .../test/test_prepared_statements.py          | 30 +++++++++++++++++++
 4 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 src/test/regress/citus_tests/test/test_prepared_statements.py

diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c
index bc2854161..a0a399fa9 100644
--- a/src/backend/distributed/planner/distributed_planner.c
+++ b/src/backend/distributed/planner/distributed_planner.c
@@ -703,6 +703,7 @@ DissuadePlannerFromUsingPlan(PlannedStmt *plan)
 	 * Arbitrarily high cost, but low enough that it can be added up
 	 * without overflowing by choose_custom_plan().
 	 */
+	Assert(plan != NULL);
 	plan->planTree->total_cost = FLT_MAX / 100000000;
 }
 
diff --git a/src/backend/distributed/planner/function_call_delegation.c b/src/backend/distributed/planner/function_call_delegation.c
index 43f7d8f7f..4a79dc25a 100644
--- a/src/backend/distributed/planner/function_call_delegation.c
+++ b/src/backend/distributed/planner/function_call_delegation.c
@@ -531,8 +531,16 @@ ShardPlacementForFunctionColocatedWithDistTable(DistObjectCacheEntry *procedure,
 
 		if (partitionParam->paramkind == PARAM_EXTERN)
 		{
-			/* Don't log a message, we should end up here again without a parameter */
-			DissuadePlannerFromUsingPlan(plan);
+			/*
+			 * Don't log a message, we should end up here again without a
+			 * parameter.
+			 * Note that "plan" can be null, for example when a CALL statement
+			 * is prepared.
+			 */
+			if (plan)
+			{
+				DissuadePlannerFromUsingPlan(plan);
+			}
 			return NULL;
 		}
 	}
diff --git a/src/test/regress/citus_tests/common.py b/src/test/regress/citus_tests/common.py
index 4d04d268c..1a9dd86bc 100644
--- a/src/test/regress/citus_tests/common.py
+++ b/src/test/regress/citus_tests/common.py
@@ -581,6 +581,14 @@ class QueryRunner(ABC):
         with self.cur(**kwargs) as cur:
             cur.execute(query, params=params)
 
+    def sql_prepared(self, query, params=None, **kwargs):
+        """Run an SQL query, with prepare=True
+
+        This opens a new connection and closes it once the query is done
+        """
+        with self.cur(**kwargs) as cur:
+            cur.execute(query, params=params, prepare=True)
+
     def sql_row(self, query, params=None, allow_empty_result=False, **kwargs):
         """Run an SQL query that returns a single row and returns this row
 
diff --git a/src/test/regress/citus_tests/test/test_prepared_statements.py b/src/test/regress/citus_tests/test/test_prepared_statements.py
new file mode 100644
index 000000000..761ecc30c
--- /dev/null
+++ b/src/test/regress/citus_tests/test/test_prepared_statements.py
@@ -0,0 +1,30 @@
+def test_call_param(cluster):
+    # create a distributed table and an associated distributed procedure
+    # to ensure parameterized CALL succeed, even when the param is the
+    # distribution key.
+    coord = cluster.coordinator
+    coord.sql("CREATE TABLE test(i int)")
+    coord.sql(
+        """
+        CREATE PROCEDURE p(_i INT) LANGUAGE plpgsql AS $$
+        BEGIN
+        INSERT INTO test(i) VALUES (_i);
+        END; $$
+        """
+    )
+    sql = "CALL p(%s)"
+
+    # prepare/exec before distributing
+    coord.sql_prepared(sql, (1,))
+
+    coord.sql("SELECT create_distributed_table('test', 'i')")
+    coord.sql(
+        "SELECT create_distributed_function('p(int)', distribution_arg_name := '_i', colocate_with := 'test')"
+    )
+
+    # prepare/exec after distribution
+    coord.sql_prepared(sql, (2,))
+
+    sum_i = coord.sql_value("select sum(i) from test;")
+
+    assert sum_i == 3