1--
2-- Test returning SETOF
3--
4CREATE FUNCTION test_setof_error() RETURNS SETOF text AS $$
5return 37
6$$ LANGUAGE plpythonu;
7SELECT test_setof_error();
8ERROR:  returned object cannot be iterated
9DETAIL:  PL/Python set-returning functions must return an iterable object.
10CONTEXT:  PL/Python function "test_setof_error"
11CREATE FUNCTION test_setof_as_list(count integer, content text) RETURNS SETOF text AS $$
12return [ content ]*count
13$$ LANGUAGE plpythonu;
14CREATE FUNCTION test_setof_as_tuple(count integer, content text) RETURNS SETOF text AS $$
15t = ()
16for i in range(count):
17	t += ( content, )
18return t
19$$ LANGUAGE plpythonu;
20CREATE FUNCTION test_setof_as_iterator(count integer, content text) RETURNS SETOF text AS $$
21class producer:
22	def __init__ (self, icount, icontent):
23		self.icontent = icontent
24		self.icount = icount
25	def __iter__ (self):
26		return self
27	def next (self):
28		if self.icount == 0:
29			raise StopIteration
30		self.icount -= 1
31		return self.icontent
32return producer(count, content)
33$$ LANGUAGE plpythonu;
34CREATE FUNCTION test_setof_spi_in_iterator() RETURNS SETOF text AS
35$$
36    for s in ('Hello', 'Brave', 'New', 'World'):
37        plpy.execute('select 1')
38        yield s
39        plpy.execute('select 2')
40$$
41LANGUAGE plpythonu;
42-- Test set returning functions
43SELECT test_setof_as_list(0, 'list');
44 test_setof_as_list
45--------------------
46(0 rows)
47
48SELECT test_setof_as_list(1, 'list');
49 test_setof_as_list
50--------------------
51 list
52(1 row)
53
54SELECT test_setof_as_list(2, 'list');
55 test_setof_as_list
56--------------------
57 list
58 list
59(2 rows)
60
61SELECT test_setof_as_list(2, null);
62 test_setof_as_list
63--------------------
64
65
66(2 rows)
67
68SELECT test_setof_as_tuple(0, 'tuple');
69 test_setof_as_tuple
70---------------------
71(0 rows)
72
73SELECT test_setof_as_tuple(1, 'tuple');
74 test_setof_as_tuple
75---------------------
76 tuple
77(1 row)
78
79SELECT test_setof_as_tuple(2, 'tuple');
80 test_setof_as_tuple
81---------------------
82 tuple
83 tuple
84(2 rows)
85
86SELECT test_setof_as_tuple(2, null);
87 test_setof_as_tuple
88---------------------
89
90
91(2 rows)
92
93SELECT test_setof_as_iterator(0, 'list');
94 test_setof_as_iterator
95------------------------
96(0 rows)
97
98SELECT test_setof_as_iterator(1, 'list');
99 test_setof_as_iterator
100------------------------
101 list
102(1 row)
103
104SELECT test_setof_as_iterator(2, 'list');
105 test_setof_as_iterator
106------------------------
107 list
108 list
109(2 rows)
110
111SELECT test_setof_as_iterator(2, null);
112 test_setof_as_iterator
113------------------------
114
115
116(2 rows)
117
118SELECT test_setof_spi_in_iterator();
119 test_setof_spi_in_iterator
120----------------------------
121 Hello
122 Brave
123 New
124 World
125(4 rows)
126
127-- set-returning function that modifies its parameters
128CREATE OR REPLACE FUNCTION ugly(x int, lim int) RETURNS SETOF int AS $$
129global x
130while x <= lim:
131    yield x
132    x = x + 1
133$$ LANGUAGE plpythonu;
134SELECT ugly(1, 5);
135 ugly
136------
137    1
138    2
139    3
140    4
141    5
142(5 rows)
143
144-- interleaved execution of such a function
145SELECT ugly(1,3), ugly(7,8);
146 ugly | ugly
147------+------
148    1 |    7
149    2 |    8
150    3 |    7
151    1 |    8
152    2 |    7
153    3 |    8
154(6 rows)
155
156-- returns set of named-composite-type tuples
157CREATE OR REPLACE FUNCTION get_user_records()
158RETURNS SETOF users
159AS $$
160    return plpy.execute("SELECT * FROM users ORDER BY username")
161$$ LANGUAGE plpythonu;
162SELECT get_user_records();
163   get_user_records
164----------------------
165 (jane,doe,j_doe,1)
166 (john,doe,johnd,2)
167 (rick,smith,slash,4)
168 (willem,doe,w_doe,3)
169(4 rows)
170
171SELECT * FROM get_user_records();
172 fname  | lname | username | userid
173--------+-------+----------+--------
174 jane   | doe   | j_doe    |      1
175 john   | doe   | johnd    |      2
176 rick   | smith | slash    |      4
177 willem | doe   | w_doe    |      3
178(4 rows)
179
180-- same, but returning set of RECORD
181CREATE OR REPLACE FUNCTION get_user_records2()
182RETURNS TABLE(fname text, lname text, username text, userid int)
183AS $$
184    return plpy.execute("SELECT * FROM users ORDER BY username")
185$$ LANGUAGE plpythonu;
186SELECT get_user_records2();
187  get_user_records2
188----------------------
189 (jane,doe,j_doe,1)
190 (john,doe,johnd,2)
191 (rick,smith,slash,4)
192 (willem,doe,w_doe,3)
193(4 rows)
194
195SELECT * FROM get_user_records2();
196 fname  | lname | username | userid
197--------+-------+----------+--------
198 jane   | doe   | j_doe    |      1
199 john   | doe   | johnd    |      2
200 rick   | smith | slash    |      4
201 willem | doe   | w_doe    |      3
202(4 rows)
203
204