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