PL/Python: Add cursor and execute methods to plan object

petere · petere · commit 70ec3f1f8f0b · 2017-03-27T11:37:22.000-04:00
Instead of

    plan = plpy.prepare(...)
    res = plpy.execute(plan, ...)

you can now write

    plan = plpy.prepare(...)
    res = plan.execute(...)

or even

    res = plpy.prepare(...).execute(...)

and similarly for the cursor() method.

This is more in object oriented style, and makes the hybrid nature of
the existing execute() function less confusing.

Reviewed-by: Andrew Dunstan &lt;andrew.dunstan@2ndquadrant.com&gt;
diff --git a/doc/src/sgml/plpython.sgml b/doc/src/sgml/plpython.sgml
@@ -1046,6 +1046,14 @@ rv = plpy.execute(plan, ["name"], 5)
       The third argument is the optional row limit as before.
      </para>
 
+     <para>
+      Alternatively, you can call the <function>execute</function> method on
+      the plan object:
+<programlisting>
+rv = plan.execute(["name"], 5)
+</programlisting>
+     </para>
+
      <para>
       Query parameters and result row fields are converted between PostgreSQL
       and Python data types as described in <xref linkend="plpython-data">.
@@ -1081,7 +1089,9 @@ $$ LANGUAGE plpythonu;
       as <literal>plpy.execute</literal> (except for the row limit) and returns
       a cursor object, which allows you to process large result sets in smaller
       chunks.  As with <literal>plpy.execute</literal>, either a query string
-      or a plan object along with a list of arguments can be used.
+      or a plan object along with a list of arguments can be used, or
+      the <function>cursor</function> function can be called as a method of
+      the plan object.
      </para>
 
      <para>
@@ -1125,7 +1135,7 @@ $$ LANGUAGE plpythonu;
 CREATE FUNCTION count_odd_prepared() RETURNS integer AS $$
 odd = 0
 plan = plpy.prepare("select num from largetable where num % $1 &lt;&gt; 0", ["integer"])
-rows = list(plpy.cursor(plan, [2]))
+rows = list(plpy.cursor(plan, [2]))  # or: = list(plan.cursor([2]))
 
 return len(rows)
 $$ LANGUAGE plpythonu;
diff --git a/src/pl/plpython/expected/plpython_spi.out b/src/pl/plpython/expected/plpython_spi.out
@@ -29,6 +29,19 @@ try:
 except Exception, ex:
 	plpy.error(str(ex))
 return None
+'
+	LANGUAGE plpythonu;
+CREATE FUNCTION spi_prepared_plan_test_two(a text) RETURNS text
+	AS
+'if "myplan" not in SD:
+	q = "SELECT count(*) FROM users WHERE lname = $1"
+	SD["myplan"] = plpy.prepare(q, [ "text" ])
+try:
+	rv = SD["myplan"].execute([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(a text) RETURNS text
@@ -80,8 +93,8 @@ select spi_prepared_plan_test_one('doe');
  there are 3 does
 (1 row)
 
-select spi_prepared_plan_test_one('smith');
- spi_prepared_plan_test_one 
+select spi_prepared_plan_test_two('smith');
+ spi_prepared_plan_test_two 
 ----------------------------
  there are 1 smiths
 (1 row)
@@ -372,7 +385,7 @@ plan = plpy.prepare(
     ["text"])
 for row in plpy.cursor(plan, ["w"]):
     yield row['fname']
-for row in plpy.cursor(plan, ["j"]):
+for row in plan.cursor(["j"]):
     yield row['fname']
 $$ LANGUAGE plpythonu;
 CREATE FUNCTION cursor_plan_wrong_args() RETURNS SETOF text AS $$
diff --git a/src/pl/plpython/plpy_cursorobject.c b/src/pl/plpython/plpy_cursorobject.c
@@ -25,7 +25,6 @@
 
 
 static PyObject *PLy_cursor_query(const char *query);
-static PyObject *PLy_cursor_plan(PyObject *ob, PyObject *args);
 static void PLy_cursor_dealloc(PyObject *arg);
 static PyObject *PLy_cursor_iternext(PyObject *self);
 static PyObject *PLy_cursor_fetch(PyObject *self, PyObject *args);
@@ -160,7 +159,7 @@ PLy_cursor_query(const char *query)
 	return (PyObject *) cursor;
 }
 
-static PyObject *
+PyObject *
 PLy_cursor_plan(PyObject *ob, PyObject *args)
 {
 	PLyCursorObject *cursor;
diff --git a/src/pl/plpython/plpy_cursorobject.h b/src/pl/plpython/plpy_cursorobject.h
@@ -19,5 +19,6 @@ typedef struct PLyCursorObject
 
 extern void PLy_cursor_init_type(void);
 extern PyObject *PLy_cursor(PyObject *self, PyObject *args);
+extern PyObject *PLy_cursor_plan(PyObject *ob, PyObject *args);
 
 #endif   /* PLPY_CURSOROBJECT_H */
diff --git a/src/pl/plpython/plpy_planobject.c b/src/pl/plpython/plpy_planobject.c
@@ -10,18 +10,24 @@
 
 #include "plpy_planobject.h"
 
+#include "plpy_cursorobject.h"
 #include "plpy_elog.h"
+#include "plpy_spi.h"
 #include "utils/memutils.h"
 
 
 static void PLy_plan_dealloc(PyObject *arg);
+static PyObject *PLy_plan_cursor(PyObject *self, PyObject *args);
+static PyObject *PLy_plan_execute(PyObject *self, PyObject *args);
 static PyObject *PLy_plan_status(PyObject *self, PyObject *args);
 
 static char PLy_plan_doc[] = {
 	"Store a PostgreSQL plan"
 };
 
 static PyMethodDef PLy_plan_methods[] = {
+	{"cursor", PLy_plan_cursor, METH_VARARGS, NULL},
+	{"execute", PLy_plan_execute, METH_VARARGS, NULL},
 	{"status", PLy_plan_status, METH_VARARGS, NULL},
 	{NULL, NULL, 0, NULL}
 };
@@ -111,6 +117,31 @@ PLy_plan_dealloc(PyObject *arg)
 }
 
 
+static PyObject *
+PLy_plan_cursor(PyObject *self, PyObject *args)
+{
+	PyObject   *planargs = NULL;
+
+	if (!PyArg_ParseTuple(args, "|O", &planargs))
+		return NULL;
+
+	return PLy_cursor_plan(self, planargs);
+}
+
+
+static PyObject *
+PLy_plan_execute(PyObject *self, PyObject *args)
+{
+	PyObject   *list = NULL;
+	long		limit = 0;
+
+	if (!PyArg_ParseTuple(args, "|Ol", &list, &limit))
+		return NULL;
+
+	return PLy_spi_execute_plan(self, list, limit);
+}
+
+
 static PyObject *
 PLy_plan_status(PyObject *self, PyObject *args)
 {
diff --git a/src/pl/plpython/plpy_spi.c b/src/pl/plpython/plpy_spi.c
@@ -30,7 +30,6 @@
 
 
 static PyObject *PLy_spi_execute_query(char *query, long limit);
-static PyObject *PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit);
 static PyObject *PLy_spi_execute_fetch_result(SPITupleTable *tuptable,
 							 uint64 rows, int status);
 static void PLy_spi_exception_set(PyObject *excclass, ErrorData *edata);
@@ -193,7 +192,7 @@ PLy_spi_execute(PyObject *self, PyObject *args)
 	return NULL;
 }
 
-static PyObject *
+PyObject *
 PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
 {
 	volatile int nargs;
diff --git a/src/pl/plpython/plpy_spi.h b/src/pl/plpython/plpy_spi.h
@@ -10,6 +10,7 @@
 
 extern PyObject *PLy_spi_prepare(PyObject *self, PyObject *args);
 extern PyObject *PLy_spi_execute(PyObject *self, PyObject *args);
+extern PyObject *PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit);
 
 typedef struct PLyExceptionEntry
 {
diff --git a/src/pl/plpython/sql/plpython_spi.sql b/src/pl/plpython/sql/plpython_spi.sql
@@ -37,6 +37,20 @@ return None
 '
 	LANGUAGE plpythonu;
 
+CREATE FUNCTION spi_prepared_plan_test_two(a text) RETURNS text
+	AS
+'if "myplan" not in SD:
+	q = "SELECT count(*) FROM users WHERE lname = $1"
+	SD["myplan"] = plpy.prepare(q, [ "text" ])
+try:
+	rv = SD["myplan"].execute([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(a text) RETURNS text
 	AS
 'if "myplan" not in SD:
@@ -79,7 +93,7 @@ return a + r
 --
 select nested_call_one('pass this along');
 select spi_prepared_plan_test_one('doe');
-select spi_prepared_plan_test_one('smith');
+select spi_prepared_plan_test_two('smith');
 select spi_prepared_plan_test_nested('smith');
 
 SELECT join_sequences(sequences) FROM sequences;
@@ -275,7 +289,7 @@ plan = plpy.prepare(
     ["text"])
 for row in plpy.cursor(plan, ["w"]):
     yield row['fname']
-for row in plpy.cursor(plan, ["j"]):
+for row in plan.cursor(["j"]):
     yield row['fname']
 $$ LANGUAGE plpythonu;
 

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,6 @@`
`25`	`25`
`26`	`26`
`27`	`27`	`static PyObject PLy_cursor_query(const char query);`
`28`		`-static PyObject PLy_cursor_plan(PyObject ob, PyObject *args);`
`29`	`28`	`static void PLy_cursor_dealloc(PyObject *arg);`
`30`	`29`	`static PyObject PLy_cursor_iternext(PyObject self);`
`31`	`30`	`static PyObject PLy_cursor_fetch(PyObject self, PyObject *args);`
`@@ -160,7 +159,7 @@ PLy_cursor_query(const char *query)`
`160`	`159`	`return (PyObject *) cursor;`
`161`	`160`	`}`
`162`	`161`
`163`		`-static PyObject *`
	`162`	`+PyObject *`
`164`	`163`	`PLy_cursor_plan(PyObject ob, PyObject args)`
`165`	`164`	`{`
`166`	`165`	`PLyCursorObject *cursor;`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,6 @@`
`30`	`30`
`31`	`31`
`32`	`32`	`static PyObject PLy_spi_execute_query(char query, long limit);`
`33`		`-static PyObject PLy_spi_execute_plan(PyObject ob, PyObject *list, long limit);`
`34`	`33`	`static PyObject PLy_spi_execute_fetch_result(SPITupleTable tuptable,`
`35`	`34`	`uint64 rows, int status);`
`36`	`35`	`static void PLy_spi_exception_set(PyObject excclass, ErrorData edata);`
`@@ -193,7 +192,7 @@ PLy_spi_execute(PyObject self, PyObject args)`
`193`	`192`	`return NULL;`
`194`	`193`	`}`
`195`	`194`
`196`		`-static PyObject *`
	`195`	`+PyObject *`
`197`	`196`	`PLy_spi_execute_plan(PyObject ob, PyObject list, long limit)`
`198`	`197`	`{`
`199`	`198`	`volatile int nargs;`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@`
`10`	`10`
`11`	`11`	`extern PyObject PLy_spi_prepare(PyObject self, PyObject *args);`
`12`	`12`	`extern PyObject PLy_spi_execute(PyObject self, PyObject *args);`
	`13`	`+extern PyObject PLy_spi_execute_plan(PyObject ob, PyObject *list, long limit);`
`13`	`14`
`14`	`15`	`typedef struct PLyExceptionEntry`
`15`	`16`	`{`