Age Owner TLA Line data Source code
1 : /*
2 : * the PLyCursor class
3 : *
4 : * src/pl/plpython/plpy_cursorobject.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 "mb/pg_wchar.h"
14 : #include "plpy_cursorobject.h"
15 : #include "plpy_elog.h"
16 : #include "plpy_main.h"
17 : #include "plpy_planobject.h"
18 : #include "plpy_procedure.h"
19 : #include "plpy_resultobject.h"
20 : #include "plpy_spi.h"
21 : #include "plpython.h"
22 : #include "utils/memutils.h"
23 :
24 : static PyObject *PLy_cursor_query(const char *query);
25 : static void PLy_cursor_dealloc(PyObject *arg);
26 : static PyObject *PLy_cursor_iternext(PyObject *self);
27 : static PyObject *PLy_cursor_fetch(PyObject *self, PyObject *args);
28 : static PyObject *PLy_cursor_close(PyObject *self, PyObject *unused);
29 :
30 : static char PLy_cursor_doc[] = "Wrapper around a PostgreSQL cursor";
31 :
32 : static PyMethodDef PLy_cursor_methods[] = {
33 : {"fetch", PLy_cursor_fetch, METH_VARARGS, NULL},
34 : {"close", PLy_cursor_close, METH_NOARGS, NULL},
35 : {NULL, NULL, 0, NULL}
36 : };
37 :
38 : static PyTypeObject PLy_CursorType = {
39 : PyVarObject_HEAD_INIT(NULL, 0)
40 : .tp_name = "PLyCursor",
41 : .tp_basicsize = sizeof(PLyCursorObject),
42 : .tp_dealloc = PLy_cursor_dealloc,
43 : .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
44 : .tp_doc = PLy_cursor_doc,
45 : .tp_iter = PyObject_SelfIter,
46 : .tp_iternext = PLy_cursor_iternext,
47 : .tp_methods = PLy_cursor_methods,
48 : };
49 :
50 : void
4130 peter_e 51 CBC 23 : PLy_cursor_init_type(void)
52 : {
53 23 : if (PyType_Ready(&PLy_CursorType) < 0)
4130 peter_e 54 UBC 0 : elog(ERROR, "could not initialize PLy_CursorType");
4130 peter_e 55 CBC 23 : }
56 :
57 : PyObject *
58 17 : PLy_cursor(PyObject *self, PyObject *args)
59 : {
60 : char *query;
61 : PyObject *plan;
62 17 : PyObject *planargs = NULL;
63 :
64 17 : if (PyArg_ParseTuple(args, "s", &query))
65 14 : return PLy_cursor_query(query);
66 :
67 3 : PyErr_Clear();
68 :
69 3 : if (PyArg_ParseTuple(args, "O|O", &plan, &planargs))
70 3 : return PLy_cursor_plan(plan, planargs);
71 :
4130 peter_e 72 UBC 0 : PLy_exception_set(PLy_exc_error, "plpy.cursor expected a query or a plan");
73 0 : return NULL;
74 : }
75 :
76 :
77 : static PyObject *
4130 peter_e 78 CBC 14 : PLy_cursor_query(const char *query)
79 : {
80 : PLyCursorObject *cursor;
1970 tgl 81 14 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
82 : volatile MemoryContext oldcontext;
83 : volatile ResourceOwner oldowner;
84 :
4130 peter_e 85 14 : if ((cursor = PyObject_New(PLyCursorObject, &PLy_CursorType)) == NULL)
4130 peter_e 86 UBC 0 : return NULL;
4130 peter_e 87 CBC 14 : cursor->portalname = NULL;
88 14 : cursor->closed = false;
2712 tgl 89 14 : cursor->mcxt = AllocSetContextCreate(TopMemoryContext,
90 : "PL/Python cursor context",
91 : ALLOCSET_DEFAULT_SIZES);
92 :
93 : /* Initialize for converting result tuples to Python */
1970 94 14 : PLy_input_setup_func(&cursor->result, cursor->mcxt,
95 : RECORDOID, -1,
96 14 : exec_ctx->curr_proc);
97 :
4130 peter_e 98 14 : oldcontext = CurrentMemoryContext;
99 14 : oldowner = CurrentResourceOwner;
100 :
101 14 : PLy_spi_subtransaction_begin(oldcontext, oldowner);
102 :
103 14 : PG_TRY();
104 : {
105 : SPIPlanPtr plan;
106 : Portal portal;
107 :
108 14 : pg_verifymbstr(query, strlen(query), false);
109 :
110 14 : plan = SPI_prepare(query, 0, NULL);
111 14 : if (plan == NULL)
4130 peter_e 112 UBC 0 : elog(ERROR, "SPI_prepare failed: %s",
113 : SPI_result_code_string(SPI_result));
114 :
4130 peter_e 115 CBC 28 : portal = SPI_cursor_open(NULL, plan, NULL, NULL,
4044 tgl 116 14 : exec_ctx->curr_proc->fn_readonly);
4130 peter_e 117 14 : SPI_freeplan(plan);
118 :
119 14 : if (portal == NULL)
4107 peter_e 120 UBC 0 : elog(ERROR, "SPI_cursor_open() failed: %s",
121 : SPI_result_code_string(SPI_result));
122 :
2712 tgl 123 CBC 14 : cursor->portalname = MemoryContextStrdup(cursor->mcxt, portal->name);
124 :
1944 peter_e 125 14 : PinPortal(portal);
126 :
4130 127 14 : PLy_spi_subtransaction_commit(oldcontext, oldowner);
128 : }
4130 peter_e 129 UBC 0 : PG_CATCH();
130 : {
131 0 : PLy_spi_subtransaction_abort(oldcontext, oldowner);
132 0 : return NULL;
133 : }
4130 peter_e 134 CBC 14 : PG_END_TRY();
135 :
136 14 : Assert(cursor->portalname != NULL);
137 14 : return (PyObject *) cursor;
138 : }
139 :
140 : PyObject *
141 4 : PLy_cursor_plan(PyObject *ob, PyObject *args)
142 : {
143 : PLyCursorObject *cursor;
144 : volatile int nargs;
145 : int i;
146 : PLyPlanObject *plan;
1970 tgl 147 4 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
148 : volatile MemoryContext oldcontext;
149 : volatile ResourceOwner oldowner;
150 :
4130 peter_e 151 4 : if (args)
152 : {
398 andres 153 3 : if (!PySequence_Check(args) || PyUnicode_Check(args))
154 : {
4130 peter_e 155 UBC 0 : PLy_exception_set(PyExc_TypeError, "plpy.cursor takes a sequence as its second argument");
156 0 : return NULL;
157 : }
4130 peter_e 158 CBC 3 : nargs = PySequence_Length(args);
159 : }
160 : else
161 1 : nargs = 0;
162 :
163 4 : plan = (PLyPlanObject *) ob;
164 :
165 4 : if (nargs != plan->nargs)
166 : {
167 : char *sv;
168 1 : PyObject *so = PyObject_Str(args);
169 :
170 1 : if (!so)
4130 peter_e 171 UBC 0 : PLy_elog(ERROR, "could not execute plan");
398 andres 172 CBC 1 : sv = PLyUnicode_AsString(so);
4130 peter_e 173 1 : PLy_exception_set_plural(PyExc_TypeError,
174 : "Expected sequence of %d argument, got %d: %s",
175 : "Expected sequence of %d arguments, got %d: %s",
176 1 : plan->nargs,
177 : plan->nargs, nargs, sv);
178 1 : Py_DECREF(so);
179 :
180 1 : return NULL;
181 : }
182 :
183 3 : if ((cursor = PyObject_New(PLyCursorObject, &PLy_CursorType)) == NULL)
4130 peter_e 184 UBC 0 : return NULL;
4130 peter_e 185 CBC 3 : cursor->portalname = NULL;
186 3 : cursor->closed = false;
2712 tgl 187 3 : cursor->mcxt = AllocSetContextCreate(TopMemoryContext,
188 : "PL/Python cursor context",
189 : ALLOCSET_DEFAULT_SIZES);
190 :
191 : /* Initialize for converting result tuples to Python */
1970 192 3 : PLy_input_setup_func(&cursor->result, cursor->mcxt,
193 : RECORDOID, -1,
194 3 : exec_ctx->curr_proc);
195 :
4130 peter_e 196 3 : oldcontext = CurrentMemoryContext;
197 3 : oldowner = CurrentResourceOwner;
198 :
199 3 : PLy_spi_subtransaction_begin(oldcontext, oldowner);
200 :
201 3 : PG_TRY();
202 : {
203 : Portal portal;
204 : char *volatile nulls;
205 : volatile int j;
206 :
207 3 : if (nargs > 0)
208 2 : nulls = palloc(nargs * sizeof(char));
209 : else
210 1 : nulls = NULL;
211 :
212 5 : for (j = 0; j < nargs; j++)
213 : {
1970 tgl 214 2 : PLyObToDatum *arg = &plan->args[j];
215 : PyObject *elem;
216 :
4130 peter_e 217 2 : elem = PySequence_GetItem(args, j);
184 drowley 218 GNC 2 : PG_TRY(2);
219 : {
220 : bool isnull;
221 :
1970 tgl 222 CBC 2 : plan->values[j] = PLy_output_convert(arg, elem, &isnull);
223 2 : nulls[j] = isnull ? 'n' : ' ';
224 : }
184 drowley 225 UNC 0 : PG_FINALLY(2);
226 : {
4130 peter_e 227 CBC 2 : Py_DECREF(elem);
228 : }
184 drowley 229 GNC 2 : PG_END_TRY(2);
230 : }
231 :
4130 peter_e 232 CBC 6 : portal = SPI_cursor_open(NULL, plan->plan, plan->values, nulls,
4044 tgl 233 3 : exec_ctx->curr_proc->fn_readonly);
4130 peter_e 234 3 : if (portal == NULL)
4107 peter_e 235 UBC 0 : elog(ERROR, "SPI_cursor_open() failed: %s",
236 : SPI_result_code_string(SPI_result));
237 :
2712 tgl 238 CBC 3 : cursor->portalname = MemoryContextStrdup(cursor->mcxt, portal->name);
239 :
1944 peter_e 240 3 : PinPortal(portal);
241 :
4130 242 3 : PLy_spi_subtransaction_commit(oldcontext, oldowner);
243 : }
4130 peter_e 244 UBC 0 : PG_CATCH();
245 : {
246 : int k;
247 :
248 : /* cleanup plan->values array */
249 0 : for (k = 0; k < nargs; k++)
250 : {
1970 tgl 251 0 : if (!plan->args[k].typbyval &&
4130 peter_e 252 0 : (plan->values[k] != PointerGetDatum(NULL)))
253 : {
254 0 : pfree(DatumGetPointer(plan->values[k]));
255 0 : plan->values[k] = PointerGetDatum(NULL);
256 : }
257 : }
258 :
259 0 : Py_DECREF(cursor);
260 :
261 0 : PLy_spi_subtransaction_abort(oldcontext, oldowner);
262 0 : return NULL;
263 : }
4130 peter_e 264 CBC 3 : PG_END_TRY();
265 :
266 5 : for (i = 0; i < nargs; i++)
267 : {
1970 tgl 268 2 : if (!plan->args[i].typbyval &&
4130 peter_e 269 2 : (plan->values[i] != PointerGetDatum(NULL)))
270 : {
271 2 : pfree(DatumGetPointer(plan->values[i]));
272 2 : plan->values[i] = PointerGetDatum(NULL);
273 : }
274 : }
275 :
276 3 : Assert(cursor->portalname != NULL);
277 3 : return (PyObject *) cursor;
278 : }
279 :
280 : static void
281 17 : PLy_cursor_dealloc(PyObject *arg)
282 : {
283 : PLyCursorObject *cursor;
284 : Portal portal;
285 :
286 17 : cursor = (PLyCursorObject *) arg;
287 :
288 17 : if (!cursor->closed)
289 : {
290 14 : portal = GetPortalByName(cursor->portalname);
291 :
292 14 : if (PortalIsValid(portal))
293 : {
1944 294 11 : UnpinPortal(portal);
4130 295 11 : SPI_cursor_close(portal);
296 : }
2712 tgl 297 14 : cursor->closed = true;
298 : }
299 17 : if (cursor->mcxt)
300 : {
301 17 : MemoryContextDelete(cursor->mcxt);
302 17 : cursor->mcxt = NULL;
303 : }
4130 peter_e 304 17 : arg->ob_type->tp_free(arg);
305 17 : }
306 :
307 : static PyObject *
308 35 : PLy_cursor_iternext(PyObject *self)
309 : {
310 : PLyCursorObject *cursor;
311 : PyObject *ret;
1970 tgl 312 35 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
313 : volatile MemoryContext oldcontext;
314 : volatile ResourceOwner oldowner;
315 : Portal portal;
316 :
4130 peter_e 317 35 : cursor = (PLyCursorObject *) self;
318 :
319 35 : if (cursor->closed)
320 : {
321 1 : PLy_exception_set(PyExc_ValueError, "iterating a closed cursor");
322 1 : return NULL;
323 : }
324 :
325 34 : portal = GetPortalByName(cursor->portalname);
326 34 : if (!PortalIsValid(portal))
327 : {
4130 peter_e 328 UBC 0 : PLy_exception_set(PyExc_ValueError,
329 : "iterating a cursor in an aborted subtransaction");
330 0 : return NULL;
331 : }
332 :
4130 peter_e 333 CBC 34 : oldcontext = CurrentMemoryContext;
334 34 : oldowner = CurrentResourceOwner;
335 :
336 34 : PLy_spi_subtransaction_begin(oldcontext, oldowner);
337 :
338 34 : PG_TRY();
339 : {
340 34 : SPI_cursor_fetch(portal, true, 1);
341 34 : if (SPI_processed == 0)
342 : {
343 8 : PyErr_SetNone(PyExc_StopIteration);
344 8 : ret = NULL;
345 : }
346 : else
347 : {
1970 tgl 348 26 : PLy_input_setup_tuple(&cursor->result, SPI_tuptable->tupdesc,
349 26 : exec_ctx->curr_proc);
350 :
351 26 : ret = PLy_input_from_tuple(&cursor->result, SPI_tuptable->vals[0],
1471 peter 352 26 : SPI_tuptable->tupdesc, true);
353 : }
354 :
4130 peter_e 355 34 : SPI_freetuptable(SPI_tuptable);
356 :
357 34 : PLy_spi_subtransaction_commit(oldcontext, oldowner);
358 : }
4130 peter_e 359 UBC 0 : PG_CATCH();
360 : {
361 0 : PLy_spi_subtransaction_abort(oldcontext, oldowner);
362 0 : return NULL;
363 : }
4130 peter_e 364 CBC 34 : PG_END_TRY();
365 :
366 34 : return ret;
367 : }
368 :
369 : static PyObject *
370 13 : PLy_cursor_fetch(PyObject *self, PyObject *args)
371 : {
372 : PLyCursorObject *cursor;
373 : int count;
374 : PLyResultObject *ret;
1970 tgl 375 13 : PLyExecutionContext *exec_ctx = PLy_current_execution_context();
376 : volatile MemoryContext oldcontext;
377 : volatile ResourceOwner oldowner;
378 : Portal portal;
379 :
2355 peter_e 380 13 : if (!PyArg_ParseTuple(args, "i:fetch", &count))
4130 peter_e 381 UBC 0 : return NULL;
382 :
4130 peter_e 383 CBC 13 : cursor = (PLyCursorObject *) self;
384 :
385 13 : if (cursor->closed)
386 : {
387 1 : PLy_exception_set(PyExc_ValueError, "fetch from a closed cursor");
388 1 : return NULL;
389 : }
390 :
391 12 : portal = GetPortalByName(cursor->portalname);
392 12 : if (!PortalIsValid(portal))
393 : {
394 2 : PLy_exception_set(PyExc_ValueError,
395 : "iterating a cursor in an aborted subtransaction");
396 2 : return NULL;
397 : }
398 :
399 10 : ret = (PLyResultObject *) PLy_result_new();
400 10 : if (ret == NULL)
4130 peter_e 401 UBC 0 : return NULL;
402 :
4130 peter_e 403 CBC 10 : oldcontext = CurrentMemoryContext;
404 10 : oldowner = CurrentResourceOwner;
405 :
406 10 : PLy_spi_subtransaction_begin(oldcontext, oldowner);
407 :
408 10 : PG_TRY();
409 : {
410 10 : SPI_cursor_fetch(portal, true, count);
411 :
412 10 : Py_DECREF(ret->status);
398 andres 413 10 : ret->status = PyLong_FromLong(SPI_OK_FETCH);
414 :
4130 peter_e 415 10 : Py_DECREF(ret->nrows);
1905 416 10 : ret->nrows = PyLong_FromUnsignedLongLong(SPI_processed);
417 :
4130 418 10 : if (SPI_processed != 0)
419 : {
420 : uint64 i;
421 :
422 : /*
423 : * PyList_New() and PyList_SetItem() use Py_ssize_t for list size
424 : * and list indices; so we cannot support a result larger than
425 : * PY_SSIZE_T_MAX.
426 : */
2584 tgl 427 7 : if (SPI_processed > (uint64) PY_SSIZE_T_MAX)
2584 tgl 428 UBC 0 : ereport(ERROR,
429 : (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
430 : errmsg("query result has too many rows to fit in a Python list")));
431 :
4130 peter_e 432 CBC 7 : Py_DECREF(ret->rows);
433 7 : ret->rows = PyList_New(SPI_processed);
1986 434 7 : if (!ret->rows)
435 : {
1986 peter_e 436 UBC 0 : Py_DECREF(ret);
437 0 : ret = NULL;
438 : }
439 : else
440 : {
1986 peter_e 441 CBC 7 : PLy_input_setup_tuple(&cursor->result, SPI_tuptable->tupdesc,
442 7 : exec_ctx->curr_proc);
443 :
444 44 : for (i = 0; i < SPI_processed; i++)
445 : {
446 74 : PyObject *row = PLy_input_from_tuple(&cursor->result,
447 37 : SPI_tuptable->vals[i],
1471 peter 448 37 : SPI_tuptable->tupdesc,
449 : true);
450 :
1986 peter_e 451 37 : PyList_SetItem(ret->rows, i, row);
452 : }
453 : }
454 : }
455 :
4130 456 10 : SPI_freetuptable(SPI_tuptable);
457 :
458 10 : PLy_spi_subtransaction_commit(oldcontext, oldowner);
459 : }
4130 peter_e 460 UBC 0 : PG_CATCH();
461 : {
462 0 : PLy_spi_subtransaction_abort(oldcontext, oldowner);
463 0 : return NULL;
464 : }
4130 peter_e 465 CBC 10 : PG_END_TRY();
466 :
467 10 : return (PyObject *) ret;
468 : }
469 :
470 : static PyObject *
471 5 : PLy_cursor_close(PyObject *self, PyObject *unused)
472 : {
473 5 : PLyCursorObject *cursor = (PLyCursorObject *) self;
474 :
475 5 : if (!cursor->closed)
476 : {
3955 bruce 477 4 : Portal portal = GetPortalByName(cursor->portalname);
478 :
4130 peter_e 479 4 : if (!PortalIsValid(portal))
480 : {
481 1 : PLy_exception_set(PyExc_ValueError,
482 : "closing a cursor in an aborted subtransaction");
483 1 : return NULL;
484 : }
485 :
1944 486 3 : UnpinPortal(portal);
4130 487 3 : SPI_cursor_close(portal);
488 3 : cursor->closed = true;
489 : }
490 :
2018 491 4 : Py_RETURN_NONE;
492 : }
|
