[CF 5692] v1 - Create DATA dump objects for partitioned tables too

This branch was automatically generated by a robot using patches from an
email thread registered at:

https://commitfest.postgresql.org/patch/5692

The branch will be overwritten each time a new patch version is posted to
the thread, and also periodically to check for bitrot caused by changes
on the master branch.

Patch(es): https://www.postgresql.org/message-id/729065.1744330102@sss.pgh.pa.us
Author(s): Tom Lane

Author: Commitfest Bot

src/bin/pg_dump/pg_backup_archiver.c

@@ -1988,12 +1988,14 @@ buildTocEntryArrays(ArchiveHandle *AH)
 		AH->tocsByDumpId[te->dumpId] = te;
 
 		/*
-		 * tableDataId provides the TABLE DATA item's dump ID for each TABLE
-		 * TOC entry that has a DATA item.  We compute this by reversing the
-		 * TABLE DATA item's dependency, knowing that a TABLE DATA item has
-		 * just one dependency and it is the TABLE item.
+		 * tableDataId provides the DATA item's dump ID for each TABLE TOC
+		 * entry that has a TABLE DATA or PARTITIONED DATA item.  We compute
+		 * this by reversing the DATA item's dependency, knowing that its
+		 * first dependency is the TABLE item.
 		 */
-		if (strcmp(te->desc, "TABLE DATA") == 0 && te->nDeps > 0)
+		if (te->nDeps > 0 &&
+			(strcmp(te->desc, "TABLE DATA") == 0 ||
+			 strcmp(te->desc, "PARTITIONED DATA") == 0))
 		{
 			DumpId		tableId = te->dependencies[0];
 
@@ -2003,7 +2005,7 @@ buildTocEntryArrays(ArchiveHandle *AH)
 			 * item's dump ID, so there should be a place for it in the array.
 			 */
 			if (tableId <= 0 || tableId > maxDumpId)
-				pg_fatal("bad table dumpId for TABLE DATA item");
+				pg_fatal("bad table dumpId for %s item", te->desc);
 
 			AH->tableDataId[tableId] = te->dumpId;
 		}
@@ -3140,6 +3142,7 @@ _tocEntryRequired(TocEntry *te, teSection curSection, ArchiveHandle *AH)
 		{
 			if (strcmp(te->desc, "TABLE") == 0 ||
 				strcmp(te->desc, "TABLE DATA") == 0 ||
+				strcmp(te->desc, "PARTITIONED DATA") == 0 ||
 				strcmp(te->desc, "VIEW") == 0 ||
 				strcmp(te->desc, "FOREIGN TABLE") == 0 ||
 				strcmp(te->desc, "MATERIALIZED VIEW") == 0 ||
@@ -3194,13 +3197,14 @@ _tocEntryRequired(TocEntry *te, teSection curSection, ArchiveHandle *AH)
 	if (!te->hadDumper)
 	{
 		/*
-		 * Special Case: If 'SEQUENCE SET' or anything to do with LOs, then it
-		 * is considered a data entry.  We don't need to check for BLOBS or
-		 * old-style BLOB COMMENTS entries, because they will have hadDumper =
-		 * true ... but we do need to check new-style BLOB ACLs, comments,
-		 * etc.
+		 * Special Case: If 'PARTITIONED DATA', 'SEQUENCE SET' or anything to
+		 * do with LOs, then it is considered a data entry.  We don't need to
+		 * check for BLOBS or old-style BLOB COMMENTS entries, because they
+		 * will have hadDumper = true ... but we do need to check new-style
+		 * BLOB ACLs, comments, etc.
 		 */
-		if (strcmp(te->desc, "SEQUENCE SET") == 0 ||
+		if (strcmp(te->desc, "PARTITIONED DATA") == 0 ||
+			strcmp(te->desc, "SEQUENCE SET") == 0 ||
 			strcmp(te->desc, "BLOB") == 0 ||
 			strcmp(te->desc, "BLOB METADATA") == 0 ||
 			(strcmp(te->desc, "ACL") == 0 &&
@@ -4996,14 +5000,14 @@ identify_locking_dependencies(ArchiveHandle *AH, TocEntry *te)
 		return;
 
 	/*
-	 * We assume the entry requires exclusive lock on each TABLE or TABLE DATA
-	 * item listed among its dependencies.  Originally all of these would have
-	 * been TABLE items, but repoint_table_dependencies would have repointed
-	 * them to the TABLE DATA items if those are present (which they might not
-	 * be, eg in a schema-only dump).  Note that all of the entries we are
-	 * processing here are POST_DATA; otherwise there might be a significant
-	 * difference between a dependency on a table and a dependency on its
-	 * data, so that closer analysis would be needed here.
+	 * We assume the entry requires exclusive lock on each TABLE, TABLE DATA,
+	 * or PARTITIONED DATA item listed among its dependencies.  Originally all
+	 * of these would have been TABLE items, but repoint_table_dependencies
+	 * would have repointed them to the DATA items if those are present (which
+	 * they might not be, eg in a schema-only dump).  Note that all of the
+	 * entries we are processing here are POST_DATA; otherwise there might be
+	 * a significant difference between a dependency on a table and a
+	 * dependency on its data, so that closer analysis would be needed here.
 	 */
 	lockids = (DumpId *) pg_malloc(te->nDeps * sizeof(DumpId));
 	nlockids = 0;
@@ -5012,8 +5016,9 @@ identify_locking_dependencies(ArchiveHandle *AH, TocEntry *te)
 		DumpId		depid = te->dependencies[i];
 
 		if (depid <= AH->maxDumpId && AH->tocsByDumpId[depid] != NULL &&
-			((strcmp(AH->tocsByDumpId[depid]->desc, "TABLE DATA") == 0) ||
-			 strcmp(AH->tocsByDumpId[depid]->desc, "TABLE") == 0))
+			(strcmp(AH->tocsByDumpId[depid]->desc, "TABLE") == 0 ||
+			 strcmp(AH->tocsByDumpId[depid]->desc, "TABLE DATA") == 0 ||
+			 strcmp(AH->tocsByDumpId[depid]->desc, "PARTITIONED DATA") == 0))
 			lockids[nlockids++] = depid;
 	}
 

src/bin/pg_dump/pg_backup_archiver.h

@@ -310,7 +310,8 @@ struct _archiveHandle
 
 	/* arrays created after the TOC list is complete: */
 	struct _tocEntry **tocsByDumpId;	/* TOCs indexed by dumpId */
-	DumpId	   *tableDataId;	/* TABLE DATA ids, indexed by table dumpId */
+	DumpId	   *tableDataId;	/* TABLE DATA and PARTITIONED DATA dumpIds,
+								 * indexed by table dumpId */
 
 	struct _tocEntry *currToc;	/* Used when dumping data */
 	pg_compress_specification compression_spec; /* Requested specification for

src/bin/pg_dump/pg_backup_custom.c

@@ -25,6 +25,8 @@
  */
 #include "postgres_fe.h"
 
+#include <limits.h>
+
 #include "common/file_utils.h"
 #include "compress_io.h"
 #include "pg_backup_utils.h"
@@ -819,17 +821,18 @@ _ReopenArchive(ArchiveHandle *AH)
 /*
  * Prepare for parallel restore.
  *
- * The main thing that needs to happen here is to fill in TABLE DATA and BLOBS
- * TOC entries' dataLength fields with appropriate values to guide the
- * ordering of restore jobs.  The source of said data is format-dependent,
- * as is the exact meaning of the values.
+ * The main thing that needs to happen here is to fill in TABLE DATA,
+ * PARTITIONED_DATA, and BLOBS TOC entries' dataLength fields with appropriate
+ * values to guide the ordering of restore jobs.  The source of said data is
+ * format-dependent, as is the exact meaning of the values.
  *
  * A format module might also choose to do other setup here.
  */
 static void
 _PrepParallelRestore(ArchiveHandle *AH)
 {
 	lclContext *ctx = (lclContext *) AH->formatData;
+	bool		have_partitioned_data = false;
 	TocEntry   *prev_te = NULL;
 	lclTocEntry *prev_tctx = NULL;
 	TocEntry   *te;
@@ -843,6 +846,10 @@ _PrepParallelRestore(ArchiveHandle *AH)
 	{
 		lclTocEntry *tctx = (lclTocEntry *) te->formatData;
 
+		/* Track whether there are any PARTITIONED_DATA items */
+		if (!have_partitioned_data && strcmp(te->desc, "PARTITIONED DATA") == 0)
+			have_partitioned_data = true;
+
 		/*
 		 * Ignore entries without a known data offset; if we were unable to
 		 * seek to rewrite the TOC when creating the archive, this'll be all
@@ -873,6 +880,38 @@ _PrepParallelRestore(ArchiveHandle *AH)
 		if (endpos > prev_tctx->dataPos)
 			prev_te->dataLength = endpos - prev_tctx->dataPos;
 	}
+
+	/*
+	 * For PARTITIONED DATA items, add up the sizes of their child objects.
+	 * (We couldn't do this earlier, since when we encounter a PARTITIONED
+	 * DATA item in the first loop we typically don't know the dataLength of
+	 * its last child yet.)
+	 */
+	if (have_partitioned_data)
+	{
+		for (te = AH->toc->next; te != AH->toc; te = te->next)
+		{
+			if (strcmp(te->desc, "PARTITIONED DATA") != 0)
+				continue;
+			for (int i = 0; i < te->nDeps; i++)
+			{
+				DumpId		depid = te->dependencies[i];
+
+				if (depid <= AH->maxDumpId && AH->tocsByDumpId[depid] != NULL)
+				{
+					pgoff_t		childLength = AH->tocsByDumpId[depid]->dataLength;
+
+					te->dataLength += childLength;
+					/* handle overflow -- unlikely except with 32-bit pgoff_t */
+					if (unlikely(te->dataLength < 0))
+					{
+						te->dataLength = INT_MAX;
+						break;
+					}
+				}
+			}
+		}
+	}
 }
 
 /*

src/bin/pg_dump/pg_backup_directory.c

@@ -37,6 +37,7 @@
 #include "postgres_fe.h"
 
 #include <dirent.h>
+#include <limits.h>
 #include <sys/stat.h>
 
 #include "common/file_utils.h"
@@ -722,16 +723,17 @@ setFilePath(ArchiveHandle *AH, char *buf, const char *relativeFilename)
 /*
  * Prepare for parallel restore.
  *
- * The main thing that needs to happen here is to fill in TABLE DATA and BLOBS
- * TOC entries' dataLength fields with appropriate values to guide the
- * ordering of restore jobs.  The source of said data is format-dependent,
- * as is the exact meaning of the values.
+ * The main thing that needs to happen here is to fill in TABLE DATA,
+ * PARTITIONED_DATA, and BLOBS TOC entries' dataLength fields with appropriate
+ * values to guide the ordering of restore jobs.  The source of said data is
+ * format-dependent, as is the exact meaning of the values.
  *
  * A format module might also choose to do other setup here.
  */
 static void
 _PrepParallelRestore(ArchiveHandle *AH)
 {
+	bool		have_partitioned_data = false;
 	TocEntry   *te;
 
 	for (te = AH->toc->next; te != AH->toc; te = te->next)
@@ -740,6 +742,10 @@ _PrepParallelRestore(ArchiveHandle *AH)
 		char		fname[MAXPGPATH];
 		struct stat st;
 
+		/* Track whether there are any PARTITIONED_DATA items */
+		if (!have_partitioned_data && strcmp(te->desc, "PARTITIONED DATA") == 0)
+			have_partitioned_data = true;
+
 		/*
 		 * A dumpable object has set tctx->filename, any other object has not.
 		 * (see _ArchiveEntry).
@@ -784,6 +790,38 @@ _PrepParallelRestore(ArchiveHandle *AH)
 		if (strcmp(te->desc, "BLOBS") == 0)
 			te->dataLength *= 1024;
 	}
+
+	/*
+	 * For PARTITIONED DATA items, add up the sizes of their child objects.
+	 * (Unlike pg_backup_custom.c, we could theoretically do this within the
+	 * previous loop, but it seems best to keep the logic looking the same in
+	 * both functions.)
+	 */
+	if (have_partitioned_data)
+	{
+		for (te = AH->toc->next; te != AH->toc; te = te->next)
+		{
+			if (strcmp(te->desc, "PARTITIONED DATA") != 0)
+				continue;
+			for (int i = 0; i < te->nDeps; i++)
+			{
+				DumpId		depid = te->dependencies[i];
+
+				if (depid <= AH->maxDumpId && AH->tocsByDumpId[depid] != NULL)
+				{
+					pgoff_t		childLength = AH->tocsByDumpId[depid]->dataLength;
+
+					te->dataLength += childLength;
+					/* handle overflow -- unlikely except with 32-bit pgoff_t */
+					if (unlikely(te->dataLength < 0))
+					{
+						te->dataLength = INT_MAX;
+						break;
+					}
+				}
+			}
+		}
+	}
 }
 
 /*