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

Generated by: LCOV version 1.13