Fix updating using Overseer node (still has errors)

Author: ildus

expected/pathman_update_node.out

@@ -15,29 +15,31 @@ SELECT create_range_partitions('test_update_node.test_range', 'val', 1, 10);
 
 /* Moving from 2st to 1st partition */
 EXPLAIN (COSTS OFF) UPDATE test_update_node.test_range SET val = 5 WHERE val = 15;
-                            QUERY PLAN                             
--------------------------------------------------------------------
- Update on test_range_2
-   ->  Custom Scan (PartitionFilter)
-         ->  Custom Scan (PartitionRouter)
-               ->  Bitmap Heap Scan on test_range_2
-                     Recheck Cond: (val = '15'::numeric)
-                     ->  Bitmap Index Scan on test_range_2_val_idx
-                           Index Cond: (val = '15'::numeric)
-(7 rows)
+                               QUERY PLAN                                
+-------------------------------------------------------------------------
+ Custom Scan (PartitionOverseer)
+   ->  Update on test_range_2
+         ->  Custom Scan (PartitionFilter)
+               ->  Custom Scan (PartitionRouter)
+                     ->  Bitmap Heap Scan on test_range_2
+                           Recheck Cond: (val = '15'::numeric)
+                           ->  Bitmap Index Scan on test_range_2_val_idx
+                                 Index Cond: (val = '15'::numeric)
+(8 rows)
 
 /* Keep same partition */
 EXPLAIN (COSTS OFF) UPDATE test_update_node.test_range SET val = 14 WHERE val = 15;
-                            QUERY PLAN                             
--------------------------------------------------------------------
- Update on test_range_2
-   ->  Custom Scan (PartitionFilter)
-         ->  Custom Scan (PartitionRouter)
-               ->  Bitmap Heap Scan on test_range_2
-                     Recheck Cond: (val = '15'::numeric)
-                     ->  Bitmap Index Scan on test_range_2_val_idx
-                           Index Cond: (val = '15'::numeric)
-(7 rows)
+                               QUERY PLAN                                
+-------------------------------------------------------------------------
+ Custom Scan (PartitionOverseer)
+   ->  Update on test_range_2
+         ->  Custom Scan (PartitionFilter)
+               ->  Custom Scan (PartitionRouter)
+                     ->  Bitmap Heap Scan on test_range_2
+                           Recheck Cond: (val = '15'::numeric)
+                           ->  Bitmap Index Scan on test_range_2_val_idx
+                                 Index Cond: (val = '15'::numeric)
+(8 rows)
 
 /* Update values in 1st partition (rows remain there) */
 UPDATE test_update_node.test_range SET val = 5 WHERE val <= 10;

src/include/partition_router.h

@@ -56,43 +56,27 @@ extern CustomScanMethods	partition_router_plan_methods;
 extern CustomExecMethods	partition_router_exec_methods;
 
 
-#define IsPartitionRouterPlan(node) \
-	( \
-		IsA((node), CustomScan) && \
-		(((CustomScan *) (node))->methods == &partition_router_plan_methods) \
-	)
-
 #define IsPartitionRouterState(node) \
 	( \
 		IsA((node), CustomScanState) && \
 		(((CustomScanState *) (node))->methods == &partition_router_exec_methods) \
 	)
 
-#define IsPartitionRouter(node) \
-	( IsPartitionRouterPlan(node) || IsPartitionRouterState(node) )
-
+/* Highlight hacks with ModifyTable's fields */
+#define MTHackField(mt_state, field) ( (mt_state)->field )
 
 void init_partition_router_static_data(void);
-
-Plan *make_partition_router(Plan *subplan, int epq_param);
-
-void prepare_modify_table_for_partition_router(PlanState *state, void *context);
-
-
-Node *partition_router_create_scan_state(CustomScan *node);
-
+void prepare_modify_table_for_partition_router(PlanState *state,
+											   void *context);
 void partition_router_begin(CustomScanState *node, EState *estate, int eflags);
-
-TupleTableSlot *partition_router_exec(CustomScanState *node);
-
 void partition_router_end(CustomScanState *node);
-
 void partition_router_rescan(CustomScanState *node);
-
 void partition_router_explain(CustomScanState *node,
 							  List *ancestors,
 							  ExplainState *es);
 
-TupleTableSlot *partition_router_run_modify_table(PlanState *state);
+Plan *make_partition_router(Plan *subplan, int epq_param);
+Node *partition_router_create_scan_state(CustomScan *node);
+TupleTableSlot *partition_router_exec(CustomScanState *node);
 
 #endif /* PARTITION_UPDATE_H */

src/partition_overseer.c

@@ -1,8 +1,9 @@
 #include "postgres.h"
 
-#include "partition_overseer.h"
 #include "partition_filter.h"
+#include "partition_overseer.h"
 #include "partition_router.h"
+#include "planner_tree_modification.h"
 
 CustomScanMethods	partition_overseer_plan_methods;
 CustomExecMethods	partition_overseer_exec_methods;
@@ -64,6 +65,30 @@ partition_overseer_create_scan_state(CustomScan *node)
 	return (Node *) state;
 }
 
+static void
+set_mt_state_for_router(PlanState *state, void *context)
+{
+    if (IsA(state, ModifyTableState))
+    {
+        ModifyTableState   *mt_state = (ModifyTableState *) state;
+        int                 i;
+
+        for (i = 0; i < mt_state->mt_nplans; i++)
+        {
+            CustomScanState        *pf_state = (CustomScanState *) mt_state->mt_plans[i];
+            PartitionRouterState   *pr_state;
+
+            /* Check if this is a PartitionFilter + PartitionRouter combo */
+            if (IsPartitionFilterState(pf_state) &&
+                IsPartitionRouterState(pr_state = linitial(pf_state->custom_ps)))
+            {
+                /* HACK: point to ModifyTable in PartitionRouter */
+                pr_state->mt_state = mt_state;
+            }
+        }
+    }
+}
+
 void
 partition_overseer_begin(CustomScanState *node,
 						 EState *estate,
@@ -74,13 +99,48 @@ partition_overseer_begin(CustomScanState *node,
 
 	/* It's convenient to store PlanState in 'custom_ps' */
 	node->custom_ps = list_make1(ExecInitNode(plan, estate, eflags));
+
+	/* Save ModifyTableState in PartitionRouterState structs */
+	state_tree_visitor((PlanState *) linitial(node->custom_ps),
+						set_mt_state_for_router,
+						NULL);
 }
 
 TupleTableSlot *
 partition_overseer_exec(CustomScanState *node)
 {
-	PlanState *state = linitial(node->custom_ps);
-	return partition_router_run_modify_table(state);
+	ModifyTableState *mt_state = linitial(node->custom_ps);
+
+	TupleTableSlot	   *slot;
+	int					mt_plans_old,
+						mt_plans_new;
+
+	/* Get initial signal */
+	mt_plans_old = mt_state->mt_nplans;
+
+restart:
+	/* Fetch next tuple */
+	slot = ExecProcNode((PlanState *) mt_state);
+
+	/* Get current signal */
+	mt_plans_new = MTHackField(mt_state, mt_nplans);
+
+	/* Did PartitionRouter ask us to restart? */
+	if (mt_plans_new != mt_plans_old)
+	{
+		/* Signal points to current plan */
+		int state_idx = -mt_plans_new;
+
+		/* HACK: partially restore ModifyTable's state */
+		MTHackField(mt_state, mt_done) = false;
+		MTHackField(mt_state, mt_nplans) = mt_plans_old;
+		MTHackField(mt_state, mt_whichplan) = state_idx;
+
+		/* Restart ModifyTable */
+		goto restart;
+	}
+
+	return slot;
 }
 
 void
@@ -101,5 +161,5 @@ partition_overseer_explain(CustomScanState *node,
 						   List *ancestors,
 						   ExplainState *es)
 {
-	/* nothing to do */
+	/* Nothing to do here now */
 }

src/partition_router.c

@@ -26,10 +26,6 @@
 #include "utils/rel.h"
 
 
-/* Highlight hacks with ModifyTable's fields */
-#define MTHackField(mt_state, field) ( (mt_state)->field )
-
-
 #define MTDisableStmtTriggers(mt_state, pr_state) \
 	do { \
 		TriggerDesc *triggers = (mt_state)->resultRelInfo->ri_TrigDesc; \
@@ -143,7 +139,6 @@ partition_router_create_scan_state(CustomScan *node)
 	state = (PartitionRouterState *) palloc0(sizeof(PartitionRouterState));
 	NodeSetTag(state, T_CustomScanState);
 
-	state = (PartitionRouterState *) makeNode(CustomScanState);
 	state->css.flags = node->flags;
 	state->css.methods = &partition_router_exec_methods;
 
@@ -246,46 +241,6 @@ partition_router_explain(CustomScanState *node,
 	/* Nothing to do here now */
 }
 
-
-/* Smart wrapper over ModifyTable */
-TupleTableSlot *
-partition_router_run_modify_table(PlanState *state)
-{
-	ModifyTableState   *mt_state;
-	TupleTableSlot	   *slot;
-	int					mt_plans_old,
-						mt_plans_new;
-
-	mt_state = (ModifyTableState *) state;
-
-	/* Get initial signal */
-	mt_plans_old = mt_state->mt_nplans;
-
-restart:
-	/* Fetch next tuple */
-	slot = ExecProcNode(state);
-
-	/* Get current signal */
-	mt_plans_new = MTHackField(mt_state, mt_nplans);
-
-	/* Did PartitionRouter ask us to restart? */
-	if (mt_plans_new != mt_plans_old)
-	{
-		/* Signal points to current plan */
-		int state_idx = -mt_plans_new;
-
-		/* HACK: partially restore ModifyTable's state */
-		MTHackField(mt_state, mt_done) = false;
-		MTHackField(mt_state, mt_nplans) = mt_plans_old;
-		MTHackField(mt_state, mt_whichplan) = state_idx;
-
-		/* Restart ModifyTable */
-		goto restart;
-	}
-
-	return slot;
-}
-
 /* Return tuple OR yield it and change ModifyTable's operation */
 static TupleTableSlot *
 router_set_slot(PartitionRouterState *state,

src/planner_tree_modification.c

@@ -765,11 +765,12 @@ partition_filter_visitor(Plan *plan, void *context)
 static Plan *
 partition_router_visitor(Plan *plan, void *context)
 {
-	List			*rtable = (List *) context;
-	ModifyTable		*modify_table = (ModifyTable *) plan;
-	ListCell		*lc1,
-					*lc2,
-					*lc3;
+	List		   *rtable = (List *) context;
+	ModifyTable	   *modify_table = (ModifyTable *) plan;
+	ListCell	   *lc1,
+				   *lc2,
+				   *lc3;
+	bool			changed = false;
 
 	/* Skip if not ModifyTable with 'UPDATE' command */
 	if (!IsA(modify_table, ModifyTable) || modify_table->operation != CMD_UPDATE)
@@ -821,10 +822,14 @@ partition_router_visitor(Plan *plan, void *context)
 											returning_list);
 
 			lfirst(lc1) = pfilter;
+			changed = true;
 		}
 	}
 
-	return make_partition_overseer(plan);
+	if (changed)
+		return make_partition_overseer(plan);
+
+	return NULL;
 }