using the same table with and without ONLY modifier in select query is prohibited

zilder · zilder · commit 4f367b78e531 · 2016-06-16T16:16:45.000+03:00
diff --git a/expected/pg_pathman.out b/expected/pg_pathman.out
@@ -116,6 +116,8 @@ SELECT COUNT(*) FROM ONLY test.num_range_rel;
      0
 (1 row)
 
+SELECT * FROM ONLY test.range_rel UNION SELECT * FROM test.range_rel;
+ERROR:  It is prohibited to query partitioned tables both with and without ONLY modifier
 SET pg_pathman.enable_runtimeappend = OFF;
 SET pg_pathman.enable_runtimemergeappend = OFF;
 VACUUM;
diff --git a/sql/pg_pathman.sql b/sql/pg_pathman.sql
@@ -46,6 +46,8 @@ INSERT INTO test.num_range_rel
 SELECT COUNT(*) FROM test.num_range_rel;
 SELECT COUNT(*) FROM ONLY test.num_range_rel;
 
+SELECT * FROM ONLY test.range_rel UNION SELECT * FROM test.range_rel;
+
 SET pg_pathman.enable_runtimeappend = OFF;
 SET pg_pathman.enable_runtimemergeappend = OFF;
 
diff --git a/src/hooks.c b/src/hooks.c
@@ -161,7 +161,8 @@ pathman_rel_pathlist_hook(PlannerInfo *root, RelOptInfo *rel, Index rti, RangeTb
 		return;
 
 	/* This works only for SELECT queries (at least for now) */
-	if (root->parse->commandType != CMD_SELECT || !inheritance_disabled)
+	if (root->parse->commandType != CMD_SELECT ||
+		!list_member_oid(inheritance_enabled_relids, rte->relid))
 		return;
 
 	/* Lookup partitioning information for parent relation */
diff --git a/src/pathman.h b/src/pathman.h
@@ -126,7 +126,16 @@ typedef struct PathmanState
 	DsmArray	databases;
 } PathmanState;
 
-extern bool				inheritance_disabled;
+/*
+ * The list of partitioned relation relids that must be handled by pg_pathman
+ */
+extern List			   *inheritance_enabled_relids;
+/*
+ * This list is used to ensure that partitioned relation isn't used both
+ * with and without ONLY modifiers
+ */
+extern List			   *inheritance_disabled_relids;
+
 extern bool 			pg_pathman_enable;
 extern PathmanState    *pmstate;
 
diff --git a/src/pg_pathman.c b/src/pg_pathman.c
@@ -46,7 +46,8 @@
 
 PG_MODULE_MAGIC;
 
-bool			inheritance_disabled;
+List		   *inheritance_disabled_relids = NIL;
+List		   *inheritance_enabled_relids = NIL;
 bool			pg_pathman_enable;
 PathmanState   *pmstate;
 
@@ -266,6 +267,11 @@ pathman_post_parse_analysis_hook(ParseState *pstate, Query *query)
 
 	if (post_parse_analyze_hook_original)
 		post_parse_analyze_hook_original(pstate, query);
+
+	// list_free(inheritance_disabled_relids);
+	// list_free(inheritance_enabled_relids);
+	inheritance_disabled_relids = NIL;
+	inheritance_enabled_relids = NIL;
 }
 
 /*
@@ -280,7 +286,7 @@ pathman_planner_hook(Query *parse, int cursorOptions, ParamListInfo boundParams)
 
 	if (pg_pathman_enable)
 	{
-		inheritance_disabled = false;
+		// inheritance_disabled = false;
 		switch(parse->commandType)
 		{
 			case CMD_SELECT:
@@ -303,6 +309,9 @@ pathman_planner_hook(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	else
 		result = standard_planner(parse, cursorOptions, boundParams);
 
+	list_free(inheritance_disabled_relids);
+	inheritance_disabled_relids = NIL;
+
 	return result;
 }
 
@@ -327,6 +336,7 @@ disable_inheritance(Query *parse)
 	foreach(lc, parse->rtable)
 	{
 		rte = (RangeTblEntry*) lfirst(lc);
+
 		switch(rte->rtekind)
 		{
 			case RTE_RELATION:
@@ -343,9 +353,28 @@ disable_inheritance(Query *parse)
 						 * when user uses ONLY statement from case when we
 						 * make rte->inh false intentionally.
 						 */
-						inheritance_disabled = true;
+						inheritance_enabled_relids = \
+							lappend_oid(inheritance_enabled_relids, rte->relid);
+
+						/*
+						 * Check if relation was already found with ONLY modifier. In
+						 * this case throw an error because we cannot handle
+						 * situations when partitioned table used both with and
+						 * without ONLY modifier in SELECT queries
+						 */
+						if (list_member_oid(inheritance_disabled_relids, rte->relid))
+							goto disable_error;
+
+						goto disable_next;
 					}
 				}
+
+				inheritance_disabled_relids = \
+					lappend_oid(inheritance_disabled_relids, rte->relid);
+
+				/* Check if relation was already found withoud ONLY modifier */
+				if (list_member_oid(inheritance_enabled_relids, rte->relid))
+						goto disable_error;
 				break;
 			case RTE_SUBQUERY:
 				/* Recursively disable inheritance for subqueries */
@@ -354,7 +383,16 @@ disable_inheritance(Query *parse)
 			default:
 				break;
 		}
+
+disable_next:
+		;
 	}
+
+	return;
+
+disable_error:
+	elog(ERROR, "It is prohibited to query partitioned tables both "
+				"with and without ONLY modifier");
 }
 
 static void

Original file line number	Diff line number	Diff line change
`@@ -116,6 +116,8 @@ SELECT COUNT(*) FROM ONLY test.num_range_rel;`
`116`	`116`	`0`
`117`	`117`	`(1 row)`
`118`	`118`
	`119`	`+SELECT * FROM ONLY test.range_rel UNION SELECT * FROM test.range_rel;`
	`120`	`+ERROR: It is prohibited to query partitioned tables both with and without ONLY modifier`
`119`	`121`	`SET pg_pathman.enable_runtimeappend = OFF;`
`120`	`122`	`SET pg_pathman.enable_runtimemergeappend = OFF;`
`121`	`123`	`VACUUM;`