From 38c5451c3c2dc24659124e8b6b464dd4624686a7 Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@clickhouse.com>
Date: Sat, 1 Oct 2022 17:47:58 +0200
Subject: [PATCH] Add more checkStackSize calls

---
 src/Interpreters/ApplyWithAliasVisitor.cpp                  | 5 +++++
 src/Interpreters/ApplyWithGlobalVisitor.cpp                 | 4 ++++
 src/Interpreters/ApplyWithSubqueryVisitor.cpp               | 5 +++++
 src/Interpreters/ApplyWithSubqueryVisitor.h                 | 1 +
 src/Interpreters/InDepthNodeVisitor.h                       | 3 +++
 src/Interpreters/InJoinSubqueriesPreprocessor.cpp           | 3 +++
 src/Interpreters/OptimizeIfChains.cpp                       | 4 ++++
 src/Interpreters/OptimizeIfWithConstantConditionVisitor.cpp | 4 ++++
 src/Interpreters/ReplaceQueryParameterVisitor.cpp           | 3 +++
 src/Interpreters/RewriteCountVariantsVisitor.cpp            | 5 +++++
 src/Interpreters/TreeRewriter.cpp                           | 1 -
 11 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/Interpreters/ApplyWithAliasVisitor.cpp b/src/Interpreters/ApplyWithAliasVisitor.cpp
index 89407f3521e..fed544598a1 100644
--- a/src/Interpreters/ApplyWithAliasVisitor.cpp
+++ b/src/Interpreters/ApplyWithAliasVisitor.cpp
@@ -2,11 +2,16 @@
 #include <Interpreters/misc.h>
 #include <Parsers/ASTSelectQuery.h>
 #include <Parsers/ASTSelectWithUnionQuery.h>
+#include <Common/checkStackSize.h>
+
 
 namespace DB
 {
+
 void ApplyWithAliasVisitor::visit(ASTPtr & ast, const Data & data)
 {
+    checkStackSize();
+
     if (auto * node_select = ast->as<ASTSelectQuery>())
     {
         std::optional<Data> new_data;
diff --git a/src/Interpreters/ApplyWithGlobalVisitor.cpp b/src/Interpreters/ApplyWithGlobalVisitor.cpp
index 55e5968b7c9..a0f256fca83 100644
--- a/src/Interpreters/ApplyWithGlobalVisitor.cpp
+++ b/src/Interpreters/ApplyWithGlobalVisitor.cpp
@@ -3,6 +3,8 @@
 #include <Parsers/ASTSelectWithUnionQuery.h>
 #include <Parsers/ASTSelectIntersectExceptQuery.h>
 #include <Parsers/ASTWithAlias.h>
+#include <Common/checkStackSize.h>
+
 
 namespace DB
 {
@@ -71,6 +73,8 @@ void ApplyWithGlobalVisitor::visit(
 
 void ApplyWithGlobalVisitor::visit(ASTPtr & ast)
 {
+    checkStackSize();
+
     if (ASTSelectWithUnionQuery * node_union = ast->as<ASTSelectWithUnionQuery>())
     {
         if (auto * first_select = typeid_cast<ASTSelectQuery *>(node_union->list_of_selects->children[0].get()))
diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.cpp b/src/Interpreters/ApplyWithSubqueryVisitor.cpp
index 6378bd6fdb2..32f5dd63916 100644
--- a/src/Interpreters/ApplyWithSubqueryVisitor.cpp
+++ b/src/Interpreters/ApplyWithSubqueryVisitor.cpp
@@ -1,3 +1,4 @@
+#include <Common/checkStackSize.h>
 #include <Interpreters/ApplyWithSubqueryVisitor.h>
 #include <Interpreters/IdentifierSemantic.h>
 #include <Interpreters/StorageID.h>
@@ -10,10 +11,14 @@
 #include <Parsers/ASTTablesInSelectQuery.h>
 #include <Parsers/ASTWithElement.h>
 
+
 namespace DB
 {
+
 void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data)
 {
+    checkStackSize();
+
     if (auto * node_select = ast->as<ASTSelectQuery>())
         visit(*node_select, data);
     else
diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.h b/src/Interpreters/ApplyWithSubqueryVisitor.h
index 3b17de1d691..baf073f12cb 100644
--- a/src/Interpreters/ApplyWithSubqueryVisitor.h
+++ b/src/Interpreters/ApplyWithSubqueryVisitor.h
@@ -4,6 +4,7 @@
 
 #include <Parsers/IAST.h>
 
+
 namespace DB
 {
 class ASTFunction;
diff --git a/src/Interpreters/InDepthNodeVisitor.h b/src/Interpreters/InDepthNodeVisitor.h
index 736a764e8e9..785c88ec77b 100644
--- a/src/Interpreters/InDepthNodeVisitor.h
+++ b/src/Interpreters/InDepthNodeVisitor.h
@@ -3,8 +3,10 @@
 #include <typeinfo>
 #include <vector>
 #include <Common/typeid_cast.h>
+#include <Common/checkStackSize.h>
 #include <Parsers/DumpASTNode.h>
 
+
 namespace DB
 {
 
@@ -24,6 +26,7 @@ public:
 
     void visit(T & ast)
     {
+        checkStackSize();
         DumpASTNode dump(*ast, ostr, visit_depth, typeid(Matcher).name());
 
         if constexpr (!_top_to_bottom)
diff --git a/src/Interpreters/InJoinSubqueriesPreprocessor.cpp b/src/Interpreters/InJoinSubqueriesPreprocessor.cpp
index 66ee32750f3..58faeb41a15 100644
--- a/src/Interpreters/InJoinSubqueriesPreprocessor.cpp
+++ b/src/Interpreters/InJoinSubqueriesPreprocessor.cpp
@@ -9,6 +9,7 @@
 #include <Parsers/ASTTablesInSelectQuery.h>
 #include <Parsers/ASTFunction.h>
 #include <Common/typeid_cast.h>
+#include <Common/checkStackSize.h>
 
 
 namespace DB
@@ -225,6 +226,8 @@ void InJoinSubqueriesPreprocessor::visit(ASTPtr & ast) const
     if (!ast)
         return;
 
+    checkStackSize();
+
     ASTSelectQuery * query = ast->as<ASTSelectQuery>();
     if (!query || !query->tables())
         return;
diff --git a/src/Interpreters/OptimizeIfChains.cpp b/src/Interpreters/OptimizeIfChains.cpp
index 5b0fab7793c..d4b4d20bb9c 100644
--- a/src/Interpreters/OptimizeIfChains.cpp
+++ b/src/Interpreters/OptimizeIfChains.cpp
@@ -1,9 +1,11 @@
 #include <Common/typeid_cast.h>
+#include <Common/checkStackSize.h>
 #include <Parsers/ASTFunction.h>
 #include <Parsers/ASTExpressionList.h>
 #include <Interpreters/OptimizeIfChains.h>
 #include <IO/WriteHelpers.h>
 
+
 namespace DB
 {
 
@@ -18,6 +20,8 @@ void OptimizeIfChainsVisitor::visit(ASTPtr & current_ast)
     if (!current_ast)
         return;
 
+    checkStackSize();
+
     for (ASTPtr & child : current_ast->children)
     {
         /// Fallthrough cases
diff --git a/src/Interpreters/OptimizeIfWithConstantConditionVisitor.cpp b/src/Interpreters/OptimizeIfWithConstantConditionVisitor.cpp
index 28b19f6670d..824ad22bb12 100644
--- a/src/Interpreters/OptimizeIfWithConstantConditionVisitor.cpp
+++ b/src/Interpreters/OptimizeIfWithConstantConditionVisitor.cpp
@@ -1,4 +1,5 @@
 #include <Common/typeid_cast.h>
+#include <Common/checkStackSize.h>
 #include <Parsers/ASTLiteral.h>
 #include <Parsers/ASTFunction.h>
 #include <Parsers/ASTHelpers.h>
@@ -6,6 +7,7 @@
 #include <Interpreters/OptimizeIfWithConstantConditionVisitor.h>
 #include <IO/WriteHelpers.h>
 
+
 namespace DB
 {
 
@@ -71,6 +73,8 @@ void OptimizeIfWithConstantConditionVisitor::visit(ASTPtr & current_ast)
     if (!current_ast)
         return;
 
+    checkStackSize();
+
     for (ASTPtr & child : current_ast->children)
     {
         auto * function_node = child->as<ASTFunction>();
diff --git a/src/Interpreters/ReplaceQueryParameterVisitor.cpp b/src/Interpreters/ReplaceQueryParameterVisitor.cpp
index 664cda74522..98e35d69ab7 100644
--- a/src/Interpreters/ReplaceQueryParameterVisitor.cpp
+++ b/src/Interpreters/ReplaceQueryParameterVisitor.cpp
@@ -12,6 +12,7 @@
 #include <Parsers/TablePropertiesQueriesASTs.h>
 #include <Common/quoteString.h>
 #include <Common/typeid_cast.h>
+#include <Common/checkStackSize.h>
 
 
 namespace DB
@@ -26,6 +27,8 @@ namespace ErrorCodes
 
 void ReplaceQueryParameterVisitor::visit(ASTPtr & ast)
 {
+    checkStackSize();
+
     if (ast->as<ASTQueryParameter>())
         visitQueryParameter(ast);
     else if (ast->as<ASTIdentifier>() || ast->as<ASTTableIdentifier>())
diff --git a/src/Interpreters/RewriteCountVariantsVisitor.cpp b/src/Interpreters/RewriteCountVariantsVisitor.cpp
index 7e7c5f14a91..741dc3e8cb7 100644
--- a/src/Interpreters/RewriteCountVariantsVisitor.cpp
+++ b/src/Interpreters/RewriteCountVariantsVisitor.cpp
@@ -5,11 +5,16 @@
 #include <Parsers/ASTTablesInSelectQuery.h>
 #include <Poco/String.h>
 #include <Common/typeid_cast.h>
+#include <Common/checkStackSize.h>
+
 
 namespace DB
 {
+
 void RewriteCountVariantsVisitor::visit(ASTPtr & node)
 {
+    checkStackSize();
+
     if (node->as<ASTSubquery>() || node->as<ASTTableExpression>() || node->as<ASTArrayJoin>())
         return;
 
diff --git a/src/Interpreters/TreeRewriter.cpp b/src/Interpreters/TreeRewriter.cpp
index d31a1ca7a8a..2f5bfd00938 100644
--- a/src/Interpreters/TreeRewriter.cpp
+++ b/src/Interpreters/TreeRewriter.cpp
@@ -3,7 +3,6 @@
 #include <Core/Settings.h>
 #include <Core/NamesAndTypes.h>
 
-#include <Common/checkStackSize.h>
 #include <Core/SettingsEnums.h>
 
 #include <Interpreters/ArrayJoinedColumnsVisitor.h>