From f885423a4a1c95118e026042a878226749d93441 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ra=C3=BAl=20Mar=C3=ADn?= <git@rmr.ninja>
Date: Mon, 26 Feb 2024 15:09:05 +0100
Subject: [PATCH] Cleanup and more tests

---
 src/Interpreters/InterpreterSelectQuery.cpp   |  3 +-
 .../02972_parallel_replicas_cte.sql           | 42 ++++++++++++++++++-
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/Interpreters/InterpreterSelectQuery.cpp b/src/Interpreters/InterpreterSelectQuery.cpp
index 135b535595c..1a9827d30f8 100644
--- a/src/Interpreters/InterpreterSelectQuery.cpp
+++ b/src/Interpreters/InterpreterSelectQuery.cpp
@@ -778,6 +778,7 @@ InterpreterSelectQuery::InterpreterSelectQuery(
     };
 
     /// This is a hack to make sure we reanalyze if GlobalSubqueriesVisitor changed allow_experimental_parallel_reading_from_replicas
+    /// inside the query context (because it doesn't have write access to the main context)
     UInt64 parallel_replicas_before_analysis
         = context->hasQueryContext() ? context->getQueryContext()->getSettingsRef().allow_experimental_parallel_reading_from_replicas : 0;
     analyze(shouldMoveToPrewhere());
@@ -787,7 +788,7 @@ InterpreterSelectQuery::InterpreterSelectQuery(
 
     if (context->hasQueryContext())
     {
-        /// No buts or ifs, if the analysis changed this setting we must reanalyze without parallel replicas
+        /// As this query can't be executed with parallel replicas, we must reanalyze it
         if (context->getQueryContext()->getSettingsRef().allow_experimental_parallel_reading_from_replicas
             != parallel_replicas_before_analysis)
         {
diff --git a/tests/queries/0_stateless/02972_parallel_replicas_cte.sql b/tests/queries/0_stateless/02972_parallel_replicas_cte.sql
index c7143b5aa93..c9ab83ff9ad 100644
--- a/tests/queries/0_stateless/02972_parallel_replicas_cte.sql
+++ b/tests/queries/0_stateless/02972_parallel_replicas_cte.sql
@@ -1,5 +1,6 @@
 DROP TABLE IF EXISTS pr_1;
 DROP TABLE IF EXISTS pr_2;
+DROP TABLE IF EXISTS numbers_1e6;
 
 CREATE TABLE pr_1 (`a` UInt32) ENGINE = MergeTree ORDER BY a PARTITION BY a % 10 AS
 SELECT 10 * intDiv(number, 10) + 1 FROM numbers(1_000_000);
@@ -28,7 +29,7 @@ SETTINGS allow_experimental_analyzer = 0, allow_experimental_parallel_reading_fr
 SELECT count() FROM pr_2 JOIN numbers(10) as pr_1 ON pr_2.a = pr_1.number
 SETTINGS allow_experimental_parallel_reading_from_replicas = 1, parallel_replicas_for_non_replicated_merge_tree = 1, cluster_for_parallel_replicas = 'test_cluster_one_shard_three_replicas_localhost', max_parallel_replicas = 3;
 
--- Being a subquery should still disable parallel replicas
+-- Parallel replicas detection should work inside subqueries
 SELECT *
 FROM
 (
@@ -37,5 +38,44 @@ FROM
 )
 SETTINGS allow_experimental_parallel_reading_from_replicas = 1, parallel_replicas_for_non_replicated_merge_tree = 1, cluster_for_parallel_replicas = 'test_cluster_one_shard_three_replicas_localhost', max_parallel_replicas = 3;
 
+-- Subquery + subquery
+SELECT count()
+FROM
+(
+    SELECT c + 1
+    FROM
+    (
+        WITH filtered_groups AS (SELECT a FROM pr_1 WHERE a >= 10000)
+        SELECT count() as c FROM pr_2 INNER JOIN filtered_groups ON pr_2.a = filtered_groups.a
+    )
+)
+SETTINGS allow_experimental_parallel_reading_from_replicas = 1, parallel_replicas_for_non_replicated_merge_tree = 1, cluster_for_parallel_replicas = 'test_cluster_one_shard_three_replicas_localhost', max_parallel_replicas = 3;
+
+CREATE TABLE numbers_1e6
+(
+    `n` UInt64
+)
+ENGINE = MergeTree
+ORDER BY n
+AS SELECT * FROM numbers(1_000_000);
+
+-- Same but nested CTE's
+WITH
+    cte1 AS
+    (
+        SELECT n
+        FROM numbers_1e6
+    ),
+    cte2 AS
+    (
+        SELECT n
+        FROM numbers_1e6
+        WHERE n IN (cte1)
+    )
+SELECT count()
+FROM cte2
+SETTINGS allow_experimental_parallel_reading_from_replicas = 1, parallel_replicas_for_non_replicated_merge_tree = 1, cluster_for_parallel_replicas = 'test_cluster_one_shard_three_replicas_localhost', max_parallel_replicas = 3;
+
+DROP TABLE IF EXISTS numbers_1e6;
 DROP TABLE IF EXISTS pr_1;
 DROP TABLE IF EXISTS pr_2;