From c4ca8cf28125c86ff11a1bd8758aaae2b31344f3 Mon Sep 17 00:00:00 2001
From: Ashutosh Bapat <ashutosh.bapat@2ndquadrant.com>
Date: Tue, 22 Sep 2020 19:00:56 +0530
Subject: [PATCH 1/3] Handle negative number of tuples passed to normal_rand()

The function converts the first argument i.e. the number of tuples to
return into an unsigned integer which turns out to be huge number when a
negative value is passed. This causes the function to take much longer
time to execute. Instead return no rows in that case.

While at it, improve SQL test to test the number of tuples returned by
this function.

Ashutosh Bapat
---
 contrib/tablefunc/expected/tablefunc.out | 15 +++++++++++----
 contrib/tablefunc/sql/tablefunc.sql      |  4 +++-
 contrib/tablefunc/tablefunc.c            |  4 +++-
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/contrib/tablefunc/expected/tablefunc.out b/contrib/tablefunc/expected/tablefunc.out
index fffadc6e1b..1c9ecef52c 100644
--- a/contrib/tablefunc/expected/tablefunc.out
+++ b/contrib/tablefunc/expected/tablefunc.out
@@ -3,10 +3,17 @@ CREATE EXTENSION tablefunc;
 -- normal_rand()
 -- no easy way to do this for regression testing
 --
-SELECT avg(normal_rand)::int FROM normal_rand(100, 250, 0.2);
- avg 
------
- 250
+SELECT avg(normal_rand)::int, count(*) FROM normal_rand(100, 250, 0.2);
+ avg | count 
+-----+-------
+ 250 |   100
+(1 row)
+
+-- negative number of tuples
+SELECT avg(normal_rand)::int, count(*) FROM normal_rand(-1, 250, 0.2);
+ avg | count 
+-----+-------
+     |     0
 (1 row)
 
 --
diff --git a/contrib/tablefunc/sql/tablefunc.sql b/contrib/tablefunc/sql/tablefunc.sql
index ec375b05c6..02e8a98c73 100644
--- a/contrib/tablefunc/sql/tablefunc.sql
+++ b/contrib/tablefunc/sql/tablefunc.sql
@@ -4,7 +4,9 @@ CREATE EXTENSION tablefunc;
 -- normal_rand()
 -- no easy way to do this for regression testing
 --
-SELECT avg(normal_rand)::int FROM normal_rand(100, 250, 0.2);
+SELECT avg(normal_rand)::int, count(*) FROM normal_rand(100, 250, 0.2);
+-- negative number of tuples
+SELECT avg(normal_rand)::int, count(*) FROM normal_rand(-1, 250, 0.2);
 
 --
 -- crosstab()
diff --git a/contrib/tablefunc/tablefunc.c b/contrib/tablefunc/tablefunc.c
index 02f02eab57..fd327b1c41 100644
--- a/contrib/tablefunc/tablefunc.c
+++ b/contrib/tablefunc/tablefunc.c
@@ -174,6 +174,7 @@ normal_rand(PG_FUNCTION_ARGS)
 	FuncCallContext *funcctx;
 	uint64		call_cntr;
 	uint64		max_calls;
+	int32		num_tuples;
 	normal_rand_fctx *fctx;
 	float8		mean;
 	float8		stddev;
@@ -193,7 +194,8 @@ normal_rand(PG_FUNCTION_ARGS)
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
 
 		/* total number of tuples to be returned */
-		funcctx->max_calls = PG_GETARG_UINT32(0);
+		num_tuples = PG_GETARG_INT32(0);
+		funcctx->max_calls = num_tuples > 0 ? num_tuples : 0;
 
 		/* allocate memory for user context */
 		fctx = (normal_rand_fctx *) palloc(sizeof(normal_rand_fctx));
-- 
2.17.1

