From 60625c15c53b0e5e0090f98b2d692f03f9d617cd Mon Sep 17 00:00:00 2001
From: Nikhil Kumar Veldanda <nikhilkv@amazon.com>
Date: Mon, 14 Apr 2025 21:53:02 +0000
Subject: [PATCH v11 6/7] pg_dump, pg_upgrade needed changes to support new
 zstd catalog

---
 src/bin/pg_dump/pg_dump.c  | 210 ++++++++++++++++++++++++++++++++-----
 src/bin/pg_upgrade/check.c |  30 ++++--
 src/bin/pg_upgrade/info.c  |   2 +-
 3 files changed, 209 insertions(+), 33 deletions(-)

diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index c6e6d3b2b86..539e046cf22 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -53,6 +53,7 @@
 #include "catalog/pg_publication_d.h"
 #include "catalog/pg_subscription_d.h"
 #include "catalog/pg_type_d.h"
+#include "catalog/pg_zstd_dictionaries_d.h"
 #include "common/connect.h"
 #include "common/int.h"
 #include "common/relpath.h"
@@ -3646,6 +3647,144 @@ dumpDatabase(Archive *fout)
 		destroyPQExpBuffer(loOutQry);
 	}
 
+	if (dopt->binary_upgrade && fout->remoteVersion >= 180000)
+	{
+		/**
+		 *  --- Process pg_zstd_dictionaries related operations ---
+		 */
+		{
+			PGresult   *zstd_res = NULL;
+			PQExpBuffer zstdFrozenQry = createPQExpBuffer();
+			PQExpBuffer zstdOutQry = createPQExpBuffer();
+			PQExpBuffer zstdHorizonQry = createPQExpBuffer();
+			int			ii_relfrozenxid,
+						ii_relfilenode,
+						ii_oid,
+						ii_relminmxid;
+
+			/* Build query to fetch pg_class information */
+			appendPQExpBuffer(zstdFrozenQry,
+							  "SELECT relfrozenxid, relminmxid, relfilenode, oid\n"
+							  "FROM pg_catalog.pg_class\n"
+							  "WHERE oid IN (%u, %u, %u, %u);\n",
+							  ZstdDictionariesRelationId,
+							  PgZstdDictionariesToastTable, /* toast table */
+							  PgZstdDictionariesToastIndex, /* toast table index */
+							  ZstdDictidIndexId);	/* index */
+
+			zstd_res = ExecuteSqlQuery(fout, zstdFrozenQry->data, PGRES_TUPLES_OK);
+
+			ii_relfrozenxid = PQfnumber(zstd_res, "relfrozenxid");
+			ii_relminmxid = PQfnumber(zstd_res, "relminmxid");
+			ii_relfilenode = PQfnumber(zstd_res, "relfilenode");
+			ii_oid = PQfnumber(zstd_res, "oid");
+
+			appendPQExpBufferStr(zstdHorizonQry,
+								 "\n-- For binary upgrade, set pg_zstd_dictionaries relfilenode, relfrozenxid, and relminmxid\n");
+			appendPQExpBufferStr(zstdOutQry,
+								 "\n-- For binary upgrade, preserve pg_zstd_dictionaries and related relfilenodes\n");
+
+			/* Loop over each result row and build update statements */
+			for (int i = 0; i < PQntuples(zstd_res); ++i)
+			{
+				Oid			oid = atooid(PQgetvalue(zstd_res, i, ii_oid));
+				Oid			relfilenumber = atooid(PQgetvalue(zstd_res, i, ii_relfilenode));
+
+				appendPQExpBuffer(zstdHorizonQry,
+								  "UPDATE pg_catalog.pg_class\n"
+								  "SET relfrozenxid = '%u', relminmxid = '%u', relfilenode = %u\n"
+								  "WHERE oid = %u;\n",
+								  atooid(PQgetvalue(zstd_res, i, ii_relfrozenxid)),
+								  atooid(PQgetvalue(zstd_res, i, ii_relminmxid)),
+								  relfilenumber,
+								  oid);
+
+				if (oid == ZstdDictionariesRelationId || oid == PgZstdDictionariesToastTable)
+					appendPQExpBuffer(zstdOutQry,
+									  "SELECT pg_catalog.binary_upgrade_set_next_heap_relfilenode('%u'::pg_catalog.oid);\n",
+									  relfilenumber);
+				else if (oid == PgZstdDictionariesToastIndex || oid == ZstdDictidIndexId)
+					appendPQExpBuffer(zstdOutQry,
+									  "SELECT pg_catalog.binary_upgrade_set_next_index_relfilenode('%u'::pg_catalog.oid);\n",
+									  relfilenumber);
+			}
+
+			appendPQExpBufferStr(zstdOutQry, zstdHorizonQry->data);
+
+			ArchiveEntry(fout, nilCatalogId, createDumpId(),
+						 ARCHIVE_OPTS(.tag = "pg_zstd_dictionaries",
+									  .description = "pg_zstd_dictionaries",
+									  .section = SECTION_PRE_DATA,
+									  .createStmt = zstdOutQry->data));
+
+			PQclear(zstd_res);
+			destroyPQExpBuffer(zstdFrozenQry);
+			destroyPQExpBuffer(zstdHorizonQry);
+			destroyPQExpBuffer(zstdOutQry);
+		}
+
+		/**
+		 * --- Process pg_depend related operations ---
+		 */
+		{
+			PGresult   *pgdep_res = NULL;
+			PQExpBuffer pgdepQry = createPQExpBuffer();
+			PQExpBuffer pgdepBuf = createPQExpBuffer();
+			int			i_classid,
+						i_objid,
+						i_objsubid,
+						i_refclassid,
+						i_refobjid,
+						i_refobjsubid,
+						i_deptype;
+
+			/*
+			 * Build query to fetch all dependency rows where refclassid
+			 * equals ZstdDictionariesRelationId
+			 */
+			appendPQExpBuffer(pgdepQry,
+							  "SELECT classid, objid, objsubid, refclassid, refobjid, refobjsubid, deptype\n"
+							  "FROM pg_catalog.pg_depend\n"
+							  "WHERE refclassid = %u;\n",
+							  ZstdDictionariesRelationId);
+
+			pgdep_res = ExecuteSqlQuery(fout, pgdepQry->data, PGRES_TUPLES_OK);
+
+			i_classid = PQfnumber(pgdep_res, "classid");
+			i_objid = PQfnumber(pgdep_res, "objid");
+			i_objsubid = PQfnumber(pgdep_res, "objsubid");
+			i_refclassid = PQfnumber(pgdep_res, "refclassid");
+			i_refobjid = PQfnumber(pgdep_res, "refobjid");
+			i_refobjsubid = PQfnumber(pgdep_res, "refobjsubid");
+			i_deptype = PQfnumber(pgdep_res, "deptype");
+
+			/* Loop over dependency rows and generate INSERT statements */
+			for (int i = 0; i < PQntuples(pgdep_res); i++)
+			{
+				appendPQExpBuffer(pgdepBuf,
+								  "INSERT INTO pg_catalog.pg_depend (classid, objid, objsubid, refclassid, refobjid, refobjsubid, deptype) "
+								  "VALUES (%s, %s, %s, %s, %s, %s, '%s');\n",
+								  PQgetvalue(pgdep_res, i, i_classid),
+								  PQgetvalue(pgdep_res, i, i_objid),
+								  PQgetvalue(pgdep_res, i, i_objsubid),
+								  PQgetvalue(pgdep_res, i, i_refclassid),
+								  PQgetvalue(pgdep_res, i, i_refobjid),
+								  PQgetvalue(pgdep_res, i, i_refobjsubid),
+								  PQgetvalue(pgdep_res, i, i_deptype));
+			}
+
+			ArchiveEntry(fout, nilCatalogId, createDumpId(),
+						 ARCHIVE_OPTS(.tag = "pg_depend_refclassid_9946",
+									  .description = "pg_depend rows for refclassid 9946",
+									  .section = SECTION_DATA,
+									  .createStmt = pgdepBuf->data));
+
+			PQclear(pgdep_res);
+			destroyPQExpBuffer(pgdepQry);
+			destroyPQExpBuffer(pgdepBuf);
+		}
+	}
+
 	PQclear(res);
 
 	free(qdatname);
@@ -9061,29 +9200,32 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 	 * collation is different from their type's default, we use a CASE here to
 	 * suppress uninteresting attcollations cheaply.
 	 */
-	appendPQExpBufferStr(q,
-						 "SELECT\n"
-						 "a.attrelid,\n"
-						 "a.attnum,\n"
-						 "a.attname,\n"
-						 "a.attstattarget,\n"
-						 "a.attstorage,\n"
-						 "t.typstorage,\n"
-						 "a.atthasdef,\n"
-						 "a.attisdropped,\n"
-						 "a.attlen,\n"
-						 "a.attalign,\n"
-						 "a.attislocal,\n"
-						 "pg_catalog.format_type(t.oid, a.atttypmod) AS atttypname,\n"
-						 "array_to_string(a.attoptions, ', ') AS attoptions,\n"
-						 "CASE WHEN a.attcollation <> t.typcollation "
-						 "THEN a.attcollation ELSE 0 END AS attcollation,\n"
-						 "pg_catalog.array_to_string(ARRAY("
-						 "SELECT pg_catalog.quote_ident(option_name) || "
-						 "' ' || pg_catalog.quote_literal(option_value) "
-						 "FROM pg_catalog.pg_options_to_table(attfdwoptions) "
-						 "ORDER BY option_name"
-						 "), E',\n    ') AS attfdwoptions,\n");
+	appendPQExpBuffer(q,
+					  "SELECT\n"
+					  "  a.attrelid,\n"
+					  "  a.attnum,\n"
+					  "  a.attname,\n"
+					  "  a.attstattarget,\n"
+					  "  a.attstorage,\n"
+					  "  t.typstorage,\n"
+					  "  a.atthasdef,\n"
+					  "  a.attisdropped,\n"
+					  "  a.attlen,\n"
+					  "  a.attalign,\n"
+					  "  a.attislocal,\n"
+					  "  pg_catalog.format_type(t.oid, a.atttypmod) AS atttypname,\n"
+					  "  array_to_string(ARRAY(\n"
+					  "    SELECT x FROM unnest(a.attoptions) AS x %s\n"
+					  "  ), ', ') AS attoptions,\n"
+					  "  CASE WHEN a.attcollation <> t.typcollation"
+					  "       THEN a.attcollation ELSE 0 END AS attcollation,\n"
+					  "  pg_catalog.array_to_string(ARRAY("
+					  "    SELECT pg_catalog.quote_ident(option_name) || ' ' || "
+					  "           pg_catalog.quote_literal(option_value)"
+					  "    FROM pg_catalog.pg_options_to_table(attfdwoptions)"
+					  "    ORDER BY option_name"
+					  "  ), E',\n    ') AS attfdwoptions,\n",
+					  dopt->binary_upgrade ? "" : "WHERE x NOT LIKE 'dictid=%'");
 
 	/*
 	 * Find out any NOT NULL markings for each column.  In 18 and up we read
@@ -12158,12 +12300,14 @@ dumpBaseType(Archive *fout, const TypeInfo *tyinfo)
 	char	   *typmodout;
 	char	   *typanalyze;
 	char	   *typsubscript;
+	char	   *typzstdsampling;
 	Oid			typreceiveoid;
 	Oid			typsendoid;
 	Oid			typmodinoid;
 	Oid			typmodoutoid;
 	Oid			typanalyzeoid;
 	Oid			typsubscriptoid;
+	Oid			typzstdsamplingoid;
 	char	   *typcategory;
 	char	   *typispreferred;
 	char	   *typdelim;
@@ -12196,10 +12340,18 @@ dumpBaseType(Archive *fout, const TypeInfo *tyinfo)
 		if (fout->remoteVersion >= 140000)
 			appendPQExpBufferStr(query,
 								 "typsubscript, "
-								 "typsubscript::pg_catalog.oid AS typsubscriptoid ");
+								 "typsubscript::pg_catalog.oid AS typsubscriptoid, ");
+		else
+			appendPQExpBufferStr(query,
+								 "'-' AS typsubscript, 0 AS typsubscriptoid, ");
+
+		if (fout->remoteVersion >= 180000)
+			appendPQExpBufferStr(query,
+								 "typzstdsampling, "
+								 "typzstdsampling::pg_catalog.oid AS typzstdsamplingoid ");
 		else
 			appendPQExpBufferStr(query,
-								 "'-' AS typsubscript, 0 AS typsubscriptoid ");
+								 "'-' AS typzstdsampling, 0 AS typzstdsamplingoid ");
 
 		appendPQExpBufferStr(query, "FROM pg_catalog.pg_type "
 							 "WHERE oid = $1");
@@ -12224,12 +12376,14 @@ dumpBaseType(Archive *fout, const TypeInfo *tyinfo)
 	typmodout = PQgetvalue(res, 0, PQfnumber(res, "typmodout"));
 	typanalyze = PQgetvalue(res, 0, PQfnumber(res, "typanalyze"));
 	typsubscript = PQgetvalue(res, 0, PQfnumber(res, "typsubscript"));
+	typzstdsampling = PQgetvalue(res, 0, PQfnumber(res, "typzstdsampling"));
 	typreceiveoid = atooid(PQgetvalue(res, 0, PQfnumber(res, "typreceiveoid")));
 	typsendoid = atooid(PQgetvalue(res, 0, PQfnumber(res, "typsendoid")));
 	typmodinoid = atooid(PQgetvalue(res, 0, PQfnumber(res, "typmodinoid")));
 	typmodoutoid = atooid(PQgetvalue(res, 0, PQfnumber(res, "typmodoutoid")));
 	typanalyzeoid = atooid(PQgetvalue(res, 0, PQfnumber(res, "typanalyzeoid")));
 	typsubscriptoid = atooid(PQgetvalue(res, 0, PQfnumber(res, "typsubscriptoid")));
+	typzstdsamplingoid = atooid(PQgetvalue(res, 0, PQfnumber(res, "typzstdsamplingoid")));
 	typcategory = PQgetvalue(res, 0, PQfnumber(res, "typcategory"));
 	typispreferred = PQgetvalue(res, 0, PQfnumber(res, "typispreferred"));
 	typdelim = PQgetvalue(res, 0, PQfnumber(res, "typdelim"));
@@ -12285,7 +12439,8 @@ dumpBaseType(Archive *fout, const TypeInfo *tyinfo)
 		appendPQExpBuffer(q, ",\n    TYPMOD_OUT = %s", typmodout);
 	if (OidIsValid(typanalyzeoid))
 		appendPQExpBuffer(q, ",\n    ANALYZE = %s", typanalyze);
-
+	if (OidIsValid(typzstdsamplingoid))
+		appendPQExpBuffer(q, ",\n    ZSTD_SAMPLING = %s", typzstdsampling);
 	if (strcmp(typcollatable, "t") == 0)
 		appendPQExpBufferStr(q, ",\n    COLLATABLE = true");
 
@@ -17542,6 +17697,9 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 					case 'l':
 						cmname = "lz4";
 						break;
+					case 'z':
+						cmname = "zstd";
+						break;
 					default:
 						cmname = NULL;
 						break;
diff --git a/src/bin/pg_upgrade/check.c b/src/bin/pg_upgrade/check.c
index 18c2d652bb6..38cae3ab0bd 100644
--- a/src/bin/pg_upgrade/check.c
+++ b/src/bin/pg_upgrade/check.c
@@ -14,6 +14,7 @@
 #include "fe_utils/string_utils.h"
 #include "pg_upgrade.h"
 #include "common/unicode_version.h"
+#include "catalog/pg_zstd_dictionaries_d.h"
 
 static void check_new_cluster_is_empty(void);
 static void check_is_install_user(ClusterInfo *cluster);
@@ -916,12 +917,29 @@ check_new_cluster_is_empty(void)
 		for (relnum = 0; relnum < rel_arr->nrels;
 			 relnum++)
 		{
-			/* pg_largeobject and its index should be skipped */
-			if (strcmp(rel_arr->rels[relnum].nspname, "pg_catalog") != 0)
-				pg_fatal("New cluster database \"%s\" is not empty: found relation \"%s.%s\"",
-						 new_cluster.dbarr.dbs[dbnum].db_name,
-						 rel_arr->rels[relnum].nspname,
-						 rel_arr->rels[relnum].relname);
+			const char *nspname = rel_arr->rels[relnum].nspname;
+			const char *relname = rel_arr->rels[relnum].relname;
+			Oid			relOid = rel_arr->rels[relnum].reloid;
+
+			/**
+			 * Allow all objects in pg_catalog
+			 * pg_largeobject, pg_zstd_dictionaries and its index should be skipped.
+			 */
+
+			if (strcmp(nspname, "pg_catalog") == 0)
+				continue;
+
+			/**
+			 * Allow the specific toast objects for pg_zstd_dictionaries:
+			 * fixed OIDs should be 9947 (toast table) or 9948 (toast index).
+			 */
+			if (strcmp(nspname, "pg_toast") == 0 && (relOid == PgZstdDictionariesToastTable || relOid == PgZstdDictionariesToastIndex))
+				continue;
+
+			pg_fatal("New cluster database \"%s\" is not empty: found relation \"%s.%s\"",
+					 new_cluster.dbarr.dbs[dbnum].db_name,
+					 nspname,
+					 relname);
 		}
 	}
 }
diff --git a/src/bin/pg_upgrade/info.c b/src/bin/pg_upgrade/info.c
index 4b7a56f5b3b..b38c125c77a 100644
--- a/src/bin/pg_upgrade/info.c
+++ b/src/bin/pg_upgrade/info.c
@@ -498,7 +498,7 @@ get_rel_infos_query(void)
 					  "                        'binary_upgrade', 'pg_toast') AND "
 					  "      c.oid >= %u::pg_catalog.oid) OR "
 					  "     (n.nspname = 'pg_catalog' AND "
-					  "      relname IN ('pg_largeobject') ))), ",
+					  "      relname IN ('pg_largeobject', 'pg_zstd_dictionaries') ))), ",
 					  (user_opts.transfer_mode == TRANSFER_MODE_SWAP) ?
 					  ", " CppAsString2(RELKIND_SEQUENCE) : "",
 					  FirstNormalObjectId);
-- 
2.47.1

