plpy_cursorobject.c

Go to the documentation of this file.

/*
 * the PLyCursor class
 *
 * src/pl/plpython/plpy_cursorobject.c
 */
 
#include "postgres.h"
 
#include <limits.h>
 
#include "catalog/pg_type.h"
#include "mb/pg_wchar.h"
#include "plpy_cursorobject.h"
#include "plpy_elog.h"
#include "plpy_main.h"
#include "plpy_planobject.h"
#include "plpy_resultobject.h"
#include "plpy_spi.h"
#include "plpy_util.h"
#include "utils/memutils.h"
 
static PyObject *PLy_cursor_query(const char *query);
static void PLy_cursor_dealloc(PLyCursorObject *self);
static PyObject *PLy_cursor_iternext(PyObject *self);
static PyObject *PLy_cursor_fetch(PyObject *self, PyObject *args);
static PyObject *PLy_cursor_close(PyObject *self, PyObject *unused);
 
static const char PLy_cursor_doc[] = "Wrapper around a PostgreSQL cursor";
 
static PyMethodDef PLy_cursor_methods[] = {
    {"fetch", PLy_cursor_fetch, METH_VARARGS, NULL},
    {"close", PLy_cursor_close, METH_NOARGS, NULL},
    {NULL, NULL, 0, NULL}
};
 
static PyType_Slot PLyCursor_slots[] =
{
    {
        Py_tp_dealloc, PLy_cursor_dealloc
    },
    {
        Py_tp_doc, (char *) PLy_cursor_doc
    },
    {
        Py_tp_iter, PyObject_SelfIter
    },
    {
        Py_tp_iternext, PLy_cursor_iternext
    },
    {
        Py_tp_methods, PLy_cursor_methods
    },
    {
        0, NULL
    }
};
 
static PyType_Spec PLyCursor_spec =
{
    .name = "PLyCursor",
    .basicsize = sizeof(PLyCursorObject),
    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
    .slots = PLyCursor_slots,
};
 
static PyTypeObject *PLy_CursorType;
 
void
PLy_cursor_init_type(void)
{
    PLy_CursorType = (PyTypeObject *) PyType_FromSpec(&PLyCursor_spec);
    if (!PLy_CursorType)
        elog(ERROR, "could not initialize PLy_CursorType");
}
 
PyObject *
PLy_cursor(PyObject *self, PyObject *args)
{
    char       *query;
    PyObject   *plan;
    PyObject   *planargs = NULL;
 
    if (PyArg_ParseTuple(args, "s", &query))
        return PLy_cursor_query(query);
 
    PyErr_Clear();
 
    if (PyArg_ParseTuple(args, "O|O", &plan, &planargs))
        return PLy_cursor_plan(plan, planargs);
 
    PLy_exception_set(PLy_exc_error, "plpy.cursor expected a query or a plan");
    return NULL;
}
 
 
static PyObject *
PLy_cursor_query(const char *query)
{
    PLyCursorObject *cursor;
    PLyExecutionContext *exec_ctx = PLy_current_execution_context();
    volatile MemoryContext oldcontext;
    volatile ResourceOwner oldowner;
 
    if ((cursor = PyObject_New(PLyCursorObject, PLy_CursorType)) == NULL)
        return NULL;
#if PY_VERSION_HEX < 0x03080000
    /* Workaround for Python issue 35810; no longer necessary in Python 3.8 */
    Py_INCREF(PLy_CursorType);
#endif
    cursor->portalname = NULL;
    cursor->closed = false;
    cursor->mcxt = AllocSetContextCreate(TopMemoryContext,
                                         "PL/Python cursor context",
                                         ALLOCSET_DEFAULT_SIZES);
 
    /* Initialize for converting result tuples to Python */
    PLy_input_setup_func(&cursor->result, cursor->mcxt,
                         RECORDOID, -1,
                         exec_ctx->curr_proc);
 
    oldcontext = CurrentMemoryContext;
    oldowner = CurrentResourceOwner;
 
    PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
    PG_TRY();
    {
        SPIPlanPtr  plan;
        Portal      portal;
 
        pg_verifymbstr(query, strlen(query), false);
 
        plan = SPI_prepare(query, 0, NULL);
        if (plan == NULL)
            elog(ERROR, "SPI_prepare failed: %s",
                 SPI_result_code_string(SPI_result));
 
        portal = SPI_cursor_open(NULL, plan, NULL, NULL,
                                 exec_ctx->curr_proc->fn_readonly);
        SPI_freeplan(plan);
 
        if (portal == NULL)
            elog(ERROR, "SPI_cursor_open() failed: %s",
                 SPI_result_code_string(SPI_result));
 
        cursor->portalname = MemoryContextStrdup(cursor->mcxt, portal->name);
 
        PinPortal(portal);
 
        PLy_spi_subtransaction_commit(oldcontext, oldowner);
    }
    PG_CATCH();
    {
        PLy_spi_subtransaction_abort(oldcontext, oldowner);
        return NULL;
    }
    PG_END_TRY();
 
    Assert(cursor->portalname != NULL);
    return (PyObject *) cursor;
}
 
PyObject *
PLy_cursor_plan(PyObject *ob, PyObject *args)
{
    PLyCursorObject *cursor;
    volatile int nargs;
    PLyPlanObject *plan;
    PLyExecutionContext *exec_ctx = PLy_current_execution_context();
    volatile MemoryContext oldcontext;
    volatile ResourceOwner oldowner;
 
    if (args)
    {
        if (!PySequence_Check(args) || PyUnicode_Check(args))
        {
            PLy_exception_set(PyExc_TypeError, "plpy.cursor takes a sequence as its second argument");
            return NULL;
        }
        nargs = PySequence_Length(args);
    }
    else
        nargs = 0;
 
    plan = (PLyPlanObject *) ob;
 
    if (nargs != plan->nargs)
    {
        char       *sv;
        PyObject   *so = PyObject_Str(args);
 
        if (!so)
            PLy_elog(ERROR, "could not execute plan");
        sv = PLyUnicode_AsString(so);
        PLy_exception_set_plural(PyExc_TypeError,
                                 "Expected sequence of %d argument, got %d: %s",
                                 "Expected sequence of %d arguments, got %d: %s",
                                 plan->nargs,
                                 plan->nargs, nargs, sv);
        Py_DECREF(so);
 
        return NULL;
    }
 
    if ((cursor = PyObject_New(PLyCursorObject, PLy_CursorType)) == NULL)
        return NULL;
#if PY_VERSION_HEX < 0x03080000
    /* Workaround for Python issue 35810; no longer necessary in Python 3.8 */
    Py_INCREF(PLy_CursorType);
#endif
    cursor->portalname = NULL;
    cursor->closed = false;
    cursor->mcxt = AllocSetContextCreate(TopMemoryContext,
                                         "PL/Python cursor context",
                                         ALLOCSET_DEFAULT_SIZES);
 
    /* Initialize for converting result tuples to Python */
    PLy_input_setup_func(&cursor->result, cursor->mcxt,
                         RECORDOID, -1,
                         exec_ctx->curr_proc);
 
    oldcontext = CurrentMemoryContext;
    oldowner = CurrentResourceOwner;
 
    PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
    PG_TRY();
    {
        Portal      portal;
        MemoryContext tmpcontext;
        Datum      *volatile values;
        char       *volatile nulls;
        volatile int j;
 
        /*
         * Converted arguments and associated cruft will be in this context,
         * which is local to our subtransaction.
         */
        tmpcontext = AllocSetContextCreate(CurTransactionContext,
                                           "PL/Python temporary context",
                                           ALLOCSET_SMALL_SIZES);
        MemoryContextSwitchTo(tmpcontext);
 
        if (nargs > 0)
        {
            values = (Datum *) palloc(nargs * sizeof(Datum));
            nulls = (char *) palloc(nargs * sizeof(char));
        }
        else
        {
            values = NULL;
            nulls = NULL;
        }
 
        for (j = 0; j < nargs; j++)
        {
            PLyObToDatum *arg = &plan->args[j];
            PyObject   *elem;
 
            elem = PySequence_GetItem(args, j);
            PG_TRY(2);
            {
                bool        isnull;
 
                values[j] = PLy_output_convert(arg, elem, &isnull);
                nulls[j] = isnull ? 'n' : ' ';
            }
            PG_FINALLY(2);
            {
                Py_DECREF(elem);
            }
            PG_END_TRY(2);
        }
 
        MemoryContextSwitchTo(oldcontext);
 
        portal = SPI_cursor_open(NULL, plan->plan, values, nulls,
                                 exec_ctx->curr_proc->fn_readonly);
        if (portal == NULL)
            elog(ERROR, "SPI_cursor_open() failed: %s",
                 SPI_result_code_string(SPI_result));
 
        cursor->portalname = MemoryContextStrdup(cursor->mcxt, portal->name);
 
        PinPortal(portal);
 
        MemoryContextDelete(tmpcontext);
        PLy_spi_subtransaction_commit(oldcontext, oldowner);
    }
    PG_CATCH();
    {
        Py_DECREF(cursor);
        /* Subtransaction abort will remove the tmpcontext */
        PLy_spi_subtransaction_abort(oldcontext, oldowner);
        return NULL;
    }
    PG_END_TRY();
 
    Assert(cursor->portalname != NULL);
    return (PyObject *) cursor;
}
 
static void
PLy_cursor_dealloc(PLyCursorObject *self)
{
#if PY_VERSION_HEX >= 0x03080000
    PyTypeObject *tp = Py_TYPE(self);
#endif
    Portal      portal;
 
    if (!self->closed)
    {
        portal = GetPortalByName(self->portalname);
 
        if (PortalIsValid(portal))
        {
            UnpinPortal(portal);
            SPI_cursor_close(portal);
        }
        self->closed = true;
    }
    if (self->mcxt)
    {
        MemoryContextDelete(self->mcxt);
        self->mcxt = NULL;
    }
 
    PyObject_Free(self);
#if PY_VERSION_HEX >= 0x03080000
    /* This was not needed before Python 3.8 (Python issue 35810) */
    Py_DECREF(tp);
#endif
}
 
static PyObject *
PLy_cursor_iternext(PyObject *self)
{
    PLyCursorObject *cursor;
    PyObject   *ret;
    PLyExecutionContext *exec_ctx = PLy_current_execution_context();
    volatile MemoryContext oldcontext;
    volatile ResourceOwner oldowner;
    Portal      portal;
 
    cursor = (PLyCursorObject *) self;
 
    if (cursor->closed)
    {
        PLy_exception_set(PyExc_ValueError, "iterating a closed cursor");
        return NULL;
    }
 
    portal = GetPortalByName(cursor->portalname);
    if (!PortalIsValid(portal))
    {
        PLy_exception_set(PyExc_ValueError,
                          "iterating a cursor in an aborted subtransaction");
        return NULL;
    }
 
    oldcontext = CurrentMemoryContext;
    oldowner = CurrentResourceOwner;
 
    PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
    PG_TRY();
    {
        SPI_cursor_fetch(portal, true, 1);
        if (SPI_processed == 0)
        {
            PyErr_SetNone(PyExc_StopIteration);
            ret = NULL;
        }
        else
        {
            PLy_input_setup_tuple(&cursor->result, SPI_tuptable->tupdesc,
                                  exec_ctx->curr_proc);
 
            ret = PLy_input_from_tuple(&cursor->result, SPI_tuptable->vals[0],
                                       SPI_tuptable->tupdesc, true);
        }
 
        SPI_freetuptable(SPI_tuptable);
 
        PLy_spi_subtransaction_commit(oldcontext, oldowner);
    }
    PG_CATCH();
    {
        PLy_spi_subtransaction_abort(oldcontext, oldowner);
        return NULL;
    }
    PG_END_TRY();
 
    return ret;
}
 
static PyObject *
PLy_cursor_fetch(PyObject *self, PyObject *args)
{
    PLyCursorObject *cursor;
    int         count;
    PLyResultObject *ret;
    PLyExecutionContext *exec_ctx = PLy_current_execution_context();
    volatile MemoryContext oldcontext;
    volatile ResourceOwner oldowner;
    Portal      portal;
 
    if (!PyArg_ParseTuple(args, "i:fetch", &count))
        return NULL;
 
    cursor = (PLyCursorObject *) self;
 
    if (cursor->closed)
    {
        PLy_exception_set(PyExc_ValueError, "fetch from a closed cursor");
        return NULL;
    }
 
    portal = GetPortalByName(cursor->portalname);
    if (!PortalIsValid(portal))
    {
        PLy_exception_set(PyExc_ValueError,
                          "iterating a cursor in an aborted subtransaction");
        return NULL;
    }
 
    ret = (PLyResultObject *) PLy_result_new();
    if (ret == NULL)
        return NULL;
 
    oldcontext = CurrentMemoryContext;
    oldowner = CurrentResourceOwner;
 
    PLy_spi_subtransaction_begin(oldcontext, oldowner);
 
    PG_TRY();
    {
        SPI_cursor_fetch(portal, true, count);
 
        Py_DECREF(ret->status);
        ret->status = PyLong_FromLong(SPI_OK_FETCH);
 
        Py_DECREF(ret->nrows);
        ret->nrows = PyLong_FromUnsignedLongLong(SPI_processed);
 
        if (SPI_processed != 0)
        {
            uint64      i;
 
            /*
             * PyList_New() and PyList_SetItem() use Py_ssize_t for list size
             * and list indices; so we cannot support a result larger than
             * PY_SSIZE_T_MAX.
             */
            if (SPI_processed > (uint64) PY_SSIZE_T_MAX)
                ereport(ERROR,
                        (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
                         errmsg("query result has too many rows to fit in a Python list")));
 
            Py_DECREF(ret->rows);
            ret->rows = PyList_New(SPI_processed);
            if (!ret->rows)
            {
                Py_DECREF(ret);
                ret = NULL;
            }
            else
            {
                PLy_input_setup_tuple(&cursor->result, SPI_tuptable->tupdesc,
                                      exec_ctx->curr_proc);
 
                for (i = 0; i < SPI_processed; i++)
                {
                    PyObject   *row = PLy_input_from_tuple(&cursor->result,
                                                           SPI_tuptable->vals[i],
                                                           SPI_tuptable->tupdesc,
                                                           true);
 
                    PyList_SetItem(ret->rows, i, row);
                }
            }
        }
 
        SPI_freetuptable(SPI_tuptable);
 
        PLy_spi_subtransaction_commit(oldcontext, oldowner);
    }
    PG_CATCH();
    {
        PLy_spi_subtransaction_abort(oldcontext, oldowner);
        return NULL;
    }
    PG_END_TRY();
 
    return (PyObject *) ret;
}
 
static PyObject *
PLy_cursor_close(PyObject *self, PyObject *unused)
{
    PLyCursorObject *cursor = (PLyCursorObject *) self;
 
    if (!cursor->closed)
    {
        Portal      portal = GetPortalByName(cursor->portalname);
 
        if (!PortalIsValid(portal))
        {
            PLy_exception_set(PyExc_ValueError,
                              "closing a cursor in an aborted subtransaction");
            return NULL;
        }
 
        UnpinPortal(portal);
        SPI_cursor_close(portal);
        cursor->closed = true;
    }
 
    Py_RETURN_NONE;
}