From f5430989edcbd02c27f21fbe1543c24f1375eed3 Mon Sep 17 00:00:00 2001
From: Andres Freund <andres@anarazel.de>
Date: Thu, 8 Sep 2016 12:27:33 -0700
Subject: [PATCH] Support NoMovement direction in router executor

This is mainly interesting because it allows to use RETURN QUERY/RETURN
QUERY EXECUTE and FOR ... IN .. LOOPs in plpgsql.
---
 .../executor/multi_router_executor.c          | 22 ++++++++++-----
 .../regress/expected/multi_router_planner.out | 27 ++++++++++++++-----
 src/test/regress/sql/multi_router_planner.sql |  5 ++--
 3 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/src/backend/distributed/executor/multi_router_executor.c b/src/backend/distributed/executor/multi_router_executor.c
index ab8c239b6..19763dca7 100644
--- a/src/backend/distributed/executor/multi_router_executor.c
+++ b/src/backend/distributed/executor/multi_router_executor.c
@@ -353,13 +353,6 @@ RouterExecutorRun(QueryDesc *queryDesc, ScanDirection direction, long count)
 	Assert(!(estate->es_top_eflags & EXEC_FLAG_EXPLAIN_ONLY));
 	Assert(task != NULL);
 
-	/* we only support default scan direction and row fetch count */
-	if (!ScanDirectionIsForward(direction))
-	{
-		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						errmsg("scan directions other than forward scans "
-							   "are unsupported")));
-	}
 
 	oldcontext = MemoryContextSwitchTo(estate->es_query_cxt);
 
@@ -376,6 +369,19 @@ RouterExecutorRun(QueryDesc *queryDesc, ScanDirection direction, long count)
 		(*destination->rStartup)(destination, operation, queryDesc->tupDesc);
 	}
 
+	/* we only support returning nothing or scanning forward */
+	if (ScanDirectionIsNoMovement(direction))
+	{
+		/* comments in PortalRunSelect() explain the reason for this case */
+		goto out;
+	}
+	else if (!ScanDirectionIsForward(direction))
+	{
+		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("scan directions other than forward scans "
+							   "are unsupported")));
+	}
+
 	/*
 	 * If query has not yet been executed, do so now. The main reason why the
 	 * query might already have been executed is cursors.
@@ -430,6 +436,8 @@ RouterExecutorRun(QueryDesc *queryDesc, ScanDirection direction, long count)
 		}
 	}
 
+out:
+
 	/* shutdown tuple receiver, if we started it */
 	if (sendTuples)
 	{
diff --git a/src/test/regress/expected/multi_router_planner.out b/src/test/regress/expected/multi_router_planner.out
index e6053d972..26a255189 100644
--- a/src/test/regress/expected/multi_router_planner.out
+++ b/src/test/regress/expected/multi_router_planner.out
@@ -1363,11 +1363,19 @@ FETCH test_cursor;
   1 |         1 | arsenous |       9572
 (1 row)
 
-FETCH test_cursor;
+FETCH ALL test_cursor;
+ id | author_id |    title     | word_count 
+----+-----------+--------------+------------
+ 11 |         1 | alamo        |       1347
+ 21 |         1 | arcading     |       5890
+ 31 |         1 | athwartships |       7271
+ 41 |         1 | aznavour     |      11814
+(4 rows)
+
+FETCH test_cursor; -- fetch one row after the last
  id | author_id | title | word_count 
 ----+-----------+-------+------------
- 11 |         1 | alamo |       1347
-(1 row)
+(0 rows)
 
 END;
 -- queries inside copy can be router plannable
@@ -1488,7 +1496,7 @@ PL/pgSQL function author_articles_max_id() line 5 at SQL statement
                      41
 (1 row)
 
--- plpgsql function that return query results are not router plannable
+-- check that function returning setof query are router plannable
 CREATE OR REPLACE FUNCTION author_articles_id_word_count() RETURNS TABLE(id bigint, word_count int) AS $$
 DECLARE
 BEGIN
@@ -1512,8 +1520,15 @@ CONTEXT:  SQL statement "SELECT ah.id, ah.word_count
 PL/pgSQL function author_articles_id_word_count() line 4 at RETURN QUERY
 DEBUG:  Plan is router executable
 CONTEXT:  PL/pgSQL function author_articles_id_word_count() line 4 at RETURN QUERY
-ERROR:  scan directions other than forward scans are unsupported
-CONTEXT:  PL/pgSQL function author_articles_id_word_count() line 4 at RETURN QUERY
+ id | word_count 
+----+------------
+  1 |       9572
+ 11 |       1347
+ 21 |       5890
+ 31 |       7271
+ 41 |      11814
+(5 rows)
+
 -- materialized views can be created for router plannable queries
 CREATE MATERIALIZED VIEW mv_articles_hash AS
 	SELECT * FROM articles_hash WHERE author_id = 1;
diff --git a/src/test/regress/sql/multi_router_planner.sql b/src/test/regress/sql/multi_router_planner.sql
index c84dbee5b..b2f7c2c57 100644
--- a/src/test/regress/sql/multi_router_planner.sql
+++ b/src/test/regress/sql/multi_router_planner.sql
@@ -577,7 +577,8 @@ DECLARE test_cursor CURSOR FOR
 		WHERE author_id = 1
 		ORDER BY id;
 FETCH test_cursor;
-FETCH test_cursor;
+FETCH ALL test_cursor;
+FETCH test_cursor; -- fetch one row after the last
 END;
 
 -- queries inside copy can be router plannable
@@ -634,7 +635,7 @@ $$ LANGUAGE plpgsql;
 
 SELECT author_articles_max_id();
 
--- plpgsql function that return query results are not router plannable
+-- check that function returning setof query are router plannable
 CREATE OR REPLACE FUNCTION author_articles_id_word_count() RETURNS TABLE(id bigint, word_count int) AS $$
 DECLARE
 BEGIN