C++ SRF_RETURN_NEXT函数代码示例

/*
 * Logic for set-returning functions.
 *
 * Currently it uses the simplest, return
 * one value/tuple per call mechanism.
 */
static Datum
handle_ret_set(FunctionCallInfo fcinfo)
{
	ProxyFunction *func;
	FuncCallContext *ret_ctx;

	if (SRF_IS_FIRSTCALL())
	{
		func = compile_and_execute(fcinfo);
		ret_ctx = SRF_FIRSTCALL_INIT();
		ret_ctx->user_fctx = func;
	}

	ret_ctx = SRF_PERCALL_SETUP();
	func = ret_ctx->user_fctx;

	if (func->cur_cluster->ret_total > 0)
	{
		SRF_RETURN_NEXT(ret_ctx, plproxy_result(func, fcinfo));
	}
	else
	{
		plproxy_clean_results(func->cur_cluster);
		SRF_RETURN_DONE(ret_ctx);
	}
}

Datum
skeys(PG_FUNCTION_ARGS)
{
	FuncCallContext *funcctx;
	AKStore    *st;

	if (SRF_IS_FIRSTCALL())
	{
		HStore	   *hs = PG_GETARG_HS(0);

		funcctx = SRF_FIRSTCALL_INIT();
		setup_firstcall(funcctx, hs);
		PG_FREE_IF_COPY(hs, 0);
	}

	funcctx = SRF_PERCALL_SETUP();
	st = (AKStore *) funcctx->user_fctx;

	if (st->i < st->hs->size)
	{
		HEntry	   *ptr = &(ARRPTR(st->hs)[st->i]);
		text	   *item = (text *) palloc(VARHDRSZ + ptr->keylen);

		SET_VARSIZE(item, VARHDRSZ + ptr->keylen);
		memcpy(VARDATA(item), STRPTR(st->hs) + ptr->pos, ptr->keylen);
		st->i++;

		SRF_RETURN_NEXT(funcctx, PointerGetDatum(item));
	}

	pfree(st->hs);
	pfree(st);

	SRF_RETURN_DONE(funcctx);
}

Datum
hstore_skeys(PG_FUNCTION_ARGS)
{
    FuncCallContext *funcctx;
    HStore	   *hs;
    int			i;

    if (SRF_IS_FIRSTCALL())
    {
        hs = PG_GETARG_HS(0);
        funcctx = SRF_FIRSTCALL_INIT();
        setup_firstcall(funcctx, hs, NULL);
    }

    funcctx = SRF_PERCALL_SETUP();
    hs = (HStore *) funcctx->user_fctx;
    i = funcctx->call_cntr;

    if (i < HS_COUNT(hs))
    {
        HEntry	   *entries = ARRPTR(hs);
        text	   *item;

        item = cstring_to_text_with_len(HS_KEY(entries, STRPTR(hs), i),
                                        HS_KEYLEN(entries, i));

        SRF_RETURN_NEXT(funcctx, PointerGetDatum(item));
    }

    SRF_RETURN_DONE(funcctx);
}

/*
 * txid_snapshot_xip(txid_snapshot) returns setof int8
 *
 *		return in-progress TXIDs in snapshot.
 */
Datum
txid_snapshot_xip(PG_FUNCTION_ARGS)
{
	FuncCallContext *fctx;
	TxidSnapshot *snap;
	txid		value;

	/* on first call initialize snap_state and get copy of snapshot */
	if (SRF_IS_FIRSTCALL())
	{
		TxidSnapshot *arg = (TxidSnapshot *) PG_GETARG_VARLENA_P(0);

		fctx = SRF_FIRSTCALL_INIT();

		/* make a copy of user snapshot */
		snap = MemoryContextAlloc(fctx->multi_call_memory_ctx, VARSIZE(arg));
		memcpy(snap, arg, VARSIZE(arg));

		fctx->user_fctx = snap;
	}

	/* return values one-by-one */
	fctx = SRF_PERCALL_SETUP();
	snap = fctx->user_fctx;
	if (fctx->call_cntr < snap->nxip)
	{
		value = snap->xip[fctx->call_cntr];
		SRF_RETURN_NEXT(fctx, Int64GetDatum(value));
	}
	else
	{
		SRF_RETURN_DONE(fctx);
	}
}

/*
 * Return the TIDs of not-all-visible tuples in pages marked all-visible
 * in the visibility map.  We hope no one will ever find any, but there could
 * be bugs, database corruption, etc.
 */
Datum
pg_check_visible(PG_FUNCTION_ARGS)
{
    FuncCallContext *funcctx;
    corrupt_items *items;

    if (SRF_IS_FIRSTCALL())
    {
        Oid			relid = PG_GETARG_OID(0);
        MemoryContext oldcontext;

        funcctx = SRF_FIRSTCALL_INIT();
        oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
        funcctx->user_fctx = collect_corrupt_items(relid, true, false);
        MemoryContextSwitchTo(oldcontext);
    }

    funcctx = SRF_PERCALL_SETUP();
    items = (corrupt_items *) funcctx->user_fctx;

    if (items->next < items->count)
        SRF_RETURN_NEXT(funcctx, PointerGetDatum(&items->tids[items->next++]));

    SRF_RETURN_DONE(funcctx);
}

Datum
get_next_row(FunctionCallInfo fcinfo)
{
    PlxResult       *plx_result;
    FuncCallContext *funcctx;
    int              call_cntr;

    plx_result = get_plx_result_or_due(fcinfo);
    /* code below will be never executed if pg_result not find */
    if (SRF_IS_FIRSTCALL())
    {
        funcctx = SRF_FIRSTCALL_INIT();
        funcctx->max_calls = PQntuples(plx_result->pg_result);
    }
    funcctx = SRF_PERCALL_SETUP();
    call_cntr = funcctx->call_cntr;
    if (call_cntr < funcctx->max_calls)
        SRF_RETURN_NEXT(funcctx, get_row(fcinfo, plx_result, call_cntr));
    else
    {
        plx_result_cache_delete(fcinfo);
        PQclear(plx_result->pg_result);
        pfree(plx_result);
        SRF_RETURN_DONE(funcctx);
    }
}

/*
 * master_get_table_ddl_events takes in a relation name, and returns the set of
 * DDL commands needed to reconstruct the relation. The returned DDL commands
 * are similar in flavor to schema definitions that pgdump returns. The function
 * errors if given relation does not exist.
 */
Datum
master_get_table_ddl_events(PG_FUNCTION_ARGS)
{
	FuncCallContext *functionContext = NULL;
	ListCell *tableDDLEventCell = NULL;

	/*
	 * On the very first call to this function, we first use the given relation
	 * name to get to the relation. We then recreate the list of DDL statements
	 * issued for this relation, and save the first statement's position in the
	 * function context.
	 */
	if (SRF_IS_FIRSTCALL())
	{
		text *relationName = PG_GETARG_TEXT_P(0);
		Oid relationId = ResolveRelationId(relationName);

		MemoryContext oldContext = NULL;
		List *tableDDLEventList = NIL;

		/* create a function context for cross-call persistence */
		functionContext = SRF_FIRSTCALL_INIT();

		/* switch to memory context appropriate for multiple function calls */
		oldContext = MemoryContextSwitchTo(functionContext->multi_call_memory_ctx);

		/* allocate DDL statements, and then save position in DDL statements */
		tableDDLEventList = GetTableDDLEvents(relationId);
		tableDDLEventCell = list_head(tableDDLEventList);

		functionContext->user_fctx = tableDDLEventCell;

		MemoryContextSwitchTo(oldContext);
	}

	/*
	 * On every call to this function, we get the current position in the
	 * statement list. We then iterate to the next position in the list and
	 * return the current statement, if we have not yet reached the end of
	 * list.
	 */
	functionContext = SRF_PERCALL_SETUP();

	tableDDLEventCell = (ListCell *) functionContext->user_fctx;
	if (tableDDLEventCell != NULL)
	{
		char *ddlStatement = (char *) lfirst(tableDDLEventCell);
		text *ddlStatementText = cstring_to_text(ddlStatement);

		functionContext->user_fctx = lnext(tableDDLEventCell);

		SRF_RETURN_NEXT(functionContext, PointerGetDatum(ddlStatementText));
	}
	else
	{
		SRF_RETURN_DONE(functionContext);
	}
}

/*
 * gp_partition_inverse
 *   Returns all child partition oids with their constraints for a given parent oid.
 *
 * Currently, this function assumes that the parent partition is the root partition.
 *
 * This function is a set-returning function.
 */
Datum
gp_partition_inverse(PG_FUNCTION_ARGS)
{
	FuncCallContext *funcCallContext = NULL;
	InverseContext *inverseContext = NULL;
	
	/*
	 * Setup the function call context for set-returning functions.
	 * At the first time of calling this function, we create and initialize
	 * necessary context data in inverseContext, such as finding the partition
	 * metadata for the given parent oid.
	 */
	if (SRF_IS_FIRSTCALL())
	{
		funcCallContext = SRF_FIRSTCALL_INIT();

		Oid parentOid = PG_GETARG_OID(0);

		MemoryContext oldContext = MemoryContextSwitchTo(funcCallContext->multi_call_memory_ctx);

		funcCallContext->user_fctx = createInverseContext(parentOid);
		inverseContext = (InverseContext *)funcCallContext->user_fctx;

		Assert(NULL != inverseContext);
		Assert(NULL != inverseContext->partitionIterator);
		Assert(NULL != inverseContext->partitionIterator->partsAndRules);

		Partition *part = inverseContext->partitionIterator->partsAndRules->part;
		Assert(NULL != part);

		Oid typeOid = 0;
		int32 typeMod = 0;
		findPartitionKeyType(parentOid, part->paratts[0], &typeOid, &typeMod);

		TupleDesc tupleDesc = createInverseTupleDesc(typeOid, typeMod);
		funcCallContext->tuple_desc = BlessTupleDesc(tupleDesc);

		MemoryContextSwitchTo(oldContext);
	}
	funcCallContext = SRF_PERCALL_SETUP();

	inverseContext = (InverseContext *)funcCallContext->user_fctx;
	Assert(inverseContext != NULL &&
		   inverseContext->partitionIterator != NULL);

	if (inverseContext->findNextRecord(inverseContext))
	{
		HeapTuple tuple = heap_form_tuple(funcCallContext->tuple_desc,
										  inverseContext->values,
										  inverseContext->nulls);
		Datum result = HeapTupleGetDatum(tuple);
		SRF_RETURN_NEXT(funcCallContext, result);
	}
	
	freeInverseContext(inverseContext);

	SRF_RETURN_DONE(funcCallContext);
}

/*
 * List a directory (returns the filenames only)
 */
Datum
pg_ls_dir(PG_FUNCTION_ARGS)
{
	FuncCallContext *funcctx;
	struct dirent *de;
	directory_fctx *fctx;

	if (!superuser())
		ereport(ERROR,
				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
				 (errmsg("must be superuser to get directory listings"))));

	if (SRF_IS_FIRSTCALL())
	{
		MemoryContext oldcontext;

		funcctx = SRF_FIRSTCALL_INIT();
		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

		fctx = palloc(sizeof(directory_fctx));
		fctx->location = convert_and_check_filename(PG_GETARG_TEXT_P(0));

		fctx->dirdesc = AllocateDir(fctx->location);

		if (!fctx->dirdesc)
			ereport(ERROR,
					(errcode_for_file_access(),
					 errmsg("could not open directory \"%s\": %m",
							fctx->location)));

		funcctx->user_fctx = fctx;
		MemoryContextSwitchTo(oldcontext);
	}

	funcctx = SRF_PERCALL_SETUP();
	fctx = (directory_fctx *) funcctx->user_fctx;

	while ((de = ReadDir(fctx->dirdesc, fctx->location)) != NULL)
	{
		int			len = strlen(de->d_name);
		text	   *result;

		if (strcmp(de->d_name, ".") == 0 ||
			strcmp(de->d_name, "..") == 0)
			continue;

		result = palloc(len + VARHDRSZ);
		SET_VARSIZE(result, len + VARHDRSZ);
		memcpy(VARDATA(result), de->d_name, len);

		SRF_RETURN_NEXT(funcctx, PointerGetDatum(result));
	}

	FreeDir(fctx->dirdesc);

	SRF_RETURN_DONE(funcctx);
}

Datum
hstore_each(PG_FUNCTION_ARGS)
{
    FuncCallContext *funcctx;
    HStore	   *hs;
    int			i;

    if (SRF_IS_FIRSTCALL())
    {
        hs = PG_GETARG_HS(0);
        funcctx = SRF_FIRSTCALL_INIT();
        setup_firstcall(funcctx, hs, fcinfo);
    }

    funcctx = SRF_PERCALL_SETUP();
    hs = (HStore *) funcctx->user_fctx;
    i = funcctx->call_cntr;

    if (i < HS_COUNT(hs))
    {
        HEntry	   *entries = ARRPTR(hs);
        char	   *ptr = STRPTR(hs);
        Datum		res,
                    dvalues[2];
        bool		nulls[2] = {false, false};
        text	   *item;
        HeapTuple	tuple;

        item = cstring_to_text_with_len(HS_KEY(entries, ptr, i),
                                        HS_KEYLEN(entries, i));
        dvalues[0] = PointerGetDatum(item);

        if (HS_VALISNULL(entries, i))
        {
            dvalues[1] = (Datum) 0;
            nulls[1] = true;
        }
        else
        {
            item = cstring_to_text_with_len(HS_VAL(entries, ptr, i),
                                            HS_VALLEN(entries, i));
            dvalues[1] = PointerGetDatum(item);
        }

        tuple = heap_form_tuple(funcctx->tuple_desc, dvalues, nulls);
        res = HeapTupleGetDatum(tuple);

        SRF_RETURN_NEXT(funcctx, PointerGetDatum(res));
    }

    SRF_RETURN_DONE(funcctx);
}

/*
 * pgstrom_debug_info
 *
 * shows user's debug information
 */
Datum
pgstrom_debug_info(PG_FUNCTION_ARGS)
{
	FuncCallContext	   *fncxt;
	MemoryContext		oldcxt;
	ListCell   *cell;
	DefElem	   *defel;
	Datum		values[2];
	bool		isnull[2];
	HeapTuple	tuple;

	if (SRF_IS_FIRSTCALL())
	{
		TupleDesc	tupdesc;
		List	   *debug_info_list = NIL;

		fncxt = SRF_FIRSTCALL_INIT();

		oldcxt = MemoryContextSwitchTo(fncxt->multi_call_memory_ctx);

		tupdesc = CreateTemplateTupleDesc(2, false);
		TupleDescInitEntry(tupdesc, (AttrNumber) 1, "key",
						   TEXTOID, -1, 0);
		TupleDescInitEntry(tupdesc, (AttrNumber) 2, "value",
						   TEXTOID, -1, 0);

		debug_info_list = pgstrom_scan_debug_info(debug_info_list);

		fncxt->user_fctx = (void *) debug_info_list;
		fncxt->tuple_desc = BlessTupleDesc(tupdesc);

		MemoryContextSwitchTo(oldcxt);
	}
	fncxt = SRF_PERCALL_SETUP();
	cell = list_head((List *)fncxt->user_fctx);
	if (!cell)
		SRF_RETURN_DONE(fncxt);

	defel = lfirst(cell);
	Assert(IsA(defel, DefElem));

	memset(isnull, false, sizeof(isnull));
	values[0] = CStringGetTextDatum(defel->defname);
	values[1] = CStringGetTextDatum(strVal(defel->arg));

	tuple = heap_form_tuple(fncxt->tuple_desc, values, isnull);

	fncxt->user_fctx = list_delete_ptr((List *)fncxt->user_fctx,
									   lfirst(cell));
	SRF_RETURN_NEXT(fncxt, HeapTupleGetDatum(tuple));
}

Datum
svals(PG_FUNCTION_ARGS)
{
	FuncCallContext *funcctx;
	AKStore    *st;

	if (SRF_IS_FIRSTCALL())
	{
		HStore	   *hs = PG_GETARG_HS(0);

		funcctx = SRF_FIRSTCALL_INIT();
		setup_firstcall(funcctx, hs);
		PG_FREE_IF_COPY(hs, 0);
	}

	funcctx = SRF_PERCALL_SETUP();
	st = (AKStore *) funcctx->user_fctx;

	if (st->i < st->hs->size)
	{
		HEntry	   *ptr = &(ARRPTR(st->hs)[st->i]);

		if (ptr->valisnull)
		{
			ReturnSetInfo *rsi;

			st->i++;
			(funcctx)->call_cntr++;
			rsi = (ReturnSetInfo *) fcinfo->resultinfo;
			rsi->isDone = ExprMultipleResult;
			PG_RETURN_NULL();
		}
		else
		{
			int			vallen = ptr->vallen;
			text	   *item = (text *) palloc(VARHDRSZ + vallen);

			SET_VARSIZE(item, VARHDRSZ + vallen);
			memcpy(VARDATA(item), STRPTR(st->hs) + ptr->pos + ptr->keylen, vallen);
			st->i++;

			SRF_RETURN_NEXT(funcctx, PointerGetDatum(item));
		}
	}

	pfree(st->hs);
	pfree(st);

	SRF_RETURN_DONE(funcctx);
}

/*
 * Return the TID array stored in a BRIN revmap page
 */
Datum
brin_revmap_data(PG_FUNCTION_ARGS)
{
    struct
    {
        ItemPointerData *tids;
        int			idx;
    }		   *state;
    FuncCallContext *fctx;

    if (!superuser())
        ereport(ERROR,
                (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
                 (errmsg("must be superuser to use raw page functions"))));

    if (SRF_IS_FIRSTCALL())
    {
        bytea	   *raw_page = PG_GETARG_BYTEA_P(0);
        MemoryContext mctx;
        Page		page;

        /* minimally verify the page we got */
        page = verify_brin_page(raw_page, BRIN_PAGETYPE_REVMAP, "revmap");

        /* create a function context for cross-call persistence */
        fctx = SRF_FIRSTCALL_INIT();

        /* switch to memory context appropriate for multiple function calls */
        mctx = MemoryContextSwitchTo(fctx->multi_call_memory_ctx);

        state = palloc(sizeof(*state));
        state->tids = ((RevmapContents *) PageGetContents(page))->rm_tids;
        state->idx = 0;

        fctx->user_fctx = state;

        MemoryContextSwitchTo(mctx);
    }

    fctx = SRF_PERCALL_SETUP();
    state = fctx->user_fctx;

    if (state->idx < REVMAP_PAGE_MAXITEMS)
        SRF_RETURN_NEXT(fctx, PointerGetDatum(&state->tids[state->idx++]));

    SRF_RETURN_DONE(fctx);
}

/*
 * Returns Oids of tables in a publication.
 */
Datum
pg_get_publication_tables(PG_FUNCTION_ARGS)
{
	FuncCallContext *funcctx;
	char	   *pubname = text_to_cstring(PG_GETARG_TEXT_PP(0));
	Publication *publication;
	List	   *tables;
	ListCell  **lcp;

	/* stuff done only on the first call of the function */
	if (SRF_IS_FIRSTCALL())
	{
		MemoryContext oldcontext;

		/* create a function context for cross-call persistence */
		funcctx = SRF_FIRSTCALL_INIT();

		/* switch to memory context appropriate for multiple function calls */
		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

		publication = GetPublicationByName(pubname, false);
		if (publication->alltables)
			tables = GetAllTablesPublicationRelations();
		else
			tables = GetPublicationRelations(publication->oid);
		lcp = (ListCell **) palloc(sizeof(ListCell *));
		*lcp = list_head(tables);
		funcctx->user_fctx = (void *) lcp;

		MemoryContextSwitchTo(oldcontext);
	}

	/* stuff done on every call of the function */
	funcctx = SRF_PERCALL_SETUP();
	lcp = (ListCell **) funcctx->user_fctx;

	while (*lcp != NULL)
	{
		Oid			relid = lfirst_oid(*lcp);

		*lcp = lnext(*lcp);
		SRF_RETURN_NEXT(funcctx, ObjectIdGetDatum(relid));
	}

	SRF_RETURN_DONE(funcctx);
}

/*
 * regexp_matches()
 *		Return a table of matches of a pattern within a string.
 */
Datum
regexp_matches(PG_FUNCTION_ARGS)
{
	FuncCallContext *funcctx;
	regexp_matches_ctx *matchctx;

	if (SRF_IS_FIRSTCALL())
	{
		text	   *pattern = PG_GETARG_TEXT_PP(1);
		text	   *flags = PG_GETARG_TEXT_PP_IF_EXISTS(2);
		MemoryContext oldcontext;

		funcctx = SRF_FIRSTCALL_INIT();
		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

		/* be sure to copy the input string into the multi-call ctx */
		matchctx = setup_regexp_matches(PG_GETARG_TEXT_P_COPY(0), pattern,
										flags,
										PG_GET_COLLATION(),
										false, true, false);

		/* Pre-create workspace that build_regexp_matches_result needs */
		matchctx->elems = (Datum *) palloc(sizeof(Datum) * matchctx->npatterns);
		matchctx->nulls = (bool *) palloc(sizeof(bool) * matchctx->npatterns);

		MemoryContextSwitchTo(oldcontext);
		funcctx->user_fctx = (void *) matchctx;
	}

	funcctx = SRF_PERCALL_SETUP();
	matchctx = (regexp_matches_ctx *) funcctx->user_fctx;

	if (matchctx->next_match < matchctx->nmatches)
	{
		ArrayType  *result_ary;

		result_ary = build_regexp_matches_result(matchctx);
		matchctx->next_match++;
		SRF_RETURN_NEXT(funcctx, PointerGetDatum(result_ary));
	}

	/* release space in multi-call ctx to avoid intraquery memory leak */
	cleanup_regexp_matches(matchctx);

	SRF_RETURN_DONE(funcctx);
}