Backport ClickHouse#89527 to 25.8: Fix possible "Context has expired" with analyzer and PK IN (subquery) (v2)

Author: robot-clickhouse

src/Planner/Planner.cpp

@@ -1256,8 +1256,14 @@ void addBuildSubqueriesForSetsStepIfNeeded(
             std::make_shared<GlobalPlannerContext>(nullptr, nullptr, FiltersForTableExpressionMap{}));
         subquery_planner.buildQueryPlanIfNeeded();
 
-        query_plan.addInterpreterContext(subquery_planner.getPlannerContext()->getQueryContext());
-        subquery->setQueryPlan(std::make_unique<QueryPlan>(std::move(subquery_planner).extractQueryPlan()));
+        auto subquery_plan = std::move(subquery_planner).extractQueryPlan();
+        /// Contexts should be copied into the root query plan, because some functions may
+        /// be created using them while this subquery plan will be destroyed after
+        /// FutureSetFromSubquery::buildSetInplace(). Otherwise, function execution may fail
+        /// with a "Context has expired" exception.
+        for (const auto & context : subquery_plan.getInterpretersContexts())
+            query_plan.addInterpreterContext(context);
+        subquery->setQueryPlan(std::make_unique<QueryPlan>(std::move(subquery_plan)));
     }
 
     if (!subqueries.empty())

src/Processors/QueryPlan/QueryPlan.h

@@ -109,6 +109,7 @@ class QueryPlan
     /// Do not allow to change the table while the pipeline alive.
     void addTableLock(TableLockHolder lock) { resources.table_locks.emplace_back(std::move(lock)); }
     void addInterpreterContext(std::shared_ptr<const Context> context) { resources.interpreter_context.emplace_back(std::move(context)); }
+    auto getInterpretersContexts() const { return resources.interpreter_context; }
     void addStorageHolder(StoragePtr storage) { resources.storage_holders.emplace_back(std::move(storage)); }
 
     void addResources(QueryPlanResourceHolder resources_) { resources = std::move(resources_); }

tests/queries/0_stateless/03671_pk_in_subquery_context_expired.reference

@@ -0,0 +1,5 @@
+1
+1
+Testing w/o relying on enable_global_with_statement...
+1
+1

tests/queries/0_stateless/03671_pk_in_subquery_context_expired.sql

@@ -0,0 +1,82 @@
+-- Issue: https://github.com/ClickHouse/ClickHouse/issues/89433
+
+DROP TABLE IF EXISTS tbl;
+DROP TABLE IF EXISTS join_engine;
+
+CREATE TABLE tbl
+(
+    `id1` LowCardinality(String),
+    `id2` LowCardinality(String),
+    `v` Int64
+)
+ENGINE = MergeTree
+ORDER BY (id1, id2, v);
+INSERT INTO tbl VALUES ('a', 'b', 1);
+CREATE TABLE join_engine
+(
+    `id1` LowCardinality(String),
+    `id2` LowCardinality(String),
+    `v` Int64
+)
+ENGINE = Join(ANY, LEFT, id1, id2);
+INSERT INTO join_engine VALUES ('a', 'b', 1);
+
+WITH cte AS
+    (
+        SELECT id2
+        FROM tbl
+        WHERE joinGet(currentDatabase() || '.join_engine', 'v', id1, id2) = tbl.v
+    )
+SELECT uniq(id2) AS count
+FROM
+(
+    -- NOTE: the bug is reproduced only because due to
+    -- enable_global_with_statement adds "cte" here, but likely it will be
+    -- fixed one day... so I've added another test below that does not rely
+    -- on this fact
+    SELECT *
+    FROM tbl AS e
+    WHERE joinGet(currentDatabase() || '.join_engine', 'v', id1, id2) = e.v
+)
+WHERE id2 IN (
+    SELECT id2
+    FROM cte
+)
+UNION ALL
+SELECT uniq(id2) AS count
+FROM cte;
+
+SELECT 'Testing w/o relying on enable_global_with_statement...';
+--
+-- The same as before, but without relying on enable_global_with_statement
+--
+SELECT uniq(id2) AS count
+FROM
+(
+    WITH cte AS
+        (
+            SELECT id2
+            FROM tbl
+            WHERE joinGet(currentDatabase() || '.join_engine', 'v', id1, id2) = tbl.v
+        )
+    SELECT *
+    FROM tbl AS e
+    WHERE joinGet(currentDatabase() || '.join_engine', 'v', id1, id2) = e.v
+)
+WHERE id2 IN (
+    SELECT id2
+    FROM
+    (
+        SELECT id2
+        FROM tbl
+        WHERE joinGet(currentDatabase() || '.join_engine', 'v', id1, id2) = tbl.v
+    )
+)
+UNION ALL
+SELECT uniq(id2) AS count
+FROM
+(
+    SELECT id2
+    FROM tbl
+    WHERE joinGet(currentDatabase() || '.join_engine', 'v', id1, id2) = tbl.v
+);