diff --git a/src/backend/utils/adt/arrayfuncs.c b/src/backend/utils/adt/arrayfuncs.c index 3ae1556..7f1d729 100644 --- a/src/backend/utils/adt/arrayfuncs.c +++ b/src/backend/utils/adt/arrayfuncs.c @@ -18,7 +18,9 @@ #include #include "access/htup_details.h" +#include "access/xact.h" #include "catalog/pg_type.h" +#include "executor/spi.h" #include "funcapi.h" #include "libpq/pqformat.h" #include "nodes/nodeFuncs.h" @@ -5939,9 +5941,227 @@ array_fill_internal(ArrayType *dims, ArrayType *lbs, return result; } +/* + * UNNEST (REFCURSOR) + */ +Datum +refcursor_unnest(PG_FUNCTION_ARGS) +{ + typedef struct + { + Portal portal; + } refcursor_unnest_fctx; + + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; + FuncCallContext *funcctx; + refcursor_unnest_fctx *fctx; + bool connected = false; + + /* stuff done only on the first call of the function */ + bool first_call = SRF_IS_FIRSTCALL(); + if (first_call) + { + /* create a function context for cross-call persistence */ + funcctx = SRF_FIRSTCALL_INIT(); + + /* Check to see if caller supports us returning a tuplestore */ + if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo)) + ereport(ERROR, + (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), + errmsg("set-valued function called in context that cannot accept a set"))); + if (!(rsinfo->allowedModes & SFRM_Materialize)) + ereport(ERROR, + (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), + errmsg("materialize mode required, but it is not " \ + "allowed in this context"))); + + /* + * switch to memory context appropriate for multiple function calls + */ + MemoryContext oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + + /* allocate memory for user context */ + fctx = (refcursor_unnest_fctx *) palloc(sizeof(refcursor_unnest_fctx)); + + MemoryContextSwitchTo(oldcontext); + + char *portal_name = text_to_cstring(PG_GETARG_TEXT_PP(0)); + + if (SPI_connect() != SPI_OK_CONNECT) + elog(ERROR, "could not connect to SPI manager"); + + /* remember to disconnect later... */ + connected = true; + + Portal portal = SPI_cursor_find(portal_name); + + if (portal == NULL) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_CURSOR), + errmsg("cursor \"%s\" does not exist", portal_name))); + + // FIXME: verify it's only a SELECT + + /* initialize state */ + fctx->portal = portal; + + funcctx->user_fctx = fctx; + + // Reset cursor position to top of set. + // FIXME: what if the cursor isn't scrollable? + SPI_scroll_cursor_move (portal, FETCH_ABSOLUTE, 0); + } + + /* stuff done on every call of the function */ + funcctx = SRF_PERCALL_SETUP(); + fctx = funcctx->user_fctx; + + if (!connected) /* if we have not connected above ... */ + if (SPI_connect() != SPI_OK_CONNECT) + elog(ERROR, "could not connect to SPI manager"); + + // Retrieve a single row... + SPI_cursor_fetch(fctx->portal, true, 1); + + // Initialise the Tuple Desriptor. (This can't be done until we have done our first fetch.) + if (first_call) + { + MemoryContext oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + + /* Build a tuplestore to return our results in */ + rsinfo->setDesc = CreateTupleDescCopy (SPI_tuptable->tupdesc); + + MemoryContextSwitchTo(oldcontext); + rsinfo->returnMode = SFRM_ValuePerCall; + } + + bool next; + Datum result; + + Assert (SPI_processed <= 1); + + if (SPI_processed == 1) + { + result = PointerGetDatum (SPI_returntuple (SPI_tuptable->vals[0], SPI_tuptable->tupdesc)); + + next = true; + } + else // no rows retrieved + { + next = false; + } + + if (SPI_finish() != SPI_OK_FINISH) + elog(ERROR, "could not disconnect from SPI manager"); + connected = false; + + if (next) + SRF_RETURN_NEXT (funcctx, result); + else + SRF_RETURN_DONE(funcctx); +} + + +/* + * Planner support function for UNNEST (REFCURSOR) + */ +Datum +refcursor_unnest_support(PG_FUNCTION_ARGS) +{ + Node *rawreq = (Node *) PG_GETARG_POINTER(0); + Node *ret = NULL; + + if (!IsA(rawreq, SupportRequestRows) + && !IsA (rawreq, SupportRequestCost)) + PG_RETURN_POINTER(NULL); + + Node *req_node; + if (IsA(rawreq, SupportRequestRows)) + { + SupportRequestRows *req = (SupportRequestRows *) rawreq; + + req_node = req->node; + } + else if (IsA (rawreq, SupportRequestCost)) + { + SupportRequestCost *req = (SupportRequestCost *) rawreq; + + req_node = req->node; + } + + // The call to UNNEST should be in a FuncExpr node. + if (!is_funcclause(req_node)) + PG_RETURN_POINTER(NULL); + + List *args = ((FuncExpr *) req_node)->args; + Node *arg1 = linitial(args); + + if (arg1 == NULL) + PG_RETURN_POINTER(NULL); + + // We can only estimate the cost if the REFCURSOR is + // already simplified to a Const. + if (!IsA (arg1, Const)) + PG_RETURN_POINTER(NULL); + + Const *cexpr = (Const *) arg1; + + if (cexpr->constisnull) + PG_RETURN_POINTER(NULL); + + if (cexpr->consttype != REFCURSOROID) + PG_RETURN_POINTER(NULL); + + // We can ignore a check on the collation because we are not + // interested in sorting, and typemod because REFCURSOR has + // no modifyable attributes. + + Oid typoutput; + bool typIsVarlena; + getTypeOutputInfo(cexpr->consttype, &typoutput, &typIsVarlena); + + char *portal_name = OidOutputFunctionCall(typoutput, cexpr->constvalue); + + if (SPI_connect() != SPI_OK_CONNECT) + elog(ERROR, "could not connect to SPI manager"); + + Portal portal = SPI_cursor_find(portal_name); + + if (portal == NULL) + PG_RETURN_POINTER(NULL); + + QueryDesc *qdesc = portal->queryDesc; + if (qdesc == NULL) + PG_RETURN_POINTER(NULL); + + PlanState *planstate = qdesc->planstate; + if (planstate == NULL) + PG_RETURN_POINTER(NULL); + + if (IsA(rawreq, SupportRequestRows)) + { + SupportRequestRows *req = (SupportRequestRows *) rawreq; + + req->rows = planstate->plan->plan_rows; + } + else if (IsA (rawreq, SupportRequestCost)) + { + SupportRequestCost *req = (SupportRequestCost *) rawreq; + + req->startup = planstate->plan->startup_cost; + req->per_tuple = (planstate->plan->total_cost - planstate->plan->startup_cost) / planstate->plan->plan_rows; + } + + if (SPI_finish() != SPI_OK_FINISH) + elog(ERROR, "could not disconnect from SPI manager"); + + ret = (Node *) rawreq; + + PG_RETURN_POINTER(ret); +} /* - * UNNEST + * UNNEST (array) */ Datum array_unnest(PG_FUNCTION_ARGS) diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat index 8733524..311d74f 100644 --- a/src/include/catalog/pg_proc.dat +++ b/src/include/catalog/pg_proc.dat @@ -1550,6 +1550,13 @@ { oid => '3996', descr => 'planner support for array_unnest', proname => 'array_unnest_support', prorettype => 'internal', proargtypes => 'internal', prosrc => 'array_unnest_support' }, +{ oid => '12921', descr => 'expand refcursor to set of rows', + proname => 'unnest', prorows => '100', prosupport => 'refcursor_unnest_support', + proretset => 't', prorettype => 'record', proargtypes => 'refcursor', + prosrc => 'refcursor_unnest' }, +{ oid => '12923', descr => 'planner support for refcursor_unnest', + proname => 'refcursor_unnest_support', prorettype => 'internal', + proargtypes => 'internal', prosrc => 'refcursor_unnest_support' }, { oid => '3167', descr => 'remove any occurrences of an element from an array', proname => 'array_remove', proisstrict => 'f', prorettype => 'anyarray',