Add psql option:

	-1 or --single-transaction

Simon Riggs

9 files changed, 67 insertions, 13 deletions

diff --git a/doc/src/sgml/ref/pg_restore.sgml b/doc/src/sgml/ref/pg_restore.sgml
index 7497e149df..269d00cf99 100644
--- a/doc/src/sgml/ref/pg_restore.sgml
+++ b/doc/src/sgml/ref/pg_restore.sgml
@@ -448,6 +448,19 @@
        </para>
       </listitem>
      </varlistentry>
+
+     <varlistentry>
+      <term><option>-1</option></term>
+      <term><option>--single-transaction</option></term>
+      <listitem>
+       <para>
+        Force the restore to execute as a single transaction. Either all
+        SQL statements complete successfully, or no changes are applied. This
+        option also forces --exit-on-error.
+       </para>
+      </listitem>
+     </varlistentry>
+
     </variablelist>
    </para>
  </refsect1>
diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index f1fa0dda1a..a77e147603 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -463,6 +463,18 @@ PostgreSQL documentation
       </listitem>
     </varlistentry>
 
+     <varlistentry>
+      <term><option>-1</option></term>
+      <term><option>--single-transaction</option></term>
+      <listitem>
+       <para>
+        When psql executes a script with the -f option, this additional option
+        will force the script to execute as a single transaction. Either all
+        SQL statements complete successfully, or no changes are applied. 
+       </para>
+      </listitem>
+     </varlistentry>
+
     <varlistentry>
       <term><option>-?</></term>
       <term><option>--help</></term>
diff --git a/src/bin/pg_dump/pg_backup.h b/src/bin/pg_dump/pg_backup.h
index 6d0ee550d1..f0ae1e2d25 100644
--- a/src/bin/pg_dump/pg_backup.h
+++ b/src/bin/pg_dump/pg_backup.h
@@ -115,6 +115,8 @@ typedef struct _restoreOptions
 
 	int			suppressDumpWarnings;	/* Suppress output of WARNING entries
 										 * to stderr */
+    bool        single_txn;
+
 } RestoreOptions;
 
 /*
diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
index 7f52491fbc..3ca70a977b 100644
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -217,6 +217,9 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
 
 	AH->stage = STAGE_PROCESSING;
 
+    if (ropt->single_txn)
+		ahprintf(AH, "BEGIN;\n\n");
+
 	/*
 	 * Drop the items at the start, in reverse order
 	 */
@@ -370,6 +373,9 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
 		}
 	}
 
+    if (ropt->single_txn)
+		ahprintf(AH, "COMMIT;\n\n");
+
 	if (AH->public.verbose)
 		dumpTimestamp(AH, "Completed on", time(NULL));
 
diff --git a/src/bin/pg_dump/pg_restore.c b/src/bin/pg_dump/pg_restore.c
index 050518fb01..d4acda30f2 100644
--- a/src/bin/pg_dump/pg_restore.c
+++ b/src/bin/pg_dump/pg_restore.c
@@ -111,6 +111,7 @@ main(int argc, char **argv)
 		{"use-list", 1, NULL, 'L'},
 		{"username", 1, NULL, 'U'},
 		{"verbose", 0, NULL, 'v'},
+		{"single-transaction", 0, NULL, '1'},
 
 		/*
 		 * the following options don't have an equivalent short option letter,
@@ -142,7 +143,7 @@ main(int argc, char **argv)
 		}
 	}
 
-	while ((c = getopt_long(argc, argv, "acCd:ef:F:h:iI:lL:n:Op:P:RsS:t:T:uU:vWxX:",
+	while ((c = getopt_long(argc, argv, "acCd:ef:F:h:iI:lL:n:Op:P:RsS:t:T:uU:vWxX:1",
 							cmdopts, NULL)) != -1)
 	{
 		switch (c)
@@ -185,9 +186,15 @@ main(int argc, char **argv)
 				opts->tocFile = strdup(optarg);
 				break;
 
+			case 'n':			/* Dump data for this schema only */
+				opts->selTypes = 1;
+				opts->schemaNames = strdup(optarg);
+				break;
+
 			case 'O':
 				opts->noOwner = 1;
 				break;
+
 			case 'p':
 				if (strlen(optarg) != 0)
 					opts->pgport = strdup(optarg);
@@ -223,11 +230,6 @@ main(int argc, char **argv)
 				opts->tableNames = strdup(optarg);
 				break;
 
-			case 'n':			/* Dump data for this schema only */
-				opts->selTypes = 1;
-				opts->schemaNames = strdup(optarg);
-				break;
-
 			case 'u':
 				opts->requirePassword = true;
 				opts->username = simple_prompt("User name: ", 100, true);
@@ -268,6 +270,11 @@ main(int argc, char **argv)
 			case 0:
 				break;
 
+			case '1':			/* Restore data in a single transaction */
+				opts->single_txn = true;
+				opts->exit_on_error = true;
+				break;
+
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
@@ -394,6 +401,7 @@ usage(const char *progname)
 	printf(_("  -X use-set-session-authorization, --use-set-session-authorization\n"
 			 "                           use SESSION AUTHORIZATION commands instead of\n"
 			 "                           OWNER TO commands\n"));
+	printf(_("  -1, --single-transaction restore as a single transaction\n"));
 
 	printf(_("\nConnection options:\n"));
 	printf(_("  -h, --host=HOSTNAME      database server host or socket directory\n"));
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 7b9274e639..dd2ca4f48c 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -563,7 +563,7 @@ exec_command(const char *cmd,
 		else
 		{
 			expand_tilde(&fname);
-			success = (process_file(fname) == EXIT_SUCCESS);
+			success = (process_file(fname, false) == EXIT_SUCCESS);
 			free(fname);
 		}
 	}
@@ -1435,11 +1435,12 @@ do_edit(const char *filename_arg, PQExpBuffer query_buf)
  * MainLoop() error code.
  */
 int
-process_file(char *filename)
+process_file(char *filename, bool single_txn)
 {
 	FILE	   *fd;
 	int			result;
 	char	   *oldfilename;
+	PGresult   *res; 
 
 	if (!filename)
 		return EXIT_FAILURE;
@@ -1455,7 +1456,13 @@ process_file(char *filename)
 
 	oldfilename = pset.inputfile;
 	pset.inputfile = filename;
+
+    if (single_txn)
+        res = PSQLexec("BEGIN", false);
 	result = MainLoop(fd);
+    if (single_txn)
+        res = PSQLexec("COMMIT", false);
+
 	fclose(fd);
 	pset.inputfile = oldfilename;
 	return result;
diff --git a/src/bin/psql/command.h b/src/bin/psql/command.h
index 7dae5d2dcc..47bd2dbe08 100644
--- a/src/bin/psql/command.h
+++ b/src/bin/psql/command.h
@@ -28,7 +28,7 @@ typedef enum _backslashResult
 extern backslashResult HandleSlashCmds(PsqlScanState scan_state,
 				PQExpBuffer query_buf);
 
-extern int	process_file(char *filename);
+extern int	process_file(char *filename, bool single_txn);
 
 extern bool do_pset(const char *param,
 		const char *value,
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index 2d72bdf85a..63555cdb6f 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -93,6 +93,7 @@ usage(void)
 	printf(_("  -d DBNAME       specify database name to connect to (default: \"%s\")\n"), env);
 	puts(_("  -c COMMAND      run only single command (SQL or internal) and exit"));
 	puts(_("  -f FILENAME     execute commands from file, then exit"));
+	puts(_("  -1 (numeral)    execute command file as a single transaction"));
 	puts(_("  -l              list available databases, then exit"));
 	puts(_("  -v NAME=VALUE   set psql variable NAME to VALUE"));
 	puts(_("  -X              do not read startup file (~/.psqlrc)"));
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index a20d5e4792..0eff9ed2a0 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -76,6 +76,7 @@ struct adhoc_opts
 	char	   *action_string;
 	bool		no_readline;
 	bool		no_psqlrc;
+    bool        single_txn;
 };
 
 static int	parse_version(const char *versionString);
@@ -268,7 +269,7 @@ main(int argc, char *argv[])
 		if (!options.no_psqlrc)
 			process_psqlrc(argv[0]);
 
-		successResult = process_file(options.action_string);
+		successResult = process_file(options.action_string, options.single_txn);
 	}
 
 	/*
@@ -425,6 +426,7 @@ parse_psql_options(int argc, char *argv[], struct adhoc_opts * options)
 		{"list", no_argument, NULL, 'l'},
 		{"log-file", required_argument, NULL, 'L'},
 		{"no-readline", no_argument, NULL, 'n'},
+		{"single-transaction", no_argument, NULL, '1'},
 		{"output", required_argument, NULL, 'o'},
 		{"port", required_argument, NULL, 'p'},
 		{"pset", required_argument, NULL, 'P'},
@@ -453,7 +455,7 @@ parse_psql_options(int argc, char *argv[], struct adhoc_opts * options)
 
 	memset(options, 0, sizeof *options);
 
-	while ((c = getopt_long(argc, argv, "aAc:d:eEf:F:h:HlL:no:p:P:qR:sStT:uU:v:VWxX?",
+	while ((c = getopt_long(argc, argv, "aAc:d:eEf:F:h:HlL:no:p:P:qR:sStT:uU:v:VWxX?1",
 							long_options, &optindex)) != -1)
 	{
 		switch (c)
@@ -606,6 +608,9 @@ parse_psql_options(int argc, char *argv[], struct adhoc_opts * options)
 			case 'X':
 				options->no_psqlrc = true;
 				break;
+			case '1':
+				options->single_txn = true;
+				break;
 			case '?':
 				/* Actual help option given */
 				if (strcmp(argv[optind - 1], "-?") == 0 || strcmp(argv[optind - 1], "--help") == 0)
@@ -690,9 +695,9 @@ process_psqlrc_file(char *filename)
 	sprintf(psqlrc, "%s-%s", filename, PG_VERSION);
 
 	if (access(psqlrc, R_OK) == 0)
-		(void) process_file(psqlrc);
+		(void) process_file(psqlrc, false);
 	else if (access(filename, R_OK) == 0)
-		(void) process_file(filename);
+		(void) process_file(filename, false);
 	free(psqlrc);
 }