diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c index 62c810a..fb2e10e 100644 --- a/contrib/dblink/dblink.c +++ b/contrib/dblink/dblink.c @@ -63,11 +63,23 @@ typedef struct remoteConn bool newXactForCursor; /* Opened a transaction for a cursor */ } remoteConn; +typedef struct storeInfo +{ + Tuplestorestate *tuplestore; + int nattrs; + AttInMetadata *attinmeta; + MemoryContext oldcontext; + char *attrvalbuf; + void **valbuf; + size_t *valbufsize; + bool error_occurred; + bool nummismatch; +} storeInfo; + /* * Internal declarations */ static Datum dblink_record_internal(FunctionCallInfo fcinfo, bool is_async); -static void materializeResult(FunctionCallInfo fcinfo, PGresult *res); static remoteConn *getConnectionByName(const char *name); static HTAB *createConnHash(void); static void createNewConnection(const char *name, remoteConn *rconn); @@ -90,6 +102,10 @@ static char *escape_param_str(const char *from); static void validate_pkattnums(Relation rel, int2vector *pkattnums_arg, int32 pknumatts_arg, int **pkattnums, int *pknumatts); +static void initStoreInfo(storeInfo *sinfo, FunctionCallInfo fcinfo); +static void finishStoreInfo(storeInfo *sinfo); +static void *addTuple(PGresult *res, AddTupFunc func, int id, size_t size); + /* Global */ static remoteConn *pconn = NULL; @@ -503,6 +519,7 @@ dblink_fetch(PG_FUNCTION_ARGS) char *curname = NULL; int howmany = 0; bool fail = true; /* default to backward compatible */ + storeInfo storeinfo; DBLINK_INIT; @@ -559,15 +576,30 @@ dblink_fetch(PG_FUNCTION_ARGS) appendStringInfo(&buf, "FETCH %d FROM %s", howmany, curname); /* + * Result is stored into storeinfo.tuplestore instead of + * res->result retuned by PQexec below + */ + initStoreInfo(&storeinfo, fcinfo); + PQregisterTupleAdder(conn, addTuple, &storeinfo); + + /* * Try to execute the query. Note that since libpq uses malloc, the * PGresult will be long-lived even though we are still in a short-lived * memory context. */ res = PQexec(conn, buf.data); + finishStoreInfo(&storeinfo); + if (!res || (PQresultStatus(res) != PGRES_COMMAND_OK && PQresultStatus(res) != PGRES_TUPLES_OK)) { + /* This is only for backward compatibility */ + if (storeinfo.nummismatch) + ereport(ERROR, + (errcode(ERRCODE_DATATYPE_MISMATCH), + errmsg("remote query result rowtype does not match " + "the specified FROM clause rowtype"))); dblink_res_error(conname, res, "could not fetch from cursor", fail); return (Datum) 0; } @@ -580,7 +612,6 @@ dblink_fetch(PG_FUNCTION_ARGS) errmsg("cursor \"%s\" does not exist", curname))); } - materializeResult(fcinfo, res); return (Datum) 0; } @@ -640,6 +671,7 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async) remoteConn *rconn = NULL; bool fail = true; /* default to backward compatible */ bool freeconn = false; + storeInfo storeinfo; /* check to see if caller supports us returning a tuplestore */ if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo)) @@ -715,164 +747,206 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async) rsinfo->setResult = NULL; rsinfo->setDesc = NULL; + + /* + * Result is stored into storeinfo.tuplestore instead of + * res->result retuned by PQexec/PQgetResult below + */ + initStoreInfo(&storeinfo, fcinfo); + PQregisterTupleAdder(conn, addTuple, &storeinfo); + /* synchronous query, or async result retrieval */ if (!is_async) res = PQexec(conn, sql); else - { res = PQgetResult(conn); - /* NULL means we're all done with the async results */ - if (!res) - return (Datum) 0; - } - /* if needed, close the connection to the database and cleanup */ - if (freeconn) - PQfinish(conn); + finishStoreInfo(&storeinfo); - if (!res || - (PQresultStatus(res) != PGRES_COMMAND_OK && - PQresultStatus(res) != PGRES_TUPLES_OK)) + /* NULL res from async get means we're all done with the results */ + if (res || !is_async) { - dblink_res_error(conname, res, "could not execute query", fail); - return (Datum) 0; + if (freeconn) + PQfinish(conn); + + if (!res || + (PQresultStatus(res) != PGRES_COMMAND_OK && + PQresultStatus(res) != PGRES_TUPLES_OK)) + { + /* This is only for backward compatibility */ + if (storeinfo.nummismatch) + { + ereport(ERROR, + (errcode(ERRCODE_DATATYPE_MISMATCH), + errmsg("remote query result rowtype does not match " + "the specified FROM clause rowtype"))); + } + dblink_res_error(conname, res, "could not execute query", fail); + return (Datum) 0; + } } - materializeResult(fcinfo, res); return (Datum) 0; } -/* - * Materialize the PGresult to return them as the function result. - * The res will be released in this function. - */ static void -materializeResult(FunctionCallInfo fcinfo, PGresult *res) +initStoreInfo(storeInfo *sinfo, FunctionCallInfo fcinfo) { ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; - - Assert(rsinfo->returnMode == SFRM_Materialize); - - PG_TRY(); + TupleDesc tupdesc; + int i; + + switch (get_call_result_type(fcinfo, NULL, &tupdesc)) { - TupleDesc tupdesc; - bool is_sql_cmd = false; - int ntuples; - int nfields; + case TYPEFUNC_COMPOSITE: + /* success */ + break; + case TYPEFUNC_RECORD: + /* failed to determine actual type of RECORD */ + ereport(ERROR, + (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), + errmsg("function returning record called in context " + "that cannot accept type record"))); + break; + default: + /* result type isn't composite */ + elog(ERROR, "return type must be a row type"); + break; + } + + sinfo->oldcontext = MemoryContextSwitchTo( + rsinfo->econtext->ecxt_per_query_memory); + + /* make sure we have a persistent copy of the tupdesc */ + tupdesc = CreateTupleDescCopy(tupdesc); + + sinfo->error_occurred = FALSE; + sinfo->nummismatch = FALSE; + sinfo->nattrs = tupdesc->natts; + sinfo->tuplestore = tuplestore_begin_heap(true, false, work_mem); + sinfo->attinmeta = TupleDescGetAttInMetadata(tupdesc); + sinfo->valbuf = (void **)malloc(sinfo->nattrs * sizeof(void *)); + sinfo->valbufsize = (size_t *)malloc(sinfo->nattrs * sizeof(size_t)); + for (i = 0 ; i < sinfo->nattrs ; i++) + { + sinfo->valbuf[i] = NULL; + sinfo->valbufsize[i] = 0; + } - if (PQresultStatus(res) == PGRES_COMMAND_OK) - { - is_sql_cmd = true; - - /* - * need a tuple descriptor representing one TEXT column to return - * the command status string as our result tuple - */ - tupdesc = CreateTemplateTupleDesc(1, false); - TupleDescInitEntry(tupdesc, (AttrNumber) 1, "status", - TEXTOID, -1, 0); - ntuples = 1; - nfields = 1; - } - else - { - Assert(PQresultStatus(res) == PGRES_TUPLES_OK); + /* Preallocate memory of same size with PGresAttDesc array for values. */ + sinfo->attrvalbuf = (char *) malloc(sinfo->nattrs * sizeof(PGresAttValue)); - is_sql_cmd = false; + rsinfo->setResult = sinfo->tuplestore; + rsinfo->setDesc = tupdesc; +} - /* get a tuple descriptor for our result type */ - switch (get_call_result_type(fcinfo, NULL, &tupdesc)) - { - case TYPEFUNC_COMPOSITE: - /* success */ - break; - case TYPEFUNC_RECORD: - /* failed to determine actual type of RECORD */ - ereport(ERROR, - (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), - errmsg("function returning record called in context " - "that cannot accept type record"))); - break; - default: - /* result type isn't composite */ - elog(ERROR, "return type must be a row type"); - break; - } +static void +finishStoreInfo(storeInfo *sinfo) +{ + int i; - /* make sure we have a persistent copy of the tupdesc */ - tupdesc = CreateTupleDescCopy(tupdesc); - ntuples = PQntuples(res); - nfields = PQnfields(res); + for (i = 0 ; i < sinfo->nattrs ; i++) + { + if (sinfo->valbuf[i]) + { + free(sinfo->valbuf[i]); + sinfo->valbuf[i] = NULL; } + } + if (sinfo->attrvalbuf) + free(sinfo->attrvalbuf); + sinfo->attrvalbuf = NULL; + MemoryContextSwitchTo(sinfo->oldcontext); +} - /* - * check result and tuple descriptor have the same number of columns - */ - if (nfields != tupdesc->natts) - ereport(ERROR, - (errcode(ERRCODE_DATATYPE_MISMATCH), - errmsg("remote query result rowtype does not match " - "the specified FROM clause rowtype"))); +static void * +addTuple(PGresult *res, AddTupFunc func, int id, size_t size) +{ + storeInfo *sinfo = (storeInfo *)PQgetAddTupleParam(res); + HeapTuple tuple; + int fields = PQnfields(res); + int i; + PGresAttValue *attval; + char **cstrs; - if (ntuples > 0) - { - AttInMetadata *attinmeta; - Tuplestorestate *tupstore; - MemoryContext oldcontext; - int row; - char **values; - - attinmeta = TupleDescGetAttInMetadata(tupdesc); - - oldcontext = MemoryContextSwitchTo( - rsinfo->econtext->ecxt_per_query_memory); - tupstore = tuplestore_begin_heap(true, false, work_mem); - rsinfo->setResult = tupstore; - rsinfo->setDesc = tupdesc; - MemoryContextSwitchTo(oldcontext); + if (sinfo->error_occurred) + return NULL; - values = (char **) palloc(nfields * sizeof(char *)); + switch (func) + { + case ADDTUP_ALLOC_TEXT: + case ADDTUP_ALLOC_BINARY: + if (id == -1) + return sinfo->attrvalbuf; + + if (id < 0 || id >= sinfo->nattrs) + return NULL; - /* put all tuples into the tuplestore */ - for (row = 0; row < ntuples; row++) + if (sinfo->valbufsize[id] < size) { - HeapTuple tuple; + if (sinfo->valbuf[id] == NULL) + sinfo->valbuf[id] = malloc(size); + else + sinfo->valbuf[id] = realloc(sinfo->valbuf[id], size); + sinfo->valbufsize[id] = size; + } + return sinfo->valbuf[id]; - if (!is_sql_cmd) - { - int i; + case ADDTUP_ADD_TUPLE: + break; /* Go through */ + default: + /* Ignore */ + break; + } - for (i = 0; i < nfields; i++) - { - if (PQgetisnull(res, row, i)) - values[i] = NULL; - else - values[i] = PQgetvalue(res, row, i); - } - } - else - { - values[0] = PQcmdStatus(res); - } + if (sinfo->nattrs != fields) + { + sinfo->error_occurred = TRUE; + sinfo->nummismatch = TRUE; + finishStoreInfo(sinfo); - /* build the tuple and put it into the tuplestore. */ - tuple = BuildTupleFromCStrings(attinmeta, values); - tuplestore_puttuple(tupstore, tuple); - } + PQsetAddTupleErrMes(res, + strdup("function returning record called in " + "context that cannot accept type record")); + return NULL; + } - /* clean up and return the tuplestore */ - tuplestore_donestoring(tupstore); - } + /* + * Rewrite PGresAttDesc[] to char(*)[] in-place. + */ + Assert(sizeof(char*) <= sizeof(PGresAttValue)); + attval = (PGresAttValue *)sinfo->attrvalbuf; + cstrs = (char **)sinfo->attrvalbuf; + for(i = 0 ; i < fields ; i++) + cstrs[i] = PQgetAsCstring(attval++); - PQclear(res); + PG_TRY(); + { + tuple = BuildTupleFromCStrings(sinfo->attinmeta, cstrs); + tuplestore_puttuple(sinfo->tuplestore, tuple); } PG_CATCH(); { - /* be sure to release the libpq result */ - PQclear(res); - PG_RE_THROW(); + /* + * Return the error message in the exception to the caller and + * cancel the exception. + */ + ErrorData *edata; + + sinfo->error_occurred = TRUE; + sinfo->nummismatch = TRUE; + + finishStoreInfo(sinfo); + + edata = CopyErrorData(); + FlushErrorState(); + + PQsetAddTupleErrMes(res, strdup(edata->message)); + return NULL; } PG_END_TRY(); + + return sinfo->attrvalbuf; } /*