static PyObject * PLy_cursor_iternext(PyObject *self) { PLyCursorObject *cursor; PyObject *ret; volatile MemoryContext oldcontext; volatile ResourceOwner oldowner; Portal portal; cursor = (PLyCursorObject *) self; if (cursor->closed) { PLy_exception_set(PyExc_ValueError, "iterating a closed cursor"); return NULL; } portal = GetPortalByName(cursor->portalname); if (!PortalIsValid(portal)) { PLy_exception_set(PyExc_ValueError, "iterating a cursor in an aborted subtransaction"); return NULL; } oldcontext = CurrentMemoryContext; oldowner = CurrentResourceOwner; PLy_spi_subtransaction_begin(oldcontext, oldowner); PG_TRY(); { SPI_cursor_fetch(portal, true, 1); if (SPI_processed == 0) { PyErr_SetNone(PyExc_StopIteration); ret = NULL; } else { if (cursor->result.is_rowtype != 1) PLy_input_tuple_funcs(&cursor->result, SPI_tuptable->tupdesc); ret = PLyDict_FromTuple(&cursor->result, SPI_tuptable->vals[0], SPI_tuptable->tupdesc); } SPI_freetuptable(SPI_tuptable); PLy_spi_subtransaction_commit(oldcontext, oldowner); } PG_CATCH(); { PLy_spi_subtransaction_abort(oldcontext, oldowner); return NULL; } PG_END_TRY(); return ret; }
static PyObject * PLy_cursor_query(const char *query) { PLyCursorObject *cursor; volatile MemoryContext oldcontext; volatile ResourceOwner oldowner; if ((cursor = PyObject_New(PLyCursorObject, &PLy_CursorType)) == NULL) return NULL; cursor->portalname = NULL; cursor->closed = false; cursor->mcxt = AllocSetContextCreate(TopMemoryContext, "PL/Python cursor context", ALLOCSET_DEFAULT_MINSIZE, ALLOCSET_DEFAULT_INITSIZE, ALLOCSET_DEFAULT_MAXSIZE); PLy_typeinfo_init(&cursor->result, cursor->mcxt); oldcontext = CurrentMemoryContext; oldowner = CurrentResourceOwner; PLy_spi_subtransaction_begin(oldcontext, oldowner); PG_TRY(); { PLyExecutionContext *exec_ctx = PLy_current_execution_context(); SPIPlanPtr plan; Portal portal; pg_verifymbstr(query, strlen(query), false); plan = SPI_prepare(query, 0, NULL); if (plan == NULL) elog(ERROR, "SPI_prepare failed: %s", SPI_result_code_string(SPI_result)); portal = SPI_cursor_open(NULL, plan, NULL, NULL, exec_ctx->curr_proc->fn_readonly); SPI_freeplan(plan); if (portal == NULL) elog(ERROR, "SPI_cursor_open() failed: %s", SPI_result_code_string(SPI_result)); cursor->portalname = MemoryContextStrdup(cursor->mcxt, portal->name); PLy_spi_subtransaction_commit(oldcontext, oldowner); } PG_CATCH(); { PLy_spi_subtransaction_abort(oldcontext, oldowner); return NULL; } PG_END_TRY(); Assert(cursor->portalname != NULL); return (PyObject *) cursor; }
static PyObject * PLy_spi_execute_query(char *query, long limit) { int rv; volatile MemoryContext oldcontext; volatile ResourceOwner oldowner; PyObject *ret = NULL; oldcontext = CurrentMemoryContext; oldowner = CurrentResourceOwner; PLy_spi_subtransaction_begin(oldcontext, oldowner); PG_TRY(); { PLyExecutionContext *exec_ctx = PLy_current_execution_context(); pg_verifymbstr(query, strlen(query), false); rv = SPI_execute(query, exec_ctx->curr_proc->fn_readonly, limit); ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv); PLy_spi_subtransaction_commit(oldcontext, oldowner); } PG_CATCH(); { PLy_spi_subtransaction_abort(oldcontext, oldowner); return NULL; } PG_END_TRY(); if (rv < 0) { Py_XDECREF(ret); PLy_exception_set(PLy_exc_spi_error, "SPI_execute failed: %s", SPI_result_code_string(rv)); return NULL; } return ret; }
/* prepare(query="select * from foo") * prepare(query="select * from foo where bar = $1", params=["text"]) * prepare(query="select * from foo where bar = $1", params=["text"], limit=5) */ PyObject * PLy_spi_prepare(PyObject *self, PyObject *args) { PLyPlanObject *plan; PyObject *list = NULL; PyObject *volatile optr = NULL; char *query; volatile MemoryContext oldcontext; volatile ResourceOwner oldowner; volatile int nargs; if (!PyArg_ParseTuple(args, "s|O", &query, &list)) return NULL; if (list && (!PySequence_Check(list))) { PLy_exception_set(PyExc_TypeError, "second argument of plpy.prepare must be a sequence"); return NULL; } if ((plan = (PLyPlanObject *) PLy_plan_new()) == NULL) return NULL; nargs = list ? PySequence_Length(list) : 0; plan->nargs = nargs; plan->types = nargs ? PLy_malloc(sizeof(Oid) * nargs) : NULL; plan->values = nargs ? PLy_malloc(sizeof(Datum) * nargs) : NULL; plan->args = nargs ? PLy_malloc(sizeof(PLyTypeInfo) * nargs) : NULL; oldcontext = CurrentMemoryContext; oldowner = CurrentResourceOwner; PLy_spi_subtransaction_begin(oldcontext, oldowner); PG_TRY(); { int i; /* * the other loop might throw an exception, if PLyTypeInfo member * isn't properly initialized the Py_DECREF(plan) will go boom */ for (i = 0; i < nargs; i++) { PLy_typeinfo_init(&plan->args[i]); plan->values[i] = PointerGetDatum(NULL); } for (i = 0; i < nargs; i++) { char *sptr; HeapTuple typeTup; Oid typeId; int32 typmod; Form_pg_type typeStruct; optr = PySequence_GetItem(list, i); if (PyString_Check(optr)) sptr = PyString_AsString(optr); else if (PyUnicode_Check(optr)) sptr = PLyUnicode_AsString(optr); else { ereport(ERROR, (errmsg("plpy.prepare: type name at ordinal position %d is not a string", i))); sptr = NULL; /* keep compiler quiet */ } /******************************************************** * Resolve argument type names and then look them up by * oid in the system cache, and remember the required *information for input conversion. ********************************************************/ parseTypeString(sptr, &typeId, &typmod); typeTup = SearchSysCache1(TYPEOID, ObjectIdGetDatum(typeId)); if (!HeapTupleIsValid(typeTup)) elog(ERROR, "cache lookup failed for type %u", typeId); Py_DECREF(optr); /* * set optr to NULL, so we won't try to unref it again in case of * an error */ optr = NULL; plan->types[i] = typeId; typeStruct = (Form_pg_type) GETSTRUCT(typeTup); if (typeStruct->typtype != TYPTYPE_COMPOSITE) PLy_output_datum_func(&plan->args[i], typeTup); else ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), errmsg("plpy.prepare does not support composite types"))); ReleaseSysCache(typeTup); } pg_verifymbstr(query, strlen(query), false); plan->plan = SPI_prepare(query, plan->nargs, plan->types); if (plan->plan == NULL) elog(ERROR, "SPI_prepare failed: %s", SPI_result_code_string(SPI_result)); /* transfer plan from procCxt to topCxt */ if (SPI_keepplan(plan->plan)) elog(ERROR, "SPI_keepplan failed"); PLy_spi_subtransaction_commit(oldcontext, oldowner); } PG_CATCH(); { Py_DECREF(plan); Py_XDECREF(optr); PLy_spi_subtransaction_abort(oldcontext, oldowner); return NULL; } PG_END_TRY(); Assert(plan->plan != NULL); return (PyObject *) plan; }
static PyObject * PLy_spi_execute_plan(PyObject *ob, PyObject *list, long limit) { volatile int nargs; int i, rv; PLyPlanObject *plan; volatile MemoryContext oldcontext; volatile ResourceOwner oldowner; PyObject *ret; if (list != NULL) { if (!PySequence_Check(list) || PyString_Check(list) || PyUnicode_Check(list)) { PLy_exception_set(PyExc_TypeError, "plpy.execute takes a sequence as its second argument"); return NULL; } nargs = PySequence_Length(list); } else nargs = 0; plan = (PLyPlanObject *) ob; if (nargs != plan->nargs) { char *sv; PyObject *so = PyObject_Str(list); if (!so) PLy_elog(ERROR, "could not execute plan"); sv = PyString_AsString(so); PLy_exception_set_plural(PyExc_TypeError, "Expected sequence of %d argument, got %d: %s", "Expected sequence of %d arguments, got %d: %s", plan->nargs, plan->nargs, nargs, sv); Py_DECREF(so); return NULL; } oldcontext = CurrentMemoryContext; oldowner = CurrentResourceOwner; PLy_spi_subtransaction_begin(oldcontext, oldowner); PG_TRY(); { char *volatile nulls; volatile int j; if (nargs > 0) nulls = palloc(nargs * sizeof(char)); else nulls = NULL; for (j = 0; j < nargs; j++) { PyObject *elem; elem = PySequence_GetItem(list, j); if (elem != Py_None) { PG_TRY(); { plan->values[j] = plan->args[j].out.d.func(&(plan->args[j].out.d), -1, elem); } PG_CATCH(); { Py_DECREF(elem); PG_RE_THROW(); } PG_END_TRY(); Py_DECREF(elem); nulls[j] = ' '; } else { Py_DECREF(elem); plan->values[j] = InputFunctionCall(&(plan->args[j].out.d.typfunc), NULL, plan->args[j].out.d.typioparam, -1); nulls[j] = 'n'; } } rv = SPI_execute_plan(plan->plan, plan->values, nulls, PLy_curr_procedure->fn_readonly, limit); ret = PLy_spi_execute_fetch_result(SPI_tuptable, SPI_processed, rv); if (nargs > 0) pfree(nulls); PLy_spi_subtransaction_commit(oldcontext, oldowner); } PG_CATCH(); { int k; /* * cleanup plan->values array */ for (k = 0; k < nargs; k++) { if (!plan->args[k].out.d.typbyval && (plan->values[k] != PointerGetDatum(NULL))) { pfree(DatumGetPointer(plan->values[k])); plan->values[k] = PointerGetDatum(NULL); } } PLy_spi_subtransaction_abort(oldcontext, oldowner); return NULL; } PG_END_TRY(); for (i = 0; i < nargs; i++) { if (!plan->args[i].out.d.typbyval && (plan->values[i] != PointerGetDatum(NULL))) { pfree(DatumGetPointer(plan->values[i])); plan->values[i] = PointerGetDatum(NULL); } } if (rv < 0) { PLy_exception_set(PLy_exc_spi_error, "SPI_execute_plan failed: %s", SPI_result_code_string(rv)); return NULL; } return ret; }
static PyObject * PLy_cursor_fetch(PyObject *self, PyObject *args) { PLyCursorObject *cursor; int count; PLyResultObject *ret; volatile MemoryContext oldcontext; volatile ResourceOwner oldowner; Portal portal; if (!PyArg_ParseTuple(args, "i", &count)) return NULL; cursor = (PLyCursorObject *) self; if (cursor->closed) { PLy_exception_set(PyExc_ValueError, "fetch from a closed cursor"); return NULL; } portal = GetPortalByName(cursor->portalname); if (!PortalIsValid(portal)) { PLy_exception_set(PyExc_ValueError, "iterating a cursor in an aborted subtransaction"); return NULL; } ret = (PLyResultObject *) PLy_result_new(); if (ret == NULL) return NULL; oldcontext = CurrentMemoryContext; oldowner = CurrentResourceOwner; PLy_spi_subtransaction_begin(oldcontext, oldowner); PG_TRY(); { SPI_cursor_fetch(portal, true, count); if (cursor->result.is_rowtype != 1) PLy_input_tuple_funcs(&cursor->result, SPI_tuptable->tupdesc); Py_DECREF(ret->status); ret->status = PyInt_FromLong(SPI_OK_FETCH); Py_DECREF(ret->nrows); ret->nrows = PyInt_FromLong(SPI_processed); if (SPI_processed != 0) { int i; Py_DECREF(ret->rows); ret->rows = PyList_New(SPI_processed); for (i = 0; i < SPI_processed; i++) { PyObject *row = PLyDict_FromTuple(&cursor->result, SPI_tuptable->vals[i], SPI_tuptable->tupdesc); PyList_SetItem(ret->rows, i, row); } } SPI_freetuptable(SPI_tuptable); PLy_spi_subtransaction_commit(oldcontext, oldowner); } PG_CATCH(); { PLy_spi_subtransaction_abort(oldcontext, oldowner); return NULL; } PG_END_TRY(); return (PyObject *) ret; }
static PyObject * PLy_cursor_plan(PyObject *ob, PyObject *args) { PLyCursorObject *cursor; volatile int nargs; int i; PLyPlanObject *plan; volatile MemoryContext oldcontext; volatile ResourceOwner oldowner; if (args) { if (!PySequence_Check(args) || PyString_Check(args) || PyUnicode_Check(args)) { PLy_exception_set(PyExc_TypeError, "plpy.cursor takes a sequence as its second argument"); return NULL; } nargs = PySequence_Length(args); } else nargs = 0; plan = (PLyPlanObject *) ob; if (nargs != plan->nargs) { char *sv; PyObject *so = PyObject_Str(args); if (!so) PLy_elog(ERROR, "could not execute plan"); sv = PyString_AsString(so); PLy_exception_set_plural(PyExc_TypeError, "Expected sequence of %d argument, got %d: %s", "Expected sequence of %d arguments, got %d: %s", plan->nargs, plan->nargs, nargs, sv); Py_DECREF(so); return NULL; } if ((cursor = PyObject_New(PLyCursorObject, &PLy_CursorType)) == NULL) return NULL; cursor->portalname = NULL; cursor->closed = false; cursor->mcxt = AllocSetContextCreate(TopMemoryContext, "PL/Python cursor context", ALLOCSET_DEFAULT_MINSIZE, ALLOCSET_DEFAULT_INITSIZE, ALLOCSET_DEFAULT_MAXSIZE); PLy_typeinfo_init(&cursor->result, cursor->mcxt); oldcontext = CurrentMemoryContext; oldowner = CurrentResourceOwner; PLy_spi_subtransaction_begin(oldcontext, oldowner); PG_TRY(); { PLyExecutionContext *exec_ctx = PLy_current_execution_context(); Portal portal; char *volatile nulls; volatile int j; if (nargs > 0) nulls = palloc(nargs * sizeof(char)); else nulls = NULL; for (j = 0; j < nargs; j++) { PyObject *elem; elem = PySequence_GetItem(args, j); if (elem != Py_None) { PG_TRY(); { plan->values[j] = plan->args[j].out.d.func(&(plan->args[j].out.d), -1, elem); } PG_CATCH(); { Py_DECREF(elem); PG_RE_THROW(); } PG_END_TRY(); Py_DECREF(elem); nulls[j] = ' '; } else { Py_DECREF(elem); plan->values[j] = InputFunctionCall(&(plan->args[j].out.d.typfunc), NULL, plan->args[j].out.d.typioparam, -1); nulls[j] = 'n'; } } portal = SPI_cursor_open(NULL, plan->plan, plan->values, nulls, exec_ctx->curr_proc->fn_readonly); if (portal == NULL) elog(ERROR, "SPI_cursor_open() failed: %s", SPI_result_code_string(SPI_result)); cursor->portalname = MemoryContextStrdup(cursor->mcxt, portal->name); PLy_spi_subtransaction_commit(oldcontext, oldowner); } PG_CATCH(); { int k; /* cleanup plan->values array */ for (k = 0; k < nargs; k++) { if (!plan->args[k].out.d.typbyval && (plan->values[k] != PointerGetDatum(NULL))) { pfree(DatumGetPointer(plan->values[k])); plan->values[k] = PointerGetDatum(NULL); } } Py_DECREF(cursor); PLy_spi_subtransaction_abort(oldcontext, oldowner); return NULL; } PG_END_TRY(); for (i = 0; i < nargs; i++) { if (!plan->args[i].out.d.typbyval && (plan->values[i] != PointerGetDatum(NULL))) { pfree(DatumGetPointer(plan->values[i])); plan->values[i] = PointerGetDatum(NULL); } } Assert(cursor->portalname != NULL); return (PyObject *) cursor; }