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