1--
2-- nested calls
3--
4CREATE FUNCTION nested_call_one(a text) RETURNS text
5	AS
6'q = "SELECT nested_call_two(''%s'')" % a
7r = plpy.execute(q)
8return r[0]'
9	LANGUAGE plpythonu ;
10CREATE FUNCTION nested_call_two(a text) RETURNS text
11	AS
12'q = "SELECT nested_call_three(''%s'')" % a
13r = plpy.execute(q)
14return r[0]'
15	LANGUAGE plpythonu ;
16CREATE FUNCTION nested_call_three(a text) RETURNS text
17	AS
18'return a'
19	LANGUAGE plpythonu ;
20-- some spi stuff
21CREATE FUNCTION spi_prepared_plan_test_one(a text) RETURNS text
22	AS
23'if "myplan" not in SD:
24	q = "SELECT count(*) FROM users WHERE lname = $1"
25	SD["myplan"] = plpy.prepare(q, [ "text" ])
26try:
27	rv = plpy.execute(SD["myplan"], [a])
28	return "there are " + str(rv[0]["count"]) + " " + str(a) + "s"
29except Exception, ex:
30	plpy.error(str(ex))
31return None
32'
33	LANGUAGE plpythonu;
34CREATE FUNCTION spi_prepared_plan_test_nested(a text) RETURNS text
35	AS
36'if "myplan" not in SD:
37	q = "SELECT spi_prepared_plan_test_one(''%s'') as count" % a
38	SD["myplan"] = plpy.prepare(q)
39try:
40	rv = plpy.execute(SD["myplan"])
41	if len(rv):
42		return rv[0]["count"]
43except Exception, ex:
44	plpy.error(str(ex))
45return None
46'
47	LANGUAGE plpythonu;
48CREATE FUNCTION join_sequences(s sequences) RETURNS text
49	AS
50'if not s["multipart"]:
51	return s["sequence"]
52q = "SELECT sequence FROM xsequences WHERE pid = ''%s''" % s["pid"]
53rv = plpy.execute(q)
54seq = s["sequence"]
55for r in rv:
56	seq = seq + r["sequence"]
57return seq
58'
59	LANGUAGE plpythonu;
60CREATE FUNCTION spi_recursive_sum(a int) RETURNS int
61	AS
62'r = 0
63if a > 1:
64    r = plpy.execute("SELECT spi_recursive_sum(%d) as a" % (a-1))[0]["a"]
65return a + r
66'
67	LANGUAGE plpythonu;
68--
69-- spi and nested calls
70--
71select nested_call_one('pass this along');
72                         nested_call_one
73-----------------------------------------------------------------
74 {'nested_call_two': "{'nested_call_three': 'pass this along'}"}
75(1 row)
76
77select spi_prepared_plan_test_one('doe');
78 spi_prepared_plan_test_one
79----------------------------
80 there are 3 does
81(1 row)
82
83select spi_prepared_plan_test_one('smith');
84 spi_prepared_plan_test_one
85----------------------------
86 there are 1 smiths
87(1 row)
88
89select spi_prepared_plan_test_nested('smith');
90 spi_prepared_plan_test_nested
91-------------------------------
92 there are 1 smiths
93(1 row)
94
95SELECT join_sequences(sequences) FROM sequences;
96 join_sequences
97----------------
98 ABCDEFGHIJKL
99 ABCDEF
100 ABCDEF
101 ABCDEF
102 ABCDEF
103 ABCDEF
104(6 rows)
105
106SELECT join_sequences(sequences) FROM sequences
107	WHERE join_sequences(sequences) ~* '^A';
108 join_sequences
109----------------
110 ABCDEFGHIJKL
111 ABCDEF
112 ABCDEF
113 ABCDEF
114 ABCDEF
115 ABCDEF
116(6 rows)
117
118SELECT join_sequences(sequences) FROM sequences
119	WHERE join_sequences(sequences) ~* '^B';
120 join_sequences
121----------------
122(0 rows)
123
124SELECT spi_recursive_sum(10);
125 spi_recursive_sum
126-------------------
127                55
128(1 row)
129
130--
131-- plan and result objects
132--
133CREATE FUNCTION result_metadata_test(cmd text) RETURNS int
134AS $$
135plan = plpy.prepare(cmd)
136plpy.info(plan.status()) # not really documented or useful
137result = plpy.execute(plan)
138if result.status() > 0:
139   plpy.info(result.colnames())
140   plpy.info(result.coltypes())
141   plpy.info(result.coltypmods())
142   return result.nrows()
143else:
144   return None
145$$ LANGUAGE plpythonu;
146SELECT result_metadata_test($$SELECT 1 AS foo, '11'::text AS bar UNION SELECT 2, '22'$$);
147INFO:  True
148INFO:  ['foo', 'bar']
149INFO:  [23, 25]
150INFO:  [-1, -1]
151 result_metadata_test
152----------------------
153                    2
154(1 row)
155
156SELECT result_metadata_test($$CREATE TEMPORARY TABLE foo1 (a int, b text)$$);
157INFO:  True
158ERROR:  plpy.Error: command did not produce a result set
159CONTEXT:  Traceback (most recent call last):
160  PL/Python function "result_metadata_test", line 6, in <module>
161    plpy.info(result.colnames())
162PL/Python function "result_metadata_test"
163CREATE FUNCTION result_nrows_test(cmd text) RETURNS int
164AS $$
165result = plpy.execute(cmd)
166return result.nrows()
167$$ LANGUAGE plpythonu;
168SELECT result_nrows_test($$SELECT 1$$);
169 result_nrows_test
170-------------------
171                 1
172(1 row)
173
174SELECT result_nrows_test($$CREATE TEMPORARY TABLE foo2 (a int, b text)$$);
175 result_nrows_test
176-------------------
177                 0
178(1 row)
179
180SELECT result_nrows_test($$INSERT INTO foo2 VALUES (1, 'one'), (2, 'two')$$);
181 result_nrows_test
182-------------------
183                 2
184(1 row)
185
186SELECT result_nrows_test($$UPDATE foo2 SET b = '' WHERE a = 2$$);
187 result_nrows_test
188-------------------
189                 1
190(1 row)
191
192CREATE FUNCTION result_len_test(cmd text) RETURNS int
193AS $$
194result = plpy.execute(cmd)
195return len(result)
196$$ LANGUAGE plpythonu;
197SELECT result_len_test($$SELECT 1$$);
198 result_len_test
199-----------------
200               1
201(1 row)
202
203SELECT result_len_test($$CREATE TEMPORARY TABLE foo3 (a int, b text)$$);
204 result_len_test
205-----------------
206               0
207(1 row)
208
209SELECT result_len_test($$INSERT INTO foo3 VALUES (1, 'one'), (2, 'two')$$);
210 result_len_test
211-----------------
212               0
213(1 row)
214
215SELECT result_len_test($$UPDATE foo3 SET b= '' WHERE a = 2$$);
216 result_len_test
217-----------------
218               0
219(1 row)
220
221CREATE FUNCTION result_subscript_test() RETURNS void
222AS $$
223result = plpy.execute("SELECT 1 AS c UNION SELECT 2 "
224                      "UNION SELECT 3 UNION SELECT 4")
225
226plpy.info(result[1]['c'])
227plpy.info(result[-1]['c'])
228
229plpy.info([item['c'] for item in result[1:3]])
230plpy.info([item['c'] for item in result[::2]])
231
232result[-1] = {'c': 1000}
233result[:2] = [{'c': 10}, {'c': 100}]
234plpy.info([item['c'] for item in result[:]])
235
236# raises TypeError, but the message differs on Python 2.6, so silence it
237try:
238    plpy.info(result['foo'])
239except TypeError:
240    pass
241else:
242    assert False, "TypeError not raised"
243
244$$ LANGUAGE plpythonu;
245SELECT result_subscript_test();
246INFO:  2
247INFO:  4
248INFO:  [2, 3]
249INFO:  [1, 3]
250INFO:  [10, 100, 3, 1000]
251 result_subscript_test
252-----------------------
253
254(1 row)
255
256CREATE FUNCTION result_empty_test() RETURNS void
257AS $$
258result = plpy.execute("select 1 where false")
259
260plpy.info(result[:])
261
262$$ LANGUAGE plpythonu;
263SELECT result_empty_test();
264INFO:  []
265 result_empty_test
266-------------------
267
268(1 row)
269
270CREATE FUNCTION result_str_test(cmd text) RETURNS text
271AS $$
272plan = plpy.prepare(cmd)
273result = plpy.execute(plan)
274return str(result)
275$$ LANGUAGE plpythonu;
276SELECT result_str_test($$SELECT 1 AS foo UNION SELECT 2$$);
277                      result_str_test
278------------------------------------------------------------
279 <PLyResult status=5 nrows=2 rows=[{'foo': 1}, {'foo': 2}]>
280(1 row)
281
282SELECT result_str_test($$CREATE TEMPORARY TABLE foo1 (a int, b text)$$);
283           result_str_test
284--------------------------------------
285 <PLyResult status=4 nrows=0 rows=[]>
286(1 row)
287
288-- cursor objects
289CREATE FUNCTION simple_cursor_test() RETURNS int AS $$
290res = plpy.cursor("select fname, lname from users")
291does = 0
292for row in res:
293    if row['lname'] == 'doe':
294        does += 1
295return does
296$$ LANGUAGE plpythonu;
297CREATE FUNCTION double_cursor_close() RETURNS int AS $$
298res = plpy.cursor("select fname, lname from users")
299res.close()
300res.close()
301$$ LANGUAGE plpythonu;
302CREATE FUNCTION cursor_fetch() RETURNS int AS $$
303res = plpy.cursor("select fname, lname from users")
304assert len(res.fetch(3)) == 3
305assert len(res.fetch(3)) == 1
306assert len(res.fetch(3)) == 0
307assert len(res.fetch(3)) == 0
308try:
309    # use next() or __next__(), the method name changed in
310    # http://www.python.org/dev/peps/pep-3114/
311    try:
312        res.next()
313    except AttributeError:
314        res.__next__()
315except StopIteration:
316    pass
317else:
318    assert False, "StopIteration not raised"
319$$ LANGUAGE plpythonu;
320CREATE FUNCTION cursor_mix_next_and_fetch() RETURNS int AS $$
321res = plpy.cursor("select fname, lname from users order by fname")
322assert len(res.fetch(2)) == 2
323
324item = None
325try:
326    item = res.next()
327except AttributeError:
328    item = res.__next__()
329assert item['fname'] == 'rick'
330
331assert len(res.fetch(2)) == 1
332$$ LANGUAGE plpythonu;
333CREATE FUNCTION fetch_after_close() RETURNS int AS $$
334res = plpy.cursor("select fname, lname from users")
335res.close()
336try:
337    res.fetch(1)
338except ValueError:
339    pass
340else:
341    assert False, "ValueError not raised"
342$$ LANGUAGE plpythonu;
343CREATE FUNCTION next_after_close() RETURNS int AS $$
344res = plpy.cursor("select fname, lname from users")
345res.close()
346try:
347    try:
348        res.next()
349    except AttributeError:
350        res.__next__()
351except ValueError:
352    pass
353else:
354    assert False, "ValueError not raised"
355$$ LANGUAGE plpythonu;
356CREATE FUNCTION cursor_fetch_next_empty() RETURNS int AS $$
357res = plpy.cursor("select fname, lname from users where false")
358assert len(res.fetch(1)) == 0
359try:
360    try:
361        res.next()
362    except AttributeError:
363        res.__next__()
364except StopIteration:
365    pass
366else:
367    assert False, "StopIteration not raised"
368$$ LANGUAGE plpythonu;
369CREATE FUNCTION cursor_plan() RETURNS SETOF text AS $$
370plan = plpy.prepare(
371    "select fname, lname from users where fname like $1 || '%' order by fname",
372    ["text"])
373for row in plpy.cursor(plan, ["w"]):
374    yield row['fname']
375for row in plpy.cursor(plan, ["j"]):
376    yield row['fname']
377$$ LANGUAGE plpythonu;
378CREATE FUNCTION cursor_plan_wrong_args() RETURNS SETOF text AS $$
379plan = plpy.prepare("select fname, lname from users where fname like $1 || '%'",
380                    ["text"])
381c = plpy.cursor(plan, ["a", "b"])
382$$ LANGUAGE plpythonu;
383CREATE TYPE test_composite_type AS (
384  a1 int,
385  a2 varchar
386);
387CREATE OR REPLACE FUNCTION plan_composite_args() RETURNS test_composite_type AS $$
388plan = plpy.prepare("select $1 as c1", ["test_composite_type"])
389res = plpy.execute(plan, [{"a1": 3, "a2": "label"}])
390return res[0]["c1"]
391$$ LANGUAGE plpythonu;
392SELECT simple_cursor_test();
393 simple_cursor_test
394--------------------
395                  3
396(1 row)
397
398SELECT double_cursor_close();
399 double_cursor_close
400---------------------
401
402(1 row)
403
404SELECT cursor_fetch();
405 cursor_fetch
406--------------
407
408(1 row)
409
410SELECT cursor_mix_next_and_fetch();
411 cursor_mix_next_and_fetch
412---------------------------
413
414(1 row)
415
416SELECT fetch_after_close();
417 fetch_after_close
418-------------------
419
420(1 row)
421
422SELECT next_after_close();
423 next_after_close
424------------------
425
426(1 row)
427
428SELECT cursor_fetch_next_empty();
429 cursor_fetch_next_empty
430-------------------------
431
432(1 row)
433
434SELECT cursor_plan();
435 cursor_plan
436-------------
437 willem
438 jane
439 john
440(3 rows)
441
442SELECT cursor_plan_wrong_args();
443ERROR:  TypeError: Expected sequence of 1 argument, got 2: ['a', 'b']
444CONTEXT:  Traceback (most recent call last):
445  PL/Python function "cursor_plan_wrong_args", line 4, in <module>
446    c = plpy.cursor(plan, ["a", "b"])
447PL/Python function "cursor_plan_wrong_args"
448SELECT plan_composite_args();
449 plan_composite_args
450---------------------
451 (3,label)
452(1 row)
453
454