Skip to content

Commit 6e09df9

Browse files
mhagandermhagander
committed
Add cancel handlers so it's possible to Ctrl-C clusterdb, reindexdb
and vacuumdb. ITAGAKI Takahiro, with minor fixes from me.
1 parent bbed5ba commit 6e09df9

File tree

5 files changed

+188
-39
lines changed

5 files changed

+188
-39
lines changed

‎src/bin/scripts/clusterdb.c

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*
55
* Portions Copyright (c) 2002-2007, PostgreSQL Global Development Group
66
*
7-
* $PostgreSQL: pgsql/src/bin/scripts/clusterdb.c,v 1.16 2007/02/13 18:06:18 momjian Exp $
7+
* $PostgreSQL: pgsql/src/bin/scripts/clusterdb.c,v 1.17 2007/04/09 18:21:22 mha Exp $
88
*
99
*-------------------------------------------------------------------------
1010
*/
@@ -111,6 +111,8 @@ main(int argc, char *argv[])
111111
exit(1);
112112
}
113113

114+
setup_cancel_handler();
115+
114116
if (alldb)
115117
{
116118
if (dbname)
@@ -159,7 +161,6 @@ cluster_one_database(const char *dbname, const char *table,
159161
PQExpBufferData sql;
160162

161163
PGconn *conn;
162-
PGresult *result;
163164

164165
initPQExpBuffer(&sql);
165166

@@ -169,12 +170,7 @@ cluster_one_database(const char *dbname, const char *table,
169170
appendPQExpBuffer(&sql, ";\n");
170171

171172
conn = connectDatabase(dbname, host, port, username, password, progname);
172-
173-
if (echo)
174-
printf("%s", sql.data);
175-
result = PQexec(conn, sql.data);
176-
177-
if (PQresultStatus(result) != PGRES_COMMAND_OK)
173+
if (!executeMaintenanceCommand(conn, sql.data, echo))
178174
{
179175
if (table)
180176
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,
185181
PQfinish(conn);
186182
exit(1);
187183
}
188-
189-
PQclear(result);
190184
PQfinish(conn);
191185
termPQExpBuffer(&sql);
192186

‎src/bin/scripts/common.c

Lines changed: 169 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,22 +7,31 @@
77
* Portions Copyright (c) 1996-2007, PostgreSQL Global Development Group
88
* Portions Copyright (c) 1994, Regents of the University of California
99
*
10-
* $PostgreSQL: pgsql/src/bin/scripts/common.c,v 1.25 2007/01/05 22:19:50 momjian Exp $
10+
* $PostgreSQL: pgsql/src/bin/scripts/common.c,v 1.26 2007/04/09 18:21:22 mha Exp $
1111
*
1212
*-------------------------------------------------------------------------
1313
*/
1414

1515
#include "postgres_fe.h"
1616

1717
#include <pwd.h>
18+
#include <signal.h>
1819
#include <unistd.h>
1920

2021
#include "common.h"
22+
#include "libpq/pqsignal.h"
23+
24+
static void SetCancelConn(PGconn *conn);
25+
static void ResetCancelConn(void);
2126

2227
#ifndef HAVE_INT_OPTRESET
2328
int optreset;
2429
#endif
2530

31+
static PGcancel *volatile cancelConn = NULL;
32+
#ifdef WIN32
33+
static CRITICAL_SECTION cancelConnLock;
34+
#endif
2635

2736
/*
2837
* Returns the current user name.
@@ -194,6 +203,33 @@ executeCommand(PGconn *conn, const char *query,
194203
}
195204

196205

206+
/*
207+
* As above for a SQL maintenance command (returns command success).
208+
* Command is executed with a cancel handler set, so Ctrl-C can
209+
* interrupt it.
210+
*/
211+
bool
212+
executeMaintenanceCommand(PGconn *conn, const char *query, bool echo)
213+
{
214+
PGresult *res;
215+
bool r;
216+
217+
if (echo)
218+
printf("%s\n", query);
219+
220+
SetCancelConn(conn);
221+
res = PQexec(conn, query);
222+
ResetCancelConn();
223+
224+
r = (res && PQresultStatus(res) == PGRES_COMMAND_OK);
225+
226+
if (res)
227+
PQclear(res);
228+
229+
return r;
230+
}
231+
232+
197233
/*
198234
* Check yes/no answer in a localized way. 1=yes, 0=no, -1=neither.
199235
*/
@@ -237,3 +273,135 @@ yesno_prompt(const char *question)
237273
_(PG_YESLETTER), _(PG_NOLETTER));
238274
}
239275
}
276+
277+
278+
/*
279+
* SetCancelConn
280+
*
281+
* Set cancelConn to point to the current database connection.
282+
*/
283+
static void
284+
SetCancelConn(PGconn *conn)
285+
{
286+
PGcancel *oldCancelConn;
287+
288+
#ifdef WIN32
289+
EnterCriticalSection(&cancelConnLock);
290+
#endif
291+
292+
/* Free the old one if we have one */
293+
oldCancelConn = cancelConn;
294+
295+
/* be sure handle_sigint doesn't use pointer while freeing */
296+
cancelConn = NULL;
297+
298+
if (oldCancelConn != NULL)
299+
PQfreeCancel(oldCancelConn);
300+
301+
cancelConn = PQgetCancel(conn);
302+
303+
#ifdef WIN32
304+
LeaveCriticalSection(&cancelConnLock);
305+
#endif
306+
}
307+
308+
/*
309+
* ResetCancelConn
310+
*
311+
* Free the current cancel connection, if any, and set to NULL.
312+
*/
313+
static void
314+
ResetCancelConn(void)
315+
{
316+
PGcancel *oldCancelConn;
317+
318+
#ifdef WIN32
319+
EnterCriticalSection(&cancelConnLock);
320+
#endif
321+
322+
oldCancelConn = cancelConn;
323+
324+
/* be sure handle_sigint doesn't use pointer while freeing */
325+
cancelConn = NULL;
326+
327+
if (oldCancelConn != NULL)
328+
PQfreeCancel(oldCancelConn);
329+
330+
#ifdef WIN32
331+
LeaveCriticalSection(&cancelConnLock);
332+
#endif
333+
}
334+
335+
#ifndef WIN32
336+
/*
337+
* Handle interrupt signals by cancelling the current command,
338+
* if it's being executed through executeMaintenanceCommand(),
339+
* and thus has a cancelConn set.
340+
*/
341+
static void
342+
handle_sigint(SIGNAL_ARGS)
343+
{
344+
int save_errno = errno;
345+
char errbuf[256];
346+
347+
/* Send QueryCancel if we are processing a database query */
348+
if (cancelConn != NULL)
349+
{
350+
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
351+
fprintf(stderr, _("Cancel request sent\n"));
352+
else
353+
fprintf(stderr, _("Could not send cancel request: %s\n"), errbuf);
354+
}
355+
356+
errno = save_errno; /* just in case the write changed it */
357+
}
358+
359+
void
360+
setup_cancel_handler(void)
361+
{
362+
pqsignal(SIGINT, handle_sigint);
363+
}
364+
365+
#else /* WIN32 */
366+
367+
/*
368+
* Console control handler for Win32. Note that the control handler will
369+
* execute on a *different thread* than the main one, so we need to do
370+
* proper locking around those structures.
371+
*/
372+
static BOOL WINAPI
373+
consoleHandler(DWORD dwCtrlType)
374+
{
375+
char errbuf[256];
376+
377+
if (dwCtrlType == CTRL_C_EVENT ||
378+
dwCtrlType == CTRL_BREAK_EVENT)
379+
{
380+
/* Send QueryCancel if we are processing a database query */
381+
EnterCriticalSection(&cancelConnLock);
382+
if (cancelConn != NULL)
383+
{
384+
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
385+
fprintf(stderr, _("Cancel request sent\n"));
386+
else
387+
fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
388+
}
389+
LeaveCriticalSection(&cancelConnLock);
390+
391+
return TRUE;
392+
}
393+
else
394+
/* Return FALSE for any signals not being handled */
395+
return FALSE;
396+
}
397+
398+
void
399+
setup_cancel_handler(void)
400+
{
401+
InitializeCriticalSection(&cancelConnLock);
402+
403+
SetConsoleCtrlHandler(consoleHandler, TRUE);
404+
}
405+
406+
#endif /* WIN32 */
407+

‎src/bin/scripts/common.h

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*
55
* Copyright (c) 2003-2007, PostgreSQL Global Development Group
66
*
7-
* $PostgreSQL: pgsql/src/bin/scripts/common.h,v 1.16 2007/01/05 22:19:50 momjian Exp $
7+
* $PostgreSQL: pgsql/src/bin/scripts/common.h,v 1.17 2007/04/09 18:21:22 mha Exp $
88
*/
99
#ifndef COMMON_H
1010
#define COMMON_H
@@ -35,6 +35,11 @@ extern PGresult *executeQuery(PGconn *conn, const char *query,
3535
extern void executeCommand(PGconn *conn, const char *query,
3636
const char *progname, bool echo);
3737

38+
extern bool executeMaintenanceCommand(PGconn *conn, const char *query,
39+
bool echo);
40+
3841
extern bool yesno_prompt(const char *question);
3942

43+
extern void setup_cancel_handler(void);
44+
4045
#endif /* COMMON_H */

‎src/bin/scripts/reindexdb.c

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*
55
* Portions Copyright (c) 1996-2007, PostgreSQL Global Development Group
66
*
7-
* $PostgreSQL: pgsql/src/bin/scripts/reindexdb.c,v 1.9 2007/02/13 18:06:18 momjian Exp $
7+
* $PostgreSQL: pgsql/src/bin/scripts/reindexdb.c,v 1.10 2007/04/09 18:21:22 mha Exp $
88
*
99
*-------------------------------------------------------------------------
1010
*/
@@ -126,6 +126,8 @@ main(int argc, char *argv[])
126126
exit(1);
127127
}
128128

129+
setup_cancel_handler();
130+
129131
if (alldb)
130132
{
131133
if (dbname)
@@ -214,7 +216,6 @@ reindex_one_database(const char *name, const char *dbname, const char *type,
214216
PQExpBufferData sql;
215217

216218
PGconn *conn;
217-
PGresult *result;
218219

219220
initPQExpBuffer(&sql);
220221

@@ -229,11 +230,7 @@ reindex_one_database(const char *name, const char *dbname, const char *type,
229230

230231
conn = connectDatabase(dbname, host, port, username, password, progname);
231232

232-
if (echo)
233-
printf("%s", sql.data);
234-
result = PQexec(conn, sql.data);
235-
236-
if (PQresultStatus(result) != PGRES_COMMAND_OK)
233+
if (!executeMaintenanceCommand(conn, sql.data, echo))
237234
{
238235
if (strcmp(type, "TABLE") == 0)
239236
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,
248245
exit(1);
249246
}
250247

251-
PQclear(result);
252248
PQfinish(conn);
253249
termPQExpBuffer(&sql);
254250

@@ -294,27 +290,19 @@ reindex_system_catalogs(const char *dbname, const char *host, const char *port,
294290
PQExpBufferData sql;
295291

296292
PGconn *conn;
297-
PGresult *result;
298293

299294
initPQExpBuffer(&sql);
300295

301296
appendPQExpBuffer(&sql, "REINDEX SYSTEM %s;\n", dbname);
302297

303298
conn = connectDatabase(dbname, host, port, username, password, progname);
304-
305-
if (echo)
306-
printf("%s", sql.data);
307-
result = PQexec(conn, sql.data);
308-
309-
if (PQresultStatus(result) != PGRES_COMMAND_OK)
299+
if (!executeMaintenanceCommand(conn, sql.data, echo))
310300
{
311301
fprintf(stderr, _("%s: reindexing of system catalogs failed: %s"),
312302
progname, PQerrorMessage(conn));
313303
PQfinish(conn);
314304
exit(1);
315305
}
316-
317-
PQclear(result);
318306
PQfinish(conn);
319307
termPQExpBuffer(&sql);
320308

‎src/bin/scripts/vacuumdb.c

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
* Portions Copyright (c) 1996-2007, PostgreSQL Global Development Group
66
* Portions Copyright (c) 1994, Regents of the University of California
77
*
8-
* $PostgreSQL: pgsql/src/bin/scripts/vacuumdb.c,v 1.16 2007/02/13 17:39:39 momjian Exp $
8+
* $PostgreSQL: pgsql/src/bin/scripts/vacuumdb.c,v 1.17 2007/04/09 18:21:22 mha Exp $
99
*
1010
*-------------------------------------------------------------------------
1111
*/
@@ -128,6 +128,8 @@ main(int argc, char *argv[])
128128
exit(1);
129129
}
130130

131+
setup_cancel_handler();
132+
131133
if (alldb)
132134
{
133135
if (dbname)
@@ -178,7 +180,6 @@ vacuum_one_database(const char *dbname, bool full, bool verbose, bool analyze,
178180
PQExpBufferData sql;
179181

180182
PGconn *conn;
181-
PGresult *result;
182183

183184
initPQExpBuffer(&sql);
184185

@@ -194,12 +195,7 @@ vacuum_one_database(const char *dbname, bool full, bool verbose, bool analyze,
194195
appendPQExpBuffer(&sql, ";\n");
195196

196197
conn = connectDatabase(dbname, host, port, username, password, progname);
197-
198-
if (echo)
199-
printf("%s", sql.data);
200-
result = PQexec(conn, sql.data);
201-
202-
if (PQresultStatus(result) != PGRES_COMMAND_OK)
198+
if (!executeMaintenanceCommand(conn, sql.data, echo))
203199
{
204200
if (table)
205201
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,
210206
PQfinish(conn);
211207
exit(1);
212208
}
213-
214-
PQclear(result);
215209
PQfinish(conn);
216210
termPQExpBuffer(&sql);
217211

0 commit comments

Comments
 (0)