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