tablefunc: Reject negative number of tuples passed to normal_rand()

The function converted 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, reject a negative value.

(If someone really wants to generate many more result rows, they
should consider adding a bigint or numeric variant.)

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

Author: Ashutosh Bapat <[email protected]>
Discussion: https://www.postgresql.org/message-id/CAG-ACPW3PUUmSnM6cLa9Rw4BEC5cEMKjX8Gogc8gvQcT3cYA1A@mail.gmail.com

diff --git a/contrib/tablefunc/expected/tablefunc.out b/contrib/tablefunc/expected/tablefunc.out
index fffadc6e1b405533a5ec0cd47250abe80ac8c095..464c210f42fd11b379c6e9adff5802df6de21830 100644 (file)
--- a/contrib/tablefunc/expected/tablefunc.out
+++ b/contrib/tablefunc/expected/tablefunc.out
@@ -3,12 +3,15 @@ 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);
+ERROR:  number of rows cannot be negative
 --
 -- crosstab()
 --

diff --git a/contrib/tablefunc/sql/tablefunc.sql b/contrib/tablefunc/sql/tablefunc.sql
index ec375b05c63c9df90fa437be88ad192bf0ab24c4..02e8a98c73e0b0f7acdb6e2c0e4301b584322908 100644 (file)
--- 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 02f02eab57422958aa52df3a8232de3230a873bb..85986ec24adfae0c227d2a2868ca3b050338a773 100644 (file)
--- a/contrib/tablefunc/tablefunc.c
+++ b/contrib/tablefunc/tablefunc.c
@@ -184,6 +184,8 @@ normal_rand(PG_FUNCTION_ARGS)
    /* stuff done only on the first call of the function */
    if (SRF_IS_FIRSTCALL())
    {
+       int32       num_tuples;
+
        /* create a function context for cross-call persistence */
        funcctx = SRF_FIRSTCALL_INIT();
 
@@ -193,7 +195,12 @@ 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);
+       if (num_tuples < 0)
+           ereport(ERROR,
+                   (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+                    errmsg("number of rows cannot be negative")));
+       funcctx->max_calls = num_tuples;
 
        /* allocate memory for user context */
        fctx = (normal_rand_fctx *) palloc(sizeof(normal_rand_fctx));