From 3e256c72c5d18743d6158bfeee51a2cb2c67e95d Mon Sep 17 00:00:00 2001
From: "Andrey V. Lepikhov" <a.lepikhov@postgrespro.ru>
Date: Mon, 11 Sep 2023 10:02:03 +0700
Subject: [PATCH] CustomScan node as a test module

---
 src/test/modules/Makefile                     |   1 +
 src/test/modules/test_custom_nodes/Makefile   |  19 ++
 .../test_custom_nodes/expected/basic.out      |  62 +++++
 .../modules/test_custom_nodes/sql/basic.sql   |  25 ++
 .../test_custom_nodes/test_custom_nodes.c     | 237 ++++++++++++++++++
 .../test_custom_nodes.control                 |   5 +
 6 files changed, 349 insertions(+)
 create mode 100644 src/test/modules/test_custom_nodes/Makefile
 create mode 100644 src/test/modules/test_custom_nodes/expected/basic.out
 create mode 100644 src/test/modules/test_custom_nodes/sql/basic.sql
 create mode 100644 src/test/modules/test_custom_nodes/test_custom_nodes.c
 create mode 100644 src/test/modules/test_custom_nodes/test_custom_nodes.control

diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 6331c976dc..ec9c1ec00a 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -16,6 +16,7 @@ SUBDIRS = \
 		  spgist_name_ops \
 		  test_bloomfilter \
 		  test_copy_callbacks \
+		  test_custom_nodes	\
 		  test_custom_rmgrs \
 		  test_ddl_deparse \
 		  test_extensions \
diff --git a/src/test/modules/test_custom_nodes/Makefile b/src/test/modules/test_custom_nodes/Makefile
new file mode 100644
index 0000000000..2e41badeb4
--- /dev/null
+++ b/src/test/modules/test_custom_nodes/Makefile
@@ -0,0 +1,19 @@
+# src/test/modules/test_custom_nodes/Makefile
+
+MODULES = test_custom_nodes
+
+EXTENSION = test_custom_nodes
+PGFILEDESC = "test_custom_nodes - custom nodes usage examples"
+
+REGRESS = basic
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_custom_nodes
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_custom_nodes/expected/basic.out b/src/test/modules/test_custom_nodes/expected/basic.out
new file mode 100644
index 0000000000..190a7d6704
--- /dev/null
+++ b/src/test/modules/test_custom_nodes/expected/basic.out
@@ -0,0 +1,62 @@
+--
+-- Copyright (c) 2023, PostgreSQL Global Development Group
+--
+-- Set of basic regression tests on Custom node
+--
+LOAD 'test_custom_nodes';
+-- CTE issue
+create view tt24v as
+with cte as materialized (select r from (values(1,2),(3,4)) r)
+select (r).column2 as col_a, (rr).column2 as col_b from
+  cte join (select rr from (values(1,7),(3,8)) rr limit 2) ss
+  on (r).column1 = (rr).column1;
+explain (verbose, costs off) select * from tt24v;
+                                              QUERY PLAN                                              
+------------------------------------------------------------------------------------------------------
+ Result
+   Output: (cte.r).column2, (ss.rr).column2
+   CTE cte
+     ->  Values Scan on "*VALUES*_1"
+           Output: ROW("*VALUES*_1".column1, "*VALUES*_1".column2)
+   ->  Custom Scan (CustomMaterializeOverHashJoin)
+         Output: cte.r, ss.rr
+         ->  Hash Join
+               Output: cte.r, (ROW("*VALUES*".column1, "*VALUES*".column2))
+               Hash Cond: ((cte.r).column1 = ((ROW("*VALUES*".column1, "*VALUES*".column2))).column1)
+               ->  CTE Scan on cte
+                     Output: cte.r
+               ->  Hash
+                     Output: (ROW("*VALUES*".column1, "*VALUES*".column2))
+                     ->  Limit
+                           Output: (ROW("*VALUES*".column1, "*VALUES*".column2))
+                           ->  Values Scan on "*VALUES*"
+                                 Output: ROW("*VALUES*".column1, "*VALUES*".column2)
+(18 rows)
+
+DROP VIEW IF EXISTS tt24v CASCADE;
+-- Subquery issue
+create view tt24v as
+with cte as (select r from (values(1,2),(3,4)) r)
+select (r).column2 as col_a, (rr).column2 as col_b from
+  cte join (select rr from (values(1,7),(3,8)) rr limit 2) ss
+  on (r).column1 = (rr).column1;
+explain (verbose, costs off) select * from tt24v;
+                                                 QUERY PLAN                                                  
+-------------------------------------------------------------------------------------------------------------
+ Result
+   Output: "*VALUES*".column2, (ss.rr).column2
+   ->  Custom Scan (CustomMaterializeOverHashJoin)
+         Output: "*VALUES*".column2, ss.rr
+         ->  Hash Join
+               Output: "*VALUES*".column2, (ROW("*VALUES*_1".column1, "*VALUES*_1".column2))
+               Hash Cond: ("*VALUES*".column1 = ((ROW("*VALUES*_1".column1, "*VALUES*_1".column2))).column1)
+               ->  Values Scan on "*VALUES*"
+                     Output: "*VALUES*".column1, "*VALUES*".column2
+               ->  Hash
+                     Output: (ROW("*VALUES*_1".column1, "*VALUES*_1".column2))
+                     ->  Limit
+                           Output: (ROW("*VALUES*_1".column1, "*VALUES*_1".column2))
+                           ->  Values Scan on "*VALUES*_1"
+                                 Output: ROW("*VALUES*_1".column1, "*VALUES*_1".column2)
+(15 rows)
+
diff --git a/src/test/modules/test_custom_nodes/sql/basic.sql b/src/test/modules/test_custom_nodes/sql/basic.sql
new file mode 100644
index 0000000000..024c5f0420
--- /dev/null
+++ b/src/test/modules/test_custom_nodes/sql/basic.sql
@@ -0,0 +1,25 @@
+--
+-- Copyright (c) 2023, PostgreSQL Global Development Group
+--
+-- Set of basic regression tests on Custom node
+--
+
+LOAD 'test_custom_nodes';
+
+-- CTE issue
+create view tt24v as
+with cte as materialized (select r from (values(1,2),(3,4)) r)
+select (r).column2 as col_a, (rr).column2 as col_b from
+  cte join (select rr from (values(1,7),(3,8)) rr limit 2) ss
+  on (r).column1 = (rr).column1;
+explain (verbose, costs off) select * from tt24v;
+
+DROP VIEW IF EXISTS tt24v CASCADE;
+
+-- Subquery issue
+create view tt24v as
+with cte as (select r from (values(1,2),(3,4)) r)
+select (r).column2 as col_a, (rr).column2 as col_b from
+  cte join (select rr from (values(1,7),(3,8)) rr limit 2) ss
+  on (r).column1 = (rr).column1;
+explain (verbose, costs off) select * from tt24v;
diff --git a/src/test/modules/test_custom_nodes/test_custom_nodes.c b/src/test/modules/test_custom_nodes/test_custom_nodes.c
new file mode 100644
index 0000000000..68146d2030
--- /dev/null
+++ b/src/test/modules/test_custom_nodes/test_custom_nodes.c
@@ -0,0 +1,237 @@
+/* -------------------------------------------------------------------------
+ *
+ * custom_scan.c
+ *		Sample custom scan code that demonstrates various coding ...
+ *
+ * Copyright (c) 2013-2023, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		src/test/modules/worker_spi/custom_scan.c
+ *
+ * -------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "executor/executor.h"
+#include "nodes/extensible.h"
+#include "optimizer/pathnode.h"
+#include "optimizer/paths.h"
+
+PG_MODULE_MAGIC;
+
+#define CUSTOM_NODE_NAME "CustomMaterializeOverHashJoin"
+
+static Plan *CreateCNPlan(PlannerInfo *root, RelOptInfo *rel,
+							CustomPath *best_path, List *tlist,
+							List *scan_clauses, List *custom_plans);
+static Node *CreateCNScanState(CustomScan *cscan);
+
+/* Exec methods */
+static void BeginCNScan(CustomScanState *node, EState *estate, int eflags);
+static TupleTableSlot *ExecCNScan(CustomScanState *node);
+static void EndCNScan(CustomScanState *node);
+static void ReScanCNScan(CustomScanState *node);
+
+static CustomPathMethods path_methods =
+{
+	.CustomName = CUSTOM_NODE_NAME,
+	.PlanCustomPath = CreateCNPlan,
+	.ReparameterizeCustomPathByChild = NULL
+};
+static CustomScanMethods plan_methods =
+{
+	.CustomName = CUSTOM_NODE_NAME,
+	.CreateCustomScanState = CreateCNScanState
+};
+static CustomExecMethods exec_methods =
+{
+	.CustomName = CUSTOM_NODE_NAME,
+
+	/* Required executor methods */
+	.BeginCustomScan = BeginCNScan,
+	.ExecCustomScan = ExecCNScan,
+	.EndCustomScan = EndCNScan,
+	.ReScanCustomScan = ReScanCNScan,
+
+	/* Optional methods: needed if mark/restore is supported */
+	.MarkPosCustomScan = NULL,
+	.RestrPosCustomScan = NULL,
+
+	/* Optional methods: needed if parallel execution is supported */
+	.EstimateDSMCustomScan = NULL,
+	.InitializeDSMCustomScan = NULL,
+	.ReInitializeDSMCustomScan = NULL,
+	.InitializeWorkerCustomScan = NULL,
+	.ShutdownCustomScan = NULL,
+	.ExplainCustomScan = NULL
+};
+
+static set_join_pathlist_hook_type prev_set_join_pathlist_hook = NULL;
+
+
+void _PG_init(void);
+static void custom_hook(PlannerInfo *root, RelOptInfo *joinrel,
+							 RelOptInfo *outerrel, RelOptInfo *innerrel,
+							 JoinType jointype, JoinPathExtraData *extra);
+
+
+void
+_PG_init(void)
+{
+	prev_set_join_pathlist_hook = set_join_pathlist_hook;
+	set_join_pathlist_hook = custom_hook;
+}
+
+static void
+custom_hook(PlannerInfo *root, RelOptInfo *joinrel, RelOptInfo *outerrel,
+			RelOptInfo *innerrel, JoinType jointype,
+			JoinPathExtraData *extra)
+{
+	CustomPath *cpath;
+	Path	   *pathnode;
+	Path	   *src_path = NULL;
+	ListCell   *lc;
+
+	/*
+	 * Some extension intercept this hook earlier. Allow it to do a work
+	 * before us.
+	 */
+	if (prev_set_join_pathlist_hook)
+		(*prev_set_join_pathlist_hook)(root, joinrel, outerrel, innerrel,
+									   jointype, extra);
+
+	if (jointype != JOIN_INNER || is_dummy_rel(joinrel))
+		return;
+
+	/*
+	 * XXX: read up on the ROWID_VAR entity and what is the reason to suppress
+	 * this case.
+	 */
+	foreach(lc, joinrel->reltarget->exprs)
+	{
+		Expr *expr = lfirst(lc);
+
+		if (IsA(expr, Var) && ((Var *)expr)->varno == ROWID_VAR)
+			return;
+	}
+
+	foreach (lc, joinrel->pathlist)
+	{
+		Path *p = (Path *) lfirst(lc);
+
+		if (!IsA(p, HashPath) || p->param_info != NULL)
+			continue;
+		src_path = p;
+		break;
+	}
+
+	if (!src_path)
+		return;
+
+	cpath = makeNode(CustomPath);
+	cpath->path.type = T_CustomPath;
+	pathnode = &cpath->path;
+	pathnode->pathtype = T_CustomScan;
+	pathnode->parent = joinrel;
+	pathnode->pathtarget = joinrel->reltarget;
+	pathnode->param_info = NULL;
+	pathnode->parallel_aware = false;
+	pathnode->parallel_safe = false;
+	pathnode->parallel_workers = 0;
+	pathnode->pathkeys = NIL;
+
+	pathnode->rows = joinrel->rows;
+	pathnode->startup_cost = src_path->startup_cost;
+	pathnode->total_cost = src_path->total_cost;
+
+	cpath->methods = &path_methods;
+	cpath->custom_paths = list_make1((void *) src_path);
+	cpath->custom_private = NIL;
+
+	joinrel->pathlist = list_delete_ptr(joinrel->pathlist, src_path);
+	add_path(joinrel, (Path *) cpath);
+}
+
+static Plan *
+CreateCNPlan(PlannerInfo *root, RelOptInfo *rel, CustomPath *best_path,
+			   List *tlist,List *scan_clauses, List *custom_plans)
+{
+	CustomScan *cscan;
+
+	cscan = makeNode(CustomScan);
+	cscan->scan.plan.targetlist = cscan->custom_scan_tlist = tlist;
+	cscan->scan.scanrelid = 0;
+	cscan->custom_exprs = NIL;
+	cscan->custom_plans = custom_plans;
+	cscan->methods = &plan_methods;
+	cscan->flags = best_path->flags;
+	cscan->custom_private = best_path->custom_private;
+
+	return &cscan->scan.plan;
+}
+
+static Node *
+CreateCNScanState(CustomScan *cscan)
+{
+	CustomScanState	   *cstate = makeNode(CustomScanState);
+
+	cstate->methods = &exec_methods;
+
+	return (Node *) cstate;
+}
+
+static void
+BeginCNScan(CustomScanState *node, EState *estate, int eflags)
+{
+	CustomScan		   *cscan = (CustomScan *) node->ss.ps.plan;
+	ListCell		   *lc;
+
+	foreach (lc, cscan->custom_plans)
+	{
+		Plan	   *child = (Plan *) lfirst(lc);
+		PlanState  *child_state;
+
+		child_state = ExecInitNode(child, estate, eflags);
+		node->custom_ps = lappend(node->custom_ps, (void *) child_state);
+	}
+}
+
+static TupleTableSlot *
+ExecCNScan(CustomScanState *node)
+{
+	return ExecProcNode((PlanState *) list_nth(node->custom_ps, 0));
+}
+
+static void
+EndCNScan(CustomScanState *node)
+{
+	ListCell		   *lc;
+
+	ExecClearTuple(node->ss.ss_ScanTupleSlot);
+	foreach (lc, node->custom_ps)
+	{
+		ExecEndNode((PlanState *) lfirst(lc));
+	}
+}
+
+static void
+ReScanCNScan(CustomScanState *node)
+{
+	ListCell		   *lc;
+
+	ExecClearTuple(node->ss.ps.ps_ResultTupleSlot);
+
+	foreach (lc, node->custom_ps)
+	{
+		PlanState *child = (PlanState *) lfirst(lc);
+
+		/*
+		 * ExecReScan doesn't know about my subplans, so I have to do
+		 * changed-parameter signaling myself.
+		 */
+		if (node->ss.ps.chgParam != NULL)
+			UpdateChangedParamSet(child, node->ss.ps.chgParam);
+
+		ExecReScan(child);
+	}
+}
diff --git a/src/test/modules/test_custom_nodes/test_custom_nodes.control b/src/test/modules/test_custom_nodes/test_custom_nodes.control
new file mode 100644
index 0000000000..69d6997dd9
--- /dev/null
+++ b/src/test/modules/test_custom_nodes/test_custom_nodes.control
@@ -0,0 +1,5 @@
+# test_custom_nodes
+comment = 'Sample custom nodes'
+default_version = '1.0'
+module_pathname = '$libdir/test_custom_nodes'
+relocatable = true
-- 
2.42.0

