plpython: Add test for returning Python set from SETOF function

This is claimed in the documentation but there was a no test case for
it.

Reported-by: Bogdan Grigorenko <gri.bogdan.2020@gmail.com>
Discussion: https://www.postgresql.org/message-id/flat/173543330569.680.6706329879058172623%40wrigleys.postgresql.org

2 files changed, 42 insertions, 0 deletions

diff --git a/src/pl/plpython/expected/plpython_setof.out b/src/pl/plpython/expected/plpython_setof.out
index 39409400290..c4461ac2762 100644
--- a/src/pl/plpython/expected/plpython_setof.out
+++ b/src/pl/plpython/expected/plpython_setof.out
@@ -17,6 +17,12 @@ for i in range(count):
 	t += ( content, )
 return t
 $$ LANGUAGE plpython3u;
+CREATE FUNCTION test_setof_as_set(count integer, content text) RETURNS SETOF text AS $$
+s = set()
+for i in range(count):
+	s.add(content * (i + 1) if content is not None else None)
+return s
+$$ LANGUAGE plpython3u;
 CREATE FUNCTION test_setof_as_iterator(count integer, content text) RETURNS SETOF text AS $$
 class producer:
 	def __init__ (self, icount, icontent):
@@ -90,6 +96,30 @@ SELECT test_setof_as_tuple(2, null);
  
 (2 rows)
 
+SELECT * FROM test_setof_as_set(0, 'set') ORDER BY 1;
+ test_setof_as_set 
+-------------------
+(0 rows)
+
+SELECT * FROM test_setof_as_set(1, 'set') ORDER BY 1;
+ test_setof_as_set 
+-------------------
+ set
+(1 row)
+
+SELECT * FROM test_setof_as_set(2, 'set') ORDER BY 1;
+ test_setof_as_set 
+-------------------
+ set
+ setset
+(2 rows)
+
+SELECT * FROM test_setof_as_set(2, null) ORDER BY 1;
+ test_setof_as_set 
+-------------------
+ 
+(1 row)
+
 SELECT test_setof_as_iterator(0, 'list');
  test_setof_as_iterator 
 ------------------------
diff --git a/src/pl/plpython/sql/plpython_setof.sql b/src/pl/plpython/sql/plpython_setof.sql
index 4cfb10192c0..1a0472b7a3b 100644
--- a/src/pl/plpython/sql/plpython_setof.sql
+++ b/src/pl/plpython/sql/plpython_setof.sql
@@ -20,6 +20,13 @@ for i in range(count):
 return t
 $$ LANGUAGE plpython3u;
 
+CREATE FUNCTION test_setof_as_set(count integer, content text) RETURNS SETOF text AS $$
+s = set()
+for i in range(count):
+	s.add(content * (i + 1) if content is not None else None)
+return s
+$$ LANGUAGE plpython3u;
+
 CREATE FUNCTION test_setof_as_iterator(count integer, content text) RETURNS SETOF text AS $$
 class producer:
 	def __init__ (self, icount, icontent):
@@ -56,6 +63,11 @@ SELECT test_setof_as_tuple(1, 'tuple');
 SELECT test_setof_as_tuple(2, 'tuple');
 SELECT test_setof_as_tuple(2, null);
 
+SELECT * FROM test_setof_as_set(0, 'set') ORDER BY 1;
+SELECT * FROM test_setof_as_set(1, 'set') ORDER BY 1;
+SELECT * FROM test_setof_as_set(2, 'set') ORDER BY 1;
+SELECT * FROM test_setof_as_set(2, null) ORDER BY 1;
+
 SELECT test_setof_as_iterator(0, 'list');
 SELECT test_setof_as_iterator(1, 'list');
 SELECT test_setof_as_iterator(2, 'list');