From a4b2197450852335559ba48bbd44c865ad0e0b92 Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Tue, 31 Mar 2020 19:24:07 +0200
Subject: [PATCH] Correctly handle non-constant LIMIT/OFFSET clauses

---
 .../planner/multi_logical_optimizer.c         | 27 ++++++++++----
 .../planner/multi_logical_planner.c           | 11 ++++++
 .../regress/expected/multi_limit_clause.out   | 37 +++++++++++++++++++
 src/test/regress/sql/multi_limit_clause.sql   | 17 +++++++++
 4 files changed, 84 insertions(+), 8 deletions(-)

diff --git a/src/backend/distributed/planner/multi_logical_optimizer.c b/src/backend/distributed/planner/multi_logical_optimizer.c
index 8264a7b8b..9dd0b12a6 100644
--- a/src/backend/distributed/planner/multi_logical_optimizer.c
+++ b/src/backend/distributed/planner/multi_logical_optimizer.c
@@ -4320,19 +4320,30 @@ WorkerLimitCount(Node *limitCount, Node *limitOffset, OrderByLimitReference
 	Node *workerLimitNode = NULL;
 	LimitPushdownable canPushDownLimit = LIMIT_CANNOT_PUSHDOWN;
 
-	/* no limit node to push down */
 	if (limitCount == NULL)
 	{
+		/* no limit node to push down */
+		return NULL;
+	}
+
+	if (!IsA(limitCount, Const))
+	{
+		/*
+		 * We only push down constant LIMIT clauses to make sure we get back
+		 * the minimum number of rows.
+		 */
+		return NULL;
+	}
+
+	if (limitOffset != NULL && !IsA(limitOffset, Const))
+	{
+		/*
+		 * If OFFSET is not a constant then we cannot calculate the LIMIT to
+		 * push down.
+		 */
 		return NULL;
 	}
 
-	/*
-	 * During subquery pushdown planning original query is used. In that case,
-	 * certain expressions such as parameters are not evaluated and converted
-	 * into Consts on the op node.
-	 */
-	Assert(IsA(limitCount, Const));
-	Assert(limitOffset == NULL || IsA(limitOffset, Const));
 
 	/*
 	 * If window functions are computed on coordinator, we cannot push down LIMIT.
diff --git a/src/backend/distributed/planner/multi_logical_planner.c b/src/backend/distributed/planner/multi_logical_planner.c
index cb0897f80..f6397a4aa 100644
--- a/src/backend/distributed/planner/multi_logical_planner.c
+++ b/src/backend/distributed/planner/multi_logical_planner.c
@@ -1024,6 +1024,17 @@ DeferErrorIfQueryNotSupported(Query *queryTree)
 		errorHint = filterHint;
 	}
 
+	if (FindNodeCheck((Node *) queryTree->limitCount, IsNodeSubquery))
+	{
+		preconditionsSatisfied = false;
+		errorMessage = "subquery in LIMIT is not supported in multi-shard queries";
+	}
+
+	if (FindNodeCheck((Node *) queryTree->limitOffset, IsNodeSubquery))
+	{
+		preconditionsSatisfied = false;
+		errorMessage = "subquery in OFFSET is not supported in multi-shard queries";
+	}
 
 	/* finally check and error out if not satisfied */
 	if (!preconditionsSatisfied)
diff --git a/src/test/regress/expected/multi_limit_clause.out b/src/test/regress/expected/multi_limit_clause.out
index 6d891ddde..b3e0367a3 100644
--- a/src/test/regress/expected/multi_limit_clause.out
+++ b/src/test/regress/expected/multi_limit_clause.out
@@ -520,4 +520,41 @@ SELECT
 (1 row)
 
 SET client_min_messages TO NOTICE;
+-- non constants should not push down
+CREATE OR REPLACE FUNCTION my_limit()
+RETURNS INT AS $$
+BEGIN
+  RETURN 5;
+END; $$ language plpgsql VOLATILE;
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT my_limit();
+ l_orderkey
+---------------------------------------------------------------------
+          1
+          1
+          1
+          1
+          1
+(5 rows)
+
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET my_limit();
+ l_orderkey
+---------------------------------------------------------------------
+          1
+          2
+          3
+          3
+          3
+          3
+          3
+          3
+          4
+          5
+(10 rows)
+
+DROP FUNCTION my_limit();
+-- subqueries should error out
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT (SELECT 10);
+ERROR:  subquery in LIMIT is not supported in multi-shard queries
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET (SELECT 10);
+ERROR:  subquery in OFFSET is not supported in multi-shard queries
 DROP TABLE lineitem_hash;
diff --git a/src/test/regress/sql/multi_limit_clause.sql b/src/test/regress/sql/multi_limit_clause.sql
index 90601ed92..8d14bbbc8 100644
--- a/src/test/regress/sql/multi_limit_clause.sql
+++ b/src/test/regress/sql/multi_limit_clause.sql
@@ -223,4 +223,21 @@ SELECT
 	LIMIT 5;
 
 SET client_min_messages TO NOTICE;
+
+-- non constants should not push down
+CREATE OR REPLACE FUNCTION my_limit()
+RETURNS INT AS $$
+BEGIN
+  RETURN 5;
+END; $$ language plpgsql VOLATILE;
+
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT my_limit();
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET my_limit();
+
+DROP FUNCTION my_limit();
+
+-- subqueries should error out
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT (SELECT 10);
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET (SELECT 10);
+
 DROP TABLE lineitem_hash;