Merge pull request #1106 from citusdata/error_out_cte_with_modify

Error out on CTEs with data modifying statement

src/backend/distributed/planner/multi_router_planner.c

@@ -127,7 +127,7 @@ static void ErrorIfInsertPartitionColumnDoesNotMatchSelect(Query *query,
 														   Oid *
 														   selectPartitionColumnTableId);
 static void AddUninstantiatedEqualityQual(Query *query, Var *targetPartitionColumnVar);
-
+static void ErrorIfQueryHasModifyingCTE(Query *queryTree);
 
 /*
  * MultiRouterPlanCreate creates a multi plan for the queries
@@ -202,8 +202,7 @@ CreateSingleTaskRouterPlan(Query *originalQuery, Query *query,
 	}
 	else
 	{
-		Assert(commandType == CMD_SELECT);
+		ErrorIfQueryHasModifyingCTE(query);
-
 		task = RouterSelectTask(originalQuery, restrictionContext, &placementList);
 	}
 
@@ -2850,3 +2849,43 @@ InstantiatePartitionQual(Node *node, void *context)
 
 	return expression_tree_mutator(node, InstantiatePartitionQual, context);
 }
+
+
+/*
+ * ErrorIfQueryHasModifyingCTE checks if the query contains modifying common table
+ * expressions and errors out if it does.
+ */
+static void
+ErrorIfQueryHasModifyingCTE(Query *queryTree)
+{
+	ListCell *cteCell = NULL;
+
+	Assert(queryTree->commandType == CMD_SELECT);
+
+	/* we do not need to do anything if there are no CTEs */
+	if (queryTree->cteList == NIL)
+	{
+		return;
+	}
+
+	foreach(cteCell, queryTree->cteList)
+	{
+		CommonTableExpr *cte = (CommonTableExpr *) lfirst(cteCell);
+		Query *cteQuery = (Query *) cte->ctequery;
+
+		/*
+		 * Here we only check for command type of top level query. Normally there can be
+		 * nested CTE, however PostgreSQL dictates that data-modifying statements must
+		 * be at top level of CTE. Therefore it is OK to just check for top level.
+		 * Similarly, we do not need to check for subqueries.
+		 */
+		if (cteQuery->commandType != CMD_SELECT)
+		{
+			ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							errmsg("cannot perform distributed planning for the given "
+								   "modification"),
+							errdetail("Data-modifying statements are not supported in "
+									  "the WITH clauses of distributed queries.")));
+		}
+	}
+}

src/test/regress/expected/multi_router_planner.out

@@ -439,6 +439,34 @@ DEBUG:  predicate pruning for shardId 840004
 DEBUG:  predicate pruning for shardId 840005
 ERROR:  cannot perform distributed planning on this query
 DETAIL:  Complex table expressions are currently unsupported
+-- CTE with queries other than SELECT is not supported
+WITH new_article AS (
+    INSERT INTO articles_hash VALUES (1,  1, 'arsenous', 9572) RETURNING *
+)
+SELECT * FROM new_article;
+ERROR:  cannot perform distributed planning for the given modification
+DETAIL:  Data-modifying statements are not supported in the WITH clauses of distributed queries.
+-- Modifying statement in nested CTE case is covered by PostgreSQL itself
+WITH new_article AS (
+    WITH nested_cte AS (
+        INSERT INTO articles_hash VALUES (1,  1, 'arsenous', 9572) RETURNING *
+    )
+    SELECT * FROM nested_cte
+)
+SELECT * FROM new_article;
+ERROR:  WITH clause containing a data-modifying statement must be at the top level
+LINE 2:     WITH nested_cte AS (
+                 ^
+-- Modifying statement in a CTE in subquwey is also covered by PostgreSQL
+SELECT * FROM (
+    WITH new_article AS (
+        INSERT INTO articles_hash VALUES (1,  1, 'arsenous', 9572) RETURNING *
+    )
+    SELECT * FROM new_article
+) AS subquery_cte;
+ERROR:  WITH clause containing a data-modifying statement must be at the top level
+LINE 2:     WITH new_article AS (
+                 ^
 -- grouping sets are supported on single shard
 SELECT
 	id, substring(title, 2, 1) AS subtitle, count(*)

src/test/regress/sql/multi_router_planner.sql

@@ -226,6 +226,29 @@ WITH RECURSIVE hierarchy as (
 				ce.company_id = 2))
 SELECT * FROM hierarchy WHERE LEVEL <= 2;
 
+-- CTE with queries other than SELECT is not supported
+WITH new_article AS (
+    INSERT INTO articles_hash VALUES (1,  1, 'arsenous', 9572) RETURNING *
+)
+SELECT * FROM new_article;
+
+-- Modifying statement in nested CTE case is covered by PostgreSQL itself
+WITH new_article AS (
+    WITH nested_cte AS (
+        INSERT INTO articles_hash VALUES (1,  1, 'arsenous', 9572) RETURNING *
+    )
+    SELECT * FROM nested_cte
+)
+SELECT * FROM new_article;
+
+-- Modifying statement in a CTE in subquwey is also covered by PostgreSQL
+SELECT * FROM (
+    WITH new_article AS (
+        INSERT INTO articles_hash VALUES (1,  1, 'arsenous', 9572) RETURNING *
+    )
+    SELECT * FROM new_article
+) AS subquery_cte;
+
 -- grouping sets are supported on single shard
 SELECT
 	id, substring(title, 2, 1) AS subtitle, count(*)