YQL-19405: fix formatter validation issues

commit_hash:8eb4c1a07911920db3521527715df0f7935127f5

2 files changed, 45 insertions, 29 deletions

diff --git a/yql/essentials/sql/v1/format/sql_format.cpp b/yql/essentials/sql/v1/format/sql_format.cpp
index 7c6d70ca09..3b3a14d850 100644
--- a/yql/essentials/sql/v1/format/sql_format.cpp
+++ b/yql/essentials/sql/v1/format/sql_format.cpp
@@ -60,20 +60,6 @@ void SkipForValidate(
     in = SkipWS(in, query.end());
     out = SkipWS(out, formattedQuery.end());
 
-    auto skipDeletedToken = [&](const TString& deletedToken, const TVector<TString>& afterTokens) {
-        if (
-            in != query.end() && in->Name == deletedToken &&
-            (out == formattedQuery.end() || out->Name != deletedToken) &&
-            in != query.begin() && IsIn(afterTokens, SkipWSOrCommentBackward(in - 1, query.begin())->Name)
-        ) {
-            in = SkipWS(++in, query.end());
-            return true;
-        }
-        return false;
-    };
-
-    while (skipDeletedToken("SEMICOLON", {"AS", "BEGIN", "LBRACE_CURLY", "SEMICOLON"})) {}
-
     auto inSkippedComments = SkipWSOrComment(in, query.end());
 
     auto skipAddedToken = [&](const TString& addedToken, const TVector<TString>& beforeTokens, const TVector<TString>& afterTokens) {
@@ -95,6 +81,20 @@ void SkipForValidate(
     skipAddedToken("LPAREN", {}, {"EQUALS"});
     skipAddedToken("RPAREN", {"END", "EOF", "SEMICOLON"}, {});
     skipAddedToken("SEMICOLON", {"END", "RBRACE_CURLY"}, {});
+
+    auto skipDeletedToken = [&](const TString& deletedToken, const TVector<TString>& afterTokens) {
+        if (
+            in != query.end() && in->Name == deletedToken &&
+            (out == formattedQuery.end() || out->Name != deletedToken) &&
+            in != query.begin() && IsIn(afterTokens, SkipWSOrCommentBackward(in - 1, query.begin())->Name)
+        ) {
+            in = SkipWS(++in, query.end());
+            return true;
+        }
+        return false;
+    };
+
+    while (skipDeletedToken("SEMICOLON", {"AS", "BEGIN", "LBRACE_CURLY", "SEMICOLON"})) {}
 }
 
 TParsedToken TransformTokenForValidate(TParsedToken token) {
@@ -141,6 +141,19 @@ bool Validate(const TParsedTokenList& query, const TParsedTokenList& formattedQu
                 if (SkipQuotes(inToken.Content) != SkipQuotes(outToken.Content)) {
                     return false;
                 }
+            } else if (inToken.Name == "COMMENT") {
+                TStringBuf inContent = inToken.Content;
+                TStringBuf outContent = outToken.Content;
+
+                auto inNextToken = SkipWS(in + 1, inEnd);
+                if (inNextToken != inEnd && inNextToken->Name == "EOF") {
+                    inContent.ChopSuffix("\n");
+                    outContent.ChopSuffix("\n");
+                }
+
+                if (inContent != outContent) {
+                    return false;
+                }
             } else {
                 if (inToken.Content != outToken.Content) {
                     return false;
@@ -600,10 +613,6 @@ private:
             WriteComments(true);
         }
 
-        if (AfterComment && Comments[LastComment - 1].Content.StartsWith("--")) {
-            return;
-        }
-
         if (OutColumn) {
             Out('\n');
         }
@@ -625,6 +634,10 @@ private:
 
         Out(text);
 
+        if (text.StartsWith("--") && !text.EndsWith("\n")) {
+            Out('\n');
+        }
+
         if (!text.StartsWith("--") &&
             TokenIndex < ParsedTokens.size() &&
             Comments[LastComment].Line < ParsedTokens[TokenIndex].Line &&
@@ -3211,15 +3224,12 @@ public:
 
             TVector<NSQLTranslation::TParsedToken> comments;
             TParsedTokenList parsedTokens, stmtTokens;
-            bool hasTrailingComments = false;
             auto onNextRawToken = [&](NSQLTranslation::TParsedToken&& token) {
                 stmtTokens.push_back(token);
                 if (token.Name == "COMMENT") {
                     comments.emplace_back(std::move(token));
-                    hasTrailingComments = true;
                 } else if (token.Name != "WS" && token.Name != "EOF") {
                     parsedTokens.emplace_back(std::move(token));
-                    hasTrailingComments = false;
                 }
             };
 
@@ -3267,11 +3277,6 @@ public:
 
             finalFormattedQuery << currentFormattedQuery;
             if (parsedTokens.back().Name != "SEMICOLON") {
-                if (hasTrailingComments
-                     && !comments.back().Content.EndsWith("\n")
-                     && comments.back().Content.StartsWith("--")) {
-                    finalFormattedQuery << "\n";
-                }
                 finalFormattedQuery << ";\n";
             }
         }
diff --git a/yql/essentials/sql/v1/format/sql_format_ut.h b/yql/essentials/sql/v1/format/sql_format_ut.h
index 2cf4894d08..079bcfced6 100644
--- a/yql/essentials/sql/v1/format/sql_format_ut.h
+++ b/yql/essentials/sql/v1/format/sql_format_ut.h
@@ -199,9 +199,16 @@ Y_UNIT_TEST(Declare) {
 
 Y_UNIT_TEST(NamedNode) {
     TCases cases = {
-        {"$x=1","$x = 1;\n"},
-        {"$x,$y=(2,3)","$x, $y = (2, 3);\n"},
-        {"$a = select 1 union all select 2","$a = (\n\tSELECT\n\t\t1\n\tUNION ALL\n\tSELECT\n\t\t2\n);\n"},
+        {"$x=1",
+            "$x = 1;\n"},
+        {"$x,$y=(2,3)",
+            "$x, $y = (2, 3);\n"},
+        {"$a = select 1 union all select 2",
+            "$a = (\n\tSELECT\n\t\t1\n\tUNION ALL\n\tSELECT\n\t\t2\n);\n"},
+        {"$a = select 1 from $as;",
+            "$a = (\n\tSELECT\n\t\t1\n\tFROM\n\t\t$as\n);\n"},
+        {"$a = select * from $t -- comment",
+            "$a = (\n\tSELECT\n\t\t*\n\tFROM\n\t\t$t -- comment\n);\n"},
     };
 
     TSetup setup;
@@ -606,6 +613,8 @@ Y_UNIT_TEST(If) {
             "SELECT\n\t\t2\n\t;\n\n\tSELECT\n\t\t3\n\t;\nEND DO;\n"},
         {"evaluate if 1=1 do begin (select 1) end do",
             "EVALUATE IF 1 == 1 DO BEGIN\n\t(\n\t\tSELECT\n\t\t\t1\n\t);\nEND DO;\n"},
+        {"evaluate if 1=1 do begin $a = select * from $begin; $end = 1; end do",
+            "EVALUATE IF 1 == 1 DO BEGIN\n\t$a = (\n\t\tSELECT\n\t\t\t*\n\t\tFROM\n\t\t\t$begin\n\t);\n\t$end = 1;\nEND DO;\n"},
     };
 
     TSetup setup;
@@ -628,6 +637,8 @@ Y_UNIT_TEST(For) {
             "EVALUATE FOR $x IN [] DO BEGIN\n\tSELECT\n\t\t$x\n\t;\n\n\tSELECT\n\t\t$y\n\t;\nEND DO;\n"},
         {"evaluate for $x in [] do begin (select 1) end do",
             "EVALUATE FOR $x IN [] DO BEGIN\n\t(\n\t\tSELECT\n\t\t\t1\n\t);\nEND DO;\n"},
+        {"evaluate for $x in [] do begin $a = select * from $begin; $end = 1; end do",
+            "EVALUATE FOR $x IN [] DO BEGIN\n\t$a = (\n\t\tSELECT\n\t\t\t*\n\t\tFROM\n\t\t\t$begin\n\t);\n\t$end = 1;\nEND DO;\n"},
     };
 
     TSetup setup;