Add tests for modifying CTE and SELECT without FROM

src/test/regress/expected/multi_mx_router_planner.out

@@ -288,6 +288,31 @@ DEBUG:  query has a single distribution column value: 1
 ---------------------------------------------------------------------
 (0 rows)
 
+WITH update_article AS (
+    UPDATE articles_hash_mx SET word_count = 11 WHERE id = 1 AND word_count = 10 RETURNING *
+)
+SELECT coalesce(1,random());
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  generating subplan XXX_1 for CTE update_article: UPDATE public.articles_hash_mx SET word_count = 11 WHERE ((id OPERATOR(pg_catalog.=) 1) AND (word_count OPERATOR(pg_catalog.=) 10)) RETURNING id, author_id, title, word_count
+DEBUG:  Creating router plan
+DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT COALESCE((1)::double precision, random()) AS "coalesce"
+DEBUG:  Creating router plan
+ coalesce
+---------------------------------------------------------------------
+        1
+(1 row)
+
+WITH update_article AS (
+    UPDATE articles_hash_mx SET word_count = 10 WHERE author_id = 1 AND id = 1 AND word_count = 11 RETURNING *
+)
+SELECT coalesce(1,random());
+DEBUG:  Creating router plan
+DEBUG:  query has a single distribution column value: 1
+ coalesce
+---------------------------------------------------------------------
+        1
+(1 row)
+
 -- recursive CTEs are supported when filtered on partition column
 INSERT INTO company_employees_mx values(1, 1, 0);
 DEBUG:  Creating router plan

src/test/regress/expected/multi_router_planner.out

@@ -507,6 +507,45 @@ DEBUG:  Creating router plan
   1 |         1 | arsenous |         10
 (1 row)
 
+WITH update_article AS (
+    UPDATE articles_hash SET word_count = 11 WHERE id = 1 AND word_count = 10 RETURNING *
+)
+SELECT coalesce(1,random());
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  generating subplan XXX_1 for CTE update_article: UPDATE public.articles_hash SET word_count = 11 WHERE ((id OPERATOR(pg_catalog.=) 1) AND (word_count OPERATOR(pg_catalog.=) 10)) RETURNING id, author_id, title, word_count
+DEBUG:  Creating router plan
+DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT COALESCE((1)::double precision, random()) AS "coalesce"
+DEBUG:  Creating router plan
+ coalesce
+---------------------------------------------------------------------
+        1
+(1 row)
+
+WITH update_article AS (
+    UPDATE articles_hash SET word_count = 10 WHERE author_id = 1 AND id = 1 AND word_count = 11 RETURNING *
+)
+SELECT coalesce(1,random());
+DEBUG:  Creating router plan
+DEBUG:  query has a single distribution column value: 1
+ coalesce
+---------------------------------------------------------------------
+        1
+(1 row)
+
+WITH update_article AS (
+    UPDATE authors_reference SET name = '' WHERE id = 0 RETURNING *
+)
+SELECT coalesce(1,random());
+DEBUG:  cannot router plan modification of a non-distributed table
+DEBUG:  generating subplan XXX_1 for CTE update_article: UPDATE public.authors_reference SET name = ''::character varying WHERE (id OPERATOR(pg_catalog.=) 0) RETURNING name, id
+DEBUG:  Creating router plan
+DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT COALESCE((1)::double precision, random()) AS "coalesce"
+DEBUG:  Creating router plan
+ coalesce
+---------------------------------------------------------------------
+        1
+(1 row)
+
 WITH delete_article AS (
     DELETE FROM articles_hash WHERE id = 1 AND word_count = 10 RETURNING *
 )

src/test/regress/sql/multi_mx_router_planner.sql

@@ -147,6 +147,16 @@ WITH id_author AS ( SELECT id, author_id FROM articles_hash_mx WHERE author_id =
 id_title AS (SELECT id, title from articles_hash_mx WHERE author_id = 2)
 SELECT * FROM id_author, id_title WHERE id_author.id = id_title.id;
 
+WITH update_article AS (
+    UPDATE articles_hash_mx SET word_count = 11 WHERE id = 1 AND word_count = 10 RETURNING *
+)
+SELECT coalesce(1,random());
+
+WITH update_article AS (
+    UPDATE articles_hash_mx SET word_count = 10 WHERE author_id = 1 AND id = 1 AND word_count = 11 RETURNING *
+)
+SELECT coalesce(1,random());
+
 -- recursive CTEs are supported when filtered on partition column
 
 INSERT INTO company_employees_mx values(1, 1, 0);

src/test/regress/sql/multi_router_planner.sql

@@ -280,6 +280,21 @@ WITH update_article AS (
 )
 SELECT * FROM update_article;
 
+WITH update_article AS (
+    UPDATE articles_hash SET word_count = 11 WHERE id = 1 AND word_count = 10 RETURNING *
+)
+SELECT coalesce(1,random());
+
+WITH update_article AS (
+    UPDATE articles_hash SET word_count = 10 WHERE author_id = 1 AND id = 1 AND word_count = 11 RETURNING *
+)
+SELECT coalesce(1,random());
+
+WITH update_article AS (
+    UPDATE authors_reference SET name = '' WHERE id = 0 RETURNING *
+)
+SELECT coalesce(1,random());
+
 WITH delete_article AS (
     DELETE FROM articles_hash WHERE id = 1 AND word_count = 10 RETURNING *
 )