v20240729-0001-Introduce-pg_sequence_state-and-SetSequenc.patch

From baa1158bb8222d36ffc8ad2b0e1bc265ae68f85e Mon Sep 17 00:00:00 2001
From: Vignesh C <vignesh21@gmail.com>
Date: Tue, 11 Jun 2024 12:15:16 +0530
Subject: [PATCH v20240729 1/3] Introduce pg_sequence_state and
 SetSequenceLastValue functions for enhanced sequence management

This patch introduces a couple of new functions: pg_sequence_state function
allows retrieval of sequence values including LSN. The SetSequenceLastValue
function enables updating sequences with specified values.
---
 doc/src/sgml/func.sgml                 |  27 ++++
 src/backend/commands/sequence.c        | 169 +++++++++++++++++++++++--
 src/include/catalog/pg_proc.dat        |   8 ++
 src/include/commands/sequence.h        |   1 +
 src/test/regress/expected/sequence.out |  12 ++
 src/test/regress/sql/sequence.sql      |   2 +
 6 files changed, 211 insertions(+), 8 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index b669ab7f97..4ff403ee01 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -19582,6 +19582,33 @@ SELECT setval('myseq', 42, false);    <lineannotation>Next <function>nextval</fu
         or <literal>SELECT</literal> privilege on the last used sequence.
        </para></entry>
       </row>
+
+      <row>
+       <entry role="func_table_entry"><para role="func_signature">
+        <indexterm>
+         <primary>pg_sequence_state</primary>
+        </indexterm>
+        <function>pg_sequence_state</function> ()
+        <returnvalue>record</returnvalue>
+        ( <parameter>lsn</parameter> <type>pg_lsn</type>,
+        <parameter>last_value</parameter> <type>bigint</type>,
+        <parameter>log_cnt</parameter> <type>bigint</type>,
+        <parameter>is_called</parameter> <type>bool</type> )
+       </para>
+       <para>
+        Returns information about the sequence. <literal>lsn</literal> is the
+        page LSN of the sequence, <literal>last_value</literal> is the most
+        recent value returned by <function>nextval</function> in the current
+        session, <literal>log_cnt</literal> shows how many fetches remain
+        before a new WAL record has to be written, and
+        <literal>is_called</literal> indicates whether the sequence has been
+        used.
+       </para>
+       <para>
+        This function requires <literal>USAGE</literal>
+        or <literal>SELECT</literal> privilege on the sequence.
+       </para></entry>
+      </row>
      </tbody>
     </tgroup>
    </table>
diff --git a/src/backend/commands/sequence.c b/src/backend/commands/sequence.c
index 9f28d40466..bff990afa7 100644
--- a/src/backend/commands/sequence.c
+++ b/src/backend/commands/sequence.c
@@ -45,6 +45,7 @@
 #include "utils/acl.h"
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
+#include "utils/pg_lsn.h"
 #include "utils/resowner.h"
 #include "utils/syscache.h"
 #include "utils/varlena.h"
@@ -102,7 +103,8 @@ static Relation lock_and_open_sequence(SeqTable seq);
 static void create_seq_hashtable(void);
 static void init_sequence(Oid relid, SeqTable *p_elm, Relation *p_rel);
 static Form_pg_sequence_data read_seq_tuple(Relation rel,
-											Buffer *buf, HeapTuple seqdatatuple);
+											Buffer *buf, HeapTuple seqdatatuple,
+											XLogRecPtr *lsn_ret);
 static void init_params(ParseState *pstate, List *options, bool for_identity,
 						bool isInit,
 						Form_pg_sequence seqform,
@@ -277,7 +279,7 @@ ResetSequence(Oid seq_relid)
 	 * indeed a sequence.
 	 */
 	init_sequence(seq_relid, &elm, &seq_rel);
-	(void) read_seq_tuple(seq_rel, &buf, &seqdatatuple);
+	(void) read_seq_tuple(seq_rel, &buf, &seqdatatuple, NULL);
 
 	pgstuple = SearchSysCache1(SEQRELID, ObjectIdGetDatum(seq_relid));
 	if (!HeapTupleIsValid(pgstuple))
@@ -328,6 +330,83 @@ ResetSequence(Oid seq_relid)
 	sequence_close(seq_rel, NoLock);
 }
 
+/*
+ * Set a sequence to a specified internal state.
+ *
+ * Caller is assumed to have acquired AccessExclusiveLock on the sequence,
+ * which must not be released until end of transaction.  Caller is also
+ * responsible for permissions checking.
+ *
+ * Note: This function resembles do_setval but does not include the locking and
+ * verification steps, as those are managed in a slightly different manner for
+ * logical replication.
+ */
+void
+SetSequenceLastValue(Oid seq_relid, int64 new_last_value)
+{
+	SeqTable        elm;
+	Relation        seqrel;
+	Buffer          buf;
+	HeapTupleData seqdatatuple;
+	Form_pg_sequence_data seq;
+
+	/* open and lock sequence */
+	init_sequence(seq_relid, &elm, &seqrel);
+
+	/* lock page buffer and read tuple */
+	seq = read_seq_tuple(seqrel, &buf, &seqdatatuple, NULL);
+
+	/* check the comment above nextval_internal()'s equivalent call. */
+	if (RelationNeedsWAL(seqrel))
+	{
+		GetTopTransactionId();
+
+		if (XLogLogicalInfoActive())
+			GetCurrentTransactionId();
+	}
+
+	/* ready to change the on-disk (or really, in-buffer) tuple */
+	START_CRIT_SECTION();
+
+	seq->last_value = new_last_value;
+	seq->is_called = true;
+	seq->log_cnt = 0;
+
+	MarkBufferDirty(buf);
+
+	/* XLOG stuff */
+	if (RelationNeedsWAL(seqrel))
+	{
+		xl_seq_rec      xlrec;
+		XLogRecPtr      recptr;
+		Page            page = BufferGetPage(buf);
+
+		XLogBeginInsert();
+		XLogRegisterBuffer(0, buf, REGBUF_WILL_INIT);
+
+		xlrec.locator = seqrel->rd_locator;
+
+		XLogRegisterData((char *) &xlrec, sizeof(xl_seq_rec));
+		XLogRegisterData((char *) seqdatatuple.t_data, seqdatatuple.t_len);
+
+		recptr = XLogInsert(RM_SEQ_ID, XLOG_SEQ_LOG);
+
+		PageSetLSN(page, recptr);
+	}
+
+	END_CRIT_SECTION();
+
+	UnlockReleaseBuffer(buf);
+
+	/*
+	 * Clear local cache so that we don't think we have cached numbers.
+	 * Note that we do not change the currval() state.
+	 */
+	elm->cached = elm->last;
+
+	relation_close(seqrel, NoLock);
+}
+
 /*
  * Initialize a sequence's relation with the specified tuple as content
  *
@@ -476,7 +555,7 @@ AlterSequence(ParseState *pstate, AlterSeqStmt *stmt)
 	seqform = (Form_pg_sequence) GETSTRUCT(seqtuple);
 
 	/* lock page buffer and read tuple into new sequence structure */
-	(void) read_seq_tuple(seqrel, &buf, &datatuple);
+	(void) read_seq_tuple(seqrel, &buf, &datatuple, NULL);
 
 	/* copy the existing sequence data tuple, so it can be modified locally */
 	newdatatuple = heap_copytuple(&datatuple);
@@ -558,7 +637,7 @@ SequenceChangePersistence(Oid relid, char newrelpersistence)
 	if (RelationNeedsWAL(seqrel))
 		GetTopTransactionId();
 
-	(void) read_seq_tuple(seqrel, &buf, &seqdatatuple);
+	(void) read_seq_tuple(seqrel, &buf, &seqdatatuple, NULL);
 	RelationSetNewRelfilenumber(seqrel, newrelpersistence);
 	fill_seq_with_data(seqrel, &seqdatatuple);
 	UnlockReleaseBuffer(buf);
@@ -687,7 +766,7 @@ nextval_internal(Oid relid, bool check_permissions)
 	ReleaseSysCache(pgstuple);
 
 	/* lock page buffer and read tuple */
-	seq = read_seq_tuple(seqrel, &buf, &seqdatatuple);
+	seq = read_seq_tuple(seqrel, &buf, &seqdatatuple, NULL);
 	page = BufferGetPage(buf);
 
 	last = next = result = seq->last_value;
@@ -983,7 +1062,7 @@ do_setval(Oid relid, int64 next, bool iscalled)
 	PreventCommandIfParallelMode("setval()");
 
 	/* lock page buffer and read tuple */
-	seq = read_seq_tuple(seqrel, &buf, &seqdatatuple);
+	seq = read_seq_tuple(seqrel, &buf, &seqdatatuple, NULL);
 
 	if ((next < minv) || (next > maxv))
 		ereport(ERROR,
@@ -1183,11 +1262,15 @@ init_sequence(Oid relid, SeqTable *p_elm, Relation *p_rel)
  * *buf receives the reference to the pinned-and-ex-locked buffer
  * *seqdatatuple receives the reference to the sequence tuple proper
  *		(this arg should point to a local variable of type HeapTupleData)
+ * *lsn_ret will be set to the page LSN if the caller requested it.
+ *		This allows the caller to determine which sequence changes are
+ *		before/after the returned sequence state.
  *
  * Function's return value points to the data payload of the tuple
  */
 static Form_pg_sequence_data
-read_seq_tuple(Relation rel, Buffer *buf, HeapTuple seqdatatuple)
+read_seq_tuple(Relation rel, Buffer *buf, HeapTuple seqdatatuple,
+			   XLogRecPtr *lsn_ret)
 {
 	Page		page;
 	ItemId		lp;
@@ -1204,6 +1287,10 @@ read_seq_tuple(Relation rel, Buffer *buf, HeapTuple seqdatatuple)
 		elog(ERROR, "bad magic number in sequence \"%s\": %08X",
 			 RelationGetRelationName(rel), sm->magic);
 
+	/* If the caller requested it, return the page LSN. */
+	if (lsn_ret)
+		*lsn_ret = PageGetLSN(page);
+
 	lp = PageGetItemId(page, FirstOffsetNumber);
 	Assert(ItemIdIsNormal(lp));
 
@@ -1807,7 +1894,7 @@ pg_sequence_last_value(PG_FUNCTION_ARGS)
 		HeapTupleData seqtuple;
 		Form_pg_sequence_data seq;
 
-		seq = read_seq_tuple(seqrel, &buf, &seqtuple);
+		seq = read_seq_tuple(seqrel, &buf, &seqtuple, NULL);
 
 		is_called = seq->is_called;
 		result = seq->last_value;
@@ -1822,6 +1909,72 @@ pg_sequence_last_value(PG_FUNCTION_ARGS)
 		PG_RETURN_NULL();
 }
 
+/*
+ * Return the current on-disk state of the sequence.
+ *
+ * Note: This is roughly equivalent to selecting the data from the sequence,
+ * except that it also returns the page LSN.
+ */
+Datum
+pg_sequence_state(PG_FUNCTION_ARGS)
+{
+	Oid			seq_relid = PG_GETARG_OID(0);
+	SeqTable	elm;
+	Relation	seqrel;
+	Buffer		buf;
+	HeapTupleData seqtuple;
+	Form_pg_sequence_data seq;
+	Datum		result;
+
+	XLogRecPtr	lsn;
+	int64		last_value;
+	int64		log_cnt;
+	bool		is_called;
+
+	TupleDesc	tupdesc;
+	HeapTuple	tuple;
+	Datum		values[4];
+	bool		nulls[4] = {false, false, false, false};
+
+	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
+		elog(ERROR, "return type must be a row type");
+
+	/* open and lock sequence */
+	init_sequence(seq_relid, &elm, &seqrel);
+
+	if (pg_class_aclcheck(elm->relid, GetUserId(),
+						  ACL_SELECT | ACL_USAGE) != ACLCHECK_OK)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				 errmsg("permission denied for sequence %s",
+						RelationGetRelationName(seqrel))));
+
+	seq = read_seq_tuple(seqrel, &buf, &seqtuple, &lsn);
+
+	last_value = seq->last_value;
+	log_cnt = seq->log_cnt;
+	is_called = seq->is_called;
+
+	UnlockReleaseBuffer(buf);
+	relation_close(seqrel, NoLock);
+
+	/* Page LSN for the sequence */
+	values[0] = LSNGetDatum(lsn);
+
+	/* The value most recently returned by nextval in the current session */
+	values[1] = Int64GetDatum(last_value);
+
+	/* How many fetches remain before a new WAL record has to be written */
+	values[2] = Int64GetDatum(log_cnt);
+
+	/* Indicates whether the sequence has been used */
+	values[3] = BoolGetDatum(is_called);
+
+	tuple = heap_form_tuple(tupdesc, values, nulls);
+	result = HeapTupleGetDatum(tuple);
+
+	PG_RETURN_DATUM(result);
+}
 
 void
 seq_redo(XLogReaderState *record)
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 06b2f4ba66..5b250bad78 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -3329,6 +3329,14 @@
   proname => 'pg_sequence_last_value', provolatile => 'v', proparallel => 'u',
   prorettype => 'int8', proargtypes => 'regclass',
   prosrc => 'pg_sequence_last_value' },
+{ oid => '6313',
+  descr => 'current on-disk sequence state',
+  proname => 'pg_sequence_state', provolatile => 'v',
+  prorettype => 'record', proargtypes => 'regclass',
+  proallargtypes => '{regclass,pg_lsn,int8,int8,bool}',
+  proargmodes => '{i,o,o,o,o}',
+  proargnames => '{seq_oid,page_lsn,last_value,log_cnt,is_called}',
+  prosrc => 'pg_sequence_state' },
 
 { oid => '275', descr => 'return the next oid for a system table',
   proname => 'pg_nextoid', provolatile => 'v', proparallel => 'u',
diff --git a/src/include/commands/sequence.h b/src/include/commands/sequence.h
index e88cbee3b5..003f2e3413 100644
--- a/src/include/commands/sequence.h
+++ b/src/include/commands/sequence.h
@@ -60,6 +60,7 @@ extern ObjectAddress AlterSequence(ParseState *pstate, AlterSeqStmt *stmt);
 extern void SequenceChangePersistence(Oid relid, char newrelpersistence);
 extern void DeleteSequenceTuple(Oid relid);
 extern void ResetSequence(Oid seq_relid);
+extern void SetSequenceLastValue(Oid seq_relid, int64 new_last_value);
 extern void ResetSequenceCaches(void);
 
 extern void seq_redo(XLogReaderState *record);
diff --git a/src/test/regress/expected/sequence.out b/src/test/regress/expected/sequence.out
index 2b47b7796b..cbcd65f499 100644
--- a/src/test/regress/expected/sequence.out
+++ b/src/test/regress/expected/sequence.out
@@ -161,6 +161,12 @@ SELECT nextval('serialTest2_f6_seq');
 CREATE SEQUENCE sequence_test;
 CREATE SEQUENCE IF NOT EXISTS sequence_test;
 NOTICE:  relation "sequence_test" already exists, skipping
+SELECT last_value, log_cnt, is_called  FROM  pg_sequence_state('sequence_test');
+ last_value | log_cnt | is_called 
+------------+---------+-----------
+          1 |       0 | f
+(1 row)
+
 SELECT nextval('sequence_test'::text);
  nextval 
 ---------
@@ -233,6 +239,12 @@ SELECT nextval('sequence_test'::text);
       99
 (1 row)
 
+SELECT last_value, log_cnt, is_called  FROM  pg_sequence_state('sequence_test');
+ last_value | log_cnt | is_called 
+------------+---------+-----------
+         99 |      32 | t
+(1 row)
+
 DISCARD SEQUENCES;
 SELECT currval('sequence_test'::regclass);
 ERROR:  currval of sequence "sequence_test" is not yet defined in this session
diff --git a/src/test/regress/sql/sequence.sql b/src/test/regress/sql/sequence.sql
index 674f5f1f66..5fcb36341d 100644
--- a/src/test/regress/sql/sequence.sql
+++ b/src/test/regress/sql/sequence.sql
@@ -112,6 +112,7 @@ SELECT nextval('serialTest2_f6_seq');
 CREATE SEQUENCE sequence_test;
 CREATE SEQUENCE IF NOT EXISTS sequence_test;
 
+SELECT last_value, log_cnt, is_called  FROM  pg_sequence_state('sequence_test');
 SELECT nextval('sequence_test'::text);
 SELECT nextval('sequence_test'::regclass);
 SELECT currval('sequence_test'::text);
@@ -124,6 +125,7 @@ SELECT setval('sequence_test'::regclass, 32);
 SELECT nextval('sequence_test'::text);
 SELECT setval('sequence_test'::regclass, 99, false);
 SELECT nextval('sequence_test'::text);
+SELECT last_value, log_cnt, is_called  FROM  pg_sequence_state('sequence_test');
 DISCARD SEQUENCES;
 SELECT currval('sequence_test'::regclass);
 
-- 
2.34.1