Add a --role option to pg_dump, pg_dumpall, and pg_restore. This allows

performing dumps and restores in accordance with a security policy that
forbids logging in directly as superuser, but instead specifies that you
should log into an admin account and then SET ROLE to the superuser.

In passing, clean up some ugly and mostly-broken code for quoting shell
arguments in pg_dumpall.

Benedek László, with some help from Tom Lane

5 files changed, 115 insertions, 86 deletions

diff --git a/src/bin/pg_dump/pg_backup.h b/src/bin/pg_dump/pg_backup.h
index ff4099fe4e9..fa7e2d55f13 100644
--- a/src/bin/pg_dump/pg_backup.h
+++ b/src/bin/pg_dump/pg_backup.h
@@ -15,7 +15,7 @@
  *
  *
  * IDENTIFICATION
- *		$PostgreSQL: pgsql/src/bin/pg_dump/pg_backup.h,v 1.47 2008/04/13 03:49:21 tgl Exp $
+ *		$PostgreSQL: pgsql/src/bin/pg_dump/pg_backup.h,v 1.48 2009/01/05 16:54:36 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -89,6 +89,7 @@ typedef struct _restoreOptions
 	int			use_setsessauth;/* Use SET SESSION AUTHORIZATION commands
 								 * instead of OWNER TO */
 	char	   *superuser;		/* Username to use as superuser */
+	char	   *use_role;		/* Issue SET ROLE to this */
 	int			dataOnly;
 	int			dropSchema;
 	char	   *filename;
diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
index ec14f1b9426..9153604edc6 100644
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -15,7 +15,7 @@
  *
  *
  * IDENTIFICATION
- *		$PostgreSQL: pgsql/src/bin/pg_dump/pg_backup_archiver.c,v 1.159 2008/12/19 16:25:17 petere Exp $
+ *		$PostgreSQL: pgsql/src/bin/pg_dump/pg_backup_archiver.c,v 1.160 2009/01/05 16:54:36 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -462,9 +462,8 @@ NewRestoreOptions(void)
 
 	opts = (RestoreOptions *) calloc(1, sizeof(RestoreOptions));
 
+	/* set any fields that shouldn't default to zeroes */
 	opts->format = archUnknown;
-	opts->suppressDumpWarnings = false;
-	opts->exit_on_error = false;
 
 	return opts;
 }
@@ -2146,6 +2145,10 @@ _doSetFixedOutputState(ArchiveHandle *AH)
 	ahprintf(AH, "SET standard_conforming_strings = %s;\n",
 			 AH->public.std_strings ? "on" : "off");
 
+	/* Select the role to be used during restore */
+	if (AH->ropt && AH->ropt->use_role)
+		ahprintf(AH, "SET ROLE %s;\n", fmtId(AH->ropt->use_role));
+
 	/* Make sure function checking is disabled */
 	ahprintf(AH, "SET check_function_bodies = false;\n");
 
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index 75ece39cc42..1935958dd3e 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -12,7 +12,7 @@
  *	by PostgreSQL
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/bin/pg_dump/pg_dump.c,v 1.511 2009/01/01 17:23:54 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/bin/pg_dump/pg_dump.c,v 1.512 2009/01/05 16:54:37 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -227,6 +227,7 @@ main(int argc, char **argv)
 	bool		outputBlobs = false;
 	int			outputNoOwner = 0;
 	char	   *outputSuperuser = NULL;
+	char	   *use_role = NULL;
 	int			my_version;
 	int			optindex;
 	RestoreOptions *ropt;
@@ -274,6 +275,7 @@ main(int argc, char **argv)
 		{"disable-triggers", no_argument, &disable_triggers, 1},
 		{"lock-wait-timeout", required_argument, NULL, 2},
 		{"no-tablespaces", no_argument, &outputNoTablespaces, 1},
+		{"role", required_argument, NULL, 3},
 		{"use-set-session-authorization", no_argument, &use_setsessauth, 1},
 
 		{NULL, 0, NULL, 0}
@@ -447,11 +449,14 @@ main(int argc, char **argv)
 				/* This covers the long options equivalent to -X xxx. */
 				break;
 
-			case 2:
-				/* lock-wait-timeout */
+			case 2:				/* lock-wait-timeout */
 				lockWaitTimeout = optarg;
 				break;
 
+			case 3:				/* SET ROLE */
+				use_role = optarg;
+				break;
+
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
@@ -570,6 +575,16 @@ main(int argc, char **argv)
 	std_strings = PQparameterStatus(g_conn, "standard_conforming_strings");
 	g_fout->std_strings = (std_strings && strcmp(std_strings, "on") == 0);
 
+	/* Set the role if requested */
+	if (use_role && g_fout->remoteVersion >= 80100)
+	{
+		PQExpBuffer query = createPQExpBuffer();
+
+		appendPQExpBuffer(query, "SET ROLE %s", fmtId(use_role));
+		do_sql_command(g_conn, query->data);
+		destroyPQExpBuffer(query);
+	}
+
 	/* Set the datestyle to ISO to ensure the dump's portability */
 	do_sql_command(g_conn, "SET DATESTYLE = ISO");
 
@@ -807,6 +822,7 @@ help(const char *progname)
 	printf(_("  --disable-dollar-quoting    disable dollar quoting, use SQL standard quoting\n"));
 	printf(_("  --disable-triggers          disable triggers during data-only restore\n"));
 	printf(_("  --no-tablespaces            do not dump tablespace assignments\n"));
+	printf(_("  --role=ROLENAME             do SET ROLE before dump\n"));
 	printf(_("  --use-set-session-authorization\n"
 			 "                              use SESSION AUTHORIZATION commands instead of\n"
 	"                              ALTER OWNER commands to set ownership\n"));
diff --git a/src/bin/pg_dump/pg_dumpall.c b/src/bin/pg_dump/pg_dumpall.c
index e7b684a5115..3ee8ad1cc77 100644
--- a/src/bin/pg_dump/pg_dumpall.c
+++ b/src/bin/pg_dump/pg_dumpall.c
@@ -6,7 +6,7 @@
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  *
- * $PostgreSQL: pgsql/src/bin/pg_dump/pg_dumpall.c,v 1.110 2009/01/01 17:23:54 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/pg_dump/pg_dumpall.c,v 1.111 2009/01/05 16:54:37 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -48,6 +48,7 @@ static void makeAlterConfigCommand(PGconn *conn, const char *arrayitem,
 					   const char *type, const char *name);
 static void dumpDatabases(PGconn *conn);
 static void dumpTimestamp(char *msg);
+static void doShellQuoting(PQExpBuffer buf, const char *str);
 
 static int	runPgDump(const char *dbname);
 static PGconn *connectDatabase(const char *dbname, const char *pghost, const char *pgport,
@@ -77,6 +78,7 @@ main(int argc, char *argv[])
 	char	   *pgport = NULL;
 	char	   *pguser = NULL;
 	char	   *pgdb = NULL;
+	char	   *use_role = NULL;
 	bool		force_password = false;
 	bool		data_only = false;
 	bool		globals_only = false;
@@ -118,9 +120,10 @@ main(int argc, char *argv[])
 		 */
 		{"disable-dollar-quoting", no_argument, &disable_dollar_quoting, 1},
 		{"disable-triggers", no_argument, &disable_triggers, 1},
+		{"lock-wait-timeout", required_argument, NULL, 2},
 		{"no-tablespaces", no_argument, &no_tablespaces, 1},
+		{"role", required_argument, NULL, 3},
 		{"use-set-session-authorization", no_argument, &use_setsessauth, 1},
-		{"lock-wait-timeout", required_argument, NULL, 2},
 
 		{NULL, 0, NULL, 0}
 	};
@@ -191,12 +194,8 @@ main(int argc, char *argv[])
 
 			case 'f':
 				filename = optarg;
-#ifndef WIN32
-				appendPQExpBuffer(pgdumpopts, " -f '%s'", filename);
-#else
-				appendPQExpBuffer(pgdumpopts, " -f \"%s\"", filename);
-#endif
-
+				appendPQExpBuffer(pgdumpopts, " -f ");
+				doShellQuoting(pgdumpopts, filename);
 				break;
 
 			case 'g':
@@ -205,12 +204,8 @@ main(int argc, char *argv[])
 
 			case 'h':
 				pghost = optarg;
-#ifndef WIN32
-				appendPQExpBuffer(pgdumpopts, " -h '%s'", pghost);
-#else
-				appendPQExpBuffer(pgdumpopts, " -h \"%s\"", pghost);
-#endif
-
+				appendPQExpBuffer(pgdumpopts, " -h ");
+				doShellQuoting(pgdumpopts, pghost);
 				break;
 
 			case 'i':
@@ -231,11 +226,8 @@ main(int argc, char *argv[])
 
 			case 'p':
 				pgport = optarg;
-#ifndef WIN32
-				appendPQExpBuffer(pgdumpopts, " -p '%s'", pgport);
-#else
-				appendPQExpBuffer(pgdumpopts, " -p \"%s\"", pgport);
-#endif
+				appendPQExpBuffer(pgdumpopts, " -p ");
+				doShellQuoting(pgdumpopts, pgport);
 				break;
 
 			case 'r':
@@ -248,11 +240,8 @@ main(int argc, char *argv[])
 				break;
 
 			case 'S':
-#ifndef WIN32
-				appendPQExpBuffer(pgdumpopts, " -S '%s'", optarg);
-#else
-				appendPQExpBuffer(pgdumpopts, " -S \"%s\"", optarg);
-#endif
+				appendPQExpBuffer(pgdumpopts, " -S ");
+				doShellQuoting(pgdumpopts, optarg);
 				break;
 
 			case 't':
@@ -261,11 +250,8 @@ main(int argc, char *argv[])
 
 			case 'U':
 				pguser = optarg;
-#ifndef WIN32
-				appendPQExpBuffer(pgdumpopts, " -U '%s'", pguser);
-#else
-				appendPQExpBuffer(pgdumpopts, " -U \"%s\"", pguser);
-#endif
+				appendPQExpBuffer(pgdumpopts, " -U ");
+				doShellQuoting(pgdumpopts, pguser);
 				break;
 
 			case 'v':
@@ -307,8 +293,14 @@ main(int argc, char *argv[])
 				break;
 
 			case 2:
-				appendPQExpBuffer(pgdumpopts, " --lock-wait-timeout=");
-				appendPQExpBuffer(pgdumpopts, "%s", optarg);
+				appendPQExpBuffer(pgdumpopts, " --lock-wait-timeout ");
+				doShellQuoting(pgdumpopts, optarg);
+				break;
+
+			case 3:
+				use_role = optarg;
+				appendPQExpBuffer(pgdumpopts, " --role ");
+				doShellQuoting(pgdumpopts, use_role);
 				break;
 
 			default:
@@ -426,6 +418,16 @@ main(int argc, char *argv[])
 	if (!std_strings)
 		std_strings = "off";
 
+	/* Set the role if requested */
+	if (use_role && server_version >= 80100)
+	{
+		PQExpBuffer query = createPQExpBuffer();
+
+		appendPQExpBuffer(query, "SET ROLE %s", fmtId(use_role));
+		executeCommand(conn, query->data);
+		destroyPQExpBuffer(query);
+	}
+
 	fprintf(OPF, "--\n-- PostgreSQL database cluster dump\n--\n\n");
 	if (verbose)
 		dumpTimestamp("Started on");
@@ -513,6 +515,7 @@ help(void)
 			 "                           disable dollar quoting, use SQL standard quoting\n"));
 	printf(_("  --disable-triggers       disable triggers during data-only restore\n"));
 	printf(_("  --no-tablespaces         do not dump tablespace assignments\n"));
+	printf(_("  --role=ROLENAME          do SET ROLE before dump\n"));
 	printf(_("  --use-set-session-authorization\n"
 			 "                           use SESSION AUTHORIZATION commands instead of\n"
 			 "                           OWNER TO commands\n"));
@@ -1271,56 +1274,21 @@ static int
 runPgDump(const char *dbname)
 {
 	PQExpBuffer cmd = createPQExpBuffer();
-	const char *p;
 	int			ret;
 
+	appendPQExpBuffer(cmd, SYSTEMQUOTE "\"%s\" %s", pg_dump_bin,
+					  pgdumpopts->data);
+
 	/*
-	 * Win32 has to use double-quotes for args, rather than single quotes.
-	 * Strangely enough, this is the only place we pass a database name on the
-	 * command line, except "postgres" which doesn't need quoting.
-	 *
 	 * If we have a filename, use the undocumented plain-append pg_dump
 	 * format.
 	 */
 	if (filename)
-	{
-#ifndef WIN32
-		appendPQExpBuffer(cmd, SYSTEMQUOTE"\"%s\" %s -Fa '", pg_dump_bin,
-#else
-		appendPQExpBuffer(cmd, SYSTEMQUOTE"\"%s\" %s -Fa \"", pg_dump_bin,
-#endif
-						  pgdumpopts->data);
-	}
+		appendPQExpBuffer(cmd, " -Fa ");
 	else
-	{
-#ifndef WIN32
-		appendPQExpBuffer(cmd, SYSTEMQUOTE "\"%s\" %s -Fp '", pg_dump_bin,
-#else
-		appendPQExpBuffer(cmd, SYSTEMQUOTE "\"%s\" %s -Fp \"", pg_dump_bin,
-#endif
-						  pgdumpopts->data);
-	}
-
-
-	/* Shell quoting is not quite like SQL quoting, so can't use fmtId */
-	for (p = dbname; *p; p++)
-	{
-#ifndef WIN32
-		if (*p == '\'')
-			appendPQExpBuffer(cmd, "'\"'\"'");
-#else
-		if (*p == '"')
-			appendPQExpBuffer(cmd, "\\\"");
-#endif
-		else
-			appendPQExpBufferChar(cmd, *p);
-	}
+		appendPQExpBuffer(cmd, " -Fp ");
 
-#ifndef WIN32
-	appendPQExpBufferChar(cmd, '\'');
-#else
-	appendPQExpBufferChar(cmd, '"');
-#endif
+	doShellQuoting(cmd, dbname);
 
 	appendPQExpBuffer(cmd, "%s", SYSTEMQUOTE);
 
@@ -1338,7 +1306,6 @@ runPgDump(const char *dbname)
 }
 
 
-
 /*
  * Make a database connection with the given parameters.  An
  * interactive password prompt is automatically issued if required.
@@ -1527,3 +1494,38 @@ dumpTimestamp(char *msg)
 				 localtime(&now)) != 0)
 		fprintf(OPF, "-- %s %s\n\n", msg, buf);
 }
+
+
+/*
+ * Append the given string to the shell command being built in the buffer,
+ * with suitable shell-style quoting.
+ */
+static void
+doShellQuoting(PQExpBuffer buf, const char *str)
+{
+	const char *p;
+
+#ifndef WIN32
+	appendPQExpBufferChar(buf, '\'');
+	for (p = str; *p; p++)
+	{
+		if (*p == '\'')
+			appendPQExpBuffer(buf, "'\"'\"'");
+		else
+			appendPQExpBufferChar(buf, *p);
+	}
+	appendPQExpBufferChar(buf, '\'');
+
+#else /* WIN32 */
+
+	appendPQExpBufferChar(buf, '"');
+	for (p = str; *p; p++)
+	{
+		if (*p == '"')
+			appendPQExpBuffer(buf, "\\\"");
+		else
+			appendPQExpBufferChar(buf, *p);
+	}
+	appendPQExpBufferChar(buf, '"');
+#endif /* WIN32 */
+}
diff --git a/src/bin/pg_dump/pg_restore.c b/src/bin/pg_dump/pg_restore.c
index 95bb61b72fe..c0171fe8928 100644
--- a/src/bin/pg_dump/pg_restore.c
+++ b/src/bin/pg_dump/pg_restore.c
@@ -34,7 +34,7 @@
  *
  *
  * IDENTIFICATION
- *		$PostgreSQL: pgsql/src/bin/pg_dump/pg_restore.c,v 1.89 2008/12/11 07:34:08 petere Exp $
+ *		$PostgreSQL: pgsql/src/bin/pg_dump/pg_restore.c,v 1.90 2009/01/05 16:54:37 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -51,6 +51,9 @@
 
 #include "getopt_long.h"
 
+extern char *optarg;
+extern int	optind;
+
 #ifndef HAVE_INT_OPTRESET
 int			optreset;
 #endif
@@ -72,8 +75,6 @@ main(int argc, char **argv)
 	int			exit_code;
 	Archive    *AH;
 	char	   *inputFileSpec;
-	extern int	optind;
-	extern char *optarg;
 	static int	disable_triggers = 0;
 	static int	no_data_for_failed_tables = 0;
 	static int  outputNoTablespaces = 0;
@@ -114,6 +115,7 @@ main(int argc, char **argv)
 		{"disable-triggers", no_argument, &disable_triggers, 1},
 		{"no-data-for-failed-tables", no_argument, &no_data_for_failed_tables, 1},
 		{"no-tablespaces", no_argument, &outputNoTablespaces, 1},
+		{"role", required_argument, NULL, 2},
 		{"use-set-session-authorization", no_argument, &use_setsessauth, 1},
 
 		{NULL, 0, NULL, 0}
@@ -261,13 +263,17 @@ main(int argc, char **argv)
 				}
 				break;
 
+			case '1':			/* Restore data in a single transaction */
+				opts->single_txn = true;
+				opts->exit_on_error = true;
+				break;
+
 			case 0:
 				/* This covers the long options equivalent to -X xxx. */
 				break;
 
-			case '1':			/* Restore data in a single transaction */
-				opts->single_txn = true;
-				opts->exit_on_error = true;
+			case 2:				/* SET ROLE */
+				opts->use_role = optarg;
 				break;
 
 			default:
@@ -405,6 +411,7 @@ usage(const char *progname)
 			 "                           do not restore data of tables that could not be\n"
 			 "                           created\n"));
 	printf(_("  --no-tablespaces         do not dump tablespace assignments\n"));
+	printf(_("  --role=ROLENAME          do SET ROLE before restore\n"));
 	printf(_("  --use-set-session-authorization\n"
 			 "                           use SESSION AUTHORIZATION commands instead of\n"
 			 "                           OWNER TO commands\n"));