granicus.if.org Git - postgresql/blob - src/pl/plpython/plpy_spi.c

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