Add PL/PgSQL FOUND and GET DIAGNOSTICS support for RETURN QUERY

bmomjian · bmomjian · commit 8c78f8e65c2f · 2009-02-05T15:25:49.000Z
statement

Pavel Stehule
diff --git a/doc/src/sgml/plpgsql.sgml b/doc/src/sgml/plpgsql.sgml
@@ -1,4 +1,4 @@
-<!-- $PostgreSQL: pgsql/doc/src/sgml/plpgsql.sgml,v 1.137 2009/02/04 21:30:41 alvherre Exp $ -->
+<!-- $PostgreSQL: pgsql/doc/src/sgml/plpgsql.sgml,v 1.138 2009/02/05 15:25:49 momjian Exp $ -->
 
 <chapter id="plpgsql">
   <title><application>PL/pgSQL</application> - <acronym>SQL</acronym> Procedural Language</title>
@@ -1356,6 +1356,14 @@ GET DIAGNOSTICS integer_var = ROW_COUNT;
             execution of other statements within the loop body.
            </para>
           </listitem>
+          <listitem>
+           <para>
+            A <command>RETURN QUERY</command> and <command>RETURN QUERY
+            EXECUTE</command> statements set <literal>FOUND</literal>
+            true if the query returns at least one row, false if no row
+            is returned.
+           </para>
+          </listitem>
          </itemizedlist>
 
      <literal>FOUND</literal> is a local variable within each
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.231 2009/01/21 11:13:14 heikki Exp $
+ *	  $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.232 2009/02/05 15:25:49 momjian Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -2286,6 +2286,7 @@ exec_stmt_return_query(PLpgSQL_execstate *estate,
 					   PLpgSQL_stmt_return_query *stmt)
 {
 	Portal		portal;
+	uint32			processed = 0;
 
 	if (!estate->retisset)
 		ereport(ERROR,
@@ -2327,6 +2328,7 @@ exec_stmt_return_query(PLpgSQL_execstate *estate,
 			HeapTuple	tuple = SPI_tuptable->vals[i];
 
 			tuplestore_puttuple(estate->tuple_store, tuple);
+			processed++;
 		}
 		MemoryContextSwitchTo(old_cxt);
 
@@ -2336,6 +2338,9 @@ exec_stmt_return_query(PLpgSQL_execstate *estate,
 	SPI_freetuptable(SPI_tuptable);
 	SPI_cursor_close(portal);
 
+	estate->eval_processed = processed;
+	exec_set_found(estate, processed != 0);
+
 	return PLPGSQL_RC_OK;
 }
 
diff --git a/src/test/regress/expected/plpgsql.out b/src/test/regress/expected/plpgsql.out
@@ -3666,3 +3666,35 @@ select * from tftest(10);
 (2 rows)
 
 drop function tftest(int);
+create or replace function rttest()
+returns setof int as $$
+declare rc int;
+begin
+  return query values(10),(20);
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+  return query select * from (values(10),(20)) f(a) where false;
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+  return query execute 'values(10),(20)';
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+  return query execute 'select * from (values(10),(20)) f(a) where false';
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+end;
+$$ language plpgsql;
+select * from rttest();
+NOTICE:  t 2
+NOTICE:  f 0
+NOTICE:  t 2
+NOTICE:  f 0
+ rttest 
+--------
+     10
+     20
+     10
+     20
+(4 rows)
+
+drop function rttest();
diff --git a/src/test/regress/sql/plpgsql.sql b/src/test/regress/sql/plpgsql.sql
@@ -2948,3 +2948,27 @@ $$ language plpgsql immutable strict;
 select * from tftest(10);
 
 drop function tftest(int);
+
+create or replace function rttest()
+returns setof int as $$
+declare rc int;
+begin
+  return query values(10),(20);
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+  return query select * from (values(10),(20)) f(a) where false;
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+  return query execute 'values(10),(20)';
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+  return query execute 'select * from (values(10),(20)) f(a) where false';
+  get diagnostics rc = row_count;
+  raise notice '% %', found, rc;
+end;
+$$ language plpgsql;
+
+select * from rttest();
+
+drop function rttest();
+

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@`
`8`	`8`	`*`
`9`	`9`	`*`
`10`	`10`	`* IDENTIFICATION`
`11`		`- * $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.231 2009/01/21 11:13:14 heikki Exp $`
	`11`	`+ * $PostgreSQL: pgsql/src/pl/plpgsql/src/pl_exec.c,v 1.232 2009/02/05 15:25:49 momjian Exp $`
`12`	`12`	`*`
`13`	`13`	`*-------------------------------------------------------------------------`
`14`	`14`	`*/`
`@@ -2286,6 +2286,7 @@ exec_stmt_return_query(PLpgSQL_execstate *estate,`
`2286`	`2286`	`PLpgSQL_stmt_return_query *stmt)`
`2287`	`2287`	`{`
`2288`	`2288`	`Portal portal;`
	`2289`	`+ uint32 processed = 0;`
`2289`	`2290`
`2290`	`2291`	`if (!estate->retisset)`
`2291`	`2292`	`ereport(ERROR,`
`@@ -2327,6 +2328,7 @@ exec_stmt_return_query(PLpgSQL_execstate *estate,`
`2327`	`2328`	`HeapTuple tuple = SPI_tuptable->vals[i];`
`2328`	`2329`
`2329`	`2330`	`tuplestore_puttuple(estate->tuple_store, tuple);`
	`2331`	`+ processed++;`
`2330`	`2332`	`}`
`2331`	`2333`	`MemoryContextSwitchTo(old_cxt);`
`2332`	`2334`
`@@ -2336,6 +2338,9 @@ exec_stmt_return_query(PLpgSQL_execstate *estate,`
`2336`	`2338`	`SPI_freetuptable(SPI_tuptable);`
`2337`	`2339`	`SPI_cursor_close(portal);`
`2338`	`2340`
	`2341`	`+ estate->eval_processed = processed;`
	`2342`	`+ exec_set_found(estate, processed != 0);`
	`2343`	`+`
`2339`	`2344`	`return PLPGSQL_RC_OK;`
`2340`	`2345`	`}`
`2341`	`2346`