From de045402f301594ba923c7c93ecc736f7cb40105 Mon Sep 17 00:00:00 2001
From: Naisila Puka <37271756+naisila@users.noreply.github.com>
Date: Fri, 26 Sep 2025 18:04:34 +0300
Subject: [PATCH] PG18 - register snapshot where needed (#8196)

Register and push snapshots as needed per the relevant PG18 commits

https://github.com/postgres/postgres/commit/8076c00592e40e8dbd1fce7a98b20d4bf075e4ba
https://github.com/postgres/postgres/commit/706054b

`citus_split_shard_columnar_partitioned`, `multi_partitioning` tests are
handled.

Fixes #8195
---
 src/backend/columnar/columnar_metadata.c | 16 +++++++++++++++-
 src/backend/distributed/commands/index.c |  6 ++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/backend/columnar/columnar_metadata.c b/src/backend/columnar/columnar_metadata.c
index e3264311c..f699553b6 100644
--- a/src/backend/columnar/columnar_metadata.c
+++ b/src/backend/columnar/columnar_metadata.c
@@ -1267,7 +1267,21 @@ StripesForRelfilelocator(RelFileLocator relfilelocator)
 {
 	uint64 storageId = LookupStorageId(relfilelocator);
 
-	return ReadDataFileStripeList(storageId, GetTransactionSnapshot());
+	/*
+	 * PG18 requires snapshot to be active or registered before it's used
+	 * Without this, we hit
+	 * Assert(snapshot->regd_count > 0 || snapshot->active_count > 0);
+	 * when reading columnar stripes.
+	 * Relevant PG18 commit:
+	 * 8076c00592e40e8dbd1fce7a98b20d4bf075e4ba
+	 */
+	Snapshot snapshot = RegisterSnapshot(GetTransactionSnapshot());
+
+	List *readDataFileStripeList = ReadDataFileStripeList(storageId, snapshot);
+
+	UnregisterSnapshot(snapshot);
+
+	return readDataFileStripeList;
 }
 
 
diff --git a/src/backend/distributed/commands/index.c b/src/backend/distributed/commands/index.c
index 1401da0a6..d95c53fb5 100644
--- a/src/backend/distributed/commands/index.c
+++ b/src/backend/distributed/commands/index.c
@@ -854,8 +854,11 @@ PostprocessIndexStmt(Node *node, const char *queryString)
 	table_close(relation, NoLock);
 	index_close(indexRelation, NoLock);
 
+	PushActiveSnapshot(GetTransactionSnapshot());
+
 	/* mark index as invalid, in-place (cannot be rolled back) */
 	index_set_state_flags(indexRelationId, INDEX_DROP_CLEAR_VALID);
+	PopActiveSnapshot();
 
 	/* re-open a transaction command from here on out */
 	CommitTransactionCommand();
@@ -1370,8 +1373,11 @@ MarkIndexValid(IndexStmt *indexStmt)
 											schemaId);
 	Relation indexRelation = index_open(indexRelationId, RowExclusiveLock);
 
+	PushActiveSnapshot(GetTransactionSnapshot());
+
 	/* mark index as valid, in-place (cannot be rolled back) */
 	index_set_state_flags(indexRelationId, INDEX_CREATE_SET_VALID);
+	PopActiveSnapshot();
 
 	table_close(relation, NoLock);
 	index_close(indexRelation, NoLock);