Allow rescan in DECLARE .. WITH HOLD

(cherry picked from commit e3b7e74f43)

src/backend/distributed/executor/citus_custom_scan.c

@@ -301,13 +301,20 @@ CitusEndScan(CustomScanState *node)
 
 
 /*
- * CitusReScan is just a place holder for rescan callback. Currently, we don't
+ * CitusReScan is not normally called, except in certain cases of
- * support rescan given that there is not any way to reach this code path.
+ * DECLARE .. CURSOR WITH HOLD ..
  */
 static void
 CitusReScan(CustomScanState *node)
+{
+	CitusScanState *scanState = (CitusScanState *) node;
+	EState *executorState = scanState->customScanState.ss.ps.state;
+	ParamListInfo paramListInfo = executorState->es_param_list_info;
+
+	if (paramListInfo != NULL)
 	{
 		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					errmsg("rescan is unsupported"),
+						errmsg("Cursors for queries on distributed tables with "
-					errdetail("We don't expect this code path to be executed.")));
+							   "parameters are currently unsupported")));
+	}
 }

src/test/regress/expected/multi_utility_statements.out

@@ -210,7 +210,100 @@ SELECT customer_key, c_name, c_address
 ERROR:  relation "customer_few" does not exist
 LINE 2:        FROM customer_few ORDER BY customer_key LIMIT 5;
                     ^
--- Test DECLARE CURSOR statements
+-- Test DECLARE CURSOR .. WITH HOLD without parameters that calls ReScan on the top-level CustomScan
+CREATE TABLE cursor_me (x int, y int);
+SELECT create_distributed_table('cursor_me', 'x');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+INSERT INTO cursor_me SELECT s/10, s FROM generate_series(1, 100) s;
+DECLARE holdCursor CURSOR WITH HOLD FOR
+		SELECT * FROM cursor_me WHERE x = 1 ORDER BY y;
+FETCH NEXT FROM holdCursor;
+ x | y  
+---+----
+ 1 | 10
+(1 row)
+
+FETCH FORWARD 3 FROM holdCursor;
+ x | y  
+---+----
+ 1 | 11
+ 1 | 12
+ 1 | 13
+(3 rows)
+
+FETCH LAST FROM holdCursor;
+ x | y  
+---+----
+ 1 | 19
+(1 row)
+
+FETCH BACKWARD 3 FROM holdCursor;
+ x | y  
+---+----
+ 1 | 18
+ 1 | 17
+ 1 | 16
+(3 rows)
+
+FETCH FORWARD 3 FROM holdCursor;
+ x | y  
+---+----
+ 1 | 17
+ 1 | 18
+ 1 | 19
+(3 rows)
+
+CLOSE holdCursor;
+-- Test DECLARE CURSOR .. WITH HOLD with parameter
+CREATE OR REPLACE FUNCTION declares_cursor(p int)
+RETURNS void AS $$
+	DECLARE c CURSOR WITH HOLD FOR SELECT * FROM cursor_me WHERE x = $1;
+$$ LANGUAGE SQL;
+SELECT declares_cursor(5);
+ERROR:  Cursors for queries on distributed tables with parameters are currently unsupported
+CREATE OR REPLACE FUNCTION cursor_plpgsql(p int)
+RETURNS SETOF int AS $$                                                                                                                                                            
+DECLARE
+  val int;
+  my_cursor CURSOR (a INTEGER) FOR SELECT y FROM cursor_me WHERE x = $1 ORDER BY y;
+BEGIN
+   -- Open the cursor
+   OPEN my_cursor(p);
+
+   LOOP
+      FETCH my_cursor INTO val;
+      EXIT WHEN NOT FOUND;
+
+      RETURN NEXT val;
+   END LOOP;
+
+   -- Close the cursor
+   CLOSE my_cursor;
+END; $$
+LANGUAGE plpgsql;
+SELECT cursor_plpgsql(4);
+ cursor_plpgsql 
+----------------
+             40
+             41
+             42
+             43
+             44
+             45
+             46
+             47
+             48
+             49
+(10 rows)
+
+DROP FUNCTION declares_cursor(int);
+DROP FUNCTION cursor_plpgsql(int);
+DROP TABLE cursor_me;
+-- Test DECLARE CURSOR statement with SCROLL
 DECLARE holdCursor SCROLL CURSOR WITH HOLD FOR
         SELECT l_orderkey, l_linenumber, l_quantity, l_discount
         FROM lineitem 

src/test/regress/sql/multi_utility_statements.sql

@@ -122,8 +122,58 @@ COMMIT;
 SELECT customer_key, c_name, c_address 
        FROM customer_few ORDER BY customer_key LIMIT 5;
 
--- Test DECLARE CURSOR statements
+-- Test DECLARE CURSOR .. WITH HOLD without parameters that calls ReScan on the top-level CustomScan
+CREATE TABLE cursor_me (x int, y int);
+SELECT create_distributed_table('cursor_me', 'x');
+INSERT INTO cursor_me SELECT s/10, s FROM generate_series(1, 100) s;
 
+DECLARE holdCursor CURSOR WITH HOLD FOR
+		SELECT * FROM cursor_me WHERE x = 1 ORDER BY y;
+
+FETCH NEXT FROM holdCursor;
+FETCH FORWARD 3 FROM holdCursor;
+FETCH LAST FROM holdCursor;
+FETCH BACKWARD 3 FROM holdCursor;
+FETCH FORWARD 3 FROM holdCursor;
+
+CLOSE holdCursor;
+
+-- Test DECLARE CURSOR .. WITH HOLD with parameter
+CREATE OR REPLACE FUNCTION declares_cursor(p int)
+RETURNS void AS $$
+	DECLARE c CURSOR WITH HOLD FOR SELECT * FROM cursor_me WHERE x = $1;
+$$ LANGUAGE SQL;
+
+SELECT declares_cursor(5);
+
+CREATE OR REPLACE FUNCTION cursor_plpgsql(p int)
+RETURNS SETOF int AS $$                                                                                                                                                            
+DECLARE
+  val int;
+  my_cursor CURSOR (a INTEGER) FOR SELECT y FROM cursor_me WHERE x = $1 ORDER BY y;
+BEGIN
+   -- Open the cursor
+   OPEN my_cursor(p);
+
+   LOOP
+      FETCH my_cursor INTO val;
+      EXIT WHEN NOT FOUND;
+
+      RETURN NEXT val;
+   END LOOP;
+
+   -- Close the cursor
+   CLOSE my_cursor;
+END; $$
+LANGUAGE plpgsql;
+
+SELECT cursor_plpgsql(4);
+
+DROP FUNCTION declares_cursor(int);
+DROP FUNCTION cursor_plpgsql(int);
+DROP TABLE cursor_me;
+
+-- Test DECLARE CURSOR statement with SCROLL
 DECLARE holdCursor SCROLL CURSOR WITH HOLD FOR
         SELECT l_orderkey, l_linenumber, l_quantity, l_discount
         FROM lineitem