From 9ae75f693adf246ad551e258198606fd63190a20 Mon Sep 17 00:00:00 2001 From: Justin Pryzby Date: Wed, 11 Mar 2020 10:09:18 -0500 Subject: [PATCH v2 1/2] SRF: avoid leaking resources if not run to completion Change to return a tuplestore populated immediately and returned in full. Discussion: https://www.postgresql.org/message-id/20200308173103.GC1357%40telsasoft.com See also: 9cb7db3f0 and 085b6b66 --- contrib/adminpack/adminpack.c | 82 +++++----- contrib/pgrowlocks/pgrowlocks.c | 162 +++++++++---------- doc/src/sgml/xfunc.sgml | 17 +- src/backend/utils/adt/datetime.c | 99 +++++------- src/backend/utils/adt/genfile.c | 111 +++++++------ src/backend/utils/adt/misc.c | 116 +++++++------ src/include/funcapi.h | 7 +- src/test/regress/expected/misc_functions.out | 18 +++ src/test/regress/sql/misc_functions.sql | 6 + 9 files changed, 325 insertions(+), 293 deletions(-) diff --git a/contrib/adminpack/adminpack.c b/contrib/adminpack/adminpack.c index bc45e79895..070025f5a8 100644 --- a/contrib/adminpack/adminpack.c +++ b/contrib/adminpack/adminpack.c @@ -504,50 +504,56 @@ pg_logdir_ls_v1_1(PG_FUNCTION_ARGS) static Datum pg_logdir_ls_internal(FunctionCallInfo fcinfo) { - FuncCallContext *funcctx; struct dirent *de; - directory_fctx *fctx; + + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; + MemoryContext oldcontext; + TupleDesc tupdesc; + Tuplestorestate *tupstore; + bool randomAccess; + DIR *dirdesc; + + /* 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_SYNTAX_ERROR), + errmsg("materialize mode required, but it is not allowed in this context"))); if (strcmp(Log_filename, "postgresql-%Y-%m-%d_%H%M%S.log") != 0) ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE), errmsg("the log_filename parameter must equal 'postgresql-%%Y-%%m-%%d_%%H%%M%%S.log'"))); - if (SRF_IS_FIRSTCALL()) - { - MemoryContext oldcontext; - TupleDesc tupdesc; - - funcctx = SRF_FIRSTCALL_INIT(); - oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + /* The Tuplestore and TupleDesc should be created in ecxt_per_query_memory */ + oldcontext = MemoryContextSwitchTo(rsinfo->econtext->ecxt_per_query_memory); + randomAccess = (rsinfo->allowedModes&SFRM_Materialize_Random) != 0; + tupstore = tuplestore_begin_heap(randomAccess, false, work_mem); - fctx = palloc(sizeof(directory_fctx)); + /* Is this right ?? */ + tupdesc = CreateTemplateTupleDesc(2); + TupleDescInitEntry(tupdesc, (AttrNumber) 1, "starttime", + TIMESTAMPOID, -1, 0); + TupleDescInitEntry(tupdesc, (AttrNumber) 2, "filename", + TEXTOID, -1, 0); - tupdesc = CreateTemplateTupleDesc(2); - TupleDescInitEntry(tupdesc, (AttrNumber) 1, "starttime", - TIMESTAMPOID, -1, 0); - TupleDescInitEntry(tupdesc, (AttrNumber) 2, "filename", - TEXTOID, -1, 0); + rsinfo->returnMode = SFRM_Materialize; + rsinfo->setResult = tupstore; + rsinfo->setDesc = tupdesc; - funcctx->attinmeta = TupleDescGetAttInMetadata(tupdesc); + MemoryContextSwitchTo(oldcontext); - fctx->location = pstrdup(Log_directory); - 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; + dirdesc = AllocateDir(Log_directory); + if (!dirdesc) + ereport(ERROR, + (errcode_for_file_access(), + errmsg("could not open directory \"%s\": %m", + Log_directory))); - while ((de = ReadDir(fctx->dirdesc, fctx->location)) != NULL) + while ((de = ReadDir(dirdesc, Log_directory)) != NULL) { char *values[2]; HeapTuple tuple; @@ -584,13 +590,13 @@ pg_logdir_ls_internal(FunctionCallInfo fcinfo) /* Seems the timestamp is OK; prepare and return tuple */ values[0] = timestampbuf; - values[1] = psprintf("%s/%s", fctx->location, de->d_name); - - tuple = BuildTupleFromCStrings(funcctx->attinmeta, values); + values[1] = psprintf("%s/%s", Log_directory, de->d_name); - SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple)); + tuple = BuildTupleFromCStrings(TupleDescGetAttInMetadata(tupdesc), values); + tuplestore_puttuple(tupstore, tuple); } - FreeDir(fctx->dirdesc); - SRF_RETURN_DONE(funcctx); + tuplestore_donestoring(tupstore); + FreeDir(dirdesc); + return (Datum) 0; } diff --git a/contrib/pgrowlocks/pgrowlocks.c b/contrib/pgrowlocks/pgrowlocks.c index a2c44a916c..baad83fcde 100644 --- a/contrib/pgrowlocks/pgrowlocks.c +++ b/contrib/pgrowlocks/pgrowlocks.c @@ -54,13 +54,6 @@ PG_FUNCTION_INFO_V1(pgrowlocks); #define NCHARS 32 -typedef struct -{ - Relation rel; - TableScanDesc scan; - int ncolumns; -} MyData; - #define Atnum_tid 0 #define Atnum_xmax 1 #define Atnum_ismulti 2 @@ -71,82 +64,84 @@ typedef struct Datum pgrowlocks(PG_FUNCTION_ARGS) { - FuncCallContext *funcctx; TableScanDesc scan; HeapScanDesc hscan; HeapTuple tuple; TupleDesc tupdesc; AttInMetadata *attinmeta; - Datum result; - MyData *mydata; Relation rel; - if (SRF_IS_FIRSTCALL()) - { - text *relname; - RangeVar *relrv; - MemoryContext oldcontext; - AclResult aclresult; - - funcctx = SRF_FIRSTCALL_INIT(); - oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); - - /* Build a tuple descriptor for our result type */ - if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE) - elog(ERROR, "return type must be a row type"); - - attinmeta = TupleDescGetAttInMetadata(tupdesc); - funcctx->attinmeta = attinmeta; - - relname = PG_GETARG_TEXT_PP(0); - relrv = makeRangeVarFromNameList(textToQualifiedNameList(relname)); - rel = relation_openrv(relrv, AccessShareLock); - - if (rel->rd_rel->relam != HEAP_TABLE_AM_OID) - ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), - errmsg("only heap AM is supported"))); - - if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE) - ereport(ERROR, - (errcode(ERRCODE_WRONG_OBJECT_TYPE), - errmsg("\"%s\" is a partitioned table", - RelationGetRelationName(rel)), - errdetail("Partitioned tables do not contain rows."))); - else if (rel->rd_rel->relkind != RELKIND_RELATION) - ereport(ERROR, - (errcode(ERRCODE_WRONG_OBJECT_TYPE), - errmsg("\"%s\" is not a table", - RelationGetRelationName(rel)))); + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; + Tuplestorestate *tupstore; + bool randomAccess; + + text *relname; + RangeVar *relrv; + MemoryContext oldcontext; + AclResult aclresult; + char **values; + + /* 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_SYNTAX_ERROR), + errmsg("materialize mode required, but it is not allowed in this context"))); + + /* The Tuplestore and TupleDesc should be created in ecxt_per_query_memory */ + oldcontext = MemoryContextSwitchTo(rsinfo->econtext->ecxt_per_query_memory); + randomAccess = (rsinfo->allowedModes&SFRM_Materialize_Random) != 0; + tupstore = tuplestore_begin_heap(randomAccess, false, work_mem); + if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE) + elog(ERROR, "return type must be a row type"); + rsinfo->returnMode = SFRM_Materialize; + rsinfo->setResult = tupstore; + rsinfo->setDesc = tupdesc; + + attinmeta = TupleDescGetAttInMetadata(tupdesc); + + relname = PG_GETARG_TEXT_PP(0); + relrv = makeRangeVarFromNameList(textToQualifiedNameList(relname)); + rel = relation_openrv(relrv, AccessShareLock); + + if (rel->rd_rel->relam != HEAP_TABLE_AM_OID) + ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), + errmsg("only heap AM is supported"))); + + if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE) + ereport(ERROR, + (errcode(ERRCODE_WRONG_OBJECT_TYPE), + errmsg("\"%s\" is a partitioned table", + RelationGetRelationName(rel)), + errdetail("Partitioned tables do not contain rows."))); + else if (rel->rd_rel->relkind != RELKIND_RELATION) + ereport(ERROR, + (errcode(ERRCODE_WRONG_OBJECT_TYPE), + errmsg("\"%s\" is not a table", + RelationGetRelationName(rel)))); + + /* + * check permissions: must have SELECT on table or be in + * pg_stat_scan_tables + */ + aclresult = pg_class_aclcheck(RelationGetRelid(rel), GetUserId(), + ACL_SELECT); + if (aclresult != ACLCHECK_OK) + aclresult = is_member_of_role(GetUserId(), DEFAULT_ROLE_STAT_SCAN_TABLES) ? ACLCHECK_OK : ACLCHECK_NO_PRIV; + + if (aclresult != ACLCHECK_OK) + aclcheck_error(aclresult, get_relkind_objtype(rel->rd_rel->relkind), + RelationGetRelationName(rel)); + + scan = table_beginscan(rel, GetActiveSnapshot(), 0, NULL); + hscan = (HeapScanDesc) scan; - /* - * check permissions: must have SELECT on table or be in - * pg_stat_scan_tables - */ - aclresult = pg_class_aclcheck(RelationGetRelid(rel), GetUserId(), - ACL_SELECT); - if (aclresult != ACLCHECK_OK) - aclresult = is_member_of_role(GetUserId(), DEFAULT_ROLE_STAT_SCAN_TABLES) ? ACLCHECK_OK : ACLCHECK_NO_PRIV; - - if (aclresult != ACLCHECK_OK) - aclcheck_error(aclresult, get_relkind_objtype(rel->rd_rel->relkind), - RelationGetRelationName(rel)); - - scan = table_beginscan(rel, GetActiveSnapshot(), 0, NULL); - hscan = (HeapScanDesc) scan; - mydata = palloc(sizeof(*mydata)); - mydata->rel = rel; - mydata->scan = scan; - mydata->ncolumns = tupdesc->natts; - funcctx->user_fctx = mydata; - - MemoryContextSwitchTo(oldcontext); - } + MemoryContextSwitchTo(oldcontext); - funcctx = SRF_PERCALL_SETUP(); - attinmeta = funcctx->attinmeta; - mydata = (MyData *) funcctx->user_fctx; - scan = mydata->scan; - hscan = (HeapScanDesc) scan; + values = (char **) palloc(tupdesc->natts * sizeof(char *)); /* scan the relation */ while ((tuple = heap_getnext(scan, ForwardScanDirection)) != NULL) @@ -169,9 +164,6 @@ pgrowlocks(PG_FUNCTION_ARGS) */ if (htsu == TM_BeingModified) { - char **values; - - values = (char **) palloc(mydata->ncolumns * sizeof(char *)); values[Atnum_tid] = (char *) DirectFunctionCall1(tidout, PointerGetDatum(&tuple->t_self)); @@ -297,16 +289,7 @@ pgrowlocks(PG_FUNCTION_ARGS) /* build a tuple */ tuple = BuildTupleFromCStrings(attinmeta, values); - - /* make the tuple into a datum */ - result = HeapTupleGetDatum(tuple); - - /* - * no need to pfree what we allocated; it's on a short-lived - * memory context anyway - */ - - SRF_RETURN_NEXT(funcctx, result); + tuplestore_puttuple(tupstore, tuple); } else { @@ -314,8 +297,9 @@ pgrowlocks(PG_FUNCTION_ARGS) } } + tuplestore_donestoring(tupstore); table_endscan(scan); - table_close(mydata->rel, AccessShareLock); + table_close(rel, AccessShareLock); - SRF_RETURN_DONE(funcctx); + return (Datum) 0; } diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml index a91f887119..432758e75e 100644 --- a/doc/src/sgml/xfunc.sgml +++ b/doc/src/sgml/xfunc.sgml @@ -2812,7 +2812,7 @@ HeapTupleGetDatum(HeapTuple tuple) Returning Sets - There is also a special API that provides support for returning + There is also a special, helper API that provides support for returning sets (multiple rows) from a C-language function. A set-returning function must follow the version-1 calling conventions. Also, source files must include funcapi.h, as @@ -2820,10 +2820,19 @@ HeapTupleGetDatum(HeapTuple tuple) - A set-returning function (SRF) is called - once for each item it returns. The SRF must - therefore save enough state to remember what it was doing and + The helper API allows writing a set-returning function + (SRF) using the ValuePerCall mode, in which the SRF is + called once for each item it returns. The SRF must + then save enough state to remember what it was doing and return the next item on each call. + Note, however, that an SRF will possibly not be run to completion, due to + a LIMIT or a cursor, and should avoid leaving resources like DIR*s or + tablescans opened across calls. To avoid leaking resources in those cases, + instead of using ValuePerCall mode, the SRF should populate and return a + tuplestore with SFRM_Materialize mode. See fmgr/README. + + + The structure FuncCallContext is provided to help control this process. Within a function, fcinfo->flinfo->fn_extra is used to hold a pointer to FuncCallContext diff --git a/src/backend/utils/adt/datetime.c b/src/backend/utils/adt/datetime.c index 4f109111d1..d0966c8305 100644 --- a/src/backend/utils/adt/datetime.c +++ b/src/backend/utils/adt/datetime.c @@ -4756,11 +4756,8 @@ Datum pg_timezone_names(PG_FUNCTION_ARGS) { MemoryContext oldcontext; - FuncCallContext *funcctx; pg_tzenum *tzenum; pg_tz *tz; - Datum result; - HeapTuple tuple; Datum values[4]; bool nulls[4]; int tzoff; @@ -4770,58 +4767,44 @@ pg_timezone_names(PG_FUNCTION_ARGS) Interval *resInterval; struct pg_tm itm; - /* stuff done only on the first call of the function */ - if (SRF_IS_FIRSTCALL()) - { - TupleDesc tupdesc; + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; - /* create a function context for cross-call persistence */ - funcctx = SRF_FIRSTCALL_INIT(); + TupleDesc tupdesc; + Tuplestorestate *tupstore; + bool randomAccess; - /* - * switch to memory context appropriate for multiple function calls - */ - oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + /* 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_SYNTAX_ERROR), + errmsg("materialize mode required, but it is not allowed in this context"))); - /* initialize timezone scanning code */ - tzenum = pg_tzenumerate_start(); - funcctx->user_fctx = (void *) tzenum; + /* The Tuplestore and TupleDesc should be created in ecxt_per_query_memory */ + oldcontext = MemoryContextSwitchTo(rsinfo->econtext->ecxt_per_query_memory); + randomAccess = (rsinfo->allowedModes&SFRM_Materialize_Random) != 0; + tupstore = tuplestore_begin_heap(randomAccess, false, work_mem); + if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE) + elog(ERROR, "return type must be a row type"); - /* - * build tupdesc for result tuples. This must match this function's - * pg_proc entry! - */ - tupdesc = CreateTemplateTupleDesc(4); - TupleDescInitEntry(tupdesc, (AttrNumber) 1, "name", - TEXTOID, -1, 0); - TupleDescInitEntry(tupdesc, (AttrNumber) 2, "abbrev", - TEXTOID, -1, 0); - TupleDescInitEntry(tupdesc, (AttrNumber) 3, "utc_offset", - INTERVALOID, -1, 0); - TupleDescInitEntry(tupdesc, (AttrNumber) 4, "is_dst", - BOOLOID, -1, 0); + rsinfo->returnMode = SFRM_Materialize; + rsinfo->setResult = tupstore; + rsinfo->setDesc = tupdesc; - funcctx->tuple_desc = BlessTupleDesc(tupdesc); - MemoryContextSwitchTo(oldcontext); - } + MemoryContextSwitchTo(oldcontext); - /* stuff done on every call of the function */ - funcctx = SRF_PERCALL_SETUP(); - tzenum = (pg_tzenum *) funcctx->user_fctx; + /* initialize timezone scanning code */ + tzenum = pg_tzenumerate_start(); /* search for another zone to display */ for (;;) { - oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); tz = pg_tzenumerate_next(tzenum); - MemoryContextSwitchTo(oldcontext); - if (!tz) - { - pg_tzenumerate_end(tzenum); - funcctx->user_fctx = NULL; - SRF_RETURN_DONE(funcctx); - } + break; /* Convert now() to local time in this zone */ if (timestamp2tm(GetCurrentTransactionStartTimestamp(), @@ -4840,25 +4823,23 @@ pg_timezone_names(PG_FUNCTION_ARGS) if (tzn && strlen(tzn) > 31) continue; - /* Found a displayable zone */ - break; - } - - MemSet(nulls, 0, sizeof(nulls)); - values[0] = CStringGetTextDatum(pg_get_timezone_name(tz)); - values[1] = CStringGetTextDatum(tzn ? tzn : ""); + MemSet(nulls, 0, sizeof(nulls)); + values[0] = CStringGetTextDatum(pg_get_timezone_name(tz)); + values[1] = CStringGetTextDatum(tzn ? tzn : ""); - MemSet(&itm, 0, sizeof(struct pg_tm)); - itm.tm_sec = -tzoff; - resInterval = (Interval *) palloc(sizeof(Interval)); - tm2interval(&itm, 0, resInterval); - values[2] = IntervalPGetDatum(resInterval); + MemSet(&itm, 0, sizeof(struct pg_tm)); + itm.tm_sec = -tzoff; + resInterval = (Interval *) palloc(sizeof(Interval)); + tm2interval(&itm, 0, resInterval); + values[2] = IntervalPGetDatum(resInterval); - values[3] = BoolGetDatum(tm.tm_isdst > 0); + values[3] = BoolGetDatum(tm.tm_isdst > 0); - tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls); - result = HeapTupleGetDatum(tuple); + tuplestore_putvalues(tupstore, tupdesc, values, nulls); + } - SRF_RETURN_NEXT(funcctx, result); + pg_tzenumerate_end(tzenum); + tuplestore_donestoring(tupstore); + return (Datum) 0; } diff --git a/src/backend/utils/adt/genfile.c b/src/backend/utils/adt/genfile.c index bcf9bd1b97..39a54d532f 100644 --- a/src/backend/utils/adt/genfile.c +++ b/src/backend/utils/adt/genfile.c @@ -36,13 +36,6 @@ #include "utils/syscache.h" #include "utils/timestamp.h" -typedef struct -{ - char *location; - DIR *dirdesc; - bool include_dot_dirs; -} directory_fctx; - /* * Convert a "text" filename argument to C string, and check it's allowable. @@ -447,67 +440,81 @@ pg_stat_file_1arg(PG_FUNCTION_ARGS) Datum pg_ls_dir(PG_FUNCTION_ARGS) { - FuncCallContext *funcctx; - struct dirent *de; - directory_fctx *fctx; - MemoryContext oldcontext; + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; + MemoryContext oldcontext; + TupleDesc tupdesc; + Tuplestorestate *tupstore; + bool randomAccess; + + bool missing_ok = false; + bool include_dot_dirs = false; + char *location; + DIR *dirdesc; + struct dirent *de; + + /* 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_SYNTAX_ERROR), + errmsg("materialize mode required, but it is not allowed in this context"))); - if (SRF_IS_FIRSTCALL()) + /* check the optional arguments */ + if (PG_NARGS() == 3) { - bool missing_ok = false; - bool include_dot_dirs = false; + if (!PG_ARGISNULL(1)) + missing_ok = PG_GETARG_BOOL(1); + if (!PG_ARGISNULL(2)) + include_dot_dirs = PG_GETARG_BOOL(2); + } - /* check the optional arguments */ - if (PG_NARGS() == 3) - { - if (!PG_ARGISNULL(1)) - missing_ok = PG_GETARG_BOOL(1); - if (!PG_ARGISNULL(2)) - include_dot_dirs = PG_GETARG_BOOL(2); - } + /* The Tuplestore and TupleDesc should be created in ecxt_per_query_memory */ + oldcontext = MemoryContextSwitchTo(rsinfo->econtext->ecxt_per_query_memory); + randomAccess = (rsinfo->allowedModes&SFRM_Materialize_Random) != 0; + tupstore = tuplestore_begin_heap(randomAccess, false, work_mem); - funcctx = SRF_FIRSTCALL_INIT(); - oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + tupdesc = CreateTemplateTupleDesc(1); + TupleDescInitEntry(tupdesc, (AttrNumber) 1, "column", TEXTOID, -1, 0); - fctx = palloc(sizeof(directory_fctx)); - fctx->location = convert_and_check_filename(PG_GETARG_TEXT_PP(0)); + rsinfo->returnMode = SFRM_Materialize; + rsinfo->setResult = tupstore; + rsinfo->setDesc = tupdesc; + MemoryContextSwitchTo(oldcontext); - fctx->include_dot_dirs = include_dot_dirs; - fctx->dirdesc = AllocateDir(fctx->location); + location = convert_and_check_filename(PG_GETARG_TEXT_PP(0)); + dirdesc = AllocateDir(location); - if (!fctx->dirdesc) - { - if (missing_ok && errno == ENOENT) - { - MemoryContextSwitchTo(oldcontext); - SRF_RETURN_DONE(funcctx); - } - else - ereport(ERROR, - (errcode_for_file_access(), - errmsg("could not open directory \"%s\": %m", - fctx->location))); - } - funcctx->user_fctx = fctx; - MemoryContextSwitchTo(oldcontext); + if (!dirdesc) + { + if (missing_ok && errno == ENOENT) + return (Datum) 0; + else + ereport(ERROR, + (errcode_for_file_access(), + errmsg("could not open directory \"%s\": %m", + location))); } - funcctx = SRF_PERCALL_SETUP(); - fctx = (directory_fctx *) funcctx->user_fctx; - - while ((de = ReadDir(fctx->dirdesc, fctx->location)) != NULL) + while ((de = ReadDir(dirdesc, location)) != NULL) { - if (!fctx->include_dot_dirs && + Datum values[1]; + bool nulls[1] = {false}; + + if (!include_dot_dirs && (strcmp(de->d_name, ".") == 0 || strcmp(de->d_name, "..") == 0)) continue; - SRF_RETURN_NEXT(funcctx, CStringGetTextDatum(de->d_name)); + values[0] = CStringGetTextDatum(de->d_name); + tuplestore_putvalues(tupstore, tupdesc, values, nulls); } - FreeDir(fctx->dirdesc); - - SRF_RETURN_DONE(funcctx); + tuplestore_donestoring(tupstore); + FreeDir(dirdesc); + return (Datum) 0; } /* diff --git a/src/backend/utils/adt/misc.c b/src/backend/utils/adt/misc.c index 323e36b81c..9919087335 100644 --- a/src/backend/utils/adt/misc.c +++ b/src/backend/utils/adt/misc.c @@ -194,72 +194,86 @@ current_query(PG_FUNCTION_ARGS) /* Function to find out which databases make use of a tablespace */ -typedef struct -{ - char *location; - DIR *dirdesc; -} ts_db_fctx; - Datum pg_tablespace_databases(PG_FUNCTION_ARGS) { - FuncCallContext *funcctx; struct dirent *de; - ts_db_fctx *fctx; - if (SRF_IS_FIRSTCALL()) - { - MemoryContext oldcontext; - Oid tablespaceOid = PG_GETARG_OID(0); + MemoryContext oldcontext; + Oid tablespaceOid = PG_GETARG_OID(0); - funcctx = SRF_FIRSTCALL_INIT(); - oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; + TupleDesc tupdesc; + Tuplestorestate *tupstore; + bool randomAccess; - fctx = palloc(sizeof(ts_db_fctx)); + DIR *dirdesc; + char *location; - if (tablespaceOid == GLOBALTABLESPACE_OID) - { - fctx->dirdesc = NULL; - ereport(WARNING, - (errmsg("global tablespace never has databases"))); - } + /* 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_SYNTAX_ERROR), + errmsg("materialize mode required, but it is not allowed in this context"))); + + if (tablespaceOid == GLOBALTABLESPACE_OID) + { + dirdesc = NULL; + ereport(WARNING, + (errmsg("global tablespace never has databases"))); + } + else + { + if (tablespaceOid == DEFAULTTABLESPACE_OID) + location = psprintf("base"); else - { - if (tablespaceOid == DEFAULTTABLESPACE_OID) - fctx->location = psprintf("base"); - else - fctx->location = psprintf("pg_tblspc/%u/%s", tablespaceOid, - TABLESPACE_VERSION_DIRECTORY); + location = psprintf("pg_tblspc/%u/%s", tablespaceOid, + TABLESPACE_VERSION_DIRECTORY); - fctx->dirdesc = AllocateDir(fctx->location); + dirdesc = AllocateDir(location); - if (!fctx->dirdesc) - { - /* the only expected error is ENOENT */ - if (errno != ENOENT) - ereport(ERROR, - (errcode_for_file_access(), - errmsg("could not open directory \"%s\": %m", - fctx->location))); - ereport(WARNING, - (errmsg("%u is not a tablespace OID", tablespaceOid))); - } + if (!dirdesc) + { + /* the only expected error is ENOENT */ + if (errno != ENOENT) + ereport(ERROR, + (errcode_for_file_access(), + errmsg("could not open directory \"%s\": %m", + location))); + ereport(WARNING, + (errmsg("%u is not a tablespace OID", tablespaceOid))); } - funcctx->user_fctx = fctx; - MemoryContextSwitchTo(oldcontext); } - funcctx = SRF_PERCALL_SETUP(); - fctx = (ts_db_fctx *) funcctx->user_fctx; + /* The Tuplestore and TupleDesc should be created in ecxt_per_query_memory */ + oldcontext = MemoryContextSwitchTo(rsinfo->econtext->ecxt_per_query_memory); + + randomAccess = (rsinfo->allowedModes&SFRM_Materialize_Random) != 0; + tupstore = tuplestore_begin_heap(randomAccess, false, work_mem); - if (!fctx->dirdesc) /* not a tablespace */ - SRF_RETURN_DONE(funcctx); + tupdesc = CreateTemplateTupleDesc(1); + TupleDescInitEntry(tupdesc, (AttrNumber) 1, "column", OIDOID, -1, 0); - while ((de = ReadDir(fctx->dirdesc, fctx->location)) != NULL) + rsinfo->returnMode = SFRM_Materialize; + rsinfo->setResult = tupstore; + rsinfo->setDesc = tupdesc; + + MemoryContextSwitchTo(oldcontext); + + if (!dirdesc) /* not a tablespace */ + return 0; + + while ((de = ReadDir(dirdesc, location)) != NULL) { Oid datOid = atooid(de->d_name); char *subdir; bool isempty; + Datum values[1]; + bool nulls[1] = {false}; /* this test skips . and .., but is awfully weak */ if (!datOid) @@ -267,18 +281,20 @@ pg_tablespace_databases(PG_FUNCTION_ARGS) /* if database subdir is empty, don't report tablespace as used */ - subdir = psprintf("%s/%s", fctx->location, de->d_name); + subdir = psprintf("%s/%s", location, de->d_name); isempty = directory_is_empty(subdir); pfree(subdir); if (isempty) continue; /* indeed, nothing in it */ - SRF_RETURN_NEXT(funcctx, ObjectIdGetDatum(datOid)); + values[0] = ObjectIdGetDatum(datOid); + tuplestore_putvalues(tupstore, tupdesc, values, nulls); } - FreeDir(fctx->dirdesc); - SRF_RETURN_DONE(funcctx); + tuplestore_donestoring(tupstore); + FreeDir(dirdesc); + return (Datum) 0; } diff --git a/src/include/funcapi.h b/src/include/funcapi.h index f9b75ae390..471100b7e5 100644 --- a/src/include/funcapi.h +++ b/src/include/funcapi.h @@ -234,7 +234,12 @@ extern Datum HeapTupleHeaderGetDatum(HeapTupleHeader tuple); /*---------- * Support for Set Returning Functions (SRFs) * - * The basic API for SRFs looks something like: + * The basic API for SRFs using ValuePerCall mode looks something like this. + * Note that an SRF will possibly not be run to completion, due to a LIMIT or a + * cursor, and should avoid leaving resources like DIR*s or tablescans opened + * across calls. To avoid leaking resources in those cases, instead of using + * ValuePerCall mode, the SRF should populate and return a tuplestore with + * SFRM_Materialize mode. See fmgr/README. * * Datum * my_Set_Returning_Function(PG_FUNCTION_ARGS) diff --git a/src/test/regress/expected/misc_functions.out b/src/test/regress/expected/misc_functions.out index e217b678d7..b56ce74586 100644 --- a/src/test/regress/expected/misc_functions.out +++ b/src/test/regress/expected/misc_functions.out @@ -180,6 +180,24 @@ select count(*) >= 0 as ok from pg_ls_archive_statusdir(); t (1 row) +select * from (select pg_ls_dir('.')a)a where a='base' limit 1; + a +------ + base +(1 row) + +select * from (select (pg_timezone_names()).name)ptn where name='UTC' limit 1; + name +------ + UTC +(1 row) + +select datname from (select pg_tablespace_databases(b.oid) as pts from pg_tablespace b where b.spcname!='pg_global')pts join pg_database db on pts.pts=db.oid limit 1; + datname +----------- + template1 +(1 row) + -- -- Test adding a support function to a subject function -- diff --git a/src/test/regress/sql/misc_functions.sql b/src/test/regress/sql/misc_functions.sql index 1e11eb3554..5441b8b643 100644 --- a/src/test/regress/sql/misc_functions.sql +++ b/src/test/regress/sql/misc_functions.sql @@ -51,6 +51,12 @@ from (select pg_ls_waldir() w) ss where length((w).name) = 24 limit 1; select count(*) >= 0 as ok from pg_ls_archive_statusdir(); +select * from (select pg_ls_dir('.')a)a where a='base' limit 1; + +select * from (select (pg_timezone_names()).name)ptn where name='UTC' limit 1; + +select datname from (select pg_tablespace_databases(b.oid) as pts from pg_tablespace b where b.spcname!='pg_global')pts join pg_database db on pts.pts=db.oid limit 1; + -- -- Test adding a support function to a subject function -- -- 2.17.0