C++ SetCancelConn函数代码示例

/*
 * SendQuery: send the query string to the backend
 * (and print out results)
 *
 * Note: This is the "front door" way to send a query. That is, use it to
 * send queries actually entered by the user. These queries will be subject to
 * single step mode.
 * To send "back door" queries (generated by slash commands, etc.) in a
 * controlled way, use PSQLexec().
 *
 * Returns true if the query executed successfully, false otherwise.
 */
bool
SendQuery(const char *query)
{
	PGresult   *results;
	TimevalStruct before,
				after;
	bool		OK;

	if (!pset.db)
	{
		psql_error("You are currently not connected to a database.\n");
		return false;
	}

	if (GetVariableBool(pset.vars, "SINGLESTEP"))
	{
		char		buf[3];

		printf(gettext("***(Single step mode: verify command)*******************************************\n"
					   "%s\n"
					   "***(press return to proceed or enter x and return to cancel)********************\n"),
			   query);
		fflush(stdout);
		if (fgets(buf, sizeof(buf), stdin) != NULL)
			if (buf[0] == 'x')
				return false;
	}
	else if (VariableEquals(pset.vars, "ECHO", "queries"))
	{
		puts(query);
		fflush(stdout);
	}

	SetCancelConn();

	if (PQtransactionStatus(pset.db) == PQTRANS_IDLE &&
		!GetVariableBool(pset.vars, "AUTOCOMMIT") &&
		!is_transact_command(query))
	{
		results = PQexec(pset.db, "BEGIN");
		if (PQresultStatus(results) != PGRES_COMMAND_OK)
		{
			psql_error("%s", PQerrorMessage(pset.db));
			PQclear(results);
			ResetCancelConn();
			return false;
		}
		PQclear(results);
	}

	if (pset.timing)
		GETTIMEOFDAY(&before);

	results = PQexec(pset.db, query);

	/* these operations are included in the timing result: */
	OK = (AcceptResult(results) && ProcessCopyResult(results));

	if (pset.timing)
		GETTIMEOFDAY(&after);

	/* but printing results isn't: */
	if (OK)
		OK = PrintQueryResults(results);

	PQclear(results);

	/* Possible microtiming output */
	if (OK && pset.timing)
		printf(gettext("Time: %.3f ms\n"), DIFF_MSEC(&after, &before));

	/* check for events that may occur during query execution */

	if (pset.encoding != PQclientEncoding(pset.db) &&
		PQclientEncoding(pset.db) >= 0)
	{
		/* track effects of SET CLIENT_ENCODING */
		pset.encoding = PQclientEncoding(pset.db);
		pset.popt.topt.encoding = pset.encoding;
		SetVariable(pset.vars, "ENCODING",
					pg_encoding_to_char(pset.encoding));
	}

	PrintNotifications();

	return OK;
}

/*
 * GetQueryResult
 *
 * Process the query result.  Returns true if there's no error, false
 * otherwise -- but errors about trying to vacuum a missing relation are
 * reported and subsequently ignored.
 */
static bool
GetQueryResult(PGconn *conn, const char *dbname, const char *progname)
{
	PGresult   *result;

	SetCancelConn(conn);
	while ((result = PQgetResult(conn)) != NULL)
	{
		/*
		 * If errors are found, report them.  Errors about a missing table are
		 * harmless so we continue processing; but die for other errors.
		 */
		if (PQresultStatus(result) != PGRES_COMMAND_OK)
		{
			char	   *sqlState = PQresultErrorField(result, PG_DIAG_SQLSTATE);

			fprintf(stderr, _("%s: vacuuming of database \"%s\" failed: %s"),
					progname, dbname, PQerrorMessage(conn));

			if (sqlState && strcmp(sqlState, ERRCODE_UNDEFINED_TABLE) != 0)
			{
				PQclear(result);
				return false;
			}
		}

		PQclear(result);
	}
	ResetCancelConn();

	return true;
}

/*
 * do_lo_export()
 *
 * Write a large object to a file
 */
bool
do_lo_export(const char *loid_arg, const char *filename_arg)
{
	int			status;
	bool		own_transaction;

	if (!start_lo_xact("\\lo_export", &own_transaction))
		return false;

	SetCancelConn();
	status = lo_export(pset.db, atooid(loid_arg), filename_arg);
	ResetCancelConn();

	/* of course this status is documented nowhere :( */
	if (status != 1)
	{
		fputs(PQerrorMessage(pset.db), stderr);
		return fail_lo_xact("\\lo_export", own_transaction);
	}

	if (!finish_lo_xact("\\lo_export", own_transaction))
		return false;

	fprintf(pset.queryFout, "lo_export\n");

	return true;
}

/*
 * do_lo_export()
 *
 * Write a large object to a file
 */
bool
do_lo_export(const char *loid_arg, const char *filename_arg)
{
	int			status;
	bool		own_transaction;

	if (!start_lo_xact("\\lo_export", &own_transaction))
		return false;

	SetCancelConn();
	status = lo_export(pset.db, atooid(loid_arg), filename_arg);
	ResetCancelConn();

	/* of course this status is documented nowhere :( */
	if (status != 1)
	{
		psql_error("%s", PQerrorMessage(pset.db));
		return fail_lo_xact("\\lo_export", own_transaction);
	}

	if (!finish_lo_xact("\\lo_export", own_transaction))
		return false;

	print_lo_result("lo_export");

	return true;
}

/*
 * do_lo_unlink()
 *
 * removes a large object out of the database
 */
bool
do_lo_unlink(const char *loid_arg)
{
	int			status;
	Oid			loid = atooid(loid_arg);
	bool		own_transaction;

	if (!start_lo_xact("\\lo_unlink", &own_transaction))
		return false;

	SetCancelConn();
	status = lo_unlink(pset.db, loid);
	ResetCancelConn();

	if (status == -1)
	{
		psql_error("%s", PQerrorMessage(pset.db));
		return fail_lo_xact("\\lo_unlink", own_transaction);
	}

	if (!finish_lo_xact("\\lo_unlink", own_transaction))
		return false;

	print_lo_result("lo_unlink %u", loid);

	return true;
}

/*
 * do_lo_unlink()
 *
 * removes a large object out of the database
 */
bool
do_lo_unlink(const char *loid_arg)
{
	int			status;
	Oid			loid = atooid(loid_arg);
	bool		own_transaction;

	if (!start_lo_xact("\\lo_unlink", &own_transaction))
		return false;

	SetCancelConn();
	status = lo_unlink(pset.db, loid);
	ResetCancelConn();

	if (status == -1)
	{
		fputs(PQerrorMessage(pset.db), stderr);
		return fail_lo_xact("\\lo_unlink", own_transaction);
	}

	if (!finish_lo_xact("\\lo_unlink", own_transaction))
		return false;

	fprintf(pset.queryFout, "lo_unlink %u\n", loid);

	return true;
}

/*
 * PSQLexec
 *
 * This is the way to send "backdoor" queries (those not directly entered
 * by the user). It is subject to -E but not -e.
 *
 * In autocommit-off mode, a new transaction block is started if start_xact
 * is true; nothing special is done when start_xact is false.  Typically,
 * start_xact = false is used for SELECTs and explicit BEGIN/COMMIT commands.
 *
 * Caller is responsible for handling the ensuing processing if a COPY
 * command is sent.
 *
 * Note: we don't bother to check PQclientEncoding; it is assumed that no
 * caller uses this path to issue "SET CLIENT_ENCODING".
 */
PGresult *
PSQLexec(const char *query, bool start_xact)
{
	PGresult   *res;

	if (!pset.db)
	{
		psql_error("You are currently not connected to a database.\n");
		return NULL;
	}

	if (pset.echo_hidden != PSQL_ECHO_HIDDEN_OFF)
	{
		printf(_("********* QUERY **********\n"
				 "%s\n"
				 "**************************\n\n"), query);
		fflush(stdout);
		if (pset.logfile)
		{
			fprintf(pset.logfile,
					_("********* QUERY **********\n"
					  "%s\n"
					  "**************************\n\n"), query);
			fflush(pset.logfile);
		}

		if (pset.echo_hidden == PSQL_ECHO_HIDDEN_NOEXEC)
			return NULL;
	}

	SetCancelConn();

	if (start_xact &&
		!pset.autocommit &&
		PQtransactionStatus(pset.db) == PQTRANS_IDLE)
	{
		res = PQexec(pset.db, "BEGIN");
		if (PQresultStatus(res) != PGRES_COMMAND_OK)
		{
			psql_error("%s", PQerrorMessage(pset.db));
			PQclear(res);
			ResetCancelConn();
			return NULL;
		}
		PQclear(res);
	}

	res = PQexec(pset.db, query);

	ResetCancelConn();

	if (!AcceptResult(res))
	{
		PQclear(res);
		res = NULL;
	}

	return res;
}

/*
 * do_lo_import()
 *
 * Copy large object from file to database
 */
bool
do_lo_import(const char *filename_arg, const char *comment_arg)
{
	PGresult   *res;
	Oid			loid;
	char		oidbuf[32];
	bool		own_transaction;

	if (!start_lo_xact("\\lo_import", &own_transaction))
		return false;

	SetCancelConn();
	loid = lo_import(pset.db, filename_arg);
	ResetCancelConn();

	if (loid == InvalidOid)
	{
		psql_error("%s", PQerrorMessage(pset.db));
		return fail_lo_xact("\\lo_import", own_transaction);
	}

	/* insert description if given */
	if (comment_arg)
	{
		char	   *cmdbuf;
		char	   *bufptr;
		size_t		slen = strlen(comment_arg);

		cmdbuf = malloc(slen * 2 + 256);
		if (!cmdbuf)
			return fail_lo_xact("\\lo_import", own_transaction);
		sprintf(cmdbuf, "COMMENT ON LARGE OBJECT %u IS '", loid);
		bufptr = cmdbuf + strlen(cmdbuf);
		bufptr += PQescapeStringConn(pset.db, bufptr, comment_arg, slen, NULL);
		strcpy(bufptr, "'");

		if (!(res = PSQLexec(cmdbuf, false)))
		{
			free(cmdbuf);
			return fail_lo_xact("\\lo_import", own_transaction);
		}

		PQclear(res);
		free(cmdbuf);
	}

	if (!finish_lo_xact("\\lo_import", own_transaction))
		return false;

	print_lo_result("lo_import %u", loid);

	sprintf(oidbuf, "%u", loid);
	SetVariable(pset.vars, "LASTOID", oidbuf);

	return true;
}

/*
 * ProcessCopyResult: if command was a COPY FROM STDIN/TO STDOUT, handle it
 *
 * Note: Utility function for use by SendQuery() only.
 *
 * Returns true if the query executed successfully, false otherwise.
 */
static bool
ProcessCopyResult(PGresult *results)
{
	bool		success = false;

	if (!results)
		return false;

	switch (PQresultStatus(results))
	{
		case PGRES_TUPLES_OK:
		case PGRES_COMMAND_OK:
		case PGRES_EMPTY_QUERY:
			/* nothing to do here */
			success = true;
			break;

		case PGRES_COPY_OUT:
			SetCancelConn();
			success = handleCopyOut(pset.db, pset.queryFout);
			ResetCancelConn();
			break;

		case PGRES_COPY_IN:
			SetCancelConn();
			success = handleCopyIn(pset.db, pset.cur_cmd_source,
								   PQbinaryTuples(results));
			ResetCancelConn();
			break;

		default:
			break;
	}

	/* may need this to recover from conn loss during COPY */
	if (!CheckConnection())
		return false;

	return success;
}

/*
 * PSQLexec
 *
 * This is the way to send "backdoor" queries (those not directly entered
 * by the user). It is subject to -E but not -e.
 *
 * In autocommit-off mode, a new transaction block is started if start_xact
 * is true; nothing special is done when start_xact is false.  Typically,
 * start_xact = false is used for SELECTs and explicit BEGIN/COMMIT commands.
 *
 * Note: we don't bother to check PQclientEncoding; it is assumed that no
 * caller uses this path to issue "SET CLIENT_ENCODING".
 */
PGresult *
PSQLexec(const char *query, bool start_xact)
{
	PGresult   *res;
	int			echo_hidden;

	if (!pset.db)
	{
		psql_error("You are currently not connected to a database.\n");
		return NULL;
	}

	echo_hidden = SwitchVariable(pset.vars, "ECHO_HIDDEN", "noexec", NULL);
	if (echo_hidden != VAR_NOTSET)
	{
		printf("********* QUERY **********\n"
			   "%s\n"
			   "**************************\n\n", query);
		fflush(stdout);

		if (echo_hidden == 1)	/* noexec? */
			return NULL;
	}

	SetCancelConn();

	if (start_xact && PQtransactionStatus(pset.db) == PQTRANS_IDLE &&
		!GetVariableBool(pset.vars, "AUTOCOMMIT"))
	{
		res = PQexec(pset.db, "BEGIN");
		if (PQresultStatus(res) != PGRES_COMMAND_OK)
		{
			psql_error("%s", PQerrorMessage(pset.db));
			PQclear(res);
			ResetCancelConn();
			return NULL;
		}
		PQclear(res);
	}

	res = PQexec(pset.db, query);

	if (!AcceptResult(res) && res)
	{
		PQclear(res);
		res = NULL;
	}

	return res;
}

/*
 * GetQueryResult
 *
 * Pump the conn till it's dry of results; return false if any are errors.
 * Note that this will block if the conn is busy.
 */
static bool
GetQueryResult(PGconn *conn, const char *progname)
{
	bool		ok = true;
	PGresult   *result;

	SetCancelConn(conn);
	while ((result = PQgetResult(conn)) != NULL)
	{
		if (!ProcessQueryResult(conn, result, progname))
			ok = false;
	}
	ResetCancelConn();
	return ok;
}

/*
 * PSQLexec
 *
 * This is the way to send "backdoor" queries (those not directly entered
 * by the user). It is subject to -E but not -e.
 *
 * Caller is responsible for handling the ensuing processing if a COPY
 * command is sent.
 *
 * Note: we don't bother to check PQclientEncoding; it is assumed that no
 * caller uses this path to issue "SET CLIENT_ENCODING".
 */
PGresult *
PSQLexec(const char *query)
{
	PGresult   *res;

	if (!pset.db)
	{
		psql_error("You are currently not connected to a database.\n");
		return NULL;
	}

	if (pset.echo_hidden != PSQL_ECHO_HIDDEN_OFF)
	{
		printf(_("********* QUERY **********\n"
				 "%s\n"
				 "**************************\n\n"), query);
		fflush(stdout);
		if (pset.logfile)
		{
			fprintf(pset.logfile,
					_("********* QUERY **********\n"
					  "%s\n"
					  "**************************\n\n"), query);
			fflush(pset.logfile);
		}

		if (pset.echo_hidden == PSQL_ECHO_HIDDEN_NOEXEC)
			return NULL;
	}

	SetCancelConn();

	res = PQexec(pset.db, query);

	ResetCancelConn();

	if (!AcceptResult(res))
	{
		PQclear(res);
		res = NULL;
	}

	return res;
}

/*
 * As above for a SQL maintenance command (returns command success).
 * Command is executed with a cancel handler set, so Ctrl-C can
 * interrupt it.
 */
bool
executeMaintenanceCommand(PGconn *conn, const char *query, bool echo)
{
	PGresult   *res;
	bool		r;

	if (echo)
		printf("%s\n", query);

	SetCancelConn(conn);
	res = PQexec(conn, query);
	ResetCancelConn();

	r = (res && PQresultStatus(res) == PGRES_COMMAND_OK);

	if (res)
		PQclear(res);

	return r;
}

/*
 * AcceptResult
 *
 * Checks whether a result is valid, giving an error message if necessary;
 * resets cancelConn as needed, and ensures that the connection to the backend
 * is still up.
 *
 * Returns true for valid result, false for error state.
 */
static bool
AcceptResult(const PGresult *result)
{
	bool		OK = true;

	ResetCancelConn();

	if (!result)
		OK = false;
	else
		switch (PQresultStatus(result))
		{
			case PGRES_COMMAND_OK:
			case PGRES_TUPLES_OK:
			case PGRES_EMPTY_QUERY:
			case PGRES_COPY_IN:
				/* Fine, do nothing */
				break;

			case PGRES_COPY_OUT:
				/* keep cancel connection for copy out state */
				SetCancelConn();
				break;

			default:
				OK = false;
				break;
		}

	if (!OK)
	{
		psql_error("%s", PQerrorMessage(pset.db));
		CheckConnection();
	}

	return OK;
}

/*
 * GetIdleSlot
 *		Return a connection slot that is ready to execute a command.
 *
 * We return the first slot we find that is marked isFree, if one is;
 * otherwise, we loop on select() until one socket becomes available.  When
 * this happens, we read the whole set and mark as free all sockets that become
 * available.
 *
 * Process the slot list, if any free slot is available then return the slotid
 * else perform the select on all the socket's and wait until at least one slot
 * becomes available.
 *
 * If an error occurs, NULL is returned.
 */
static ParallelSlot *
GetIdleSlot(ParallelSlot slots[], int numslots, const char *dbname,
			const char *progname)
{
	int			i;
	int			firstFree = -1;
	fd_set		slotset;
	pgsocket	maxFd;

	for (i = 0; i < numslots; i++)
		if ((slots + i)->isFree)
			return slots + i;

	FD_ZERO(&slotset);

	maxFd = slots->sock;
	for (i = 0; i < numslots; i++)
	{
		FD_SET((slots + i)->sock, &slotset);
		if ((slots + i)->sock > maxFd)
			maxFd = (slots + i)->sock;
	}

	/*
	 * No free slot found, so wait until one of the connections has finished
	 * its task and return the available slot.
	 */
	for (firstFree = -1; firstFree < 0;)
	{
		bool		aborting;

		SetCancelConn(slots->connection);
		i = select_loop(maxFd, &slotset, &aborting);
		ResetCancelConn();

		if (aborting)
		{
			/*
			 * We set the cancel-receiving connection to the one in the zeroth
			 * slot above, so fetch the error from there.
			 */
			GetQueryResult(slots->connection, dbname, progname);
			return NULL;
		}
		Assert(i != 0);

		for (i = 0; i < numslots; i++)
		{
			if (!FD_ISSET((slots + i)->sock, &slotset))
				continue;

			PQconsumeInput((slots + i)->connection);
			if (PQisBusy((slots + i)->connection))
				continue;

			(slots + i)->isFree = true;

			if (!GetQueryResult((slots + i)->connection, dbname, progname))
				return NULL;

			if (firstFree < 0)
				firstFree = i;
		}
	}

	return slots + firstFree;
}