Fix partition router running

Author: ildus

src/partition_overseer.c

@@ -68,25 +68,24 @@ partition_overseer_create_scan_state(CustomScan *node)
 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;
-            }
-        }
-    }
+    ModifyTableState   *mt_state = (ModifyTableState *) state;
+
+    if (!IsA(state, ModifyTableState))
+		return;
+
+	for (int 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
@@ -119,7 +118,7 @@ partition_overseer_exec(CustomScanState *node)
 	mt_plans_old = mt_state->mt_nplans;
 
 restart:
-	/* Fetch next tuple */
+	/* Run ModifyTable */
 	slot = ExecProcNode((PlanState *) mt_state);
 
 	/* Get current signal */
@@ -136,7 +135,7 @@ partition_overseer_exec(CustomScanState *node)
 		MTHackField(mt_state, mt_nplans) = mt_plans_old;
 		MTHackField(mt_state, mt_whichplan) = state_idx;
 
-		/* Restart ModifyTable */
+		/* Rerun ModifyTable */
 		goto restart;
 	}
 

src/partition_router.c

@@ -63,7 +63,6 @@ static TupleTableSlot *router_set_slot(PartitionRouterState *state,
 static TupleTableSlot *router_get_slot(PartitionRouterState *state,
 									   bool *should_process);
 
-static void router_lazy_init_junkfilter(PartitionRouterState *state);
 static void router_lazy_init_constraint(PartitionRouterState *state);
 
 static ItemPointerData router_extract_ctid(PartitionRouterState *state,
@@ -185,8 +184,9 @@ partition_router_exec(CustomScanState *node)
 
 		ItemPointerSetInvalid(&ctid);
 
-		/* Build new junkfilter lazily */
-		router_lazy_init_junkfilter(state);
+		/* Build new junkfilter if needed */
+		if (state->junkfilter == NULL)
+			state->junkfilter = state->current_rri->ri_junkFilter;
 
 		/* Build recheck constraint state lazily */
 		router_lazy_init_constraint(state);
@@ -257,15 +257,14 @@ router_set_slot(PartitionRouterState *state,
 	MTHackField(mt_state, mt_nplans) = -mt_state->mt_whichplan;
 	MTHackField(mt_state, operation) = operation;
 
+	/* HACK: disable AFTER STATEMENT triggers */
+	MTDisableStmtTriggers(mt_state, state);
+
 	if (!TupIsNull(slot))
 	{
 		/* We should've cached junk filter already */
 		Assert(state->junkfilter);
 
-		/* HACK: disable AFTER STATEMENT triggers */
-		MTDisableStmtTriggers(mt_state, state);
-
-
 		/* HACK: conditionally disable junk filter in result relation */
 		state->current_rri->ri_junkFilter = (operation == CMD_UPDATE) ?
 												state->junkfilter :
@@ -276,11 +275,6 @@ router_set_slot(PartitionRouterState *state,
 			slot->tts_tupleDescriptor);
 		ExecCopySlot(state->yielded_slot, slot);
 	}
-	else
-	{
-		/* HACK: enable AFTER STATEMENT triggers */
-		MTEnableStmtTriggers(mt_state, state);
-	}
 
 	/* Yield */
 	state->yielded = true;
@@ -324,13 +318,6 @@ router_get_slot(PartitionRouterState *state,
 	return slot;
 }
 
-static void
-router_lazy_init_junkfilter(PartitionRouterState *state)
-{
-	if (state->junkfilter == NULL)
-		state->junkfilter = state->current_rri->ri_junkFilter;
-}
-
 static void
 router_lazy_init_constraint(PartitionRouterState *state)
 {