fixed format on nested select

2 files changed, 15 insertions, 3 deletions

diff --git a/ydb/library/yql/sql/v1/format/sql_format.cpp b/ydb/library/yql/sql/v1/format/sql_format.cpp
index 3e41eca9f6..4aff679ede 100644
--- a/ydb/library/yql/sql/v1/format/sql_format.cpp
+++ b/ydb/library/yql/sql/v1/format/sql_format.cpp
@@ -357,14 +357,17 @@ private:
         if (descr == TToken::GetDescriptor()) {
             const auto& token = dynamic_cast<const TToken&>(msg);
             MarkToken(token);
-        }
-
-        if (descr == TRule_lambda_body::GetDescriptor()) {
+        } else if (descr == TRule_lambda_body::GetDescriptor()) {
             Y_ENSURE(TokenIndex >= 1);
             auto prevIndex = TokenIndex - 1;
             Y_ENSURE(prevIndex < ParsedTokens.size());
             Y_ENSURE(ParsedTokens[prevIndex].Content == "{");
             MarkedTokens[prevIndex].OpeningBracket = false;
+        } else if (descr == TRule_in_atom_expr::GetDescriptor()) {
+            const auto& value = dynamic_cast<const TRule_in_atom_expr&>(msg);
+            if (value.Alt_case() == TRule_in_atom_expr::kAltInAtomExpr7) {
+                AfterInAtom = true;
+            }
         }
 
         VisitAllFieldsImpl<&TVisitor::MarkTokens>(descr, msg);
@@ -400,6 +403,12 @@ private:
             PopBracket("<");
         }
 
+        if (AfterInAtom) {
+            auto& info = MarkedTokens[TokenIndex];
+            info.OpeningBracket = false;
+            AfterInAtom = false;
+        }
+
         TokenIndex++;
     }
 
@@ -1986,6 +1995,7 @@ private:
     ui32 TokenIndex = 0;
     TMarkTokenStack MarkTokenStack;
     TVector<TTokenInfo> MarkedTokens;
+    bool AfterInAtom = false;
 };
 
 template <typename T>
diff --git a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
index e6a5c24d15..c0287e1164 100644
--- a/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
+++ b/ydb/library/yql/sql/v1/format/sql_format_ut.cpp
@@ -780,6 +780,8 @@ Y_UNIT_TEST_SUITE(CheckSqlFormatter) {
              "SELECT\n\t1 IN (\n\t\tSELECT\n\t\t\t1\n\t);\n\n"},
             {"select 1 in ((select 1))",
              "SELECT\n\t1 IN (\n\t\t(\n\t\t\tSELECT\n\t\t\t\t1\n\t\t)\n\t);\n\n"},
+            {"select 1 in (\nselect 1)",
+             "SELECT\n\t1 IN (\n\t\tSELECT\n\t\t\t1\n\t);\n\n"},
         };
 
         TSetup setup;