Limit transformation count while trying to annotate single node (#7588)

3 files changed, 28 insertions, 1 deletions

diff --git a/ydb/library/yql/ast/yql_expr.h b/ydb/library/yql/ast/yql_expr.h
index a063783396..2f9b6e9f77 100644
--- a/ydb/library/yql/ast/yql_expr.h
+++ b/ydb/library/yql/ast/yql_expr.h
@@ -2525,6 +2525,7 @@ struct TExprContext : private TNonCopyable {
     ui64 StringsAllocationLimit = 100000000;
     ui64 RepeatTransformLimit = 1000000;
     ui64 RepeatTransformCounter = 0;
+    ui64 TypeAnnNodeRepeatLimit = 1000;
 
     TGcNodeConfig GcConfig;
 
diff --git a/ydb/library/yql/core/type_ann/type_ann_expr.cpp b/ydb/library/yql/core/type_ann/type_ann_expr.cpp
index 9109875f95..7681c087a8 100644
--- a/ydb/library/yql/core/type_ann/type_ann_expr.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_expr.cpp
@@ -180,7 +180,7 @@ private:
         }
 
         auto input = start;
-        for (;;) {
+        for (size_t transformCount = 0; true; ++transformCount) {
             TIssueScopeGuard issueScope(ctx.IssueManager, [this, input, &ctx]() -> TIssuePtr {
                 TStringBuilder str;
                 str << "At ";
@@ -258,6 +258,22 @@ private:
                 YQL_ENSURE(false, "Unknown state");
             }
 
+            if (transformCount >= ctx.TypeAnnNodeRepeatLimit) {
+                TConvertToAstSettings settings;
+                settings.AllowFreeArgs = true;
+                auto ast = ConvertToAst(*input, ctx, settings);
+                if (ast.Root) {
+                    TStringStream s;
+                    ast.Root->PrettyPrintTo(s, TAstPrintFlags::ShortQuote | TAstPrintFlags::AdaptArbitraryContent | TAstPrintFlags::PerLine);
+                    YQL_CLOG(INFO, Core) << "Too many transformations for node:\n" << s.Str();
+                }
+                ctx.AddError(TIssue(ctx.GetPosition(input->Pos()),
+                    TStringBuilder() << "YQL: Internal core error! Type annotation of node " << input->Content()
+                        << " with type " << input->Type() << " takes too much iterations: "
+                        << ctx.TypeAnnNodeRepeatLimit << ". You may set TypeAnnNodeRepeatLimit as flags for config provider."));
+                return TStatus::Error;
+            }
+
             input->SetState(TExprNode::EState::TypePending);
             switch (input->Type()) {
             case TExprNode::Atom:
diff --git a/ydb/library/yql/providers/config/yql_config_provider.cpp b/ydb/library/yql/providers/config/yql_config_provider.cpp
index 5bd0f8220f..160536ac7b 100644
--- a/ydb/library/yql/providers/config/yql_config_provider.cpp
+++ b/ydb/library/yql/providers/config/yql_config_provider.cpp
@@ -547,6 +547,16 @@ namespace {
                     return false;
                 }
             }
+            else if (name == "TypeAnnNodeRepeatLimit") {
+                if (args.size() != 1) {
+                    ctx.AddError(TIssue(pos, TStringBuilder() << "Expected 1 argument, but got " << args.size()));
+                    return false;
+                }
+                if (!TryFromString(args[0], ctx.TypeAnnNodeRepeatLimit)) {
+                    ctx.AddError(TIssue(pos, TStringBuilder() << "Expected integer, but got: " << args[0]));
+                    return false;
+                }
+            }
             else if (name == "PureDataSource") {
                 if (args.size() != 1) {
                     ctx.AddError(TIssue(pos, TStringBuilder() << "Expected 1 argument, but got " << args.size()));