From 457ab1ac24d5ee438f89b71e2829c593c032bb91 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Tue, 30 Apr 2024 19:33:24 -0500
Subject: [PATCH v3 3/3] cache more sequence data

---
 src/bin/pg_dump/pg_dump.c | 126 ++++++++++++++++++++++++++++++--------
 1 file changed, 102 insertions(+), 24 deletions(-)

diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index a54e32c7be..fd738f7087 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -114,6 +114,8 @@ typedef struct
 	int64		startv;			/* start value */
 	int64		incby;			/* increment value */
 	int64		cache;			/* cache size */
+	int64		last_value;		/* last value of sequence */
+	bool		is_called;		/* whether nextval advances before returning */
 } SequenceItem;
 
 typedef enum OidOptions
@@ -17231,7 +17233,7 @@ static void
 collectSequences(Archive *fout)
 {
 	PGresult   *res;
-	const char *query;
+	PQExpBuffer query;
 
 	/*
 	 * Before Postgres 10, sequence metadata is in the sequence itself.  We
@@ -17241,14 +17243,63 @@ collectSequences(Archive *fout)
 	if (fout->remoteVersion < 100000)
 		return;
 
-	query = "SELECT seqrelid, format_type(seqtypid, NULL), "
-		"seqstart, seqincrement, "
-		"seqmax, seqmin, "
-		"seqcache, seqcycle "
-		"FROM pg_catalog.pg_sequence "
-		"ORDER BY seqrelid";
+	query = createPQExpBuffer();
 
-	res = ExecuteSqlQuery(fout, query, PGRES_TUPLES_OK);
+	appendPQExpBuffer(query, "SELECT seqrelid, format_type(seqtypid, NULL), "
+					  "seqstart, seqincrement, "
+					  "seqmax, seqmin, "
+					  "seqcache, seqcycle, ");
+
+	if (fout->dopt->schemaOnly && !fout->dopt->sequence_data)
+	{
+		/*
+		 * If we aren't dumping the sequence data, just return garbage for
+		 * is_called and last_value.
+		 */
+		appendPQExpBuffer(query, "'f', NULL FROM pg_catalog.pg_sequence ");
+	}
+	else if (fout->remoteVersion >= 180000)
+	{
+		/*
+		 * Since version 18, pg_sequence_last_value() returns NULL if the
+		 * caller doesn't have privileges on the sequence (instead of
+		 * ERROR-ing), so we don't need to make any extra calls to
+		 * has_sequence_privilege() in the query.
+		 */
+		appendPQExpBuffer(query,
+						  "pg_sequence_last_value(seqrelid) IS NOT NULL, "
+						  "CASE WHEN pg_sequence_last_value(seqrelid) IS NOT NULL "
+						  "THEN pg_sequence_last_value(seqrelid) "
+						  "ELSE seqstart END "
+						  "FROM pg_catalog.pg_sequence ");
+	}
+	else
+	{
+		/*
+		 * Before v18, pg_sequence_last_value() ERRORs for sequences for which
+		 * the caller lacks privileges, so we must check privileges before
+		 * calling it.  Furthermore, some older versions have a buggy
+		 * pg_sequence_last_value() that ERRORs for unlogged sequences on
+		 * standbys (see commit 3cb2f13ac5), so we need some additional checks
+		 * to avoid that.
+		 */
+		appendPQExpBuffer(query,
+						  "CASE WHEN has_sequence_privilege(seqrelid, 'SELECT,USAGE'::text) "
+						  "AND (pg_is_in_recovery() = 'f' OR c.relpersistence = 'p') "
+						  "THEN pg_sequence_last_value(seqrelid) IS NOT NULL "
+						  "ELSE 'f' END, "
+						  "CASE WHEN has_sequence_privilege(seqrelid, 'SELECT,USAGE'::text) "
+						  "AND (pg_is_in_recovery() = 'f' OR c.relpersistence = 'p') "
+						  "AND pg_sequence_last_value(seqrelid) IS NOT NULL "
+						  "THEN pg_sequence_last_value(seqrelid) "
+						  "ELSE seqstart END "
+						  "FROM pg_catalog.pg_sequence s "
+						  "JOIN pg_class c ON s.seqrelid = c.oid ");
+	}
+
+	appendPQExpBuffer(query, "ORDER BY seqrelid");
+
+	res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
 
 	nsequences = PQntuples(res);
 	sequences = (SequenceItem *) pg_malloc(nsequences * sizeof(SequenceItem));
@@ -17269,9 +17320,12 @@ collectSequences(Archive *fout)
 		sequences[i].minv = strtoi64(PQgetvalue(res, i, 5), NULL, 10);
 		sequences[i].cache = strtoi64(PQgetvalue(res, i, 6), NULL, 10);
 		sequences[i].cycled = (strcmp(PQgetvalue(res, i, 7), "t") == 0);
+		sequences[i].is_called = (strcmp(PQgetvalue(res, i, 8), "t") == 0);
+		sequences[i].last_value = strtoi64(PQgetvalue(res, i, 9), NULL, 10);
 	}
 
 	PQclear(res);
+	destroyPQExpBuffer(query);
 }
 
 /*
@@ -17550,30 +17604,56 @@ static void
 dumpSequenceData(Archive *fout, const TableDataInfo *tdinfo)
 {
 	TableInfo  *tbinfo = tdinfo->tdtable;
-	PGresult   *res;
-	char	   *last;
+	int64		last;
 	bool		called;
 	PQExpBuffer query = createPQExpBuffer();
 
-	appendPQExpBuffer(query,
-					  "SELECT last_value, is_called FROM %s",
-					  fmtQualifiedDumpable(tbinfo));
+	/*
+	 * For versions >= 10, the sequence information is gathered in a sorted
+	 * table before any calls to dumpSequenceData().  See collectSequences()
+	 * for more information.
+	 */
+	if (fout->remoteVersion < 100000)
+	{
+		PGresult   *res;
 
-	res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
+		appendPQExpBuffer(query,
+						  "SELECT last_value, is_called FROM %s",
+						  fmtQualifiedDumpable(tbinfo));
 
-	if (PQntuples(res) != 1)
-		pg_fatal(ngettext("query to get data of sequence \"%s\" returned %d row (expected 1)",
-						  "query to get data of sequence \"%s\" returned %d rows (expected 1)",
-						  PQntuples(res)),
-				 tbinfo->dobj.name, PQntuples(res));
+		res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
+
+		if (PQntuples(res) != 1)
+			pg_fatal(ngettext("query to get data of sequence \"%s\" returned %d row (expected 1)",
+							  "query to get data of sequence \"%s\" returned %d rows (expected 1)",
+							  PQntuples(res)),
+					 tbinfo->dobj.name, PQntuples(res));
+
+		last = strtoi64(PQgetvalue(res, 0, 0), NULL, 10);
+		called = (strcmp(PQgetvalue(res, 0, 1), "t") == 0);
+
+		PQclear(res);
+	}
+	else
+	{
+		SequenceItem key = {0};
+		SequenceItem *entry;
 
-	last = PQgetvalue(res, 0, 0);
-	called = (strcmp(PQgetvalue(res, 0, 1), "t") == 0);
+		Assert(sequences);
+		Assert(tbinfo->dobj.catId.oid);
+
+		key.oid = tbinfo->dobj.catId.oid;
+		entry = bsearch(&key, sequences, nsequences,
+						sizeof(SequenceItem), SequenceItemCmp);
+
+		last = entry->last_value;
+		called = entry->is_called;
+	}
 
 	resetPQExpBuffer(query);
 	appendPQExpBufferStr(query, "SELECT pg_catalog.setval(");
 	appendStringLiteralAH(query, fmtQualifiedDumpable(tbinfo), fout);
-	appendPQExpBuffer(query, ", %s, %s);\n",
+	appendPQExpBuffer(query, ", " INT64_FORMAT ", %s);\n",
 					  last, (called ? "true" : "false"));
 
 	if (tdinfo->dobj.dump & DUMP_COMPONENT_DATA)
@@ -17587,8 +17667,6 @@ dumpSequenceData(Archive *fout, const TableDataInfo *tdinfo)
 								  .deps = &(tbinfo->dobj.dumpId),
 								  .nDeps = 1));
 
-	PQclear(res);
-
 	destroyPQExpBuffer(query);
 }
 
-- 
2.39.3 (Apple Git-146)

