plpython improvements:
authorBruce Momjian <bruce@momjian.us>
Thu, 27 Apr 2006 01:05:05 +0000 (01:05 +0000)
committerBruce Momjian <bruce@momjian.us>
Thu, 27 Apr 2006 01:05:05 +0000 (01:05 +0000)
1) named parameters additionally to args[]
2) return composite-types from plpython as dictionary
3) return result-set from plpython as list, iterator or generator

Hannu Krosing
Sven Suursoho

src/pl/plpython/plpython.c

index bf96db1776df51532ef5084e13c78cc48a5b884d..a722fce917bace91b3b796339a35793bd7c67eae 100644 (file)
@@ -1,7 +1,7 @@
 /**********************************************************************
  * plpython.c - python as a procedural language for PostgreSQL
  *
- * $PostgreSQL: pgsql/src/pl/plpython/plpython.c,v 1.77 2006/04/04 19:35:37 tgl Exp $
+ * $PostgreSQL: pgsql/src/pl/plpython/plpython.c,v 1.78 2006/04/27 01:05:05 momjian Exp $
  *
  *********************************************************************
  */
@@ -19,6 +19,7 @@
 #include "catalog/pg_type.h"
 #include "commands/trigger.h"
 #include "executor/spi.h"
+#include "funcapi.h"
 #include "fmgr.h"
 #include "nodes/makefuncs.h"
 #include "parser/parse_type.h"
@@ -108,6 +109,11 @@ typedef struct PLyProcedure
    bool        fn_readonly;
    PLyTypeInfo result;         /* also used to store info for trigger tuple
                                 * type */
+   bool        is_setof;       /* true, if procedure returns result set */
+   PyObject    *setof;     /* contents of result set. */
+   int     setof_count;    /* numbef of items to return in result set */
+   int     setof_current;  /* current item in result set */
+   char        **argnames;     /* Argument names */
    PLyTypeInfo args[FUNC_MAX_ARGS];
    int         nargs;
    PyObject   *code;           /* compiled procedure code */
@@ -184,6 +190,7 @@ static Datum PLy_function_handler(FunctionCallInfo fcinfo, PLyProcedure *);
 static HeapTuple PLy_trigger_handler(FunctionCallInfo fcinfo, PLyProcedure *);
 
 static PyObject *PLy_function_build_args(FunctionCallInfo fcinfo, PLyProcedure *);
+static void PLy_function_delete_args(PLyProcedure *);
 static PyObject *PLy_trigger_build_args(FunctionCallInfo fcinfo, PLyProcedure *,
                       HeapTuple *);
 static HeapTuple PLy_modify_tuple(PLyProcedure *, PyObject *,
@@ -218,6 +225,7 @@ static PyObject *PLyFloat_FromString(const char *);
 static PyObject *PLyInt_FromString(const char *);
 static PyObject *PLyLong_FromString(const char *);
 static PyObject *PLyString_FromString(const char *);
+static HeapTuple PLyDict_ToTuple(PLyTypeInfo *, PyObject *);
 
 
 /* global data */
@@ -726,11 +734,17 @@ PLy_function_handler(FunctionCallInfo fcinfo, PLyProcedure * proc)
 
    PG_TRY();
    {
-       plargs = PLy_function_build_args(fcinfo, proc);
-       plrv = PLy_procedure_call(proc, "args", plargs);
-
-       Assert(plrv != NULL);
-       Assert(!PLy_error_in_progress);
+       if (!proc->is_setof || proc->setof_count == -1)
+       {
+           /* python function not called yet, do it */
+           plargs = PLy_function_build_args(fcinfo, proc);
+           plrv = PLy_procedure_call(proc, "args", plargs);
+           if (!proc->is_setof)
+               /* SETOF function parameters are deleted when called last row is returned */
+               PLy_function_delete_args(proc);
+           Assert(plrv != NULL);
+           Assert(!PLy_error_in_progress);
+       }
 
        /*
         * Disconnect from SPI manager and then create the return values datum
@@ -741,6 +755,76 @@ PLy_function_handler(FunctionCallInfo fcinfo, PLyProcedure * proc)
        if (SPI_finish() != SPI_OK_FINISH)
            elog(ERROR, "SPI_finish failed");
 
+       if (proc->is_setof)
+       {
+           bool is_done = false;
+           ReturnSetInfo *rsi = (ReturnSetInfo *)fcinfo->resultinfo;
+
+           if (proc->setof_current == -1)
+           {
+               /* first time -- do checks and setup */
+               if (!rsi || !IsA(rsi, ReturnSetInfo) ||
+                       (rsi->allowedModes & SFRM_ValuePerCall) == 0)
+               {
+                   ereport(ERROR,
+                           (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+                            errmsg("only value per call is allowed")));
+               }
+               rsi->returnMode = SFRM_ValuePerCall;
+
+               /* fetch information about returned object */
+               proc->setof = plrv;
+               plrv = NULL;
+               if (PyList_Check(proc->setof))
+                   /* SETOF as list */
+                   proc->setof_count = PyList_GET_SIZE(proc->setof);
+               else if (PyIter_Check(proc->setof))
+                   /* SETOF as iterator, unknown number of items */
+                   proc->setof_current = proc->setof_count = 0;
+               else
+               {
+                   ereport(ERROR,
+                           (errcode(ERRCODE_INVALID_BINARY_REPRESENTATION),
+                            errmsg("SETOF must be returned as list or iterator")));
+               }
+           }
+
+           Assert(proc->setof != NULL);
+
+           /* Fetch next of SETOF */
+           if (PyList_Check(proc->setof))
+           {
+               is_done = ++proc->setof_current == proc->setof_count;
+               if (!is_done)
+                   plrv = PyList_GET_ITEM(proc->setof, proc->setof_current);
+           }
+           else if (PyIter_Check(proc->setof))
+           {
+               plrv = PyIter_Next(proc->setof);
+               is_done = plrv == NULL;
+           }
+
+           if (!is_done)
+           {
+               rsi->isDone = ExprMultipleResult;
+           }
+           else
+           {
+               rsi->isDone = ExprEndResult;
+               proc->setof_count = proc->setof_current = -1;
+               Py_DECREF(proc->setof);
+               proc->setof = NULL;
+
+               Py_XDECREF(plargs);
+               Py_XDECREF(plrv);
+               Py_XDECREF(plrv_so);
+
+               PLy_function_delete_args(proc);
+               fcinfo->isnull = true;
+               return (Datum)NULL;
+           }
+       }
+
        /*
         * If the function is declared to return void, the Python
         * return value must be None. For void-returning functions, we
@@ -767,6 +851,26 @@ PLy_function_handler(FunctionCallInfo fcinfo, PLyProcedure * proc)
                                   proc->result.out.d.typioparam,
                                   -1);
        }
+       else if (proc->result.is_rowtype >= 1)
+       {
+           HeapTuple   tuple;
+
+           /* returning composite type */
+           if (!PyDict_Check(plrv))
+               elog(ERROR, "tuple must be returned as dictionary");
+
+           tuple = PLyDict_ToTuple(&proc->result, plrv);
+           if (tuple != NULL)
+           {
+               fcinfo->isnull = false;
+               rv = HeapTupleGetDatum(tuple);
+           }
+           else
+           {
+               fcinfo->isnull = true;
+               rv = (Datum) NULL;
+           }
+       }
        else
        {
            fcinfo->isnull = false;
@@ -893,6 +997,7 @@ PLy_function_build_args(FunctionCallInfo fcinfo, PLyProcedure * proc)
             * FIXME -- error check this
             */
            PyList_SetItem(args, i, arg);
+           PyDict_SetItemString(proc->globals, proc->argnames[i], arg);
            arg = NULL;
        }
    }
@@ -909,6 +1014,16 @@ PLy_function_build_args(FunctionCallInfo fcinfo, PLyProcedure * proc)
 }
 
 
+static void
+PLy_function_delete_args(PLyProcedure *proc)
+{
+   int i;
+
+   for (i = 0; i < proc->nargs; i++)
+       PyDict_DelItemString(proc->globals, proc->argnames[i]);
+}
+
+
 /*
  * PLyProcedure functions
  */
@@ -979,6 +1094,9 @@ PLy_procedure_create(FunctionCallInfo fcinfo, Oid tgreloid,
    bool        isnull;
    int         i,
                rv;
+   Datum       argnames;
+   Datum           *elems;
+   int     nelems;
 
    procStruct = (Form_pg_proc) GETSTRUCT(procTup);
 
@@ -1010,6 +1128,10 @@ PLy_procedure_create(FunctionCallInfo fcinfo, Oid tgreloid,
    proc->nargs = 0;
    proc->code = proc->statics = NULL;
    proc->globals = proc->me = NULL;
+   proc->is_setof = procStruct->proretset;
+   proc->setof = NULL;
+   proc->setof_count = proc->setof_current = -1;
+   proc->argnames = NULL;
 
    PG_TRY();
    {
@@ -1046,9 +1168,11 @@ PLy_procedure_create(FunctionCallInfo fcinfo, Oid tgreloid,
            }
 
            if (rvTypeStruct->typtype == 'c')
-               ereport(ERROR,
-                       (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-                    errmsg("plpython functions cannot return tuples yet")));
+           {
+               /* Tuple: set up later, during first call to PLy_function_handler */
+               proc->result.out.d.typoid = procStruct->prorettype;
+               proc->result.is_rowtype = 2;
+           }
            else
                PLy_output_datum_func(&proc->result, rvTypeTup);
 
@@ -1071,6 +1195,21 @@ PLy_procedure_create(FunctionCallInfo fcinfo, Oid tgreloid,
         * arguments.
         */
        proc->nargs = fcinfo->nargs;
+       proc->argnames = NULL;
+       if (proc->nargs)
+       {
+           argnames = SysCacheGetAttr(PROCOID, procTup, Anum_pg_proc_proargnames, &isnull);
+           if (!isnull)
+           {
+               deconstruct_array(DatumGetArrayTypeP(argnames), TEXTOID, -1, false, 'i',
+                       &elems, NULL, &nelems);
+               if (nelems != proc->nargs)
+                   elog(ERROR,
+                           "proargnames must have the same number of elements "
+                           "as the function has arguments");
+               proc->argnames = (char **) PLy_malloc(sizeof(char *)*proc->nargs);
+           }
+       }
        for (i = 0; i < fcinfo->nargs; i++)
        {
            HeapTuple   argTypeTup;
@@ -1099,8 +1238,11 @@ PLy_procedure_create(FunctionCallInfo fcinfo, Oid tgreloid,
                proc->args[i].is_rowtype = 2;   /* still need to set I/O funcs */
 
            ReleaseSysCache(argTypeTup);
-       }
 
+           /* Fetch argument name */
+           if (proc->argnames)
+               proc->argnames[i] = PLy_strdup(DatumGetCString(DirectFunctionCall1(textout, elems[i])));
+       }
 
        /*
         * get the text of the function.
@@ -1236,6 +1378,7 @@ PLy_procedure_delete(PLyProcedure * proc)
    if (proc->pyname)
        PLy_free(proc->pyname);
    for (i = 0; i < proc->nargs; i++)
+   {
        if (proc->args[i].is_rowtype == 1)
        {
            if (proc->args[i].in.r.atts)
@@ -1243,6 +1386,11 @@ PLy_procedure_delete(PLyProcedure * proc)
            if (proc->args[i].out.r.atts)
                PLy_free(proc->args[i].out.r.atts);
        }
+       if (proc->argnames && proc->argnames[i])
+           PLy_free(proc->argnames[i]);
+   }
+   if (proc->argnames)
+       PLy_free(proc->argnames);
 }
 
 /* conversion functions.  remember output from python is
@@ -1501,6 +1649,78 @@ PLyDict_FromTuple(PLyTypeInfo * info, HeapTuple tuple, TupleDesc desc)
    return dict;
 }
 
+
+static HeapTuple
+PLyDict_ToTuple(PLyTypeInfo *info, PyObject *dict)
+{
+   TupleDesc   desc;
+   HeapTuple   tuple;
+   Datum       *values;
+   char        *nulls;
+   int     i;
+
+   desc = CreateTupleDescCopy(lookup_rowtype_tupdesc(info->out.d.typoid, -1));
+
+   /* Set up tuple type, if neccessary */
+   if (info->is_rowtype == 2)
+   {
+       PLy_output_tuple_funcs(info, desc);
+       info->is_rowtype = 1;
+   }
+   Assert(info->is_rowtype == 1);
+
+   /* Build tuple */
+   values = palloc(sizeof(Datum)*desc->natts);
+   nulls = palloc(sizeof(char)*desc->natts);
+   for (i = 0;  i < desc->natts;  ++i)
+   {
+       char        *key;
+       PyObject    *value,
+               *so;
+
+       key = NameStr(desc->attrs[i]->attname);
+       value = so = NULL;
+       PG_TRY();
+       {
+           value = PyDict_GetItemString(dict, key);
+           if (value != Py_None && value != NULL)
+           {
+               char *valuestr;
+
+               so = PyObject_Str(value);
+               valuestr = PyString_AsString(so);
+               values[i] = InputFunctionCall(&info->out.r.atts[i].typfunc
+                       , valuestr
+                       , info->out.r.atts[i].typioparam
+                       , -1);
+               Py_DECREF(so);
+               value = so = NULL;
+               nulls[i] = ' ';
+           }
+           else
+           {
+               value = NULL;
+               values[i] = (Datum) NULL;
+               nulls[i] = 'n';
+           }
+       }
+       PG_CATCH();
+       {
+           Py_XDECREF(value);
+           Py_XDECREF(so);
+           PG_RE_THROW();
+       }
+       PG_END_TRY();
+   }
+
+   tuple = heap_formtuple(desc, values, nulls);
+   FreeTupleDesc(desc);
+   pfree(values);
+   pfree(nulls);
+
+   return tuple;
+}
+
 /* initialization, some python variables function declared here */
 
 /* interface to postgresql elog */
@@ -2644,3 +2864,4 @@ PLy_free(void *ptr)
 {
    free(ptr);
 }
+/* vim: set noexpandtab nosmarttab shiftwidth=8 cinoptions=l1j1: */