diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c index 36a8e3e..a0c9bd8 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; + MemoryContext oldcontext; + AttInMetadata *attinmeta; + char** valbuf; + int *valbuflen; + char **cstrs; + 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 int storeHandler(PGresult *res, void *param, PGrowValue *columns); + /* 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,51 @@ 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); + PQsetRowProcessor(conn, storeHandler, &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); + PG_TRY(); + { + res = PQexec(conn, buf.data); + } + PG_CATCH(); + { + ErrorData *edata; + + finishStoreInfo(&storeinfo); + edata = CopyErrorData(); + FlushErrorState(); + + /* Skip remaining results when storeHandler raises exception. */ + PQskipResult(conn, FALSE); + ReThrowError(edata); + } + PG_END_TRY(); + + finishStoreInfo(&storeinfo); + if (!res || (PQresultStatus(res) != PGRES_COMMAND_OK && PQresultStatus(res) != PGRES_TUPLES_OK)) { + /* finishStoreInfo saves the fields referred to below. */ + if (storeinfo.nummismatch) + { + /* This is only for backward compatibility */ + 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; } @@ -579,8 +632,8 @@ dblink_fetch(PG_FUNCTION_ARGS) (errcode(ERRCODE_INVALID_CURSOR_NAME), errmsg("cursor \"%s\" does not exist", curname))); } + PQclear(res); - materializeResult(fcinfo, res); return (Datum) 0; } @@ -640,6 +693,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 +769,252 @@ 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); + PQsetRowProcessor(conn, storeHandler, &storeinfo); + /* synchronous query, or async result retrieval */ - if (!is_async) - res = PQexec(conn, sql); - else + PG_TRY(); { - res = PQgetResult(conn); - /* NULL means we're all done with the async results */ - if (!res) - return (Datum) 0; + if (!is_async) + res = PQexec(conn, sql); + else + res = PQgetResult(conn); } + PG_CATCH(); + { + ErrorData *edata; - /* if needed, close the connection to the database and cleanup */ - if (freeconn) - PQfinish(conn); + finishStoreInfo(&storeinfo); + edata = CopyErrorData(); + FlushErrorState(); - if (!res || - (PQresultStatus(res) != PGRES_COMMAND_OK && - PQresultStatus(res) != PGRES_TUPLES_OK)) + /* Skip remaining results when storeHandler raises exception. */ + PQskipResult(conn, FALSE); + ReThrowError(edata); + } + PG_END_TRY(); + + finishStoreInfo(&storeinfo); + + /* 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)) + { + /* finishStoreInfo saves the fields referred to below. */ + if (storeinfo.nummismatch) + { + /* This is only for backward compatibility */ + 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; + } } + PQclear(res); - 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; + TupleDesc tupdesc; + int i; + + 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; + } - Assert(rsinfo->returnMode == SFRM_Materialize); + sinfo->oldcontext = MemoryContextSwitchTo( + rsinfo->econtext->ecxt_per_query_memory); - PG_TRY(); + /* make sure we have a persistent copy of the tupdesc */ + tupdesc = CreateTupleDescCopy(tupdesc); + + sinfo->error_occurred = FALSE; + sinfo->nummismatch = FALSE; + sinfo->attinmeta = TupleDescGetAttInMetadata(tupdesc); + sinfo->nattrs = tupdesc->natts; + sinfo->tuplestore = tuplestore_begin_heap(true, false, work_mem); + sinfo->valbuf = NULL; + sinfo->valbuflen = NULL; + + /* Preallocate memory of same size with c string array for values. */ + sinfo->valbuf = (char **)malloc(sinfo->nattrs * sizeof(char*)); + if (sinfo->valbuf) + sinfo->valbuflen = (int *)malloc(sinfo->nattrs * sizeof(int)); + if (sinfo->valbuflen) + sinfo->cstrs = (char **)malloc(sinfo->nattrs * sizeof(char*)); + + if (sinfo->cstrs == NULL) { - TupleDesc tupdesc; - bool is_sql_cmd = false; - int ntuples; - int nfields; + if (sinfo->valbuf) + free(sinfo->valbuf); + if (sinfo->valbuflen) + free(sinfo->valbuflen); - if (PQresultStatus(res) == PGRES_COMMAND_OK) - { - is_sql_cmd = true; + ereport(ERROR, + (errcode(ERRCODE_OUT_OF_MEMORY), + errmsg("out of memory"))); + } - /* - * 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); + for (i = 0 ; i < sinfo->nattrs ; i++) + { + sinfo->valbuf[i] = NULL; + sinfo->valbuflen[i] = -1; + } - 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); + if (sinfo->valbuf) + { + for (i = 0 ; i < sinfo->nattrs ; i++) + { + if (sinfo->valbuf[i]) + free(sinfo->valbuf[i]); } + free(sinfo->valbuf); + sinfo->valbuf = NULL; + } - /* - * 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"))); + if (sinfo->valbuflen) + { + free(sinfo->valbuflen); + sinfo->valbuflen = NULL; + } - 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->cstrs) + { + free(sinfo->cstrs); + sinfo->cstrs = NULL; + } - values = (char **) palloc(nfields * sizeof(char *)); + MemoryContextSwitchTo(sinfo->oldcontext); +} - /* put all tuples into the tuplestore */ - for (row = 0; row < ntuples; row++) - { - HeapTuple tuple; +static int +storeHandler(PGresult *res, void *param, PGrowValue *columns) +{ + storeInfo *sinfo = (storeInfo *)param; + HeapTuple tuple; + int fields = PQnfields(res); + int i; + char **cstrs = sinfo->cstrs; - if (!is_sql_cmd) - { - int i; + if (sinfo->error_occurred) + return FALSE; - 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); + + /* This error will be processed in + * dblink_record_internal(). So do not set error message + * here. */ + return FALSE; + } - /* build the tuple and put it into the tuplestore. */ - tuple = BuildTupleFromCStrings(attinmeta, values); - tuplestore_puttuple(tupstore, tuple); + /* + * value input functions assumes that the input string is + * terminated by zero. We should make the values to be so. + */ + for(i = 0 ; i < fields ; i++) + { + int len = columns[i].len; + if (len < 0) + cstrs[i] = NULL; + else + { + char *tmp = sinfo->valbuf[i]; + int tmplen = sinfo->valbuflen[i]; + + /* + * Divide calls to malloc and realloc so that things will + * go fine even on the systems of which realloc() does not + * accept NULL as old memory block. + * + * Also try to (re)allocate in bigger steps to + * avoid flood of allocations on weird data. + */ + if (tmp == NULL) + { + tmplen = len + 1; + if (tmplen < 64) + tmplen = 64; + tmp = (char *)malloc(tmplen); + } + else if (tmplen < len + 1) + { + if (len + 1 > tmplen * 2) + tmplen = len + 1; + else + tmplen = tmplen * 2; + tmp = (char *)realloc(tmp, tmplen); } - /* clean up and return the tuplestore */ - tuplestore_donestoring(tupstore); - } + /* + * sinfo->valbuf[n] will be freed in finishStoreInfo() + * when realloc returns NULL. + */ + if (tmp == NULL) + return FALSE; - PQclear(res); - } - PG_CATCH(); - { - /* be sure to release the libpq result */ - PQclear(res); - PG_RE_THROW(); + sinfo->valbuf[i] = tmp; + sinfo->valbuflen[i] = tmplen; + + cstrs[i] = sinfo->valbuf[i]; + memcpy(cstrs[i], columns[i].value, len); + cstrs[i][len] = '\0'; + } } - PG_END_TRY(); + + /* + * These functions may throw exception. It will be caught in + * dblink_record_internal() + */ + tuple = BuildTupleFromCStrings(sinfo->attinmeta, cstrs); + tuplestore_puttuple(sinfo->tuplestore, tuple); + + return TRUE; } /*