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