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

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 <[email protected]>

8 files changed, 79 insertions, 11 deletions

diff --git a/doc/src/sgml/plpython.sgml b/doc/src/sgml/plpython.sgml
index fb5d336efc9..777a7ef780f 100644
--- a/doc/src/sgml/plpython.sgml
+++ b/doc/src/sgml/plpython.sgml
@@ -1047,6 +1047,14 @@ rv = plpy.execute(plan, ["name"], 5)
      </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">.
      </para>
@@ -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
index 0d78ca1de48..e54dca9e2ef 100644
--- a/src/pl/plpython/expected/plpython_spi.out
+++ b/src/pl/plpython/expected/plpython_spi.out
@@ -31,6 +31,19 @@ except Exception, 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
 	AS
 'if "myplan" not in SD:
@@ -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
index 7bb89921484..18e689f1415 100644
--- 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
index c73033c486b..ef23865dd2c 100644
--- 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
index 16c39a05ddf..390b4e90d45 100644
--- a/src/pl/plpython/plpy_planobject.c
+++ b/src/pl/plpython/plpy_planobject.c
@@ -10,11 +10,15 @@
 
 #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[] = {
@@ -22,6 +26,8 @@ static char PLy_plan_doc[] = {
 };
 
 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}
 };
@@ -112,6 +118,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)
 {
 	if (PyArg_ParseTuple(args, ":status"))
diff --git a/src/pl/plpython/plpy_spi.c b/src/pl/plpython/plpy_spi.c
index 07ab6a087e5..c6856ccbacc 100644
--- 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
index b0427947ef4..817a7584e79 100644
--- 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
index 7427de824b3..fcf049cb66d 100644
--- 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;