From 95c8b4834fd2276d7cbbf110950df8278526a047 Mon Sep 17 00:00:00 2001
From: Mehmet Yilmaz <mehmetyilmaz@microsoft.com>
Date: Fri, 14 Mar 2025 10:17:35 +0000
Subject: [PATCH] Enhance insert/select planner to handle nextval expressions
 correctly

---
 citus-tools                                   |  1 +
 .../planner/insert_select_planner.c           | 72 +++++++++++++++++++
 2 files changed, 73 insertions(+)
 create mode 160000 citus-tools

diff --git a/citus-tools b/citus-tools
new file mode 160000
index 000000000..3376bd684
--- /dev/null
+++ b/citus-tools
@@ -0,0 +1 @@
+Subproject commit 3376bd6845f0614908ed304f5033bd644c82d3bf
diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c
index ce61bd0ae..356de8c87 100644
--- a/src/backend/distributed/planner/insert_select_planner.c
+++ b/src/backend/distributed/planner/insert_select_planner.c
@@ -28,6 +28,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/fmgroids.h"
 
 #include "pg_version_constants.h"
 
@@ -107,6 +108,7 @@ static void ResetTargetEntryResno(List *targetList);
 static void ProcessEntryPair(TargetEntry *insertEntry, TargetEntry *selectEntry,
 							 Form_pg_attribute attr, int targetEntryIndex,
 							 List **projectedEntries, List **nonProjectedEntries);
+static bool IsNextvalExpression(Expr *expr);							 
 
 
 /* depth of current insert/select planner. */
@@ -1699,6 +1701,44 @@ ProcessEntryPair(TargetEntry *insertEntry, TargetEntry *selectEntry,
 				 Form_pg_attribute attr, int targetEntryIndex,
 				 List **projectedEntries, List **nonProjectedEntries)
 {
+	
+	/* Check if it's nextval(...) */
+	if (IsNextvalExpression((Expr *) selectEntry->expr))
+	{									
+
+		/*
+			* The real nextval(...) TLE belongs in nonProjectedEntries
+			* so that the coordinator (not the worker) will evaluate it.
+			*/
+		TargetEntry *coordinatorEntry = copyObject(selectEntry);
+
+		/*
+			* If there's a type mismatch, we might still do a cast. For now, assume
+			* nextval returns INT8. If your target column is INT4, you can do a cast
+			* as well. Or rely on subsequent code to do that if needed.
+			*/
+		if (attr->atttypid != exprType((Node *) coordinatorEntry->expr))
+		{
+			coordinatorEntry->expr =
+				CastExpr((Expr *) coordinatorEntry->expr,
+							exprType((Node *) coordinatorEntry->expr),
+							attr->atttypid,
+							attr->attcollation,
+							attr->atttypmod);
+		}
+
+		/* Optionally rename the coordinator entry so it won't conflict with placeholders */
+		coordinatorEntry->resname = pstrdup("coordinator_nextval");
+
+		/*
+			* (C) Add them to the respective lists:
+			*   - real nextval entry goes to nonProjectedEntries (the coordinator aggregator).
+			*/
+		*nonProjectedEntries = lappend(*nonProjectedEntries, coordinatorEntry);
+
+		return; /* done handling nextval */
+	}
+
 	Oid effectiveSourceType = exprType((Node *) selectEntry->expr);
 	Oid targetType = attr->atttypid;
 
@@ -1754,6 +1794,38 @@ ResetTargetEntryResno(List *targetList)
 }
 
 
+/*
+ * IsNextvalExpression returns true if `expr` is either a NextValueExpr
+ * or a FuncExpr calling nextval(...).
+ */
+static bool
+IsNextvalExpression(Expr *expr)
+{
+	if (expr == NULL)
+	{
+		return false;
+	}
+
+	/* e.g. identity column usage */
+	if (IsA(expr, NextValueExpr))
+	{
+		return true;
+	}
+
+	/* or an explicit nextval() call */
+	if (IsA(expr, FuncExpr))
+	{
+		FuncExpr *funcExpr = (FuncExpr *) expr;
+		if (funcExpr->funcid == F_NEXTVAL)  /* #include "utils/builtins.h" for F_NEXTVAL */
+		{
+			return true;
+		}
+	}
+
+	return false;
+}
+
+
 /*
  * Looks up the nextval(regclass) function in pg_proc, returning its actual
  * rettype. In a standard build, that will be INT8OID, but this is more robust.