postgres/contrib/dblink/dblink.c

/*
 * dblink.c
 *
 * Functions returning results from a remote database
 *
 * Copyright (c) Joseph Conway <joe.conway@mail.com>, 2001;
 * 
 * Permission to use, copy, modify, and distribute this software and its
 * documentation for any purpose, without fee, and without a written agreement
 * is hereby granted, provided that the above copyright notice and this
 * paragraph and the following two paragraphs appear in all copies.
 * 
 * IN NO EVENT SHALL THE AUTHOR OR DISTRIBUTORS BE LIABLE TO ANY PARTY FOR
 * DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING
 * LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS
 * DOCUMENTATION, EVEN IF THE AUTHOR OR DISTRIBUTORS HAVE BEEN ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 * 
 * THE AUTHOR AND DISTRIBUTORS SPECIFICALLY DISCLAIMS ANY WARRANTIES,
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
 * ON AN "AS IS" BASIS, AND THE AUTHOR AND DISTRIBUTORS HAS NO OBLIGATIONS TO
 * PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
 *
 */

#include "dblink.h"

PG_FUNCTION_INFO_V1(dblink);
Datum
dblink(PG_FUNCTION_ARGS)
{
	PGconn			*conn = NULL;
	PGresult		*res = NULL;
	dblink_results	*results;
	char			*optstr;
	char			*sqlstatement;
	char			*curstr = "DECLARE mycursor CURSOR FOR ";
	char			*execstatement;
	char			*msg;
	int				ntuples = 0;
	ReturnSetInfo 	*rsi;

	if (PG_ARGISNULL(0) || PG_ARGISNULL(1)) {
		elog(ERROR, "dblink: NULL arguments are not permitted");
	}

	if (fcinfo->resultinfo == NULL || ! IsA(fcinfo->resultinfo, ReturnSetInfo)) {
		elog(ERROR, "dblink: function called in context that does not accept a set result");
	}

	optstr = DatumGetCString(DirectFunctionCall1(textout, PointerGetDatum(PG_GETARG_TEXT_P(0))));
	sqlstatement = DatumGetCString(DirectFunctionCall1(textout, PointerGetDatum(PG_GETARG_TEXT_P(1))));

	if (fcinfo->flinfo->fn_extra == NULL) {

		conn = PQconnectdb(optstr);
		if (PQstatus(conn) == CONNECTION_BAD)
		{
			msg = pstrdup(PQerrorMessage(conn));
			PQfinish(conn);
			elog(ERROR, "dblink: connection error: %s", msg);
		}

		res = PQexec(conn, "BEGIN");
		if (PQresultStatus(res) != PGRES_COMMAND_OK)
		{
			msg = pstrdup(PQerrorMessage(conn));
			PQclear(res);
			PQfinish(conn);
			elog(ERROR, "dblink: begin error: %s", msg);
		}
		PQclear(res);

		execstatement = (char *) palloc(strlen(curstr) + strlen(sqlstatement) + 1);
		if (execstatement != NULL) {
			strcpy(execstatement, curstr);
			strcat(execstatement, sqlstatement);
			strcat(execstatement, "\0");
		} else {
			elog(ERROR, "dblink: insufficient memory" );
		}

		res = PQexec(conn, execstatement);
		if (!res || (PQresultStatus(res) != PGRES_COMMAND_OK && PQresultStatus(res) != PGRES_TUPLES_OK))
		{
			msg = pstrdup(PQerrorMessage(conn));
			PQclear(res);
			PQfinish(conn);
			elog(ERROR, "dblink: sql error: %s", msg);
		} else {
			/*
			 * got results, start fetching them
			 */
		    PQclear(res);

		    res = PQexec(conn, "FETCH ALL in mycursor");
		    if (!res || PQresultStatus(res) != PGRES_TUPLES_OK) {
				msg = pstrdup(PQerrorMessage(conn));
				PQclear(res);
				PQfinish(conn);
				elog(ERROR, "dblink: sql error: %s", msg);
		    }

			ntuples = PQntuples(res);

			if (ntuples > 0) {

				results = init_dblink_results(fcinfo->flinfo->fn_mcxt);
				results->tup_num = 0;
				results->res = res;
				res = NULL;

				fcinfo->flinfo->fn_extra = (void *) results;

				results = NULL;
				results = fcinfo->flinfo->fn_extra;

			    /* close the cursor */
			    res = PQexec(conn, "CLOSE mycursor");
			    PQclear(res);

			    /* commit the transaction */
			    res = PQexec(conn, "COMMIT");
			    PQclear(res);

			    /* close the connection to the database and cleanup */
			    PQfinish(conn);

				rsi = (ReturnSetInfo *)fcinfo->resultinfo;
				rsi->isDone = ExprMultipleResult;

				PG_RETURN_POINTER(results);

			} else {

			    PQclear(res);

			    /* close the cursor */
			    res = PQexec(conn, "CLOSE mycursor");
			    PQclear(res);

			    /* commit the transaction */
			    res = PQexec(conn, "COMMIT");
			    PQclear(res);

			    /* close the connection to the database and cleanup */
			    PQfinish(conn);

				rsi = (ReturnSetInfo *)fcinfo->resultinfo;
				rsi->isDone = ExprEndResult ;

				PG_RETURN_NULL();
			}
		}
	} else {
		/*
		 * check for more results
		 */

		results = fcinfo->flinfo->fn_extra;
		results->tup_num++;
		ntuples = PQntuples(results->res);

		if (results->tup_num < ntuples) {

			/*
			 * fetch them if available
			 */

			rsi = (ReturnSetInfo *)fcinfo->resultinfo;
			rsi->isDone = ExprMultipleResult;

			PG_RETURN_POINTER(results);

		} else {

			/*
			 * or if no more, clean things up
			 */

			results = fcinfo->flinfo->fn_extra;

		    PQclear(results->res);

			rsi = (ReturnSetInfo *)fcinfo->resultinfo;
			rsi->isDone = ExprEndResult ;

			PG_RETURN_NULL();
		}
	}
	PG_RETURN_NULL();
}


/*
 * dblink_tok
 * parse dblink output string
 * return fldnum item (0 based)
 * based on provided field separator
 */

PG_FUNCTION_INFO_V1(dblink_tok);
Datum
dblink_tok(PG_FUNCTION_ARGS)
{
	dblink_results	*results;
	int				fldnum;
	text			*result_text;
	char			*result;
	int				nfields = 0;
	int				text_len = 0;

	if (PG_ARGISNULL(0) || PG_ARGISNULL(1)) {
		elog(ERROR, "dblink: NULL arguments are not permitted");
	}

	results = (dblink_results *) PG_GETARG_POINTER(0);
	if (results == NULL) {
		elog(ERROR, "dblink: function called with invalid result pointer");
	}

	fldnum = PG_GETARG_INT32(1);
	if (fldnum < 0) {
		elog(ERROR, "dblink: field number < 0 not permitted");
	}

	nfields = PQnfields(results->res);
	if (fldnum > (nfields - 1)) {
		elog(ERROR, "dblink: field number %d does not exist", fldnum);
	}

	if (PQgetisnull(results->res, results->tup_num, fldnum) == 1) {

		PG_RETURN_NULL();

	} else {

		text_len = PQgetlength(results->res, results->tup_num, fldnum);

		result = (char *) palloc(text_len + 1);

		if (result != NULL) {
			strcpy(result, PQgetvalue(results->res, results->tup_num, fldnum));
			strcat(result, "\0");
		} else {
			elog(ERROR, "dblink: insufficient memory" );
		}

		result_text = DatumGetTextP(DirectFunctionCall1(textin, CStringGetDatum(result)));

		PG_RETURN_TEXT_P(result_text);

	}
}


/*
 * internal functions
 */


/*
 * init_dblink_results
 *	 - create an empty dblink_results data structure
 */
dblink_results *
init_dblink_results(MemoryContext fn_mcxt)
{
	MemoryContext	oldcontext;
	dblink_results	*retval;

	oldcontext = MemoryContextSwitchTo(fn_mcxt);

	retval = (dblink_results *) palloc(sizeof(dblink_results));
	MemSet(retval, 0, sizeof(dblink_results));

	retval->tup_num = -1;
	retval->res = NULL;

	MemoryContextSwitchTo(oldcontext);

	return retval;
}