LCOV - code coverage report
Current view: top level - src/pl/plpython - plpy_spi.c (source / functions) Hit Total Coverage
Test: PostgreSQL 13devel Lines: 190 230 82.6 %
Date: 2019-11-13 23:06:49 Functions: 9 9 100.0 %
Legend: Lines: hit not hit

          Line data    Source code
       1             : /*
       2             :  * interface to SPI functions
       3             :  *
       4             :  * src/pl/plpython/plpy_spi.c
       5             :  */
       6             : 
       7             : #include "postgres.h"
       8             : 
       9             : #include <limits.h>
      10             : 
      11             : #include "access/htup_details.h"
      12             : #include "access/xact.h"
      13             : #include "catalog/pg_type.h"
      14             : #include "executor/spi.h"
      15             : #include "mb/pg_wchar.h"
      16             : #include "parser/parse_type.h"
      17             : #include "utils/memutils.h"
      18             : #include "utils/syscache.h"
      19             : 
      20             : #include "plpython.h"
      21             : 
      22             : #include "plpy_spi.h"
      23             : 
      24             : #include "plpy_elog.h"
      25             : #include "plpy_main.h"
      26             : #include "plpy_planobject.h"
      27             : #include "plpy_plpymodule.h"
      28             : #include "plpy_procedure.h"
      29             : #include "plpy_resultobject.h"
      30             : 
      31             : 
      32             : static PyObject *PLy_spi_execute_query(char *query, long limit);
      33             : static PyObject *PLy_spi_execute_fetch_result(SPITupleTable *tuptable,
      34             :                                               uint64 rows, int status);
      35             : static void PLy_spi_exception_set(PyObject *excclass, ErrorData *edata);
      36             : 
      37             : 
      38             : /* prepare(query="select * from foo")
      39             :  * prepare(query="select * from foo where bar = $1", params=["text"])
      40             :  * prepare(query="select * from foo where bar = $1", params=["text"], limit=5)
      41             :  */
      42             : PyObject *
      43          50 : PLy_spi_prepare(PyObject *self, PyObject *args)
      44             : {
      45             :     PLyPlanObject *plan;
      46          50 :     PyObject   *list = NULL;
      47          50 :     PyObject   *volatile optr = NULL;
      48             :     char       *query;
      49          50 :     PLyExecutionContext *exec_ctx = PLy_current_execution_context();
      50             :     volatile MemoryContext oldcontext;
      51             :     volatile ResourceOwner oldowner;
      52             :     volatile int nargs;
      53             : 
      54          50 :     if (!PyArg_ParseTuple(args, "s|O:prepare", &query, &list))
      55           0 :         return NULL;
      56             : 
      57          50 :     if (list && (!PySequence_Check(list)))
      58             :     {
      59           0 :         PLy_exception_set(PyExc_TypeError,
      60             :                           "second argument of plpy.prepare must be a sequence");
      61           0 :         return NULL;
      62             :     }
      63             : 
      64          50 :     if ((plan = (PLyPlanObject *) PLy_plan_new()) == NULL)
      65           0 :         return NULL;
      66             : 
      67          50 :     plan->mcxt = AllocSetContextCreate(TopMemoryContext,
      68             :                                        "PL/Python plan context",
      69             :                                        ALLOCSET_DEFAULT_SIZES);
      70          50 :     oldcontext = MemoryContextSwitchTo(plan->mcxt);
      71             : 
      72          50 :     nargs = list ? PySequence_Length(list) : 0;
      73             : 
      74          50 :     plan->nargs = nargs;
      75          50 :     plan->types = nargs ? palloc0(sizeof(Oid) * nargs) : NULL;
      76          50 :     plan->values = nargs ? palloc0(sizeof(Datum) * nargs) : NULL;
      77          50 :     plan->args = nargs ? palloc0(sizeof(PLyObToDatum) * nargs) : NULL;
      78             : 
      79          50 :     MemoryContextSwitchTo(oldcontext);
      80             : 
      81          50 :     oldcontext = CurrentMemoryContext;
      82          50 :     oldowner = CurrentResourceOwner;
      83             : 
      84          50 :     PLy_spi_subtransaction_begin(oldcontext, oldowner);
      85             : 
      86          50 :     PG_TRY();
      87             :     {
      88             :         int         i;
      89             : 
      90          78 :         for (i = 0; i < nargs; i++)
      91             :         {
      92             :             char       *sptr;
      93             :             Oid         typeId;
      94             :             int32       typmod;
      95             : 
      96          34 :             optr = PySequence_GetItem(list, i);
      97          34 :             if (PyString_Check(optr))
      98          32 :                 sptr = PyString_AsString(optr);
      99           2 :             else if (PyUnicode_Check(optr))
     100           2 :                 sptr = PLyUnicode_AsString(optr);
     101             :             else
     102             :             {
     103           0 :                 ereport(ERROR,
     104             :                         (errmsg("plpy.prepare: type name at ordinal position %d is not a string", i)));
     105             :                 sptr = NULL;    /* keep compiler quiet */
     106             :             }
     107             : 
     108             :             /********************************************************
     109             :              * Resolve argument type names and then look them up by
     110             :              * oid in the system cache, and remember the required
     111             :              *information for input conversion.
     112             :              ********************************************************/
     113             : 
     114          34 :             parseTypeString(sptr, &typeId, &typmod, false);
     115             : 
     116          28 :             Py_DECREF(optr);
     117             : 
     118             :             /*
     119             :              * set optr to NULL, so we won't try to unref it again in case of
     120             :              * an error
     121             :              */
     122          28 :             optr = NULL;
     123             : 
     124          28 :             plan->types[i] = typeId;
     125          28 :             PLy_output_setup_func(&plan->args[i], plan->mcxt,
     126             :                                   typeId, typmod,
     127          28 :                                   exec_ctx->curr_proc);
     128             :         }
     129             : 
     130          44 :         pg_verifymbstr(query, strlen(query), false);
     131          44 :         plan->plan = SPI_prepare(query, plan->nargs, plan->types);
     132          44 :         if (plan->plan == NULL)
     133           0 :             elog(ERROR, "SPI_prepare failed: %s",
     134             :                  SPI_result_code_string(SPI_result));
     135             : 
     136             :         /* transfer plan from procCxt to topCxt */
     137          44 :         if (SPI_keepplan(plan->plan))
     138           0 :             elog(ERROR, "SPI_keepplan failed");
     139             : 
     140          44 :         PLy_spi_subtransaction_commit(oldcontext, oldowner);
     141             :     }
     142           6 :     PG_CATCH();
     143             :     {
     144           6 :         Py_DECREF(plan);
     145           6 :         Py_XDECREF(optr);
     146             : 
     147           6 :         PLy_spi_subtransaction_abort(oldcontext, oldowner);
     148           6 :         return NULL;
     149             :     }
     150          44 :     PG_END_TRY();
     151             : 
     152             :     Assert(plan->plan != NULL);
     153          44 :     return (PyObject *) plan;
     154             : }
     155             : 
     156             : /* execute(query="select * from foo", limit=5)
     157             :  * execute(plan=plan, values=(foo, bar), limit=5)
     158             :  */
     159             : PyObject *
     160         330 : PLy_spi_execute(PyObject *self, PyObject *args)
     161             : {
     162             :     char       *query;
     163             :     PyObject   *plan;
     164         330 :     PyObject   *list = NULL;
     165         330 :     long        limit = 0;
     166             : 
     167         330 :     if (PyArg_ParseTuple(args, "s|l", &query, &limit))
     168         288 :         return PLy_spi_execute_query(query, limit);
     169             : 
     170          42 :     PyErr_Clear();
     171             : 
     172          84 :     if (PyArg_ParseTuple(args, "O|Ol", &plan, &list, &limit) &&
     173          42 :         is_PLyPlanObject(plan))
     174          42 :         return PLy_spi_execute_plan(plan, list, limit);
     175             : 
     176           0 :     PLy_exception_set(PLy_exc_error, "plpy.execute expected a query or a plan");
     177           0 :     return NULL;
     178             : }
     179             : 
     180             : PyObject *
     181          44 : PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
     182             : {
     183             :     volatile int nargs;
     184             :     int         i,
     185             :                 rv;
     186             :     PLyPlanObject *plan;
     187             :     volatile MemoryContext oldcontext;
     188             :     volatile ResourceOwner oldowner;
     189             :     PyObject   *ret;
     190             : 
     191          44 :     if (list != NULL)
     192             :     {
     193          32 :         if (!PySequence_Check(list) || PyString_Check(list) || PyUnicode_Check(list))
     194             :         {
     195           0 :             PLy_exception_set(PyExc_TypeError, "plpy.execute takes a sequence as its second argument");
     196           0 :             return NULL;
     197             :         }
     198          32 :         nargs = PySequence_Length(list);
     199             :     }
     200             :     else
     201          12 :         nargs = 0;
     202             : 
     203          44 :     plan = (PLyPlanObject *) ob;
     204             : 
     205          44 :     if (nargs != plan->nargs)
     206             :     {
     207             :         char       *sv;
     208           0 :         PyObject   *so = PyObject_Str(list);
     209             : 
     210           0 :         if (!so)
     211           0 :             PLy_elog(ERROR, "could not execute plan");
     212           0 :         sv = PyString_AsString(so);
     213           0 :         PLy_exception_set_plural(PyExc_TypeError,
     214             :                                  "Expected sequence of %d argument, got %d: %s",
     215             :                                  "Expected sequence of %d arguments, got %d: %s",
     216           0 :                                  plan->nargs,
     217             :                                  plan->nargs, nargs, sv);
     218           0 :         Py_DECREF(so);
     219             : 
     220           0 :         return NULL;
     221             :     }
     222             : 
     223          44 :     oldcontext = CurrentMemoryContext;
     224          44 :     oldowner = CurrentResourceOwner;
     225             : 
     226          44 :     PLy_spi_subtransaction_begin(oldcontext, oldowner);
     227             : 
     228          44 :     PG_TRY();
     229             :     {
     230          44 :         PLyExecutionContext *exec_ctx = PLy_current_execution_context();
     231             :         char       *volatile nulls;
     232             :         volatile int j;
     233             : 
     234          44 :         if (nargs > 0)
     235          30 :             nulls = palloc(nargs * sizeof(char));
     236             :         else
     237          14 :             nulls = NULL;
     238             : 
     239          72 :         for (j = 0; j < nargs; j++)
     240             :         {
     241          32 :             PLyObToDatum *arg = &plan->args[j];
     242             :             PyObject   *elem;
     243             : 
     244          32 :             elem = PySequence_GetItem(list, j);
     245          32 :             PG_TRY();
     246             :             {
     247             :                 bool        isnull;
     248             : 
     249          32 :                 plan->values[j] = PLy_output_convert(arg, elem, &isnull);
     250          28 :                 nulls[j] = isnull ? 'n' : ' ';
     251             :             }
     252           4 :             PG_FINALLY();
     253             :             {
     254          32 :                 Py_DECREF(elem);
     255             :             }
     256          32 :             PG_END_TRY();
     257             :         }
     258             : 
     259          40 :         rv = SPI_execute_plan(plan->plan, plan->values, nulls,
     260          40 :                               exec_ctx->curr_proc->fn_readonly, limit);
     261          40 :         ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv);
     262             : 
     263          40 :         if (nargs > 0)
     264          26 :             pfree(nulls);
     265             : 
     266          40 :         PLy_spi_subtransaction_commit(oldcontext, oldowner);
     267             :     }
     268           4 :     PG_CATCH();
     269             :     {
     270             :         int         k;
     271             : 
     272             :         /*
     273             :          * cleanup plan->values array
     274             :          */
     275           8 :         for (k = 0; k < nargs; k++)
     276             :         {
     277           4 :             if (!plan->args[k].typbyval &&
     278           0 :                 (plan->values[k] != PointerGetDatum(NULL)))
     279             :             {
     280           0 :                 pfree(DatumGetPointer(plan->values[k]));
     281           0 :                 plan->values[k] = PointerGetDatum(NULL);
     282             :             }
     283             :         }
     284             : 
     285           4 :         PLy_spi_subtransaction_abort(oldcontext, oldowner);
     286           4 :         return NULL;
     287             :     }
     288          40 :     PG_END_TRY();
     289             : 
     290          68 :     for (i = 0; i < nargs; i++)
     291             :     {
     292          48 :         if (!plan->args[i].typbyval &&
     293          20 :             (plan->values[i] != PointerGetDatum(NULL)))
     294             :         {
     295          20 :             pfree(DatumGetPointer(plan->values[i]));
     296          20 :             plan->values[i] = PointerGetDatum(NULL);
     297             :         }
     298             :     }
     299             : 
     300          40 :     if (rv < 0)
     301             :     {
     302           2 :         PLy_exception_set(PLy_exc_spi_error,
     303             :                           "SPI_execute_plan failed: %s",
     304             :                           SPI_result_code_string(rv));
     305           2 :         return NULL;
     306             :     }
     307             : 
     308          38 :     return ret;
     309             : }
     310             : 
     311             : static PyObject *
     312         288 : PLy_spi_execute_query(char *query, long limit)
     313             : {
     314             :     int         rv;
     315             :     volatile MemoryContext oldcontext;
     316             :     volatile ResourceOwner oldowner;
     317         288 :     PyObject   *ret = NULL;
     318             : 
     319         288 :     oldcontext = CurrentMemoryContext;
     320         288 :     oldowner = CurrentResourceOwner;
     321             : 
     322         288 :     PLy_spi_subtransaction_begin(oldcontext, oldowner);
     323             : 
     324         288 :     PG_TRY();
     325             :     {
     326         288 :         PLyExecutionContext *exec_ctx = PLy_current_execution_context();
     327             : 
     328         288 :         pg_verifymbstr(query, strlen(query), false);
     329         288 :         rv = SPI_execute(query, exec_ctx->curr_proc->fn_readonly, limit);
     330         244 :         ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv);
     331             : 
     332         244 :         PLy_spi_subtransaction_commit(oldcontext, oldowner);
     333             :     }
     334          44 :     PG_CATCH();
     335             :     {
     336          44 :         PLy_spi_subtransaction_abort(oldcontext, oldowner);
     337          44 :         return NULL;
     338             :     }
     339         244 :     PG_END_TRY();
     340             : 
     341         244 :     if (rv < 0)
     342             :     {
     343           2 :         Py_XDECREF(ret);
     344           2 :         PLy_exception_set(PLy_exc_spi_error,
     345             :                           "SPI_execute failed: %s",
     346             :                           SPI_result_code_string(rv));
     347           2 :         return NULL;
     348             :     }
     349             : 
     350         242 :     return ret;
     351             : }
     352             : 
     353             : static PyObject *
     354         284 : PLy_spi_execute_fetch_result(SPITupleTable *tuptable, uint64 rows, int status)
     355             : {
     356             :     PLyResultObject *result;
     357         284 :     PLyExecutionContext *exec_ctx = PLy_current_execution_context();
     358             :     volatile MemoryContext oldcontext;
     359             : 
     360         284 :     result = (PLyResultObject *) PLy_result_new();
     361         284 :     if (!result)
     362             :     {
     363           0 :         SPI_freetuptable(tuptable);
     364           0 :         return NULL;
     365             :     }
     366         284 :     Py_DECREF(result->status);
     367         284 :     result->status = PyInt_FromLong(status);
     368             : 
     369         284 :     if (status > 0 && tuptable == NULL)
     370             :     {
     371         162 :         Py_DECREF(result->nrows);
     372         162 :         result->nrows = PyLong_FromUnsignedLongLong(rows);
     373             :     }
     374         122 :     else if (status > 0 && tuptable != NULL)
     375             :     {
     376             :         PLyDatumToOb ininfo;
     377             :         MemoryContext cxt;
     378             : 
     379         118 :         Py_DECREF(result->nrows);
     380         118 :         result->nrows = PyLong_FromUnsignedLongLong(rows);
     381             : 
     382         118 :         cxt = AllocSetContextCreate(CurrentMemoryContext,
     383             :                                     "PL/Python temp context",
     384             :                                     ALLOCSET_DEFAULT_SIZES);
     385             : 
     386             :         /* Initialize for converting result tuples to Python */
     387         118 :         PLy_input_setup_func(&ininfo, cxt, RECORDOID, -1,
     388         118 :                              exec_ctx->curr_proc);
     389             : 
     390         118 :         oldcontext = CurrentMemoryContext;
     391         118 :         PG_TRY();
     392             :         {
     393             :             MemoryContext oldcontext2;
     394             : 
     395         118 :             if (rows)
     396             :             {
     397             :                 uint64      i;
     398             : 
     399             :                 /*
     400             :                  * PyList_New() and PyList_SetItem() use Py_ssize_t for list
     401             :                  * size and list indices; so we cannot support a result larger
     402             :                  * than PY_SSIZE_T_MAX.
     403             :                  */
     404         114 :                 if (rows > (uint64) PY_SSIZE_T_MAX)
     405           0 :                     ereport(ERROR,
     406             :                             (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
     407             :                              errmsg("query result has too many rows to fit in a Python list")));
     408             : 
     409         114 :                 Py_DECREF(result->rows);
     410         114 :                 result->rows = PyList_New(rows);
     411         114 :                 if (result->rows)
     412             :                 {
     413         114 :                     PLy_input_setup_tuple(&ininfo, tuptable->tupdesc,
     414         114 :                                           exec_ctx->curr_proc);
     415             : 
     416         262 :                     for (i = 0; i < rows; i++)
     417             :                     {
     418         296 :                         PyObject   *row = PLy_input_from_tuple(&ininfo,
     419         148 :                                                                tuptable->vals[i],
     420             :                                                                tuptable->tupdesc,
     421             :                                                                true);
     422             : 
     423         148 :                         PyList_SetItem(result->rows, i, row);
     424             :                     }
     425             :                 }
     426             :             }
     427             : 
     428             :             /*
     429             :              * Save tuple descriptor for later use by result set metadata
     430             :              * functions.  Save it in TopMemoryContext so that it survives
     431             :              * outside of an SPI context.  We trust that PLy_result_dealloc()
     432             :              * will clean it up when the time is right.  (Do this as late as
     433             :              * possible, to minimize the number of ways the tupdesc could get
     434             :              * leaked due to errors.)
     435             :              */
     436         118 :             oldcontext2 = MemoryContextSwitchTo(TopMemoryContext);
     437         118 :             result->tupdesc = CreateTupleDescCopy(tuptable->tupdesc);
     438         118 :             MemoryContextSwitchTo(oldcontext2);
     439             :         }
     440           0 :         PG_CATCH();
     441             :         {
     442           0 :             MemoryContextSwitchTo(oldcontext);
     443           0 :             MemoryContextDelete(cxt);
     444           0 :             Py_DECREF(result);
     445           0 :             PG_RE_THROW();
     446             :         }
     447         118 :         PG_END_TRY();
     448             : 
     449         118 :         MemoryContextDelete(cxt);
     450         118 :         SPI_freetuptable(tuptable);
     451             : 
     452             :         /* in case PyList_New() failed above */
     453         118 :         if (!result->rows)
     454             :         {
     455           0 :             Py_DECREF(result);
     456           0 :             result = NULL;
     457             :         }
     458             :     }
     459             : 
     460         284 :     return (PyObject *) result;
     461             : }
     462             : 
     463             : /*
     464             :  * Utilities for running SPI functions in subtransactions.
     465             :  *
     466             :  * Usage:
     467             :  *
     468             :  *  MemoryContext oldcontext = CurrentMemoryContext;
     469             :  *  ResourceOwner oldowner = CurrentResourceOwner;
     470             :  *
     471             :  *  PLy_spi_subtransaction_begin(oldcontext, oldowner);
     472             :  *  PG_TRY();
     473             :  *  {
     474             :  *      <call SPI functions>
     475             :  *      PLy_spi_subtransaction_commit(oldcontext, oldowner);
     476             :  *  }
     477             :  *  PG_CATCH();
     478             :  *  {
     479             :  *      <do cleanup>
     480             :  *      PLy_spi_subtransaction_abort(oldcontext, oldowner);
     481             :  *      return NULL;
     482             :  *  }
     483             :  *  PG_END_TRY();
     484             :  *
     485             :  * These utilities take care of restoring connection to the SPI manager and
     486             :  * setting a Python exception in case of an abort.
     487             :  */
     488             : void
     489         504 : PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner)
     490             : {
     491         504 :     BeginInternalSubTransaction(NULL);
     492             :     /* Want to run inside function's memory context */
     493         504 :     MemoryContextSwitchTo(oldcontext);
     494         504 : }
     495             : 
     496             : void
     497         450 : PLy_spi_subtransaction_commit(MemoryContext oldcontext, ResourceOwner oldowner)
     498             : {
     499             :     /* Commit the inner transaction, return to outer xact context */
     500         450 :     ReleaseCurrentSubTransaction();
     501         450 :     MemoryContextSwitchTo(oldcontext);
     502         450 :     CurrentResourceOwner = oldowner;
     503         450 : }
     504             : 
     505             : void
     506          54 : PLy_spi_subtransaction_abort(MemoryContext oldcontext, ResourceOwner oldowner)
     507             : {
     508             :     ErrorData  *edata;
     509             :     PLyExceptionEntry *entry;
     510             :     PyObject   *exc;
     511             : 
     512             :     /* Save error info */
     513          54 :     MemoryContextSwitchTo(oldcontext);
     514          54 :     edata = CopyErrorData();
     515          54 :     FlushErrorState();
     516             : 
     517             :     /* Abort the inner transaction */
     518          54 :     RollbackAndReleaseCurrentSubTransaction();
     519          54 :     MemoryContextSwitchTo(oldcontext);
     520          54 :     CurrentResourceOwner = oldowner;
     521             : 
     522             :     /* Look up the correct exception */
     523          54 :     entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
     524             :                         HASH_FIND, NULL);
     525             : 
     526             :     /*
     527             :      * This could be a custom error code, if that's the case fallback to
     528             :      * SPIError
     529             :      */
     530          54 :     exc = entry ? entry->exc : PLy_exc_spi_error;
     531             :     /* Make Python raise the exception */
     532          54 :     PLy_spi_exception_set(exc, edata);
     533          54 :     FreeErrorData(edata);
     534          54 : }
     535             : 
     536             : /*
     537             :  * Raise a SPIError, passing in it more error details, like the
     538             :  * internal query and error position.
     539             :  */
     540             : static void
     541          54 : PLy_spi_exception_set(PyObject *excclass, ErrorData *edata)
     542             : {
     543          54 :     PyObject   *args = NULL;
     544          54 :     PyObject   *spierror = NULL;
     545          54 :     PyObject   *spidata = NULL;
     546             : 
     547          54 :     args = Py_BuildValue("(s)", edata->message);
     548          54 :     if (!args)
     549           0 :         goto failure;
     550             : 
     551             :     /* create a new SPI exception with the error message as the parameter */
     552          54 :     spierror = PyObject_CallObject(excclass, args);
     553          54 :     if (!spierror)
     554           0 :         goto failure;
     555             : 
     556          54 :     spidata = Py_BuildValue("(izzzizzzzz)", edata->sqlerrcode, edata->detail, edata->hint,
     557             :                             edata->internalquery, edata->internalpos,
     558             :                             edata->schema_name, edata->table_name, edata->column_name,
     559             :                             edata->datatype_name, edata->constraint_name);
     560          54 :     if (!spidata)
     561           0 :         goto failure;
     562             : 
     563          54 :     if (PyObject_SetAttrString(spierror, "spidata", spidata) == -1)
     564           0 :         goto failure;
     565             : 
     566          54 :     PyErr_SetObject(excclass, spierror);
     567             : 
     568          54 :     Py_DECREF(args);
     569          54 :     Py_DECREF(spierror);
     570          54 :     Py_DECREF(spidata);
     571          54 :     return;
     572             : 
     573             : failure:
     574           0 :     Py_XDECREF(args);
     575           0 :     Py_XDECREF(spierror);
     576           0 :     Py_XDECREF(spidata);
     577           0 :     elog(ERROR, "could not convert SPI error to Python exception");
     578             : }

Generated by: LCOV version 1.13