v15-0002-Enable-dumping-of-table-index-stats-in-pg_dump.patch

From b86fa7d7c3c031c7dd5f7b42e34ff4689b7ac62b Mon Sep 17 00:00:00 2001
From: Corey Huinker <corey.huinker@gmail.com>
Date: Sat, 16 Mar 2024 17:21:10 -0400
Subject: [PATCH v15 2/2] Enable dumping of table/index stats in pg_dump.

For each table/matview/index dumped, it will generate a statement that
calls pg_set_relation_stats(), and it will generate a series of
statements that call pg_set_attribute_stats(), one per attribute. These
statements will restore the statistics of the current system onto the
destination system.

As is the pattern with pg_dump options, this can be disabled with
--no-statistics.
---
 src/bin/pg_dump/pg_backup.h            |   2 +
 src/bin/pg_dump/pg_backup_archiver.c   |   5 +
 src/bin/pg_dump/pg_dump.c              | 289 ++++++++++++++++++++++++-
 src/bin/pg_dump/pg_dump.h              |   1 +
 src/bin/pg_dump/pg_dumpall.c           |   5 +
 src/bin/pg_dump/pg_restore.c           |   3 +
 src/bin/pg_upgrade/t/002_pg_upgrade.pl |   4 +-
 7 files changed, 305 insertions(+), 4 deletions(-)

diff --git a/src/bin/pg_dump/pg_backup.h b/src/bin/pg_dump/pg_backup.h
index 9ef2f2017e..1db5cf52eb 100644
--- a/src/bin/pg_dump/pg_backup.h
+++ b/src/bin/pg_dump/pg_backup.h
@@ -112,6 +112,7 @@ typedef struct _restoreOptions
 	int			no_publications;	/* Skip publication entries */
 	int			no_security_labels; /* Skip security label entries */
 	int			no_subscriptions;	/* Skip subscription entries */
+	int			no_statistics;		/* Skip statistics import */
 	int			strict_names;
 
 	const char *filename;
@@ -179,6 +180,7 @@ typedef struct _dumpOptions
 	int			no_security_labels;
 	int			no_publications;
 	int			no_subscriptions;
+	int			no_statistics;
 	int			no_toast_compression;
 	int			no_unlogged_table_data;
 	int			serializable_deferrable;
diff --git a/src/bin/pg_dump/pg_backup_archiver.c b/src/bin/pg_dump/pg_backup_archiver.c
index d97ebaff5b..d5f61399d9 100644
--- a/src/bin/pg_dump/pg_backup_archiver.c
+++ b/src/bin/pg_dump/pg_backup_archiver.c
@@ -2833,6 +2833,10 @@ _tocEntryRequired(TocEntry *te, teSection curSection, ArchiveHandle *AH)
 	if (ropt->no_subscriptions && strcmp(te->desc, "SUBSCRIPTION") == 0)
 		return 0;
 
+	/* If it's a stats dump, maybe ignore it */
+	if (ropt->no_statistics && strcmp(te->desc, "STATISTICS") == 0)
+		return 0;
+
 	/* Ignore it if section is not to be dumped/restored */
 	switch (curSection)
 	{
@@ -2862,6 +2866,7 @@ _tocEntryRequired(TocEntry *te, teSection curSection, ArchiveHandle *AH)
 	 */
 	if (strcmp(te->desc, "ACL") == 0 ||
 		strcmp(te->desc, "COMMENT") == 0 ||
+		strcmp(te->desc, "STATISTICS") == 0 ||
 		strcmp(te->desc, "SECURITY LABEL") == 0)
 	{
 		/* Database properties react to createDB, not selectivity options. */
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index b1c4c3ec7f..f7ba76217c 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -428,6 +428,7 @@ main(int argc, char **argv)
 		{"no-comments", no_argument, &dopt.no_comments, 1},
 		{"no-publications", no_argument, &dopt.no_publications, 1},
 		{"no-security-labels", no_argument, &dopt.no_security_labels, 1},
+		{"no-statistics", no_argument, &dopt.no_statistics, 1},
 		{"no-subscriptions", no_argument, &dopt.no_subscriptions, 1},
 		{"no-toast-compression", no_argument, &dopt.no_toast_compression, 1},
 		{"no-unlogged-table-data", no_argument, &dopt.no_unlogged_table_data, 1},
@@ -1144,6 +1145,7 @@ help(const char *progname)
 	printf(_("  --no-comments                do not dump comments\n"));
 	printf(_("  --no-publications            do not dump publications\n"));
 	printf(_("  --no-security-labels         do not dump security label assignments\n"));
+	printf(_("  --no-statistics              do not dump statistics\n"));
 	printf(_("  --no-subscriptions           do not dump subscriptions\n"));
 	printf(_("  --no-table-access-method     do not dump table access methods\n"));
 	printf(_("  --no-tablespaces             do not dump tablespace assignments\n"));
@@ -7001,6 +7003,7 @@ getTables(Archive *fout, int *numTables)
 
 		/* Tables have data */
 		tblinfo[i].dobj.components |= DUMP_COMPONENT_DATA;
+		tblinfo[i].dobj.components |= DUMP_COMPONENT_STATISTICS;
 
 		/* Mark whether table has an ACL */
 		if (!PQgetisnull(res, i, i_relacl))
@@ -7498,6 +7501,7 @@ getIndexes(Archive *fout, TableInfo tblinfo[], int numTables)
 			indxinfo[j].dobj.catId.tableoid = atooid(PQgetvalue(res, j, i_tableoid));
 			indxinfo[j].dobj.catId.oid = atooid(PQgetvalue(res, j, i_oid));
 			AssignDumpId(&indxinfo[j].dobj);
+			indxinfo[j].dobj.components |= DUMP_COMPONENT_STATISTICS;
 			indxinfo[j].dobj.dump = tbinfo->dobj.dump;
 			indxinfo[j].dobj.name = pg_strdup(PQgetvalue(res, j, i_indexname));
 			indxinfo[j].dobj.namespace = tbinfo->dobj.namespace;
@@ -10247,6 +10251,272 @@ dumpComment(Archive *fout, const char *type,
 						catalogId, subid, dumpId, NULL);
 }
 
+/*
+ * Tabular description of the parameters to pg_set_relation_stats()
+ * required, param_name, param_type
+ */
+static const char *rel_stats_arginfo[][3] = {
+	{"t", "relation", "regclass"},
+	{"t", "relpages", "integer"},
+	{"t", "reltuples", "real"},
+	{"t", "relallvisible", "integer"},
+};
+
+/*
+ * Tabular description of the parameters to pg_set_attribute_stats()
+ * required, param_name, param_type
+ */
+static const char *att_stats_arginfo[][3] = {
+	{"t", "relation", "regclass"},
+	{"t", "attname", "name"},
+	{"t", "inherited", "boolean"},
+	{"t", "null_frac", "float4"},
+	{"t", "avg_width", "integer"},
+	{"t", "n_distinct", "float4"},
+	{"f", "most_common_vals", "text"},
+	{"f", "most_common_freqs", "float4[]"},
+	{"f", "histogram_bounds", "text"},
+	{"f", "correlation", "float4"},
+	{"f", "most_common_elems", "text"},
+	{"f", "most_common_elem_freqs", "float4[]"},
+	{"f", "elem_count_histogram", "float4[]"},
+	{"f", "range_length_histogram", "text"},
+	{"f", "range_empty_frac", "float4"},
+	{"f", "range_bounds_histogram", "text"},
+};
+
+/*
+ * getRelStatsExportQuery --
+ *
+ * Generate a query that will fetch all relation (e.g. pg_class)
+ * stats for a given relation.
+ */
+static void
+getRelStatsExportQuery(PQExpBuffer query, Archive *fout,
+					   const char *schemaname, const char *relname)
+{
+	resetPQExpBuffer(query);
+	appendPQExpBufferStr(query,
+						 "SELECT c.oid::regclass AS relation, c.relpages, "
+						 "c.reltuples, c.relallvisible "
+						 "FROM pg_class c "
+						 "JOIN pg_namespace n "
+						 "ON n.oid = c.relnamespace "
+						 "WHERE n.nspname = ");
+	appendStringLiteralAH(query, schemaname, fout);
+	appendPQExpBufferStr(query, " AND c.relname = ");
+	appendStringLiteralAH(query, relname, fout);
+}
+
+/*
+ * getAttStatsExportQuery --
+ *
+ * Generate a query that will fetch all attribute (e.g. pg_statistic)
+ * stats for a given relation.
+ */
+static void
+getAttStatsExportQuery(PQExpBuffer query, Archive *fout,
+					   const char *schemaname, const char *relname)
+{
+	resetPQExpBuffer(query);
+	appendPQExpBufferStr(query,
+						 "SELECT c.oid::regclass AS relation, "
+						 "s.attname,"
+						 "s.inherited,"
+						 "s.null_frac,"
+						 "s.avg_width,"
+						 "s.n_distinct,"
+						 "s.most_common_vals,"
+						 "s.most_common_freqs,"
+						 "s.histogram_bounds,"
+						 "s.correlation,"
+						 "s.most_common_elems,"
+						 "s.most_common_elem_freqs,"
+						 "s.elem_count_histogram,");
+
+	if (fout->remoteVersion >= 170000)
+		appendPQExpBufferStr(query,
+							 "s.range_length_histogram,"
+							 "s.range_empty_frac,"
+							 "s.range_bounds_histogram ");
+	else
+		appendPQExpBufferStr(query,
+							 "NULL AS range_length_histogram,"
+							 "NULL AS range_empty_frac,"
+							 "NULL AS range_bounds_histogram ");
+
+	appendPQExpBufferStr(query,
+						 "FROM pg_stats s "
+						 "JOIN pg_namespace n "
+						 "ON n.nspname = s.schemaname "
+						 "JOIN pg_class c "
+						 "ON c.relname = s.tablename "
+						 "AND c.relnamespace = n.oid "
+						 "WHERE s.schemaname = ");
+	appendStringLiteralAH(query, schemaname, fout);
+	appendPQExpBufferStr(query, " AND s.tablename = ");
+	appendStringLiteralAH(query, relname, fout);
+	appendPQExpBufferStr(query, " ORDER BY s.attname, s.inherited");
+}
+
+
+/*
+ * appendNamedArgument --
+ *
+ * Convenience routine for constructing parameters of the form:
+ *  paraname => 'value'::type
+ */
+static void
+appendNamedArgument(PQExpBuffer out, Archive *fout, const char *argname,
+					const char *argval, const char *argtype)
+{
+	appendPQExpBuffer(out, "\t%s => ", argname);
+	appendStringLiteralAH(out, argval, fout);
+	appendPQExpBuffer(out, "::%s", argtype);
+}
+
+/*
+ * appendRelStatsImport --
+ *
+ * Append a formatted pg_set_relation_stats statement.
+ */
+static void
+appendRelStatsImport(PQExpBuffer out, Archive *fout, PGresult *res)
+{
+	const char *sep = "";
+
+	if (PQntuples(res) == 0)
+		return;
+
+	appendPQExpBufferStr(out, "SELECT pg_catalog.pg_set_relation_stats(\n");
+
+	for (int argno = 0; argno < lengthof(rel_stats_arginfo); argno++)
+	{
+		bool			required = (rel_stats_arginfo[argno][0][0] == 't');
+		const char	   *argname = rel_stats_arginfo[argno][1];
+		const char	   *argtype = rel_stats_arginfo[argno][2];
+		int				fieldno = PQfnumber(res, argname);
+
+		if (fieldno < 0)
+			pg_fatal("relation stats export query missing field '%s'",
+					 argname);
+
+		if (PQgetisnull(res, 0, fieldno))
+		{
+			if (required)
+				pg_fatal("relation stats export query unexpected NULL in '%s'",
+						 argname);
+			else
+				continue;
+		}
+
+		appendPQExpBufferStr(out, sep);
+		appendNamedArgument(out, fout, argname, 
+							PQgetvalue(res, 0, fieldno), argtype);
+
+		sep = ",\n";
+	}
+	appendPQExpBufferStr(out, "\n);\n");
+}
+
+/*
+ * appendAttStatsImport --
+ *
+ * Append a series of formatted pg_set_attribute_stats statements.
+ */
+static void
+appendAttStatsImport(PQExpBuffer out, Archive *fout, PGresult *res)
+{
+	for (int rownum = 0; rownum < PQntuples(res); rownum++)
+	{
+		const char *sep = "";
+		appendPQExpBufferStr(out, "SELECT pg_catalog.pg_set_attribute_stats(\n");
+		for (int argno = 0; argno < lengthof(att_stats_arginfo); argno++)
+		{
+			bool		required = (att_stats_arginfo[argno][0][0] == 't');
+			const char	*argname = att_stats_arginfo[argno][1];
+			const char	*argtype = att_stats_arginfo[argno][2];
+			int			 fieldno = PQfnumber(res, argname);
+
+			if (fieldno < 0)
+				pg_fatal("attribute stats export query missing field '%s'",
+						 argname);
+
+			if (PQgetisnull(res, rownum, fieldno))
+			{
+				if (required)
+					pg_fatal("attribute stats export query unexpected NULL in '%s'",
+							 argname);
+				else
+					continue;
+			}
+
+			appendPQExpBufferStr(out, sep);
+			appendNamedArgument(out, fout, argname,
+								PQgetvalue(res, rownum, fieldno), argtype);
+			sep = ",\n";
+		}
+		appendPQExpBufferStr(out, "\n);\n");
+	}
+}
+
+/*
+ * dumpRelationStats --
+ *
+ * Dump command to import stats into the relation on the new database.
+ */
+static void
+dumpRelationStats(Archive *fout, const DumpableObject *dobj,
+				  const char *reltypename, DumpId dumpid)
+{
+	PGresult   *res;
+	PQExpBuffer query;
+	PQExpBuffer out;
+	PQExpBuffer tag;
+
+	/* do nothing, if --no-statistics is supplied */
+	if (fout->dopt->no_statistics)
+		return;
+
+	tag = createPQExpBuffer();
+	appendPQExpBuffer(tag, "%s %s", reltypename, fmtId(dobj->name));
+
+
+	query = createPQExpBuffer();
+	out = createPQExpBuffer();
+
+	getRelStatsExportQuery(query, fout, dobj->namespace->dobj.name,
+						   dobj->name);
+	res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
+	appendRelStatsImport(out, fout, res);
+	PQclear(res);
+
+	getRelStatsExportQuery(query, fout, dobj->namespace->dobj.name,
+						   dobj->name);
+	res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
+	appendRelStatsImport(out, fout, res);
+	PQclear(res);
+
+	getAttStatsExportQuery(query, fout, dobj->namespace->dobj.name,
+							  dobj->name);
+	res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
+	appendAttStatsImport(out, fout, res);
+	PQclear(res);
+
+	ArchiveEntry(fout, nilCatalogId, createDumpId(),
+				 ARCHIVE_OPTS(.tag = tag->data,
+							  .namespace = dobj->namespace->dobj.name,
+							  .description = "STATISTICS DATA",
+							  .section = SECTION_NONE,
+							  .createStmt = out->data,
+							  .deps = &dumpid,
+							  .nDeps = 1));
+
+	destroyPQExpBuffer(query);
+	destroyPQExpBuffer(out);
+	destroyPQExpBuffer(tag);
+}
+
 /*
  * dumpTableComment --
  *
@@ -16681,6 +16951,13 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 	if (tbinfo->dobj.dump & DUMP_COMPONENT_SECLABEL)
 		dumpTableSecLabel(fout, tbinfo, reltypename);
 
+	/* Statistics are dependent on the definition, not the data */
+	/* Views don't have stats */
+	if ((tbinfo->dobj.dump & DUMP_COMPONENT_STATISTICS) &&
+		(tbinfo->relkind != RELKIND_VIEW))
+		dumpRelationStats(fout, &tbinfo->dobj, reltypename,
+						  tbinfo->dobj.dumpId);
+
 	/* Dump comments on inlined table constraints */
 	for (j = 0; j < tbinfo->ncheck; j++)
 	{
@@ -16882,6 +17159,7 @@ dumpIndex(Archive *fout, const IndxInfo *indxinfo)
 	PQExpBuffer delq;
 	char	   *qindxname;
 	char	   *qqindxname;
+	DumpId		dumpid;
 
 	/* Do nothing in data-only dump */
 	if (dopt->dataOnly)
@@ -16994,14 +17272,21 @@ dumpIndex(Archive *fout, const IndxInfo *indxinfo)
 		free(indstatvalsarray);
 	}
 
+	/* Comments and stats share same .dep */
+	dumpid = is_constraint ? indxinfo->indexconstraint :
+		indxinfo->dobj.dumpId;
+
 	/* Dump Index Comments */
 	if (indxinfo->dobj.dump & DUMP_COMPONENT_COMMENT)
 		dumpComment(fout, "INDEX", qindxname,
 					tbinfo->dobj.namespace->dobj.name,
 					tbinfo->rolname,
 					indxinfo->dobj.catId, 0,
-					is_constraint ? indxinfo->indexconstraint :
-					indxinfo->dobj.dumpId);
+					dumpid);
+
+	/* Dump Index Stats */
+	if (indxinfo->dobj.dump & DUMP_COMPONENT_STATISTICS)
+		dumpRelationStats(fout, &indxinfo->dobj, "INDEX", dumpid);
 
 	destroyPQExpBuffer(q);
 	destroyPQExpBuffer(delq);
diff --git a/src/bin/pg_dump/pg_dump.h b/src/bin/pg_dump/pg_dump.h
index 9bc93520b4..d6a071ec28 100644
--- a/src/bin/pg_dump/pg_dump.h
+++ b/src/bin/pg_dump/pg_dump.h
@@ -101,6 +101,7 @@ typedef uint32 DumpComponents;
 #define DUMP_COMPONENT_ACL			(1 << 4)
 #define DUMP_COMPONENT_POLICY		(1 << 5)
 #define DUMP_COMPONENT_USERMAP		(1 << 6)
+#define DUMP_COMPONENT_STATISTICS	(1 << 7)
 #define DUMP_COMPONENT_ALL			(0xFFFF)
 
 /*
diff --git a/src/bin/pg_dump/pg_dumpall.c b/src/bin/pg_dump/pg_dumpall.c
index 046c0dc3b3..69652aa205 100644
--- a/src/bin/pg_dump/pg_dumpall.c
+++ b/src/bin/pg_dump/pg_dumpall.c
@@ -105,6 +105,7 @@ static int	use_setsessauth = 0;
 static int	no_comments = 0;
 static int	no_publications = 0;
 static int	no_security_labels = 0;
+static int	no_statistics = 0;
 static int	no_subscriptions = 0;
 static int	no_toast_compression = 0;
 static int	no_unlogged_table_data = 0;
@@ -174,6 +175,7 @@ main(int argc, char *argv[])
 		{"no-role-passwords", no_argument, &no_role_passwords, 1},
 		{"no-security-labels", no_argument, &no_security_labels, 1},
 		{"no-subscriptions", no_argument, &no_subscriptions, 1},
+		{"no-statistics", no_argument, &no_statistics, 1},
 		{"no-sync", no_argument, NULL, 4},
 		{"no-toast-compression", no_argument, &no_toast_compression, 1},
 		{"no-unlogged-table-data", no_argument, &no_unlogged_table_data, 1},
@@ -453,6 +455,8 @@ main(int argc, char *argv[])
 		appendPQExpBufferStr(pgdumpopts, " --no-publications");
 	if (no_security_labels)
 		appendPQExpBufferStr(pgdumpopts, " --no-security-labels");
+	if (no_statistics)
+		appendPQExpBufferStr(pgdumpopts, " --no-statistics");
 	if (no_subscriptions)
 		appendPQExpBufferStr(pgdumpopts, " --no-subscriptions");
 	if (no_toast_compression)
@@ -668,6 +672,7 @@ help(void)
 	printf(_("  --no-publications            do not dump publications\n"));
 	printf(_("  --no-role-passwords          do not dump passwords for roles\n"));
 	printf(_("  --no-security-labels         do not dump security label assignments\n"));
+	printf(_("  --no-statistics              do not dump statistics\n"));
 	printf(_("  --no-subscriptions           do not dump subscriptions\n"));
 	printf(_("  --no-sync                    do not wait for changes to be written safely to disk\n"));
 	printf(_("  --no-table-access-method     do not dump table access methods\n"));
diff --git a/src/bin/pg_dump/pg_restore.c b/src/bin/pg_dump/pg_restore.c
index c3beacdec1..2d326dec72 100644
--- a/src/bin/pg_dump/pg_restore.c
+++ b/src/bin/pg_dump/pg_restore.c
@@ -75,6 +75,7 @@ main(int argc, char **argv)
 	static int	no_publications = 0;
 	static int	no_security_labels = 0;
 	static int	no_subscriptions = 0;
+	static int  no_statistics = 0;
 	static int	strict_names = 0;
 
 	struct option cmdopts[] = {
@@ -126,6 +127,7 @@ main(int argc, char **argv)
 		{"no-security-labels", no_argument, &no_security_labels, 1},
 		{"no-subscriptions", no_argument, &no_subscriptions, 1},
 		{"filter", required_argument, NULL, 4},
+		{"no-statistics", no_argument, &no_statistics, 1},
 
 		{NULL, 0, NULL, 0}
 	};
@@ -358,6 +360,7 @@ main(int argc, char **argv)
 	opts->no_publications = no_publications;
 	opts->no_security_labels = no_security_labels;
 	opts->no_subscriptions = no_subscriptions;
+	opts->no_statistics = no_statistics;
 
 	if (if_exists && !opts->dropSchema)
 		pg_fatal("option --if-exists requires option -c/--clean");
diff --git a/src/bin/pg_upgrade/t/002_pg_upgrade.pl b/src/bin/pg_upgrade/t/002_pg_upgrade.pl
index 3e67121a8d..5006223fbe 100644
--- a/src/bin/pg_upgrade/t/002_pg_upgrade.pl
+++ b/src/bin/pg_upgrade/t/002_pg_upgrade.pl
@@ -315,7 +315,7 @@ if (defined($ENV{oldinstall}))
 # that we need to use pg_dumpall from the new node here.
 my @dump_command = (
 	'pg_dumpall', '--no-sync', '-d', $oldnode->connstr('postgres'),
-	'-f', $dump1_file);
+	'-f', $dump1_file, '--no-statistics');
 # --extra-float-digits is needed when upgrading from a version older than 11.
 push(@dump_command, '--extra-float-digits', '0')
   if ($oldnode->pg_version < 12);
@@ -483,7 +483,7 @@ is( $result,
 # Second dump from the upgraded instance.
 @dump_command = (
 	'pg_dumpall', '--no-sync', '-d', $newnode->connstr('postgres'),
-	'-f', $dump2_file);
+	'-f', $dump2_file, '--no-statistics');
 # --extra-float-digits is needed when upgrading from a version older than 11.
 push(@dump_command, '--extra-float-digits', '0')
   if ($oldnode->pg_version < 12);
-- 
2.44.0