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->values = nargs ? palloc0(sizeof(Datum) * nargs) : NULL;
70 52 : plan->args = nargs ? palloc0(sizeof(PLyObToDatum) * nargs) : NULL;
71 :
72 52 : MemoryContextSwitchTo(oldcontext);
73 :
74 52 : oldcontext = CurrentMemoryContext;
75 52 : oldowner = CurrentResourceOwner;
76 :
77 52 : PLy_spi_subtransaction_begin(oldcontext, oldowner);
78 :
79 52 : PG_TRY();
80 : {
81 : int i;
82 :
83 80 : for (i = 0; i < nargs; i++)
84 : {
85 : char *sptr;
86 : Oid typeId;
87 : int32 typmod;
88 :
89 34 : optr = PySequence_GetItem(list, i);
90 34 : if (PyUnicode_Check(optr))
91 34 : 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 34 : (void) parseTypeString(sptr, &typeId, &typmod, NULL);
106 :
107 28 : 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 28 : optr = NULL;
114 :
115 28 : plan->types[i] = typeId;
116 28 : PLy_output_setup_func(&plan->args[i], plan->mcxt,
117 : typeId, typmod,
118 28 : exec_ctx->curr_proc);
119 : }
120 :
121 46 : pg_verifymbstr(query, strlen(query), false);
122 46 : plan->plan = SPI_prepare(query, plan->nargs, plan->types);
123 46 : 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 46 : if (SPI_keepplan(plan->plan))
129 0 : elog(ERROR, "SPI_keepplan failed");
130 :
131 46 : PLy_spi_subtransaction_commit(oldcontext, oldowner);
132 : }
133 6 : PG_CATCH();
134 : {
135 6 : Py_DECREF(plan);
136 6 : Py_XDECREF(optr);
137 :
138 6 : PLy_spi_subtransaction_abort(oldcontext, oldowner);
139 6 : return NULL;
140 : }
141 46 : PG_END_TRY();
142 :
143 : Assert(plan->plan != NULL);
144 46 : 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 328 : PLy_spi_execute(PyObject *self, PyObject *args)
152 : {
153 : char *query;
154 : PyObject *plan;
155 328 : PyObject *list = NULL;
156 328 : long limit = 0;
157 :
158 328 : if (PyArg_ParseTuple(args, "s|l", &query, &limit))
159 286 : return PLy_spi_execute_query(query, limit);
160 :
161 42 : PyErr_Clear();
162 :
163 84 : if (PyArg_ParseTuple(args, "O|Ol", &plan, &list, &limit) &&
164 42 : is_PLyPlanObject(plan))
165 42 : return PLy_spi_execute_plan(plan, list, limit);
166 :
167 0 : PLy_exception_set(PLy_exc_error, "plpy.execute expected a query or a plan");
168 0 : return NULL;
169 : }
170 :
171 : PyObject *
172 44 : 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 44 : if (list != NULL)
183 : {
184 32 : 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 32 : nargs = PySequence_Length(list);
190 : }
191 : else
192 12 : nargs = 0;
193 :
194 44 : plan = (PLyPlanObject *) ob;
195 :
196 44 : 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 0 : Py_DECREF(so);
210 :
211 0 : return NULL;
212 : }
213 :
214 44 : oldcontext = CurrentMemoryContext;
215 44 : oldowner = CurrentResourceOwner;
216 :
217 44 : PLy_spi_subtransaction_begin(oldcontext, oldowner);
218 :
219 44 : PG_TRY();
220 : {
221 44 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
222 : char *volatile nulls;
223 : volatile int j;
224 :
225 44 : if (nargs > 0)
226 30 : nulls = palloc(nargs * sizeof(char));
227 : else
228 14 : nulls = NULL;
229 :
230 72 : for (j = 0; j < nargs; j++)
231 : {
232 32 : PLyObToDatum *arg = &plan->args[j];
233 : PyObject *elem;
234 :
235 32 : elem = PySequence_GetItem(list, j);
236 32 : PG_TRY(2);
237 : {
238 : bool isnull;
239 :
240 32 : plan->values[j] = PLy_output_convert(arg, elem, &isnull);
241 28 : nulls[j] = isnull ? 'n' : ' ';
242 : }
243 4 : PG_FINALLY(2);
244 : {
245 32 : Py_DECREF(elem);
246 : }
247 32 : PG_END_TRY(2);
248 : }
249 :
250 80 : rv = SPI_execute_plan(plan->plan, plan->values, nulls,
251 40 : exec_ctx->curr_proc->fn_readonly, limit);
252 40 : ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv);
253 :
254 40 : if (nargs > 0)
255 26 : pfree(nulls);
256 :
257 40 : PLy_spi_subtransaction_commit(oldcontext, oldowner);
258 : }
259 4 : PG_CATCH();
260 : {
261 : int k;
262 :
263 : /*
264 : * cleanup plan->values array
265 : */
266 8 : for (k = 0; k < nargs; k++)
267 : {
268 4 : if (!plan->args[k].typbyval &&
269 0 : (plan->values[k] != PointerGetDatum(NULL)))
270 : {
271 0 : pfree(DatumGetPointer(plan->values[k]));
272 0 : plan->values[k] = PointerGetDatum(NULL);
273 : }
274 : }
275 :
276 4 : PLy_spi_subtransaction_abort(oldcontext, oldowner);
277 4 : return NULL;
278 : }
279 40 : PG_END_TRY();
280 :
281 68 : for (i = 0; i < nargs; i++)
282 : {
283 28 : if (!plan->args[i].typbyval &&
284 20 : (plan->values[i] != PointerGetDatum(NULL)))
285 : {
286 20 : pfree(DatumGetPointer(plan->values[i]));
287 20 : plan->values[i] = PointerGetDatum(NULL);
288 : }
289 : }
290 :
291 40 : if (rv < 0)
292 : {
293 2 : PLy_exception_set(PLy_exc_spi_error,
294 : "SPI_execute_plan failed: %s",
295 : SPI_result_code_string(rv));
296 2 : return NULL;
297 : }
298 :
299 38 : return ret;
300 : }
301 :
302 : static PyObject *
303 286 : PLy_spi_execute_query(char *query, long limit)
304 : {
305 : int rv;
306 : volatile MemoryContext oldcontext;
307 : volatile ResourceOwner oldowner;
308 286 : PyObject *ret = NULL;
309 :
310 286 : oldcontext = CurrentMemoryContext;
311 286 : oldowner = CurrentResourceOwner;
312 :
313 286 : PLy_spi_subtransaction_begin(oldcontext, oldowner);
314 :
315 286 : PG_TRY();
316 : {
317 286 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
318 :
319 286 : pg_verifymbstr(query, strlen(query), false);
320 286 : rv = SPI_execute(query, exec_ctx->curr_proc->fn_readonly, limit);
321 244 : ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv);
322 :
323 244 : PLy_spi_subtransaction_commit(oldcontext, oldowner);
324 : }
325 42 : PG_CATCH();
326 : {
327 42 : PLy_spi_subtransaction_abort(oldcontext, oldowner);
328 42 : return NULL;
329 : }
330 244 : PG_END_TRY();
331 :
332 244 : if (rv < 0)
333 : {
334 2 : Py_XDECREF(ret);
335 2 : PLy_exception_set(PLy_exc_spi_error,
336 : "SPI_execute failed: %s",
337 : SPI_result_code_string(rv));
338 2 : return NULL;
339 : }
340 :
341 242 : return ret;
342 : }
343 :
344 : static PyObject *
345 284 : PLy_spi_execute_fetch_result(SPITupleTable *tuptable, uint64 rows, int status)
346 : {
347 : PLyResultObject *result;
348 284 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
349 : volatile MemoryContext oldcontext;
350 :
351 284 : result = (PLyResultObject *) PLy_result_new();
352 284 : if (!result)
353 : {
354 0 : SPI_freetuptable(tuptable);
355 0 : return NULL;
356 : }
357 284 : Py_DECREF(result->status);
358 284 : result->status = PyLong_FromLong(status);
359 :
360 284 : if (status > 0 && tuptable == NULL)
361 : {
362 160 : Py_DECREF(result->nrows);
363 160 : result->nrows = PyLong_FromUnsignedLongLong(rows);
364 : }
365 124 : else if (status > 0 && tuptable != NULL)
366 : {
367 : PLyDatumToOb ininfo;
368 : MemoryContext cxt;
369 :
370 120 : Py_DECREF(result->nrows);
371 120 : result->nrows = PyLong_FromUnsignedLongLong(rows);
372 :
373 120 : cxt = AllocSetContextCreate(CurrentMemoryContext,
374 : "PL/Python temp context",
375 : ALLOCSET_DEFAULT_SIZES);
376 :
377 : /* Initialize for converting result tuples to Python */
378 120 : PLy_input_setup_func(&ininfo, cxt, RECORDOID, -1,
379 120 : exec_ctx->curr_proc);
380 :
381 120 : oldcontext = CurrentMemoryContext;
382 120 : PG_TRY();
383 : {
384 : MemoryContext oldcontext2;
385 :
386 120 : 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 116 : if (rows > (uint64) PY_SSIZE_T_MAX)
396 0 : ereport(ERROR,
397 : (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
398 : errmsg("query result has too many rows to fit in a Python list")));
399 :
400 116 : Py_DECREF(result->rows);
401 116 : result->rows = PyList_New(rows);
402 116 : if (result->rows)
403 : {
404 116 : PLy_input_setup_tuple(&ininfo, tuptable->tupdesc,
405 116 : exec_ctx->curr_proc);
406 :
407 266 : for (i = 0; i < rows; i++)
408 : {
409 300 : PyObject *row = PLy_input_from_tuple(&ininfo,
410 150 : tuptable->vals[i],
411 : tuptable->tupdesc,
412 : true);
413 :
414 150 : 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 120 : oldcontext2 = MemoryContextSwitchTo(TopMemoryContext);
428 120 : result->tupdesc = CreateTupleDescCopy(tuptable->tupdesc);
429 120 : MemoryContextSwitchTo(oldcontext2);
430 : }
431 0 : PG_CATCH();
432 : {
433 0 : MemoryContextSwitchTo(oldcontext);
434 0 : MemoryContextDelete(cxt);
435 0 : Py_DECREF(result);
436 0 : PG_RE_THROW();
437 : }
438 120 : PG_END_TRY();
439 :
440 120 : MemoryContextDelete(cxt);
441 120 : SPI_freetuptable(tuptable);
442 :
443 : /* in case PyList_New() failed above */
444 120 : if (!result->rows)
445 : {
446 0 : Py_DECREF(result);
447 0 : result = NULL;
448 : }
449 : }
450 :
451 284 : return (PyObject *) result;
452 : }
453 :
454 : PyObject *
455 52 : PLy_commit(PyObject *self, PyObject *args)
456 : {
457 52 : MemoryContext oldcontext = CurrentMemoryContext;
458 52 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
459 :
460 52 : PG_TRY();
461 : {
462 52 : SPI_commit();
463 :
464 : /* was cleared at transaction end, reset pointer */
465 40 : exec_ctx->scratch_ctx = NULL;
466 : }
467 12 : PG_CATCH();
468 : {
469 : ErrorData *edata;
470 : PLyExceptionEntry *entry;
471 : PyObject *exc;
472 :
473 : /* Save error info */
474 12 : MemoryContextSwitchTo(oldcontext);
475 12 : edata = CopyErrorData();
476 12 : FlushErrorState();
477 :
478 : /* was cleared at transaction end, reset pointer */
479 12 : exec_ctx->scratch_ctx = NULL;
480 :
481 : /* Look up the correct exception */
482 12 : 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 12 : exc = entry ? entry->exc : PLy_exc_spi_error;
490 : /* Make Python raise the exception */
491 12 : PLy_spi_exception_set(exc, edata);
492 12 : FreeErrorData(edata);
493 :
494 12 : return NULL;
495 : }
496 40 : PG_END_TRY();
497 :
498 40 : Py_RETURN_NONE;
499 : }
500 :
501 : PyObject *
502 34 : PLy_rollback(PyObject *self, PyObject *args)
503 : {
504 34 : MemoryContext oldcontext = CurrentMemoryContext;
505 34 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
506 :
507 34 : PG_TRY();
508 : {
509 34 : SPI_rollback();
510 :
511 : /* was cleared at transaction end, reset pointer */
512 34 : exec_ctx->scratch_ctx = NULL;
513 : }
514 0 : PG_CATCH();
515 : {
516 : ErrorData *edata;
517 : PLyExceptionEntry *entry;
518 : PyObject *exc;
519 :
520 : /* Save error info */
521 0 : MemoryContextSwitchTo(oldcontext);
522 0 : edata = CopyErrorData();
523 0 : FlushErrorState();
524 :
525 : /* was cleared at transaction end, reset pointer */
526 0 : exec_ctx->scratch_ctx = NULL;
527 :
528 : /* Look up the correct exception */
529 0 : 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 0 : exc = entry ? entry->exc : PLy_exc_spi_error;
537 : /* Make Python raise the exception */
538 0 : PLy_spi_exception_set(exc, edata);
539 0 : FreeErrorData(edata);
540 :
541 0 : return NULL;
542 : }
543 34 : PG_END_TRY();
544 :
545 34 : 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 508 : PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner)
575 : {
576 508 : BeginInternalSubTransaction(NULL);
577 : /* Want to run inside function's memory context */
578 508 : MemoryContextSwitchTo(oldcontext);
579 508 : }
580 :
581 : void
582 454 : PLy_spi_subtransaction_commit(MemoryContext oldcontext, ResourceOwner oldowner)
583 : {
584 : /* Commit the inner transaction, return to outer xact context */
585 454 : ReleaseCurrentSubTransaction();
586 454 : MemoryContextSwitchTo(oldcontext);
587 454 : CurrentResourceOwner = oldowner;
588 454 : }
589 :
590 : void
591 54 : PLy_spi_subtransaction_abort(MemoryContext oldcontext, ResourceOwner oldowner)
592 : {
593 : ErrorData *edata;
594 : PLyExceptionEntry *entry;
595 : PyObject *exc;
596 :
597 : /* Save error info */
598 54 : MemoryContextSwitchTo(oldcontext);
599 54 : edata = CopyErrorData();
600 54 : FlushErrorState();
601 :
602 : /* Abort the inner transaction */
603 54 : RollbackAndReleaseCurrentSubTransaction();
604 54 : MemoryContextSwitchTo(oldcontext);
605 54 : CurrentResourceOwner = oldowner;
606 :
607 : /* Look up the correct exception */
608 54 : 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 54 : exc = entry ? entry->exc : PLy_exc_spi_error;
616 : /* Make Python raise the exception */
617 54 : PLy_spi_exception_set(exc, edata);
618 54 : FreeErrorData(edata);
619 54 : }
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 66 : PLy_spi_exception_set(PyObject *excclass, ErrorData *edata)
627 : {
628 66 : PyObject *args = NULL;
629 66 : PyObject *spierror = NULL;
630 66 : PyObject *spidata = NULL;
631 :
632 66 : args = Py_BuildValue("(s)", edata->message);
633 66 : if (!args)
634 0 : goto failure;
635 :
636 : /* create a new SPI exception with the error message as the parameter */
637 66 : spierror = PyObject_CallObject(excclass, args);
638 66 : if (!spierror)
639 0 : goto failure;
640 :
641 66 : 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 66 : if (!spidata)
646 0 : goto failure;
647 :
648 66 : if (PyObject_SetAttrString(spierror, "spidata", spidata) == -1)
649 0 : goto failure;
650 :
651 66 : PyErr_SetObject(excclass, spierror);
652 :
653 66 : Py_DECREF(args);
654 66 : Py_DECREF(spierror);
655 66 : Py_DECREF(spidata);
656 66 : return;
657 :
658 0 : failure:
659 0 : Py_XDECREF(args);
660 0 : Py_XDECREF(spierror);
661 0 : Py_XDECREF(spidata);
662 0 : elog(ERROR, "could not convert SPI error to Python exception");
663 : }
|
