Allow PL/python to return composite types and result sets

Sven Suursoho
2025-06-14 18:42:34 +03:00 · 2006-09-02 12:30:01 +00:00
parent b1620c538d
commit 819f22a302
8 changed files with 1359 additions and 99 deletions
--- a/src/pl/plpython/sql/plpython_function.sql
+++ b/src/pl/plpython/sql/plpython_function.sql
@ -65,29 +65,29 @@ except Exception, ex:
 return "succeeded, as expected"'
    LANGUAGE plpythonu;

-CREATE FUNCTION import_test_one(text) RETURNS text
+CREATE FUNCTION import_test_one(p text) RETURNS text
 	AS
 'import sha
-digest = sha.new(args[0])
+digest = sha.new(p)
 return digest.hexdigest()'
 	LANGUAGE plpythonu;

-CREATE FUNCTION import_test_two(users) RETURNS text
+CREATE FUNCTION import_test_two(u users) RETURNS text
 	AS
 'import sha
-plain = args[0]["fname"] + args[0]["lname"]
+plain = u["fname"] + u["lname"]
 digest = sha.new(plain);
 return "sha hash of " + plain + " is " + digest.hexdigest()'
 	LANGUAGE plpythonu;

-CREATE FUNCTION argument_test_one(users, text, text) RETURNS text
+CREATE FUNCTION argument_test_one(u users, a1 text, a2 text) RETURNS text
 	AS
-'keys = args[0].keys()
+'keys = u.keys()
 keys.sort()
 out = []
 for key in keys:
-    out.append("%s: %s" % (key, args[0][key]))
-words = args[1] + " " + args[2] + " => {" + ", ".join(out) + "}"
+    out.append("%s: %s" % (key, u[key]))
+words = a1 + " " + a2 + " => {" + ", ".join(out) + "}"
 return words'
 	LANGUAGE plpythonu;

@ -176,45 +176,45 @@ DROP FUNCTION trigger_data();
 -- nested calls
 --

-CREATE FUNCTION nested_call_one(text) RETURNS text
+CREATE FUNCTION nested_call_one(a text) RETURNS text
 	AS
-'q = "SELECT nested_call_two(''%s'')" % args[0]
+'q = "SELECT nested_call_two(''%s'')" % a
 r = plpy.execute(q)
 return r[0]'
 	LANGUAGE plpythonu ;

-CREATE FUNCTION nested_call_two(text) RETURNS text
+CREATE FUNCTION nested_call_two(a text) RETURNS text
 	AS
-'q = "SELECT nested_call_three(''%s'')" % args[0]
+'q = "SELECT nested_call_three(''%s'')" % a
 r = plpy.execute(q)
 return r[0]'
 	LANGUAGE plpythonu ;

-CREATE FUNCTION nested_call_three(text) RETURNS text
+CREATE FUNCTION nested_call_three(a text) RETURNS text
 	AS
-'return args[0]'
+'return a'
 	LANGUAGE plpythonu ;

 -- some spi stuff

-CREATE FUNCTION spi_prepared_plan_test_one(text) RETURNS text
+CREATE FUNCTION spi_prepared_plan_test_one(a text) RETURNS text
 	AS
 'if not SD.has_key("myplan"):
 	q = "SELECT count(*) FROM users WHERE lname = $1"
 	SD["myplan"] = plpy.prepare(q, [ "text" ])
 try:
-	rv = plpy.execute(SD["myplan"], [args[0]])
-	return "there are " + str(rv[0]["count"]) + " " + str(args[0]) + "s"
+	rv = plpy.execute(SD["myplan"], [a])
+	return "there are " + str(rv[0]["count"]) + " " + str(a) + "s"
 except Exception, ex:
 	plpy.error(str(ex))
 return None
 '
 	LANGUAGE plpythonu;

-CREATE FUNCTION spi_prepared_plan_test_nested(text) RETURNS text
+CREATE FUNCTION spi_prepared_plan_test_nested(a text) RETURNS text
 	AS
 'if not SD.has_key("myplan"):
-	q = "SELECT spi_prepared_plan_test_one(''%s'') as count" % args[0]
+	q = "SELECT spi_prepared_plan_test_one(''%s'') as count" % a
 	SD["myplan"] = plpy.prepare(q)
 try:
 	rv = plpy.execute(SD["myplan"])
@ -233,12 +233,12 @@ CREATE FUNCTION stupid() RETURNS text AS 'return "zarkon"' LANGUAGE plpythonu;

 /* a typo
 */
-CREATE FUNCTION invalid_type_uncaught(text) RETURNS text
+CREATE FUNCTION invalid_type_uncaught(a text) RETURNS text
 	AS
 'if not SD.has_key("plan"):
 	q = "SELECT fname FROM users WHERE lname = $1"
 	SD["plan"] = plpy.prepare(q, [ "test" ])
-rv = plpy.execute(SD["plan"], [ args[0] ])
+rv = plpy.execute(SD["plan"], [ a ])
 if len(rv):
 	return rv[0]["fname"]
 return None
@ -248,7 +248,7 @@ return None
 /* for what it's worth catch the exception generated by
 * the typo, and return None
 */
-CREATE FUNCTION invalid_type_caught(text) RETURNS text
+CREATE FUNCTION invalid_type_caught(a text) RETURNS text
 	AS
 'if not SD.has_key("plan"):
 	q = "SELECT fname FROM users WHERE lname = $1"
@ -257,7 +257,7 @@ CREATE FUNCTION invalid_type_caught(text) RETURNS text
 	except plpy.SPIError, ex:
 		plpy.notice(str(ex))
 		return None
-rv = plpy.execute(SD["plan"], [ args[0] ])
+rv = plpy.execute(SD["plan"], [ a ])
 if len(rv):
 	return rv[0]["fname"]
 return None
@ -267,7 +267,7 @@ return None
 /* for what it's worth catch the exception generated by
 * the typo, and reraise it as a plain error
 */
-CREATE FUNCTION invalid_type_reraised(text) RETURNS text
+CREATE FUNCTION invalid_type_reraised(a text) RETURNS text
 	AS
 'if not SD.has_key("plan"):
 	q = "SELECT fname FROM users WHERE lname = $1"
@ -275,7 +275,7 @@ CREATE FUNCTION invalid_type_reraised(text) RETURNS text
 		SD["plan"] = plpy.prepare(q, [ "test" ])
 	except plpy.SPIError, ex:
 		plpy.error(str(ex))
-rv = plpy.execute(SD["plan"], [ args[0] ])
+rv = plpy.execute(SD["plan"], [ a ])
 if len(rv):
 	return rv[0]["fname"]
 return None
@ -285,11 +285,11 @@ return None

 /* no typo no messing about
 */
-CREATE FUNCTION valid_type(text) RETURNS text
+CREATE FUNCTION valid_type(a text) RETURNS text
 	AS
 'if not SD.has_key("plan"):
 	SD["plan"] = plpy.prepare("SELECT fname FROM users WHERE lname = $1", [ "text" ])
-rv = plpy.execute(SD["plan"], [ args[0] ])
+rv = plpy.execute(SD["plan"], [ a ])
 if len(rv):
 	return rv[0]["fname"]
 return None
@ -318,13 +318,13 @@ return rv[0]'
 	LANGUAGE plpythonu;


-CREATE FUNCTION join_sequences(sequences) RETURNS text
+CREATE FUNCTION join_sequences(s sequences) RETURNS text
 	AS
-'if not args[0]["multipart"]:
-	return args[0]["sequence"]
-q = "SELECT sequence FROM xsequences WHERE pid = ''%s''" % args[0]["pid"]
+'if not s["multipart"]:
+	return s["sequence"]
+q = "SELECT sequence FROM xsequences WHERE pid = ''%s''" % s["pid"]
 rv = plpy.execute(q)
-seq = args[0]["sequence"]
+seq = s["sequence"]
 for r in rv:
 	seq = seq + r["sequence"]
 return seq
@ -389,3 +389,95 @@ $$ LANGUAGE plpythonu;
 CREATE FUNCTION test_return_none() RETURNS int AS $$
 None
 $$ LANGUAGE plpythonu;
+
+
+--
+-- Test named parameters
+--
+CREATE FUNCTION test_param_names1(a0 integer, a1 text) RETURNS boolean AS $$
+assert a0 == args[0]
+assert a1 == args[1]
+return True
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION test_param_names2(u users) RETURNS text AS $$
+assert u == args[0]
+return str(u)
+$$ LANGUAGE plpythonu;
+
+-- use deliberately wrong parameter names
+CREATE FUNCTION test_param_names3(a0 integer) RETURNS boolean AS $$
+try:
+	assert a1 == args[0]
+	return False
+except NameError, e:
+	assert e.args[0].find("a1") > -1
+	return True
+$$ LANGUAGE plpythonu;
+
+
+--
+-- Test returning SETOF
+--
+CREATE FUNCTION test_setof_as_list(count integer, content text) RETURNS SETOF text AS $$
+return [ content ]*count
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION test_setof_as_tuple(count integer, content text) RETURNS SETOF text AS $$
+t = ()
+for i in xrange(count):
+	t += ( content, )
+return t
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION test_setof_as_iterator(count integer, content text) RETURNS SETOF text AS $$
+class producer:
+	def __init__ (self, icount, icontent):
+		self.icontent = icontent
+		self.icount = icount
+	def __iter__ (self):
+		return self
+	def next (self):
+		if self.icount == 0:
+			raise StopIteration
+		self.icount -= 1
+		return self.icontent
+return producer(count, content)
+$$ LANGUAGE plpythonu;
+
+
+--
+-- Test returning tuples
+--
+CREATE FUNCTION test_table_record_as(typ text, first text, second integer, retnull boolean) RETURNS table_record AS $$
+if retnull:
+	return None
+if typ == 'dict':
+	return { 'first': first, 'second': second, 'additionalfield': 'must not cause trouble' }
+elif typ == 'tuple':
+	return ( first, second )
+elif typ == 'list':
+	return [ first, second ]
+elif typ == 'obj':
+	class type_record: pass
+	type_record.first = first
+	type_record.second = second
+	return type_record
+$$ LANGUAGE plpythonu;
+
+CREATE FUNCTION test_type_record_as(typ text, first text, second integer, retnull boolean) RETURNS type_record AS $$
+if retnull:
+	return None
+if typ == 'dict':
+	return { 'first': first, 'second': second, 'additionalfield': 'must not cause trouble' }
+elif typ == 'tuple':
+	return ( first, second )
+elif typ == 'list':
+	return [ first, second ]
+elif typ == 'obj':
+	class type_record: pass
+	type_record.first = first
+	type_record.second = second
+	return type_record
+$$ LANGUAGE plpythonu;
+