Commit 56cbb611 authored by Joe Conway's avatar Joe Conway

Rewrite dblink_record_internal() and dblink_fetch() to use a tuplestore

(SFRM_Materialize mode) to return tuples. Since we don't return from the
dblink function in tuplestore mode, release the PGresult with a PG_CATCH
block on error. Also rearrange to share the same code to materialize the
tuplestore. Patch by Takahiro Itagaki.
parent 875353b9
...@@ -8,7 +8,7 @@ ...@@ -8,7 +8,7 @@
* Darko Prenosil <Darko.Prenosil@finteh.hr> * Darko Prenosil <Darko.Prenosil@finteh.hr>
* Shridhar Daithankar <shridhar_daithankar@persistent.co.in> * Shridhar Daithankar <shridhar_daithankar@persistent.co.in>
* *
* $PostgreSQL: pgsql/contrib/dblink/dblink.c,v 1.86 2010/01/02 16:57:32 momjian Exp $ * $PostgreSQL: pgsql/contrib/dblink/dblink.c,v 1.87 2010/01/24 22:19:38 joe Exp $
* Copyright (c) 2001-2010, PostgreSQL Global Development Group * Copyright (c) 2001-2010, PostgreSQL Global Development Group
* ALL RIGHTS RESERVED; * ALL RIGHTS RESERVED;
* *
...@@ -80,6 +80,7 @@ typedef struct remoteConn ...@@ -80,6 +80,7 @@ typedef struct remoteConn
* Internal declarations * Internal declarations
*/ */
static Datum dblink_record_internal(FunctionCallInfo fcinfo, bool is_async); static Datum dblink_record_internal(FunctionCallInfo fcinfo, bool is_async);
static void materializeResult(FunctionCallInfo fcinfo, PGresult *res);
static remoteConn *getConnectionByName(const char *name); static remoteConn *getConnectionByName(const char *name);
static HTAB *createConnHash(void); static HTAB *createConnHash(void);
static void createNewConnection(const char *name, remoteConn *rconn); static void createNewConnection(const char *name, remoteConn *rconn);
...@@ -504,27 +505,18 @@ PG_FUNCTION_INFO_V1(dblink_fetch); ...@@ -504,27 +505,18 @@ PG_FUNCTION_INFO_V1(dblink_fetch);
Datum Datum
dblink_fetch(PG_FUNCTION_ARGS) dblink_fetch(PG_FUNCTION_ARGS)
{ {
FuncCallContext *funcctx; ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
TupleDesc tupdesc = NULL;
int call_cntr;
int max_calls;
AttInMetadata *attinmeta;
PGresult *res = NULL; PGresult *res = NULL;
MemoryContext oldcontext;
char *conname = NULL; char *conname = NULL;
remoteConn *rconn = NULL; remoteConn *rconn = NULL;
DBLINK_INIT;
/* stuff done only on the first call of the function */
if (SRF_IS_FIRSTCALL())
{
PGconn *conn = NULL; PGconn *conn = NULL;
StringInfoData buf; StringInfoData buf;
char *curname = NULL; char *curname = NULL;
int howmany = 0; int howmany = 0;
bool fail = true; /* default to backward compatible */ bool fail = true; /* default to backward compatible */
DBLINK_INIT;
if (PG_NARGS() == 4) if (PG_NARGS() == 4)
{ {
/* text,text,int,bool */ /* text,text,int,bool */
...@@ -569,12 +561,14 @@ dblink_fetch(PG_FUNCTION_ARGS) ...@@ -569,12 +561,14 @@ dblink_fetch(PG_FUNCTION_ARGS)
if (!conn) if (!conn)
DBLINK_CONN_NOT_AVAIL; DBLINK_CONN_NOT_AVAIL;
/* let the caller know we're sending back a tuplestore */
rsinfo->returnMode = SFRM_Materialize;
rsinfo->setResult = NULL;
rsinfo->setDesc = NULL;
initStringInfo(&buf); initStringInfo(&buf);
appendStringInfo(&buf, "FETCH %d FROM %s", howmany, curname); appendStringInfo(&buf, "FETCH %d FROM %s", howmany, curname);
/* create a function context for cross-call persistence */
funcctx = SRF_FIRSTCALL_INIT();
/* /*
* Try to execute the query. Note that since libpq uses malloc, the * Try to execute the query. Note that since libpq uses malloc, the
* PGresult will be long-lived even though we are still in a * PGresult will be long-lived even though we are still in a
...@@ -586,7 +580,7 @@ dblink_fetch(PG_FUNCTION_ARGS) ...@@ -586,7 +580,7 @@ dblink_fetch(PG_FUNCTION_ARGS)
PQresultStatus(res) != PGRES_TUPLES_OK)) PQresultStatus(res) != PGRES_TUPLES_OK))
{ {
dblink_res_error(conname, res, "could not fetch from cursor", fail); dblink_res_error(conname, res, "could not fetch from cursor", fail);
SRF_RETURN_DONE(funcctx); return (Datum) 0;
} }
else if (PQresultStatus(res) == PGRES_COMMAND_OK) else if (PQresultStatus(res) == PGRES_COMMAND_OK)
{ {
...@@ -597,107 +591,8 @@ dblink_fetch(PG_FUNCTION_ARGS) ...@@ -597,107 +591,8 @@ dblink_fetch(PG_FUNCTION_ARGS)
errmsg("cursor \"%s\" does not exist", curname))); errmsg("cursor \"%s\" does not exist", curname)));
} }
funcctx->max_calls = PQntuples(res); materializeResult(fcinfo, res);
return (Datum) 0;
/* got results, keep track of them */
funcctx->user_fctx = res;
/* 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;
}
/* check result and tuple descriptor have the same number of columns */
if (PQnfields(res) != tupdesc->natts)
ereport(ERROR,
(errcode(ERRCODE_DATATYPE_MISMATCH),
errmsg("remote query result rowtype does not match "
"the specified FROM clause rowtype")));
/*
* fast track when no results. We could exit earlier, but then we'd
* not report error if the result tuple type is wrong.
*/
if (funcctx->max_calls < 1)
{
PQclear(res);
SRF_RETURN_DONE(funcctx);
}
/*
* switch to memory context appropriate for multiple function calls,
* so we can make long-lived copy of tupdesc etc
*/
oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
/* make sure we have a persistent copy of the tupdesc */
tupdesc = CreateTupleDescCopy(tupdesc);
/* store needed metadata for subsequent calls */
attinmeta = TupleDescGetAttInMetadata(tupdesc);
funcctx->attinmeta = attinmeta;
MemoryContextSwitchTo(oldcontext);
}
/* stuff done on every call of the function */
funcctx = SRF_PERCALL_SETUP();
/*
* initialize per-call variables
*/
call_cntr = funcctx->call_cntr;
max_calls = funcctx->max_calls;
res = (PGresult *) funcctx->user_fctx;
attinmeta = funcctx->attinmeta;
tupdesc = attinmeta->tupdesc;
if (call_cntr < max_calls) /* do when there is more left to send */
{
char **values;
HeapTuple tuple;
Datum result;
int i;
int nfields = PQnfields(res);
values = (char **) palloc(nfields * sizeof(char *));
for (i = 0; i < nfields; i++)
{
if (PQgetisnull(res, call_cntr, i) == 0)
values[i] = PQgetvalue(res, call_cntr, i);
else
values[i] = NULL;
}
/* build the tuple */
tuple = BuildTupleFromCStrings(attinmeta, values);
/* make the tuple into a datum */
result = HeapTupleGetDatum(tuple);
SRF_RETURN_NEXT(funcctx, result);
}
else
{
/* do when there is no more left */
PQclear(res);
SRF_RETURN_DONE(funcctx);
}
} }
/* /*
...@@ -749,37 +644,29 @@ dblink_get_result(PG_FUNCTION_ARGS) ...@@ -749,37 +644,29 @@ dblink_get_result(PG_FUNCTION_ARGS)
static Datum static Datum
dblink_record_internal(FunctionCallInfo fcinfo, bool is_async) dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
{ {
FuncCallContext *funcctx; ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
TupleDesc tupdesc = NULL;
int call_cntr;
int max_calls;
AttInMetadata *attinmeta;
char *msg; char *msg;
PGresult *res = NULL; PGresult *res = NULL;
bool is_sql_cmd = false;
char *sql_cmd_status = NULL;
MemoryContext oldcontext;
bool freeconn = false;
DBLINK_INIT;
/* stuff done only on the first call of the function */
if (SRF_IS_FIRSTCALL())
{
PGconn *conn = NULL; PGconn *conn = NULL;
char *connstr = NULL; char *connstr = NULL;
char *sql = NULL; char *sql = NULL;
char *conname = NULL; char *conname = NULL;
remoteConn *rconn = NULL; remoteConn *rconn = NULL;
bool fail = true; /* default to backward compatible */ bool fail = true; /* default to backward compatible */
bool freeconn = false;
/* create a function context for cross-call persistence */ /* check to see if caller supports us returning a tuplestore */
funcctx = SRF_FIRSTCALL_INIT(); 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")));
/* DBLINK_INIT;
* switch to memory context appropriate for multiple function calls
*/
oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
if (!is_async) if (!is_async)
{ {
...@@ -837,6 +724,11 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async) ...@@ -837,6 +724,11 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
if (!conn) if (!conn)
DBLINK_CONN_NOT_AVAIL; DBLINK_CONN_NOT_AVAIL;
/* let the caller know we're sending back a tuplestore */
rsinfo->returnMode = SFRM_Materialize;
rsinfo->setResult = NULL;
rsinfo->setDesc = NULL;
/* synchronous query, or async result retrieval */ /* synchronous query, or async result retrieval */
if (!is_async) if (!is_async)
res = PQexec(conn, sql); res = PQexec(conn, sql);
...@@ -845,51 +737,63 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async) ...@@ -845,51 +737,63 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
res = PQgetResult(conn); res = PQgetResult(conn);
/* NULL means we're all done with the async results */ /* NULL means we're all done with the async results */
if (!res) if (!res)
{ return (Datum) 0;
MemoryContextSwitchTo(oldcontext);
SRF_RETURN_DONE(funcctx);
}
} }
/* if needed, close the connection to the database and cleanup */
if (freeconn)
PQfinish(conn);
if (!res || if (!res ||
(PQresultStatus(res) != PGRES_COMMAND_OK && (PQresultStatus(res) != PGRES_COMMAND_OK &&
PQresultStatus(res) != PGRES_TUPLES_OK)) PQresultStatus(res) != PGRES_TUPLES_OK))
{ {
if (freeconn)
PQfinish(conn);
dblink_res_error(conname, res, "could not execute query", fail); dblink_res_error(conname, res, "could not execute query", fail);
MemoryContextSwitchTo(oldcontext); return (Datum) 0;
SRF_RETURN_DONE(funcctx);
} }
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)
{
ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
Assert(rsinfo->returnMode == SFRM_Materialize);
PG_TRY();
{
TupleDesc tupdesc;
bool is_sql_cmd = false;
int ntuples;
int nfields;
if (PQresultStatus(res) == PGRES_COMMAND_OK) if (PQresultStatus(res) == PGRES_COMMAND_OK)
{ {
is_sql_cmd = true; is_sql_cmd = true;
/* need a tuple descriptor representing one TEXT column */ /*
* need a tuple descriptor representing one TEXT column to
* return the command status string as our result tuple
*/
tupdesc = CreateTemplateTupleDesc(1, false); tupdesc = CreateTemplateTupleDesc(1, false);
TupleDescInitEntry(tupdesc, (AttrNumber) 1, "status", TupleDescInitEntry(tupdesc, (AttrNumber) 1, "status",
TEXTOID, -1, 0); TEXTOID, -1, 0);
ntuples = 1;
/* nfields = 1;
* and save a copy of the command status string to return as our
* result tuple
*/
sql_cmd_status = PQcmdStatus(res);
funcctx->max_calls = 1;
} }
else else
funcctx->max_calls = PQntuples(res); {
Assert(PQresultStatus(res) == PGRES_TUPLES_OK);
/* got results, keep track of them */
funcctx->user_fctx = res;
/* if needed, close the connection to the database and cleanup */ is_sql_cmd = false;
if (freeconn)
PQfinish(conn);
if (!is_sql_cmd)
{
/* get a tuple descriptor for our result type */ /* get a tuple descriptor for our result type */
switch (get_call_result_type(fcinfo, NULL, &tupdesc)) switch (get_call_result_type(fcinfo, NULL, &tupdesc))
{ {
...@@ -911,87 +815,78 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async) ...@@ -911,87 +815,78 @@ dblink_record_internal(FunctionCallInfo fcinfo, bool is_async)
/* make sure we have a persistent copy of the tupdesc */ /* make sure we have a persistent copy of the tupdesc */
tupdesc = CreateTupleDescCopy(tupdesc); tupdesc = CreateTupleDescCopy(tupdesc);
ntuples = PQntuples(res);
nfields = PQnfields(res);
} }
/* /*
* check result and tuple descriptor have the same number of columns * check result and tuple descriptor have the same number of columns
*/ */
if (PQnfields(res) != tupdesc->natts) if (nfields != tupdesc->natts)
ereport(ERROR, ereport(ERROR,
(errcode(ERRCODE_DATATYPE_MISMATCH), (errcode(ERRCODE_DATATYPE_MISMATCH),
errmsg("remote query result rowtype does not match " errmsg("remote query result rowtype does not match "
"the specified FROM clause rowtype"))); "the specified FROM clause rowtype")));
/* fast track when no results */ if (ntuples > 0)
if (funcctx->max_calls < 1)
{ {
if (res) AttInMetadata *attinmeta;
PQclear(res); Tuplestorestate *tupstore;
MemoryContextSwitchTo(oldcontext); MemoryContext oldcontext;
SRF_RETURN_DONE(funcctx); int row;
} char **values;
/* store needed metadata for subsequent calls */
attinmeta = TupleDescGetAttInMetadata(tupdesc); attinmeta = TupleDescGetAttInMetadata(tupdesc);
funcctx->attinmeta = attinmeta;
oldcontext = MemoryContextSwitchTo(
rsinfo->econtext->ecxt_per_query_memory);
tupstore = tuplestore_begin_heap(true, false, work_mem);
rsinfo->setResult = tupstore;
rsinfo->setDesc = tupdesc;
MemoryContextSwitchTo(oldcontext); MemoryContextSwitchTo(oldcontext);
} values = (char **) palloc(nfields * sizeof(char *));
/* stuff done on every call of the function */
funcctx = SRF_PERCALL_SETUP();
/*
* initialize per-call variables
*/
call_cntr = funcctx->call_cntr;
max_calls = funcctx->max_calls;
res = (PGresult *) funcctx->user_fctx;
attinmeta = funcctx->attinmeta;
tupdesc = attinmeta->tupdesc;
if (call_cntr < max_calls) /* do when there is more left to send */ /* put all tuples into the tuplestore */
for (row = 0; row < ntuples; row++)
{ {
char **values;
HeapTuple tuple; HeapTuple tuple;
Datum result;
if (!is_sql_cmd) if (!is_sql_cmd)
{ {
int i; int i;
int nfields = PQnfields(res);
values = (char **) palloc(nfields * sizeof(char *));
for (i = 0; i < nfields; i++) for (i = 0; i < nfields; i++)
{ {
if (PQgetisnull(res, call_cntr, i) == 0) if (PQgetisnull(res, row, i))
values[i] = PQgetvalue(res, call_cntr, i);
else
values[i] = NULL; values[i] = NULL;
else
values[i] = PQgetvalue(res, row, i);
} }
} }
else else
{ {
values = (char **) palloc(1 * sizeof(char *)); values[0] = PQcmdStatus(res);
values[0] = sql_cmd_status;
} }
/* build the tuple */ /* build the tuple and put it into the tuplestore. */
tuple = BuildTupleFromCStrings(attinmeta, values); tuple = BuildTupleFromCStrings(attinmeta, values);
tuplestore_puttuple(tupstore, tuple);
}
/* make the tuple into a datum */ /* clean up and return the tuplestore */
result = HeapTupleGetDatum(tuple); tuplestore_donestoring(tupstore);
}
SRF_RETURN_NEXT(funcctx, result); PQclear(res);
} }
else PG_CATCH();
{ {
/* do when there is no more left */ /* be sure to release the libpq result */
PQclear(res); PQclear(res);
SRF_RETURN_DONE(funcctx); PG_RE_THROW();
} }
PG_END_TRY();
} }
/* /*
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment