PartitionFilter now transforms INSERT statements within subplans

funbringer · funbringer · commit c63c9a16d00d · 2016-05-27T01:12:40.000+03:00
diff --git a/src/partition_filter.c b/src/partition_filter.c
@@ -40,7 +40,7 @@ init_partition_filter_static_data(void)
 }
 
 Plan *
-make_partition_filter_plan(Plan *subplan, Oid partitioned_table,
+make_partition_filter(Plan *subplan, Oid partitioned_table,
 						   OnConflictAction	conflict_action)
 {
 	CustomScan *cscan = makeNode(CustomScan);
@@ -265,27 +265,73 @@ pfilter_build_tlist(List *tlist)
 	return result_tlist;
 }
 
-/* Add proxy PartitionFilter nodes to subplans of ModifyTable node */
 void
-add_partition_filters(List *rtable, ModifyTable *modify_table)
+add_partition_filters(List *rtable, Plan *plan)
 {
-	ListCell *lc1,
-			 *lc2;
+	ListCell   *l;
 
-	Assert(IsA(modify_table, ModifyTable));
-
-	if (!pg_pathman_enable_partition_filter)
+	if (plan == NULL || !pg_pathman_enable_partition_filter)
 		return;
 
-	forboth (lc1, modify_table->plans, lc2, modify_table->resultRelations)
+	/* Plan-type-specific fixes*/
+	switch (nodeTag(plan))
 	{
-		Index				rindex = lfirst_int(lc2);
-		Oid					relid = getrelid(rindex, rtable);
-		PartRelationInfo   *prel = get_pathman_relation_info(relid, NULL);
-
-		if (prel)
-			lfirst(lc1) = make_partition_filter_plan((Plan *) lfirst(lc1),
-													 relid,
-													 modify_table->onConflictAction);
+		case T_SubqueryScan:
+			add_partition_filters(rtable, ((SubqueryScan *) plan)->subplan);
+			break;
+
+		case T_CustomScan:
+			foreach(l, ((CustomScan *) plan)->custom_plans)
+				add_partition_filters(rtable, (Plan *) lfirst(l));
+			break;
+
+		/*
+		 * Add proxy PartitionFilter nodes
+		 * to subplans of ModifyTable node
+		 */
+		case T_ModifyTable:
+			{
+				ModifyTable	   *modify_table = ((ModifyTable *) plan);
+				ListCell	   *lc1,
+							   *lc2;
+
+				forboth (lc1, modify_table->plans, lc2, modify_table->resultRelations)
+				{
+					Index				rindex = lfirst_int(lc2);
+					Oid					relid = getrelid(rindex, rtable);
+					PartRelationInfo   *prel = get_pathman_relation_info(relid, NULL);
+
+					add_partition_filters(rtable, (Plan *) lfirst(lc1));
+
+					if (prel)
+						lfirst(lc1) = make_partition_filter((Plan *) lfirst(lc1),
+															relid,
+															modify_table->onConflictAction);
+				}
+			}
+			break;
+
+		/* Since they look alike */
+		case T_MergeAppend:
+		case T_Append:
+			foreach(l, ((Append *) plan)->appendplans)
+				add_partition_filters(rtable, (Plan *) lfirst(l));
+			break;
+
+		case T_BitmapAnd:
+			foreach(l, ((BitmapAnd *) plan)->bitmapplans)
+				add_partition_filters(rtable, (Plan *) lfirst(l));
+			break;
+
+		case T_BitmapOr:
+			foreach(l, ((BitmapOr *) plan)->bitmapplans)
+				add_partition_filters(rtable, (Plan *) lfirst(l));
+			break;
+
+		default:
+			break;
 	}
+
+	add_partition_filters(rtable, plan->lefttree);
+	add_partition_filters(rtable, plan->righttree);
 }
diff --git a/src/partition_filter.h b/src/partition_filter.h
@@ -37,12 +37,12 @@ extern CustomScanMethods	partition_filter_plan_methods;
 extern CustomExecMethods	partition_filter_exec_methods;
 
 
-void add_partition_filters(List *rtable, ModifyTable *modify_table);
+void add_partition_filters(List *rtable, Plan *plan);
 
 void init_partition_filter_static_data(void);
 
-Plan * make_partition_filter_plan(Plan *subplan, Oid partitioned_table,
-								  OnConflictAction	conflict_action);
+Plan * make_partition_filter(Plan *subplan, Oid partitioned_table,
+							 OnConflictAction conflict_action);
 
 Node * partition_filter_create_scan_state(CustomScan *node);
 
diff --git a/src/pg_pathman.c b/src/pg_pathman.c
@@ -254,10 +254,17 @@ pathman_planner_hook(Query *parse, int cursorOptions, ParamListInfo boundParams)
 				handle_modification_query(parse);
 				break;
 			case CMD_INSERT:
+			{
+				ListCell *lc;
+
 				result = standard_planner(parse, cursorOptions, boundParams);
-				add_partition_filters(result->rtable,
-									  (ModifyTable *) result->planTree);
+
+				add_partition_filters(result->rtable, result->planTree);
+				foreach (lc, result->subplans)
+					add_partition_filters(result->rtable, (Plan *) lfirst(lc));
+
 				return result;
+			}
 			default:
 				break;
 		}
