Support for GiST in get_equal_strategy_number()
[pgsql.git] / src / pl / plpython / plpy_spi.c
blobbcbd07b70ae98d844efe928a7d1cbe0f6b128ac2
1 /*
2 * interface to SPI functions
4 * src/pl/plpython/plpy_spi.c
5 */
7 #include "postgres.h"
9 #include <limits.h>
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"
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);
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)
35 PyObject *
36 PLy_spi_prepare(PyObject *self, PyObject *args)
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;
47 if (!PyArg_ParseTuple(args, "s|O:prepare", &query, &list))
48 return NULL;
50 if (list && (!PySequence_Check(list)))
52 PLy_exception_set(PyExc_TypeError,
53 "second argument of plpy.prepare must be a sequence");
54 return NULL;
57 if ((plan = (PLyPlanObject *) PLy_plan_new()) == NULL)
58 return NULL;
60 plan->mcxt = AllocSetContextCreate(TopMemoryContext,
61 "PL/Python plan context",
62 ALLOCSET_DEFAULT_SIZES);
63 oldcontext = MemoryContextSwitchTo(plan->mcxt);
65 nargs = list ? PySequence_Length(list) : 0;
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;
72 MemoryContextSwitchTo(oldcontext);
74 oldcontext = CurrentMemoryContext;
75 oldowner = CurrentResourceOwner;
77 PLy_spi_subtransaction_begin(oldcontext, oldowner);
79 PG_TRY();
81 int i;
83 for (i = 0; i < nargs; i++)
85 char *sptr;
86 Oid typeId;
87 int32 typmod;
89 optr = PySequence_GetItem(list, i);
90 if (PyUnicode_Check(optr))
91 sptr = PLyUnicode_AsString(optr);
92 else
94 ereport(ERROR,
95 (errmsg("plpy.prepare: type name at ordinal position %d is not a string", i)));
96 sptr = NULL; /* keep compiler quiet */
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 ********************************************************/
105 (void) parseTypeString(sptr, &typeId, &typmod, NULL);
107 Py_DECREF(optr);
110 * set optr to NULL, so we won't try to unref it again in case of
111 * an error
113 optr = NULL;
115 plan->types[i] = typeId;
116 PLy_output_setup_func(&plan->args[i], plan->mcxt,
117 typeId, typmod,
118 exec_ctx->curr_proc);
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));
127 /* transfer plan from procCxt to topCxt */
128 if (SPI_keepplan(plan->plan))
129 elog(ERROR, "SPI_keepplan failed");
131 PLy_spi_subtransaction_commit(oldcontext, oldowner);
133 PG_CATCH();
135 Py_DECREF(plan);
136 Py_XDECREF(optr);
138 PLy_spi_subtransaction_abort(oldcontext, oldowner);
139 return NULL;
141 PG_END_TRY();
143 Assert(plan->plan != NULL);
144 return (PyObject *) plan;
147 /* execute(query="select * from foo", limit=5)
148 * execute(plan=plan, values=(foo, bar), limit=5)
150 PyObject *
151 PLy_spi_execute(PyObject *self, PyObject *args)
153 char *query;
154 PyObject *plan;
155 PyObject *list = NULL;
156 long limit = 0;
158 if (PyArg_ParseTuple(args, "s|l", &query, &limit))
159 return PLy_spi_execute_query(query, limit);
161 PyErr_Clear();
163 if (PyArg_ParseTuple(args, "O|Ol", &plan, &list, &limit) &&
164 is_PLyPlanObject(plan))
165 return PLy_spi_execute_plan(plan, list, limit);
167 PLy_exception_set(PLy_exc_error, "plpy.execute expected a query or a plan");
168 return NULL;
171 PyObject *
172 PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit)
174 volatile int nargs;
175 int i,
177 PLyPlanObject *plan;
178 volatile MemoryContext oldcontext;
179 volatile ResourceOwner oldowner;
180 PyObject *ret;
182 if (list != NULL)
184 if (!PySequence_Check(list) || PyUnicode_Check(list))
186 PLy_exception_set(PyExc_TypeError, "plpy.execute takes a sequence as its second argument");
187 return NULL;
189 nargs = PySequence_Length(list);
191 else
192 nargs = 0;
194 plan = (PLyPlanObject *) ob;
196 if (nargs != plan->nargs)
198 char *sv;
199 PyObject *so = PyObject_Str(list);
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);
211 return NULL;
214 oldcontext = CurrentMemoryContext;
215 oldowner = CurrentResourceOwner;
217 PLy_spi_subtransaction_begin(oldcontext, oldowner);
219 PG_TRY();
221 PLyExecutionContext *exec_ctx = PLy_current_execution_context();
222 char *volatile nulls;
223 volatile int j;
225 if (nargs > 0)
226 nulls = palloc(nargs * sizeof(char));
227 else
228 nulls = NULL;
230 for (j = 0; j < nargs; j++)
232 PLyObToDatum *arg = &plan->args[j];
233 PyObject *elem;
235 elem = PySequence_GetItem(list, j);
236 PG_TRY(2);
238 bool isnull;
240 plan->values[j] = PLy_output_convert(arg, elem, &isnull);
241 nulls[j] = isnull ? 'n' : ' ';
243 PG_FINALLY(2);
245 Py_DECREF(elem);
247 PG_END_TRY(2);
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);
254 if (nargs > 0)
255 pfree(nulls);
257 PLy_spi_subtransaction_commit(oldcontext, oldowner);
259 PG_CATCH();
261 int k;
264 * cleanup plan->values array
266 for (k = 0; k < nargs; k++)
268 if (!plan->args[k].typbyval &&
269 (plan->values[k] != PointerGetDatum(NULL)))
271 pfree(DatumGetPointer(plan->values[k]));
272 plan->values[k] = PointerGetDatum(NULL);
276 PLy_spi_subtransaction_abort(oldcontext, oldowner);
277 return NULL;
279 PG_END_TRY();
281 for (i = 0; i < nargs; i++)
283 if (!plan->args[i].typbyval &&
284 (plan->values[i] != PointerGetDatum(NULL)))
286 pfree(DatumGetPointer(plan->values[i]));
287 plan->values[i] = PointerGetDatum(NULL);
291 if (rv < 0)
293 PLy_exception_set(PLy_exc_spi_error,
294 "SPI_execute_plan failed: %s",
295 SPI_result_code_string(rv));
296 return NULL;
299 return ret;
302 static PyObject *
303 PLy_spi_execute_query(char *query, long limit)
305 int rv;
306 volatile MemoryContext oldcontext;
307 volatile ResourceOwner oldowner;
308 PyObject *ret = NULL;
310 oldcontext = CurrentMemoryContext;
311 oldowner = CurrentResourceOwner;
313 PLy_spi_subtransaction_begin(oldcontext, oldowner);
315 PG_TRY();
317 PLyExecutionContext *exec_ctx = PLy_current_execution_context();
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);
323 PLy_spi_subtransaction_commit(oldcontext, oldowner);
325 PG_CATCH();
327 PLy_spi_subtransaction_abort(oldcontext, oldowner);
328 return NULL;
330 PG_END_TRY();
332 if (rv < 0)
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;
341 return ret;
344 static PyObject *
345 PLy_spi_execute_fetch_result(SPITupleTable *tuptable, uint64 rows, int status)
347 PLyResultObject *result;
348 PLyExecutionContext *exec_ctx = PLy_current_execution_context();
349 volatile MemoryContext oldcontext;
351 result = (PLyResultObject *) PLy_result_new();
352 if (!result)
354 SPI_freetuptable(tuptable);
355 return NULL;
357 Py_DECREF(result->status);
358 result->status = PyLong_FromLong(status);
360 if (status > 0 && tuptable == NULL)
362 Py_DECREF(result->nrows);
363 result->nrows = PyLong_FromUnsignedLongLong(rows);
365 else if (status > 0 && tuptable != NULL)
367 PLyDatumToOb ininfo;
368 MemoryContext cxt;
370 Py_DECREF(result->nrows);
371 result->nrows = PyLong_FromUnsignedLongLong(rows);
373 cxt = AllocSetContextCreate(CurrentMemoryContext,
374 "PL/Python temp context",
375 ALLOCSET_DEFAULT_SIZES);
377 /* Initialize for converting result tuples to Python */
378 PLy_input_setup_func(&ininfo, cxt, RECORDOID, -1,
379 exec_ctx->curr_proc);
381 oldcontext = CurrentMemoryContext;
382 PG_TRY();
384 MemoryContext oldcontext2;
386 if (rows)
388 uint64 i;
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.
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")));
400 Py_DECREF(result->rows);
401 result->rows = PyList_New(rows);
402 if (result->rows)
404 PLy_input_setup_tuple(&ininfo, tuptable->tupdesc,
405 exec_ctx->curr_proc);
407 for (i = 0; i < rows; i++)
409 PyObject *row = PLy_input_from_tuple(&ininfo,
410 tuptable->vals[i],
411 tuptable->tupdesc,
412 true);
414 PyList_SetItem(result->rows, i, row);
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.)
427 oldcontext2 = MemoryContextSwitchTo(TopMemoryContext);
428 result->tupdesc = CreateTupleDescCopy(tuptable->tupdesc);
429 MemoryContextSwitchTo(oldcontext2);
431 PG_CATCH();
433 MemoryContextSwitchTo(oldcontext);
434 MemoryContextDelete(cxt);
435 Py_DECREF(result);
436 PG_RE_THROW();
438 PG_END_TRY();
440 MemoryContextDelete(cxt);
441 SPI_freetuptable(tuptable);
443 /* in case PyList_New() failed above */
444 if (!result->rows)
446 Py_DECREF(result);
447 result = NULL;
451 return (PyObject *) result;
454 PyObject *
455 PLy_commit(PyObject *self, PyObject *args)
457 MemoryContext oldcontext = CurrentMemoryContext;
458 PLyExecutionContext *exec_ctx = PLy_current_execution_context();
460 PG_TRY();
462 SPI_commit();
464 /* was cleared at transaction end, reset pointer */
465 exec_ctx->scratch_ctx = NULL;
467 PG_CATCH();
469 ErrorData *edata;
470 PLyExceptionEntry *entry;
471 PyObject *exc;
473 /* Save error info */
474 MemoryContextSwitchTo(oldcontext);
475 edata = CopyErrorData();
476 FlushErrorState();
478 /* was cleared at transaction end, reset pointer */
479 exec_ctx->scratch_ctx = NULL;
481 /* Look up the correct exception */
482 entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
483 HASH_FIND, NULL);
486 * This could be a custom error code, if that's the case fallback to
487 * SPIError
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);
494 return NULL;
496 PG_END_TRY();
498 Py_RETURN_NONE;
501 PyObject *
502 PLy_rollback(PyObject *self, PyObject *args)
504 MemoryContext oldcontext = CurrentMemoryContext;
505 PLyExecutionContext *exec_ctx = PLy_current_execution_context();
507 PG_TRY();
509 SPI_rollback();
511 /* was cleared at transaction end, reset pointer */
512 exec_ctx->scratch_ctx = NULL;
514 PG_CATCH();
516 ErrorData *edata;
517 PLyExceptionEntry *entry;
518 PyObject *exc;
520 /* Save error info */
521 MemoryContextSwitchTo(oldcontext);
522 edata = CopyErrorData();
523 FlushErrorState();
525 /* was cleared at transaction end, reset pointer */
526 exec_ctx->scratch_ctx = NULL;
528 /* Look up the correct exception */
529 entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
530 HASH_FIND, NULL);
533 * This could be a custom error code, if that's the case fallback to
534 * SPIError
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);
541 return NULL;
543 PG_END_TRY();
545 Py_RETURN_NONE;
549 * Utilities for running SPI functions in subtransactions.
551 * Usage:
553 * MemoryContext oldcontext = CurrentMemoryContext;
554 * ResourceOwner oldowner = CurrentResourceOwner;
556 * PLy_spi_subtransaction_begin(oldcontext, oldowner);
557 * PG_TRY();
559 * <call SPI functions>
560 * PLy_spi_subtransaction_commit(oldcontext, oldowner);
562 * PG_CATCH();
564 * <do cleanup>
565 * PLy_spi_subtransaction_abort(oldcontext, oldowner);
566 * return NULL;
568 * PG_END_TRY();
570 * These utilities take care of restoring connection to the SPI manager and
571 * setting a Python exception in case of an abort.
573 void
574 PLy_spi_subtransaction_begin(MemoryContext oldcontext, ResourceOwner oldowner)
576 BeginInternalSubTransaction(NULL);
577 /* Want to run inside function's memory context */
578 MemoryContextSwitchTo(oldcontext);
581 void
582 PLy_spi_subtransaction_commit(MemoryContext oldcontext, ResourceOwner oldowner)
584 /* Commit the inner transaction, return to outer xact context */
585 ReleaseCurrentSubTransaction();
586 MemoryContextSwitchTo(oldcontext);
587 CurrentResourceOwner = oldowner;
590 void
591 PLy_spi_subtransaction_abort(MemoryContext oldcontext, ResourceOwner oldowner)
593 ErrorData *edata;
594 PLyExceptionEntry *entry;
595 PyObject *exc;
597 /* Save error info */
598 MemoryContextSwitchTo(oldcontext);
599 edata = CopyErrorData();
600 FlushErrorState();
602 /* Abort the inner transaction */
603 RollbackAndReleaseCurrentSubTransaction();
604 MemoryContextSwitchTo(oldcontext);
605 CurrentResourceOwner = oldowner;
607 /* Look up the correct exception */
608 entry = hash_search(PLy_spi_exceptions, &(edata->sqlerrcode),
609 HASH_FIND, NULL);
612 * This could be a custom error code, if that's the case fallback to
613 * SPIError
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);
622 * Raise a SPIError, passing in it more error details, like the
623 * internal query and error position.
625 static void
626 PLy_spi_exception_set(PyObject *excclass, ErrorData *edata)
628 PyObject *args = NULL;
629 PyObject *spierror = NULL;
630 PyObject *spidata = NULL;
632 args = Py_BuildValue("(s)", edata->message);
633 if (!args)
634 goto failure;
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;
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;
648 if (PyObject_SetAttrString(spierror, "spidata", spidata) == -1)
649 goto failure;
651 PyErr_SetObject(excclass, spierror);
653 Py_DECREF(args);
654 Py_DECREF(spierror);
655 Py_DECREF(spidata);
656 return;
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");