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,200 +505,94 @@ PG_FUNCTION_INFO_V1(dblink_fetch); ...@@ -504,200 +505,94 @@ 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; PGresult *res = NULL;
int call_cntr; char *conname = NULL;
int max_calls; remoteConn *rconn = NULL;
AttInMetadata *attinmeta; PGconn *conn = NULL;
PGresult *res = NULL; StringInfoData buf;
MemoryContext oldcontext; char *curname = NULL;
char *conname = NULL; int howmany = 0;
remoteConn *rconn = NULL; bool fail = true; /* default to backward compatible */
DBLINK_INIT; DBLINK_INIT;
/* stuff done only on the first call of the function */ if (PG_NARGS() == 4)
if (SRF_IS_FIRSTCALL())
{ {
PGconn *conn = NULL; /* text,text,int,bool */
StringInfoData buf; conname = text_to_cstring(PG_GETARG_TEXT_PP(0));
char *curname = NULL; curname = text_to_cstring(PG_GETARG_TEXT_PP(1));
int howmany = 0; howmany = PG_GETARG_INT32(2);
bool fail = true; /* default to backward compatible */ fail = PG_GETARG_BOOL(3);
if (PG_NARGS() == 4) rconn = getConnectionByName(conname);
if (rconn)
conn = rconn->conn;
}
else if (PG_NARGS() == 3)
{
/* text,text,int or text,int,bool */
if (get_fn_expr_argtype(fcinfo->flinfo, 2) == BOOLOID)
{
curname = text_to_cstring(PG_GETARG_TEXT_PP(0));
howmany = PG_GETARG_INT32(1);
fail = PG_GETARG_BOOL(2);
conn = pconn->conn;
}
else
{ {
/* text,text,int,bool */
conname = text_to_cstring(PG_GETARG_TEXT_PP(0)); conname = text_to_cstring(PG_GETARG_TEXT_PP(0));
curname = text_to_cstring(PG_GETARG_TEXT_PP(1)); curname = text_to_cstring(PG_GETARG_TEXT_PP(1));
howmany = PG_GETARG_INT32(2); howmany = PG_GETARG_INT32(2);
fail = PG_GETARG_BOOL(3);
rconn = getConnectionByName(conname); rconn = getConnectionByName(conname);
if (rconn) if (rconn)
conn = rconn->conn; conn = rconn->conn;
} }
else if (PG_NARGS() == 3) }
{ else if (PG_NARGS() == 2)
/* text,text,int or text,int,bool */ {
if (get_fn_expr_argtype(fcinfo->flinfo, 2) == BOOLOID) /* text,int */
{ curname = text_to_cstring(PG_GETARG_TEXT_PP(0));
curname = text_to_cstring(PG_GETARG_TEXT_PP(0)); howmany = PG_GETARG_INT32(1);
howmany = PG_GETARG_INT32(1); conn = pconn->conn;
fail = PG_GETARG_BOOL(2); }
conn = pconn->conn;
}
else
{
conname = text_to_cstring(PG_GETARG_TEXT_PP(0));
curname = text_to_cstring(PG_GETARG_TEXT_PP(1));
howmany = PG_GETARG_INT32(2);
rconn = getConnectionByName(conname);
if (rconn)
conn = rconn->conn;
}
}
else if (PG_NARGS() == 2)
{
/* text,int */
curname = text_to_cstring(PG_GETARG_TEXT_PP(0));
howmany = PG_GETARG_INT32(1);
conn = pconn->conn;
}
if (!conn)
DBLINK_CONN_NOT_AVAIL;
initStringInfo(&buf);
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
* PGresult will be long-lived even though we are still in a
* short-lived memory context.
*/
res = PQexec(conn, buf.data);
if (!res ||
(PQresultStatus(res) != PGRES_COMMAND_OK &&
PQresultStatus(res) != PGRES_TUPLES_OK))
{
dblink_res_error(conname, res, "could not fetch from cursor", fail);
SRF_RETURN_DONE(funcctx);
}
else if (PQresultStatus(res) == PGRES_COMMAND_OK)
{
/* cursor does not exist - closed already or bad name */
PQclear(res);
ereport(ERROR,
(errcode(ERRCODE_INVALID_CURSOR_NAME),
errmsg("cursor \"%s\" does not exist", curname)));
}
funcctx->max_calls = PQntuples(res);
/* 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 */ if (!conn)
attinmeta = TupleDescGetAttInMetadata(tupdesc); DBLINK_CONN_NOT_AVAIL;
funcctx->attinmeta = attinmeta;
MemoryContextSwitchTo(oldcontext); /* let the caller know we're sending back a tuplestore */
} rsinfo->returnMode = SFRM_Materialize;
rsinfo->setResult = NULL;
rsinfo->setDesc = NULL;
/* stuff done on every call of the function */ initStringInfo(&buf);
funcctx = SRF_PERCALL_SETUP(); appendStringInfo(&buf, "FETCH %d FROM %s", howmany, curname);
/* /*
* initialize per-call variables * 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.
*/ */
call_cntr = funcctx->call_cntr; res = PQexec(conn, buf.data);
max_calls = funcctx->max_calls; if (!res ||
(PQresultStatus(res) != PGRES_COMMAND_OK &&
res = (PGresult *) funcctx->user_fctx; PQresultStatus(res) != PGRES_TUPLES_OK))
attinmeta = funcctx->attinmeta;
tupdesc = attinmeta->tupdesc;
if (call_cntr < max_calls) /* do when there is more left to send */
{ {
char **values; dblink_res_error(conname, res, "could not fetch from cursor", fail);
HeapTuple tuple; return (Datum) 0;
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 else if (PQresultStatus(res) == PGRES_COMMAND_OK)
{ {
/* do when there is no more left */ /* cursor does not exist - closed already or bad name */
PQclear(res); PQclear(res);
SRF_RETURN_DONE(funcctx); ereport(ERROR,
(errcode(ERRCODE_INVALID_CURSOR_NAME),
errmsg("cursor \"%s\" does not exist", curname)));
} }
materializeResult(fcinfo, res);
return (Datum) 0;
} }
/* /*
...@@ -749,147 +644,156 @@ dblink_get_result(PG_FUNCTION_ARGS) ...@@ -749,147 +644,156 @@ 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; char *msg;
int call_cntr; PGresult *res = NULL;
int max_calls; PGconn *conn = NULL;
AttInMetadata *attinmeta; char *connstr = NULL;
char *msg; char *sql = NULL;
PGresult *res = NULL; char *conname = NULL;
bool is_sql_cmd = false; remoteConn *rconn = NULL;
char *sql_cmd_status = NULL; bool fail = true; /* default to backward compatible */
MemoryContext oldcontext; bool freeconn = false;
bool freeconn = false;
/* 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_FEATURE_NOT_SUPPORTED),
errmsg("materialize mode required, but it is not " \
"allowed in this context")));
DBLINK_INIT; DBLINK_INIT;
/* stuff done only on the first call of the function */ if (!is_async)
if (SRF_IS_FIRSTCALL())
{ {
PGconn *conn = NULL; if (PG_NARGS() == 3)
char *connstr = NULL;
char *sql = NULL;
char *conname = NULL;
remoteConn *rconn = NULL;
bool fail = true; /* default to backward compatible */
/* 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);
if (!is_async)
{ {
if (PG_NARGS() == 3) /* text,text,bool */
{ DBLINK_GET_CONN;
/* text,text,bool */ sql = text_to_cstring(PG_GETARG_TEXT_PP(1));
DBLINK_GET_CONN; fail = PG_GETARG_BOOL(2);
sql = text_to_cstring(PG_GETARG_TEXT_PP(1));
fail = PG_GETARG_BOOL(2);
}
else if (PG_NARGS() == 2)
{
/* text,text or text,bool */
if (get_fn_expr_argtype(fcinfo->flinfo, 1) == BOOLOID)
{
conn = pconn->conn;
sql = text_to_cstring(PG_GETARG_TEXT_PP(0));
fail = PG_GETARG_BOOL(1);
}
else
{
DBLINK_GET_CONN;
sql = text_to_cstring(PG_GETARG_TEXT_PP(1));
}
}
else if (PG_NARGS() == 1)
{
/* text */
conn = pconn->conn;
sql = text_to_cstring(PG_GETARG_TEXT_PP(0));
}
else
/* shouldn't happen */
elog(ERROR, "wrong number of arguments");
} }
else /* is_async */ else if (PG_NARGS() == 2)
{ {
/* get async result */ /* text,text or text,bool */
if (PG_NARGS() == 2) if (get_fn_expr_argtype(fcinfo->flinfo, 1) == BOOLOID)
{ {
/* text,bool */ conn = pconn->conn;
DBLINK_GET_CONN; sql = text_to_cstring(PG_GETARG_TEXT_PP(0));
fail = PG_GETARG_BOOL(1); fail = PG_GETARG_BOOL(1);
} }
else if (PG_NARGS() == 1) else
{ {
/* text */
DBLINK_GET_CONN; DBLINK_GET_CONN;
sql = text_to_cstring(PG_GETARG_TEXT_PP(1));
} }
else
/* shouldn't happen */
elog(ERROR, "wrong number of arguments");
} }
else if (PG_NARGS() == 1)
if (!conn) {
DBLINK_CONN_NOT_AVAIL; /* text */
conn = pconn->conn;
/* synchronous query, or async result retrieval */ sql = text_to_cstring(PG_GETARG_TEXT_PP(0));
if (!is_async) }
res = PQexec(conn, sql);
else else
/* shouldn't happen */
elog(ERROR, "wrong number of arguments");
}
else /* is_async */
{
/* get async result */
if (PG_NARGS() == 2)
{ {
res = PQgetResult(conn); /* text,bool */
/* NULL means we're all done with the async results */ DBLINK_GET_CONN;
if (!res) fail = PG_GETARG_BOOL(1);
{
MemoryContextSwitchTo(oldcontext);
SRF_RETURN_DONE(funcctx);
}
} }
else if (PG_NARGS() == 1)
if (!res ||
(PQresultStatus(res) != PGRES_COMMAND_OK &&
PQresultStatus(res) != PGRES_TUPLES_OK))
{ {
if (freeconn) /* text */
PQfinish(conn); DBLINK_GET_CONN;
dblink_res_error(conname, res, "could not execute query", fail);
MemoryContextSwitchTo(oldcontext);
SRF_RETURN_DONE(funcctx);
} }
else
/* shouldn't happen */
elog(ERROR, "wrong number of arguments");
}
if (!conn)
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 */
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);
if (!res ||
(PQresultStatus(res) != PGRES_COMMAND_OK &&
PQresultStatus(res) != PGRES_TUPLES_OK))
{
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)
{
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;
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); MemoryContextSwitchTo(oldcontext);
SRF_RETURN_DONE(funcctx);
}
/* 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; values = (char **) palloc(nfields * sizeof(char *));
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, row, i))
if (PQgetisnull(res, call_cntr, i) == 0) values[i] = NULL;
values[i] = PQgetvalue(res, call_cntr, i); else
values[i] = PQgetvalue(res, row, i);
}
}
else else
values[i] = NULL; {
} values[0] = PQcmdStatus(res);
} }
else
{
values = (char **) palloc(1 * sizeof(char *));
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