>From 127efe4a19c5a6977f3989b441002f71f677476e Mon Sep 17 00:00:00 2001
From: Shubhankar Anand Kulkarni <shubhankar.ak@zohocorp.com>
Date: Wed, 14 Oct 2026 12:55:02 +0530
Subject: [PATCH v2 2/2] Add a test module for planner extension state

test_extendplan stores a tag in the extension state of every base and
join relation as the planner builds it, recording which RelOptInfo the
tag was stored in.  join_path_setup_hook then checks the joinrel and
both input rels, and complains if any of them holds a tag stored in a
different RelOptInfo.

The test plans an eager-aggregation query, which builds grouped rels by
flat-copying base and join rels, and a semijoin, which can unique-ify
its inner rel the same way.  Before the previous commit, the copies
inherited the original rel's extension state and the test reported it.

Author: Shubhankar Anand Kulkarni <shubhankar.ak@zohocorp.com>
Discussion: https://postgr.es/m/1a12fb3a08d.429ac7fc62960.6325095284115187647@zohocorp.com
Backpatch-through: 19
---
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 src/test/modules/test_extendplan/.gitignore   |   4 +
 src/test/modules/test_extendplan/Makefile     |  23 +++
 .../expected/test_extendplan.out              |  70 +++++++
 src/test/modules/test_extendplan/meson.build  |  33 ++++
 .../test_extendplan/sql/test_extendplan.sql   |  32 ++++
 .../test_extendplan/test_extendplan--1.0.sql  |  11 ++
 .../modules/test_extendplan/test_extendplan.c | 172 ++++++++++++++++++
 .../test_extendplan/test_extendplan.control   |   4 +
 src/tools/pgindent/typedefs.list              |   1 +
 11 files changed, 352 insertions(+)
 create mode 100644 src/test/modules/test_extendplan/.gitignore
 create mode 100644 src/test/modules/test_extendplan/Makefile
 create mode 100644 src/test/modules/test_extendplan/expected/test_extendplan.out
 create mode 100644 src/test/modules/test_extendplan/meson.build
 create mode 100644 src/test/modules/test_extendplan/sql/test_extendplan.sql
 create mode 100644 src/test/modules/test_extendplan/test_extendplan--1.0.sql
 create mode 100644 src/test/modules/test_extendplan/test_extendplan.c
 create mode 100644 src/test/modules/test_extendplan/test_extendplan.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6c1504ac93..ca5a1c67df 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -30,6 +30,7 @@ SUBDIRS = \
 		  test_dsa \
 		  test_dsm_registry \
 		  test_escape \
+		  test_extendplan \
 		  test_extensible \
 		  test_extensions \
 		  test_ginpostinglist \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index ef3cdb8faf..31c4c218ef 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -32,6 +32,7 @@ subdir('test_ddl_deparse')
 subdir('test_dsa')
 subdir('test_dsm_registry')
 subdir('test_escape')
+subdir('test_extendplan')
 subdir('test_extensible')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
diff --git a/src/test/modules/test_extendplan/.gitignore b/src/test/modules/test_extendplan/.gitignore
new file mode 100644
index 0000000000..5dcb3ff972
--- /dev/null
+++ b/src/test/modules/test_extendplan/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_extendplan/Makefile b/src/test/modules/test_extendplan/Makefile
new file mode 100644
index 0000000000..6b3df630f5
--- /dev/null
+++ b/src/test/modules/test_extendplan/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_extendplan/Makefile
+
+MODULE_big = test_extendplan
+OBJS = \
+	$(WIN32RES) \
+	test_extendplan.o
+PGFILEDESC = "test_extendplan - test module for planner extension state"
+
+EXTENSION = test_extendplan
+DATA = test_extendplan--1.0.sql
+
+REGRESS = test_extendplan
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_extendplan
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_extendplan/expected/test_extendplan.out b/src/test/modules/test_extendplan/expected/test_extendplan.out
new file mode 100644
index 0000000000..f4e1a2074c
--- /dev/null
+++ b/src/test/modules/test_extendplan/expected/test_extendplan.out
@@ -0,0 +1,70 @@
+--
+-- Planner extension state must not be shared between a RelOptInfo and the
+-- RelOptInfos that the planner builds by flat-copying it.
+--
+CREATE EXTENSION test_extendplan;
+-- The first call loads the module, which installs its planner hooks.
+SELECT test_extendplan_reset();
+ test_extendplan_reset 
+-----------------------
+ 
+(1 row)
+
+CREATE TABLE extendplan_t1 (a int, b int, c double precision);
+CREATE TABLE extendplan_t2 (a int, b int, c double precision);
+INSERT INTO extendplan_t1 SELECT i, i, i FROM generate_series(1, 1000) i;
+INSERT INTO extendplan_t2 SELECT i, i % 10, i FROM generate_series(1, 1000) i;
+ANALYZE extendplan_t1, extendplan_t2;
+-- Eager aggregation copies base and join rels to build grouped rels.
+SET enable_eager_aggregate = on;
+SELECT test_extendplan_reset();
+ test_extendplan_reset 
+-----------------------
+ 
+(1 row)
+
+SELECT t1.a, avg(t2.c)
+  FROM extendplan_t1 t1 JOIN extendplan_t2 t2 ON t1.b = t2.b
+ GROUP BY t1.a ORDER BY t1.a;
+ a | avg 
+---+-----
+ 1 | 496
+ 2 | 497
+ 3 | 498
+ 4 | 499
+ 5 | 500
+ 6 | 501
+ 7 | 502
+ 8 | 503
+ 9 | 504
+(9 rows)
+
+SELECT * FROM test_extendplan_summary();
+ grouped_rels_checked | unique_rels_checked | mismatches 
+----------------------+---------------------+------------
+ t                    | f                   |          0
+(1 row)
+
+RESET enable_eager_aggregate;
+-- A semijoin can unique-ify its inner rel, which copies that rel.
+SELECT test_extendplan_reset();
+ test_extendplan_reset 
+-----------------------
+ 
+(1 row)
+
+SELECT count(*) FROM extendplan_t1 t1
+ WHERE t1.b IN (SELECT t2.b FROM extendplan_t2 t2);
+ count 
+-------
+     9
+(1 row)
+
+SELECT * FROM test_extendplan_summary();
+ grouped_rels_checked | unique_rels_checked | mismatches 
+----------------------+---------------------+------------
+ f                    | t                   |          0
+(1 row)
+
+DROP TABLE extendplan_t1, extendplan_t2;
+DROP EXTENSION test_extendplan;
diff --git a/src/test/modules/test_extendplan/meson.build b/src/test/modules/test_extendplan/meson.build
new file mode 100644
index 0000000000..e1eb33ff84
--- /dev/null
+++ b/src/test/modules/test_extendplan/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) 2026, PostgreSQL Global Development Group
+
+test_extendplan_sources = files(
+  'test_extendplan.c',
+)
+
+if host_system == 'windows'
+  test_extendplan_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_extendplan',
+    '--FILEDESC', 'test_extendplan - test module for planner extension state',])
+endif
+
+test_extendplan = shared_module('test_extendplan',
+  test_extendplan_sources,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_extendplan
+
+test_install_data += files(
+  'test_extendplan.control',
+  'test_extendplan--1.0.sql',
+)
+
+tests += {
+  'name': 'test_extendplan',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'regress': {
+    'sql': [
+      'test_extendplan',
+    ],
+  },
+}
diff --git a/src/test/modules/test_extendplan/sql/test_extendplan.sql b/src/test/modules/test_extendplan/sql/test_extendplan.sql
new file mode 100644
index 0000000000..795053e83f
--- /dev/null
+++ b/src/test/modules/test_extendplan/sql/test_extendplan.sql
@@ -0,0 +1,32 @@
+--
+-- Planner extension state must not be shared between a RelOptInfo and the
+-- RelOptInfos that the planner builds by flat-copying it.
+--
+CREATE EXTENSION test_extendplan;
+
+-- The first call loads the module, which installs its planner hooks.
+SELECT test_extendplan_reset();
+
+CREATE TABLE extendplan_t1 (a int, b int, c double precision);
+CREATE TABLE extendplan_t2 (a int, b int, c double precision);
+INSERT INTO extendplan_t1 SELECT i, i, i FROM generate_series(1, 1000) i;
+INSERT INTO extendplan_t2 SELECT i, i % 10, i FROM generate_series(1, 1000) i;
+ANALYZE extendplan_t1, extendplan_t2;
+
+-- Eager aggregation copies base and join rels to build grouped rels.
+SET enable_eager_aggregate = on;
+SELECT test_extendplan_reset();
+SELECT t1.a, avg(t2.c)
+  FROM extendplan_t1 t1 JOIN extendplan_t2 t2 ON t1.b = t2.b
+ GROUP BY t1.a ORDER BY t1.a;
+SELECT * FROM test_extendplan_summary();
+RESET enable_eager_aggregate;
+
+-- A semijoin can unique-ify its inner rel, which copies that rel.
+SELECT test_extendplan_reset();
+SELECT count(*) FROM extendplan_t1 t1
+ WHERE t1.b IN (SELECT t2.b FROM extendplan_t2 t2);
+SELECT * FROM test_extendplan_summary();
+
+DROP TABLE extendplan_t1, extendplan_t2;
+DROP EXTENSION test_extendplan;
diff --git a/src/test/modules/test_extendplan/test_extendplan--1.0.sql b/src/test/modules/test_extendplan/test_extendplan--1.0.sql
new file mode 100644
index 0000000000..84c3cc07a6
--- /dev/null
+++ b/src/test/modules/test_extendplan/test_extendplan--1.0.sql
@@ -0,0 +1,11 @@
+/* src/test/modules/test_extendplan/test_extendplan--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_extendplan" to load this file. \quit
+
+CREATE FUNCTION test_extendplan_reset() RETURNS void
+  AS 'MODULE_PATHNAME', 'test_extendplan_reset' LANGUAGE C STRICT;
+
+CREATE FUNCTION test_extendplan_summary(OUT grouped_rels_checked bool,
+  OUT unique_rels_checked bool, OUT mismatches int8)
+  AS 'MODULE_PATHNAME', 'test_extendplan_summary' LANGUAGE C STRICT;
diff --git a/src/test/modules/test_extendplan/test_extendplan.c b/src/test/modules/test_extendplan/test_extendplan.c
new file mode 100644
index 0000000000..309bd29e44
--- /dev/null
+++ b/src/test/modules/test_extendplan/test_extendplan.c
@@ -0,0 +1,172 @@
+/*-------------------------------------------------------------------------
+ *
+ * test_extendplan.c
+ *		Test module for planner extension state
+ *
+ * When the planner builds a base or join relation, this module stores a tag
+ * in the relation's extension state recording which RelOptInfo it was
+ * stored in.  join_path_setup_hook then checks the joinrel and both of its
+ * input rels: any tag found there must belong to that same RelOptInfo.
+ *
+ * The planner creates some RelOptInfos by flat-copying another one, such as
+ * the grouped rels built for eager aggregation and the unique-ified rels
+ * built for semijoins.  Such a copy must not inherit, or share, the extension
+ * state of the rel it was copied from.
+ *
+ * Copyright (c) 2026, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_extendplan/test_extendplan.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/htup_details.h"
+#include "fmgr.h"
+#include "funcapi.h"
+#include "optimizer/extendplan.h"
+#include "optimizer/pathnode.h"
+#include "optimizer/paths.h"
+
+PG_MODULE_MAGIC;
+
+/* What we store in the extension state of each rel we see being built */
+typedef struct TestExtendPlanTag
+{
+	RelOptInfo *owner;			/* the RelOptInfo the tag was stored in */
+} TestExtendPlanTag;
+
+static int	extension_id = -1;
+
+/* Counters reported by test_extendplan_summary() */
+static int64 grouped_rels_checked = 0;
+static int64 unique_rels_checked = 0;
+static int64 mismatches = 0;
+
+static build_simple_rel_hook_type prev_build_simple_rel_hook = NULL;
+static joinrel_setup_hook_type prev_joinrel_setup_hook = NULL;
+static join_path_setup_hook_type prev_join_path_setup_hook = NULL;
+
+PG_FUNCTION_INFO_V1(test_extendplan_reset);
+PG_FUNCTION_INFO_V1(test_extendplan_summary);
+
+/*
+ * Store a tag naming rel in rel's extension state.
+ */
+static void
+tag_rel(RelOptInfo *rel)
+{
+	TestExtendPlanTag *tag = palloc_object(TestExtendPlanTag);
+
+	tag->owner = rel;
+	SetRelOptInfoExtensionState(rel, extension_id, tag);
+}
+
+/*
+ * Complain if rel's extension state holds a tag stored in some other rel.
+ */
+static void
+check_rel(RelOptInfo *rel, const char *role)
+{
+	TestExtendPlanTag *tag;
+
+	tag = GetRelOptInfoExtensionState(rel, extension_id);
+	if (tag != NULL && tag->owner != rel)
+	{
+		mismatches++;
+		elog(WARNING, "extension state of %s%s %s belongs to a different RelOptInfo",
+			 IS_GROUPED_REL(rel) ? "grouped " : "", role,
+			 bmsToString(rel->relids));
+	}
+}
+
+static void
+test_build_simple_rel(PlannerInfo *root, RelOptInfo *rel, RangeTblEntry *rte)
+{
+	if (prev_build_simple_rel_hook)
+		prev_build_simple_rel_hook(root, rel, rte);
+
+	tag_rel(rel);
+}
+
+static void
+test_joinrel_setup(PlannerInfo *root, RelOptInfo *joinrel,
+				   RelOptInfo *outer_rel, RelOptInfo *inner_rel,
+				   SpecialJoinInfo *sjinfo, List *restrictlist)
+{
+	if (prev_joinrel_setup_hook)
+		prev_joinrel_setup_hook(root, joinrel, outer_rel, inner_rel,
+								sjinfo, restrictlist);
+
+	tag_rel(joinrel);
+}
+
+static void
+test_join_path_setup(PlannerInfo *root, RelOptInfo *joinrel,
+					 RelOptInfo *outerrel, RelOptInfo *innerrel,
+					 JoinType jointype, JoinPathExtraData *extra)
+{
+	if (prev_join_path_setup_hook)
+		prev_join_path_setup_hook(root, joinrel, outerrel, innerrel,
+								  jointype, extra);
+
+	/* Count the flat-copied rels this call lets us check */
+	if (IS_GROUPED_REL(joinrel) || IS_GROUPED_REL(outerrel) ||
+		IS_GROUPED_REL(innerrel))
+		grouped_rels_checked++;
+	if (jointype == JOIN_UNIQUE_OUTER || jointype == JOIN_UNIQUE_INNER)
+		unique_rels_checked++;
+
+	check_rel(joinrel, "joinrel");
+	check_rel(outerrel, "outer rel");
+	check_rel(innerrel, "inner rel");
+}
+
+void
+_PG_init(void)
+{
+	extension_id = GetPlannerExtensionId("test_extendplan");
+
+	prev_build_simple_rel_hook = build_simple_rel_hook;
+	build_simple_rel_hook = test_build_simple_rel;
+	prev_joinrel_setup_hook = joinrel_setup_hook;
+	joinrel_setup_hook = test_joinrel_setup;
+	prev_join_path_setup_hook = join_path_setup_hook;
+	join_path_setup_hook = test_join_path_setup;
+}
+
+/*
+ * Reset the counters.
+ */
+Datum
+test_extendplan_reset(PG_FUNCTION_ARGS)
+{
+	grouped_rels_checked = 0;
+	unique_rels_checked = 0;
+	mismatches = 0;
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * Report whether grouped and unique-ified rels were checked since the last
+ * reset, and how many tags were found in the wrong RelOptInfo.
+ */
+Datum
+test_extendplan_summary(PG_FUNCTION_ARGS)
+{
+	TupleDesc	tupdesc;
+	Datum		values[3];
+	bool		nulls[3] = {0};
+
+	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
+		elog(ERROR, "return type must be a row type");
+
+	values[0] = BoolGetDatum(grouped_rels_checked > 0);
+	values[1] = BoolGetDatum(unique_rels_checked > 0);
+	values[2] = Int64GetDatum(mismatches);
+
+	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
+}
diff --git a/src/test/modules/test_extendplan/test_extendplan.control b/src/test/modules/test_extendplan/test_extendplan.control
new file mode 100644
index 0000000000..9f8f2ef81f
--- /dev/null
+++ b/src/test/modules/test_extendplan/test_extendplan.control
@@ -0,0 +1,4 @@
+comment = 'Test code for planner extension state'
+default_version = '1.0'
+module_pathname = '$libdir/test_extendplan'
+relocatable = false
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index a7a8dbdd23..8abe1d75d2 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -3165,6 +3165,7 @@ TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
 TestExtNode
+TestExtendPlanTag
 TestShmemData
 TestSpec
 TestValueType
-- 
2.43.0

