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 <andrew.dunstan@2ndquadrant.com>

diff --git a/doc/src/sgml/plpython.sgml b/doc/src/sgml/plpython.sgml
index fb5d336efc98d6d2bc08ebd1ca77a463457f688c..777a7ef780f9aa986606e9497403b6f9d43168ca 100644 (file)
--- 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
index 0d78ca1de481d103b1590079304e690dca292ba6..e54dca9e2ef00c85b4c08b032fdaff96c71b076c 100644 (file)
--- 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
index 7bb89921484906fa6667927e5cee52a8a2447cb2..18e689f1415accf417f87412d88398218d34abbf 100644 (file)
--- 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 c73033c486bbe60a314f7935ce6ddb234f830144..ef23865dd2cdda4321ae193457f9871528775051 100644 (file)
--- 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 16c39a05ddf0574a6b4c0e9894191b8f2f19e143..390b4e90d45bba9eb0a53bb8f45d63c4a28845e3 100644 (file)
--- 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}
 };
@@ -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
index 07ab6a087e531ccfaab13301d618cf0a8c8c7670..c6856ccbacc691f22aa5a84aeb1b98ea670a9b53 100644 (file)
--- 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 b0427947ef42541efeb5de33f3cc51f514e2c048..817a7584e79a29e0dafaffdaf880a5ac0de986de 100644 (file)
--- 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 7427de824b3c88a065a11c392ccde9499ccd7a95..fcf049cb66dc704da90798b4deb91a02aa6e6c9f 100644 (file)
--- 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;