From a4cd6abcc901c1a8009c62a27f78696717bb8fe1 Mon Sep 17 00:00:00 2001
From: Andrew Dunstan <andrew@dunslane.net>
Date: Fri, 16 Dec 2011 19:09:38 -0500
Subject: [PATCH] Add --section option to pg_dump and pg_restore.

Valid values are --pre-data, data and post-data. The option can be
given more than once. --schema-only is equivalent to
--section=pre-data --section=post-data. --data-only is equivalent
to --section=data.

Andrew Dunstan, reviewed by Joachim Wieland and Josh Berkus.
---
 doc/src/sgml/ref/pg_dump.sgml        | 24 ++++++++++++--
 doc/src/sgml/ref/pg_restore.sgml     | 23 +++++++++++++
 src/bin/pg_dump/common.c             |  2 +-
 src/bin/pg_dump/dumputils.c          | 30 +++++++++++++++++
 src/bin/pg_dump/dumputils.h          |  1 +
 src/bin/pg_dump/pg_backup.h          |  9 +++++
 src/bin/pg_dump/pg_backup_archiver.c | 19 +++++++++++
 src/bin/pg_dump/pg_backup_archiver.h |  3 ++
 src/bin/pg_dump/pg_dump.c            | 49 ++++++++++++++++++++++++++--
 src/bin/pg_dump/pg_restore.c         | 30 +++++++++++++++++
 10 files changed, 184 insertions(+), 6 deletions(-)

diff --git a/doc/src/sgml/ref/pg_dump.sgml b/doc/src/sgml/ref/pg_dump.sgml
index cf6d689918..710051a4fd 100644
--- a/doc/src/sgml/ref/pg_dump.sgml
+++ b/doc/src/sgml/ref/pg_dump.sgml
@@ -116,9 +116,7 @@ PostgreSQL documentation
        </para>
 
        <para>
-        This option is only meaningful for the plain-text format.  For
-        the archive formats, you can specify the option when you
-        call <command>pg_restore</command>.
+        This option is equivalent to specifying <option>--section=data</>.
        </para>
       </listitem>
      </varlistentry>
@@ -408,9 +406,29 @@ PostgreSQL documentation
         To exclude table data for only a subset of tables in the database, 
         see <option>--exclude-table-data</>.
        </para>
+       <para>
+        This option is equivalent to specifying 
+        <option>--section=pre-data --section=post-data</>.
+       </para>
       </listitem>
      </varlistentry>
 
+     <varlistentry>
+       <term><option>--section=<replaceable class="parameter">sectionname</replaceable></option></term>
+       <listitem>
+         <para>
+           Only dump the named section. The name can be one of <option>pre-data</>, <option>data</> 
+           and <option>post-data</>. 
+           This option can be specified more than once. The default is to dump all sections.
+         </para>
+         <para>
+           Post-data items consist of definitions of indexes, triggers, rules 
+           and constraints other than check constraints. 
+           Pre-data items consist of all other data definition items.
+         </para>
+       </listitem>
+     </varlistentry>
+
      <varlistentry>
       <term><option>-S <replaceable class="parameter">username</replaceable></option></term>
       <term><option>--superuser=<replaceable class="parameter">username</replaceable></option></term>
diff --git a/doc/src/sgml/ref/pg_restore.sgml b/doc/src/sgml/ref/pg_restore.sgml
index be11d176cd..a28faf87a4 100644
--- a/doc/src/sgml/ref/pg_restore.sgml
+++ b/doc/src/sgml/ref/pg_restore.sgml
@@ -93,6 +93,9 @@
        <para>
         Restore only the data, not the schema (data definitions).
        </para>
+       <para>
+        This option is equivalent to specifying <option>--section=data</>.
+       </para>
       </listitem>
      </varlistentry>
 
@@ -359,6 +362,10 @@
         (Do not confuse this with the <option>--schema</> option, which
         uses the word <quote>schema</> in a different meaning.)
        </para>
+       <para>
+        This option is equivalent to specifying 
+        <option>--section=pre-data --section=post-data</>.
+       </para>
       </listitem>
      </varlistentry>
 
@@ -504,6 +511,22 @@
       </listitem>
      </varlistentry>
 
+     <varlistentry>
+       <term><option>--section=<replaceable class="parameter">sectionname</replaceable></option></term>
+       <listitem>
+         <para>
+           Only restore the named section. The name can be one of <option>pre-data</>, <option>data</> 
+           and <option>post-data</>. 
+           This option can be specified more than once. The default is to restore all sections.
+         </para>
+         <para>
+           Post-data items consist of definitions of indexes, triggers, rules 
+           and constraints other than check constraints. 
+           Pre-data items consist of all other data definition items.
+         </para>
+       </listitem>
+     </varlistentry>
+
      <varlistentry>
       <term><option>--use-set-session-authorization</option></term>
       <listitem>
diff --git a/src/bin/pg_dump/common.c b/src/bin/pg_dump/common.c
index d4e906ddd0..92c73ad3c1 100644
--- a/src/bin/pg_dump/common.c
+++ b/src/bin/pg_dump/common.c
@@ -2,7 +2,7 @@
  *
  * common.c
  *	Catalog routines used by pg_dump; long ago these were shared
- *	by another dump tool, but not anymore.  
+ *	by another dump tool, but not anymore.
  *
  * Portions Copyright (c) 1996-2011, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
diff --git a/src/bin/pg_dump/dumputils.c b/src/bin/pg_dump/dumputils.c
index 5dab9675fc..659ec06dc7 100644
--- a/src/bin/pg_dump/dumputils.c
+++ b/src/bin/pg_dump/dumputils.c
@@ -17,6 +17,7 @@
 #include <ctype.h>
 
 #include "dumputils.h"
+#include "pg_backup.h"
 
 #include "parser/keywords.h"
 
@@ -1262,3 +1263,32 @@ exit_horribly(const char *modulename, const char *fmt,...)
 
 	exit(1);
 }
+
+/*
+ * Set the bitmask in dumpSections according to the first argument.
+ * dumpSections is initialised as DUMP_UNSECTIONED by pg_dump and
+ * pg_restore so they can know if this has even been called.
+ */
+
+void
+set_section (const char *arg, int *dumpSections)
+{
+	/* if this is the first, clear all the bits */
+	if (*dumpSections == DUMP_UNSECTIONED)
+		*dumpSections = 0;
+
+	if (strcmp(arg,"pre-data") == 0)
+		*dumpSections |= DUMP_PRE_DATA;
+	else if (strcmp(arg,"data") == 0)
+		*dumpSections |= DUMP_DATA;
+	else if (strcmp(arg,"post-data") == 0)
+		*dumpSections |= DUMP_POST_DATA;
+	else
+	{
+		fprintf(stderr, _("%s: unknown section name \"%s\")\n"),
+				progname, arg);
+		fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
+				progname);
+		exit(1);
+	}
+}
diff --git a/src/bin/pg_dump/dumputils.h b/src/bin/pg_dump/dumputils.h
index c857f3a098..f80dd27943 100644
--- a/src/bin/pg_dump/dumputils.h
+++ b/src/bin/pg_dump/dumputils.h
@@ -58,5 +58,6 @@ extern void vwrite_msg(const char *modulename, const char *fmt, va_list ap)
 				__attribute__((format(PG_PRINTF_ATTRIBUTE, 2, 0)));
 extern void exit_horribly(const char *modulename, const char *fmt,...)
 				__attribute__((format(PG_PRINTF_ATTRIBUTE, 2, 3)));
+extern void set_section (const char *arg, int *dumpSections);
 
 #endif   /* DUMPUTILS_H */
diff --git a/src/bin/pg_dump/pg_backup.h b/src/bin/pg_dump/pg_backup.h
index 8168cfffed..89264885b1 100644
--- a/src/bin/pg_dump/pg_backup.h
+++ b/src/bin/pg_dump/pg_backup.h
@@ -69,6 +69,14 @@ typedef enum _teSection
 	SECTION_POST_DATA			/* stuff to be processed after data */
 } teSection;
 
+typedef enum
+{
+	DUMP_PRE_DATA = 0x01,
+	DUMP_DATA = 0x02,
+	DUMP_POST_DATA = 0x04,
+	DUMP_UNSECTIONED = 0xff
+} DumpSections;
+
 /*
  *	We may want to have some more user-readable data, but in the mean
  *	time this gives us some abstraction and type checking.
@@ -111,6 +119,7 @@ typedef struct _restoreOptions
 	int			dropSchema;
 	char	   *filename;
 	int			schemaOnly;
+	int         dumpSections;
 	int			verbose;
 	int			aclsSkip;
 	int			tocSummary;
diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
index 66123094fe..7d895c46a3 100644
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -665,6 +665,7 @@ NewRestoreOptions(void)
 	/* set any fields that shouldn't default to zeroes */
 	opts->format = archUnknown;
 	opts->promptPassword = TRI_DEFAULT;
+	opts->dumpSections = DUMP_UNSECTIONED;
 
 	return opts;
 }
@@ -2120,6 +2121,7 @@ ReadToc(ArchiveHandle *AH)
 	int			depIdx;
 	int			depSize;
 	TocEntry   *te;
+	bool        in_post_data = false;
 
 	AH->tocCount = ReadInt(AH);
 	AH->maxDumpId = 0;
@@ -2185,6 +2187,12 @@ ReadToc(ArchiveHandle *AH)
 				te->section = SECTION_PRE_DATA;
 		}
 
+		/* will stay true even for SECTION_NONE items */
+		if (te->section == SECTION_POST_DATA)
+			in_post_data = true;
+
+		te->inPostData = in_post_data;
+
 		te->defn = ReadStr(AH);
 		te->dropStmt = ReadStr(AH);
 
@@ -2334,6 +2342,17 @@ _tocEntryRequired(TocEntry *te, RestoreOptions *ropt, bool include_acls)
 	if (!ropt->createDB && strcmp(te->desc, "DATABASE") == 0)
 		return 0;
 
+	/* skip (all but) post data section as required */
+	/* table data is filtered if necessary lower down */
+	if (ropt->dumpSections != DUMP_UNSECTIONED)
+	{
+		if (!(ropt->dumpSections & DUMP_POST_DATA) && te->inPostData)
+			return 0;
+		if (!(ropt->dumpSections & DUMP_PRE_DATA) && ! te->inPostData && strcmp(te->desc, "TABLE DATA") != 0)
+			return 0;
+	}
+
+
 	/* Check options for selective dump/restore */
 	if (ropt->schemaNames)
 	{
diff --git a/src/bin/pg_dump/pg_backup_archiver.h b/src/bin/pg_dump/pg_backup_archiver.h
index 07a4b6e50a..7a4fd36073 100644
--- a/src/bin/pg_dump/pg_backup_archiver.h
+++ b/src/bin/pg_dump/pg_backup_archiver.h
@@ -287,6 +287,9 @@ typedef struct _tocEntry
 	void	   *dataDumperArg;	/* Arg for above routine */
 	void	   *formatData;		/* TOC Entry data specific to file format */
 
+	/* in post data? not quite the same as section, might be SECTION_NONE */
+	bool        inPostData;
+
 	/* working state (needed only for parallel restore) */
 	struct _tocEntry *par_prev; /* list links for pending/ready items; */
 	struct _tocEntry *par_next; /* these are NULL if not in either list */
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index 2fa99dd0bb..7dfa1dd8f8 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -91,6 +91,7 @@ PGconn	   *g_conn;				/* the database connection */
 /* various user-settable parameters */
 bool		schemaOnly;
 bool		dataOnly;
+int         dumpSections; /* bitmask of chosen sections */
 bool		aclsSkip;
 const char *lockWaitTimeout;
 
@@ -250,7 +251,6 @@ static void do_sql_command(PGconn *conn, const char *query);
 static void check_sql_result(PGresult *res, PGconn *conn, const char *query,
 				 ExecStatusType expected);
 
-
 int
 main(int argc, char **argv)
 {
@@ -332,6 +332,7 @@ main(int argc, char **argv)
 		{"no-tablespaces", no_argument, &outputNoTablespaces, 1},
 		{"quote-all-identifiers", no_argument, &quote_all_identifiers, 1},
 		{"role", required_argument, NULL, 3},
+		{"section", required_argument, NULL, 5},
 		{"serializable-deferrable", no_argument, &serializable_deferrable, 1},
 		{"use-set-session-authorization", no_argument, &use_setsessauth, 1},
 		{"no-security-labels", no_argument, &no_security_labels, 1},
@@ -349,6 +350,7 @@ main(int argc, char **argv)
 	strcpy(g_opaque_type, "opaque");
 
 	dataOnly = schemaOnly = false;
+	dumpSections = DUMP_UNSECTIONED;
 	lockWaitTimeout = NULL;
 
 	progname = get_progname(argv[0]);
@@ -494,6 +496,10 @@ main(int argc, char **argv)
 				simple_string_list_append(&tabledata_exclude_patterns, optarg);
 				break;
 
+			case 5:				/* section */
+				set_section(optarg, &dumpSections);
+				break;
+
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
@@ -524,6 +530,22 @@ main(int argc, char **argv)
 		exit(1);
 	}
 
+	if ((dataOnly || schemaOnly) && dumpSections != DUMP_UNSECTIONED)
+	{
+		write_msg(NULL, "options -s/--schema-only and -a/--data-only cannot be used with --section\n");
+		exit(1);
+	}
+
+	if (dataOnly)
+		dumpSections = DUMP_DATA;
+	else if (schemaOnly)
+		dumpSections = DUMP_PRE_DATA | DUMP_POST_DATA;
+	else if ( dumpSections != DUMP_UNSECTIONED)
+	{
+		dataOnly = dumpSections == DUMP_DATA;
+		schemaOnly = !(dumpSections & DUMP_DATA);
+	}
+
 	if (dataOnly && outputClean)
 	{
 		write_msg(NULL, "options -c/--clean and -a/--data-only cannot be used together\n");
@@ -871,6 +893,7 @@ help(const char *progname)
 	printf(_("  --no-tablespaces            do not dump tablespace assignments\n"));
 	printf(_("  --no-unlogged-table-data    do not dump unlogged table data\n"));
 	printf(_("  --quote-all-identifiers     quote all identifiers, even if not key words\n"));
+	printf(_("  --section=SECTION           dump named section (pre-data, data or post-data)\n"));
 	printf(_("  --serializable-deferrable   wait until the dump can run without anomalies\n"));
 	printf(_("  --use-set-session-authorization\n"
 			 "                              use SET SESSION AUTHORIZATION commands instead of\n"
@@ -1107,7 +1130,7 @@ selectDumpableTable(TableInfo *tbinfo)
 		tbinfo->dobj.dumpdata = true;
 	else
 		tbinfo->dobj.dumpdata = false;
-		
+
 }
 
 /*
@@ -7093,6 +7116,28 @@ collectComments(Archive *fout, CommentItem **items)
 static void
 dumpDumpableObject(Archive *fout, DumpableObject *dobj)
 {
+
+	bool skip = false;
+
+	switch (dobj->objType)
+	{
+		case DO_INDEX:
+		case DO_TRIGGER:
+		case DO_CONSTRAINT:
+		case DO_FK_CONSTRAINT:
+		case DO_RULE:
+			skip = !(dumpSections & DUMP_POST_DATA);
+			break;
+		case DO_TABLE_DATA:
+			skip = !(dumpSections & DUMP_DATA);
+			break;
+		default:
+			skip = !(dumpSections & DUMP_PRE_DATA);
+	}
+
+	if (skip)
+		return;
+
 	switch (dobj->objType)
 	{
 		case DO_NAMESPACE:
diff --git a/src/bin/pg_dump/pg_restore.c b/src/bin/pg_dump/pg_restore.c
index 8d6edacd7d..6ff1ab8904 100644
--- a/src/bin/pg_dump/pg_restore.c
+++ b/src/bin/pg_dump/pg_restore.c
@@ -118,6 +118,7 @@ main(int argc, char **argv)
 		{"no-data-for-failed-tables", no_argument, &no_data_for_failed_tables, 1},
 		{"no-tablespaces", no_argument, &outputNoTablespaces, 1},
 		{"role", required_argument, NULL, 2},
+		{"section", required_argument, NULL, 3},
 		{"use-set-session-authorization", no_argument, &use_setsessauth, 1},
 		{"no-security-labels", no_argument, &no_security_labels, 1},
 
@@ -272,6 +273,10 @@ main(int argc, char **argv)
 				opts->use_role = optarg;
 				break;
 
+			case 3:				/* section */
+				set_section(optarg, &(opts->dumpSections));
+				break;
+
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
@@ -294,6 +299,30 @@ main(int argc, char **argv)
 		exit(1);
 	}
 
+	if (opts->dataOnly && opts->schemaOnly)
+	{
+		fprintf(stderr, _("%s: options -s/--schema-only and -a/--data-only cannot be used together\n"),
+			progname);
+		exit(1);
+	}
+
+	if ((opts->dataOnly || opts->schemaOnly) && (opts->dumpSections != DUMP_UNSECTIONED))
+	{
+		fprintf(stderr, _("%s: options -s/--schema-only and -a/--data-only cannot be used with --section\n"),
+			progname);
+		exit(1);
+	}
+
+	if (opts->dataOnly)
+		opts->dumpSections = DUMP_DATA;
+	else if (opts->schemaOnly)
+		opts->dumpSections = DUMP_PRE_DATA | DUMP_POST_DATA;
+	else if ( opts->dumpSections != DUMP_UNSECTIONED)
+	{
+		opts->dataOnly = opts->dumpSections == DUMP_DATA;
+		opts->schemaOnly = !(opts->dumpSections & DUMP_DATA);
+	}
+
 	/* Should get at most one of -d and -f, else user is confused */
 	if (opts->dbname)
 	{
@@ -434,6 +463,7 @@ usage(const char *progname)
 			 "                           created\n"));
 	printf(_("  --no-security-labels     do not restore security labels\n"));
 	printf(_("  --no-tablespaces         do not restore tablespace assignments\n"));
+	printf(_("  --section=SECTION        restore named section (pre-data, data or post-data)\n"));
 	printf(_("  --use-set-session-authorization\n"
 			 "                           use SET SESSION AUTHORIZATION commands instead of\n"
 	  "                           ALTER OWNER commands to set ownership\n"));