Add cancel handlers so it's possible to Ctrl-C clusterdb, reindexdb

and vacuumdb.
ITAGAKI Takahiro, with minor fixes from me.

src/bin/scripts/clusterdb.c

@@ -4,7 +4,7 @@
  *
  * Portions Copyright (c) 2002-2007, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/scripts/clusterdb.c,v 1.16 2007/02/13 18:06:18 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/scripts/clusterdb.c,v 1.17 2007/04/09 18:21:22 mha Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -111,6 +111,8 @@ main(int argc, char *argv[])
 			exit(1);
 	}
 
+	setup_cancel_handler();
+
 	if (alldb)
 	{
 		if (dbname)
@@ -159,7 +161,6 @@ cluster_one_database(const char *dbname, const char *table,
 	PQExpBufferData sql;
 
 	PGconn	   *conn;
-	PGresult   *result;
 
 	initPQExpBuffer(&sql);
 
@@ -169,12 +170,7 @@ cluster_one_database(const char *dbname, const char *table,
 	appendPQExpBuffer(&sql, ";\n");
 
 	conn = connectDatabase(dbname, host, port, username, password, progname);
-
-	if (echo)
-		printf("%s", sql.data);
-	result = PQexec(conn, sql.data);
-
-	if (PQresultStatus(result) != PGRES_COMMAND_OK)
+	if (!executeMaintenanceCommand(conn, sql.data, echo))
 	{
 		if (table)
 			fprintf(stderr, _("%s: clustering of table \"%s\" in database \"%s\" failed: %s"),
@@ -185,8 +181,6 @@ cluster_one_database(const char *dbname, const char *table,
 		PQfinish(conn);
 		exit(1);
 	}
-
-	PQclear(result);
 	PQfinish(conn);
 	termPQExpBuffer(&sql);
 

src/bin/scripts/common.c

@@ -7,7 +7,7 @@
  * Portions Copyright (c) 1996-2007, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- * $PostgreSQL: pgsql/src/bin/scripts/common.c,v 1.25 2007/01/05 22:19:50 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/scripts/common.c,v 1.26 2007/04/09 18:21:22 mha Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -15,14 +15,23 @@
 #include "postgres_fe.h"
 
 #include <pwd.h>
+#include <signal.h>
 #include <unistd.h>
 
 #include "common.h"
+#include "libpq/pqsignal.h"
+
+static void SetCancelConn(PGconn *conn);
+static void ResetCancelConn(void);
 
 #ifndef HAVE_INT_OPTRESET
 int			optreset;
 #endif
 
+static PGcancel *volatile cancelConn = NULL;
+#ifdef WIN32
+static CRITICAL_SECTION cancelConnLock;
+#endif
 
 /*
  * Returns the current user name.
@@ -194,6 +203,33 @@ executeCommand(PGconn *conn, const char *query,
 }
 
 
+/*
+ * 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;
+}
+
+
 /*
  * Check yes/no answer in a localized way.	1=yes, 0=no, -1=neither.
  */
@@ -237,3 +273,135 @@ yesno_prompt(const char *question)
 			   _(PG_YESLETTER), _(PG_NOLETTER));
 	}
 }
+
+
+/*
+ * SetCancelConn
+ *
+ * Set cancelConn to point to the current database connection.
+ */
+static void
+SetCancelConn(PGconn *conn)
+{
+	PGcancel   *oldCancelConn;
+
+#ifdef WIN32
+	EnterCriticalSection(&cancelConnLock);
+#endif
+
+	/* Free the old one if we have one */
+	oldCancelConn = cancelConn;
+
+	/* be sure handle_sigint doesn't use pointer while freeing */
+	cancelConn = NULL;
+
+	if (oldCancelConn != NULL)
+		PQfreeCancel(oldCancelConn);
+
+	cancelConn = PQgetCancel(conn);
+
+#ifdef WIN32
+	LeaveCriticalSection(&cancelConnLock);
+#endif
+}
+
+/*
+ * ResetCancelConn
+ *
+ * Free the current cancel connection, if any, and set to NULL.
+ */
+static void
+ResetCancelConn(void)
+{
+	PGcancel   *oldCancelConn;
+
+#ifdef WIN32
+	EnterCriticalSection(&cancelConnLock);
+#endif
+
+	oldCancelConn = cancelConn;
+
+	/* be sure handle_sigint doesn't use pointer while freeing */
+	cancelConn = NULL;
+
+	if (oldCancelConn != NULL)
+		PQfreeCancel(oldCancelConn);
+
+#ifdef WIN32
+	LeaveCriticalSection(&cancelConnLock);
+#endif
+}
+
+#ifndef WIN32
+/*
+ * Handle interrupt signals by cancelling the current command,
+ * if it's being executed through executeMaintenanceCommand(),
+ * and thus has a cancelConn set.
+ */
+static void
+handle_sigint(SIGNAL_ARGS)
+{
+	int			save_errno = errno;
+	char		errbuf[256];
+
+	/* Send QueryCancel if we are processing a database query */
+	if (cancelConn != NULL)
+	{
+		if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
+			fprintf(stderr, _("Cancel request sent\n"));
+		else
+			fprintf(stderr, _("Could not send cancel request: %s\n"), errbuf);
+	}
+
+	errno = save_errno;			/* just in case the write changed it */
+}
+
+void
+setup_cancel_handler(void)
+{
+	pqsignal(SIGINT, handle_sigint);
+}
+
+#else							/* WIN32 */
+
+/*
+ * Console control handler for Win32. Note that the control handler will
+ * execute on a *different thread* than the main one, so we need to do
+ * proper locking around those structures.
+ */
+static BOOL WINAPI
+consoleHandler(DWORD dwCtrlType)
+{
+	char		errbuf[256];
+
+	if (dwCtrlType == CTRL_C_EVENT ||
+		dwCtrlType == CTRL_BREAK_EVENT)
+	{
+		/* Send QueryCancel if we are processing a database query */
+		EnterCriticalSection(&cancelConnLock);
+		if (cancelConn != NULL)
+		{
+			if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
+				fprintf(stderr, _("Cancel request sent\n"));
+			else
+				fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
+		}
+		LeaveCriticalSection(&cancelConnLock);
+
+		return TRUE;
+	}
+	else
+		/* Return FALSE for any signals not being handled */
+		return FALSE;
+}
+
+void
+setup_cancel_handler(void)
+{
+	InitializeCriticalSection(&cancelConnLock);
+
+	SetConsoleCtrlHandler(consoleHandler, TRUE);
+}
+
+#endif   /* WIN32 */
+

src/bin/scripts/common.h

@@ -4,7 +4,7 @@
  *
  *	Copyright (c) 2003-2007, PostgreSQL Global Development Group
  *
- *	$PostgreSQL: pgsql/src/bin/scripts/common.h,v 1.16 2007/01/05 22:19:50 momjian Exp $
+ *	$PostgreSQL: pgsql/src/bin/scripts/common.h,v 1.17 2007/04/09 18:21:22 mha Exp $
  */
 #ifndef COMMON_H
 #define COMMON_H
@@ -35,6 +35,11 @@ extern PGresult *executeQuery(PGconn *conn, const char *query,
 extern void executeCommand(PGconn *conn, const char *query,
 			   const char *progname, bool echo);
 
+extern bool executeMaintenanceCommand(PGconn *conn, const char *query,
+							   bool echo);
+
 extern bool yesno_prompt(const char *question);
 
+extern void setup_cancel_handler(void);
+
 #endif   /* COMMON_H */

src/bin/scripts/reindexdb.c

@@ -4,7 +4,7 @@
  *
  * Portions Copyright (c) 1996-2007, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/scripts/reindexdb.c,v 1.9 2007/02/13 18:06:18 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/scripts/reindexdb.c,v 1.10 2007/04/09 18:21:22 mha Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -126,6 +126,8 @@ main(int argc, char *argv[])
 			exit(1);
 	}
 
+	setup_cancel_handler();
+
 	if (alldb)
 	{
 		if (dbname)
@@ -214,7 +216,6 @@ reindex_one_database(const char *name, const char *dbname, const char *type,
 	PQExpBufferData sql;
 
 	PGconn	   *conn;
-	PGresult   *result;
 
 	initPQExpBuffer(&sql);
 
@@ -229,11 +230,7 @@ reindex_one_database(const char *name, const char *dbname, const char *type,
 
 	conn = connectDatabase(dbname, host, port, username, password, progname);
 
-	if (echo)
-		printf("%s", sql.data);
-	result = PQexec(conn, sql.data);
-
-	if (PQresultStatus(result) != PGRES_COMMAND_OK)
+	if (!executeMaintenanceCommand(conn, sql.data, echo))
 	{
 		if (strcmp(type, "TABLE") == 0)
 			fprintf(stderr, _("%s: reindexing of table \"%s\" in database \"%s\" failed: %s"),
@@ -248,7 +245,6 @@ reindex_one_database(const char *name, const char *dbname, const char *type,
 		exit(1);
 	}
 
-	PQclear(result);
 	PQfinish(conn);
 	termPQExpBuffer(&sql);
 
@@ -294,27 +290,19 @@ reindex_system_catalogs(const char *dbname, const char *host, const char *port,
 	PQExpBufferData sql;
 
 	PGconn	   *conn;
-	PGresult   *result;
 
 	initPQExpBuffer(&sql);
 
 	appendPQExpBuffer(&sql, "REINDEX SYSTEM %s;\n", dbname);
 
 	conn = connectDatabase(dbname, host, port, username, password, progname);
-
-	if (echo)
-		printf("%s", sql.data);
-	result = PQexec(conn, sql.data);
-
-	if (PQresultStatus(result) != PGRES_COMMAND_OK)
+	if (!executeMaintenanceCommand(conn, sql.data, echo))
 	{
 		fprintf(stderr, _("%s: reindexing of system catalogs failed: %s"),
 				progname, PQerrorMessage(conn));
 		PQfinish(conn);
 		exit(1);
 	}
-
-	PQclear(result);
 	PQfinish(conn);
 	termPQExpBuffer(&sql);
 

src/bin/scripts/vacuumdb.c

@@ -5,7 +5,7 @@
  * Portions Copyright (c) 1996-2007, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- * $PostgreSQL: pgsql/src/bin/scripts/vacuumdb.c,v 1.16 2007/02/13 17:39:39 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/scripts/vacuumdb.c,v 1.17 2007/04/09 18:21:22 mha Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -128,6 +128,8 @@ main(int argc, char *argv[])
 			exit(1);
 	}
 
+	setup_cancel_handler();
+
 	if (alldb)
 	{
 		if (dbname)
@@ -178,7 +180,6 @@ vacuum_one_database(const char *dbname, bool full, bool verbose, bool analyze,
 	PQExpBufferData sql;
 
 	PGconn	   *conn;
-	PGresult   *result;
 
 	initPQExpBuffer(&sql);
 
@@ -194,12 +195,7 @@ vacuum_one_database(const char *dbname, bool full, bool verbose, bool analyze,
 	appendPQExpBuffer(&sql, ";\n");
 
 	conn = connectDatabase(dbname, host, port, username, password, progname);
-
-	if (echo)
-		printf("%s", sql.data);
-	result = PQexec(conn, sql.data);
-
-	if (PQresultStatus(result) != PGRES_COMMAND_OK)
+	if (!executeMaintenanceCommand(conn, sql.data, echo))
 	{
 		if (table)
 			fprintf(stderr, _("%s: vacuuming of table \"%s\" in database \"%s\" failed: %s"),
@@ -210,8 +206,6 @@ vacuum_one_database(const char *dbname, bool full, bool verbose, bool analyze,
 		PQfinish(conn);
 		exit(1);
 	}
-
-	PQclear(result);
 	PQfinish(conn);
 	termPQExpBuffer(&sql);