Always use Utf8 aware TextWalker for antlr4 (#9245)

vitstn · web-flow · commit c52a7898d21a · 2024-09-16T11:36:00.000+03:00
diff --git a/ydb/library/yql/parser/lexer_common/hints.cpp b/ydb/library/yql/parser/lexer_common/hints.cpp
@@ -31,9 +31,10 @@ namespace {
 
 class TTokenProcessor {
 public:
-    TTokenProcessor(const TString& queryFile, TSQLHints& hints)
+    TTokenProcessor(const TString& queryFile, TSQLHints& hints, bool utf8Aware)
         : QueryFile(queryFile)
         , Hints(hints)
+        , Utf8Aware(utf8Aware)
     {}
 
     TPosition ExtractPosition(const TParsedToken& token) const {
@@ -57,7 +58,7 @@ class TTokenProcessor {
             // skip leading comments
             return;
         }
-        TVector<TSQLHint> currentHints = NDetail::ParseSqlHints(pos, token.Content);
+        TVector<TSQLHint> currentHints = NDetail::ParseSqlHints(pos, token.Content, Utf8Aware);
         if (currentHints.empty()) {
             // no hints here
             return;
@@ -70,12 +71,14 @@ class TTokenProcessor {
     TMaybe<TPosition> PrevNonCommentPos;
     const TString QueryFile;
     TSQLHints& Hints;
+    const bool Utf8Aware;
 };
 
 }
 
-bool CollectSqlHints(ILexer& lexer, const TString& query, const TString& queryName, const TString& queryFile, TSQLHints& hints, NYql::TIssues& issues, size_t maxErrors) {
-    TTokenProcessor tp(queryFile, hints);
+bool CollectSqlHints(ILexer& lexer, const TString& query, const TString& queryName,
+    const TString& queryFile, TSQLHints& hints, NYql::TIssues& issues, size_t maxErrors, bool utf8Aware) {
+    TTokenProcessor tp(queryFile, hints, utf8Aware);
     return lexer.Tokenize(query, queryName, [&tp](TParsedToken&& token) { tp.ProcessToken(std::move(token)); }, issues, maxErrors);
 }
 
diff --git a/ydb/library/yql/parser/lexer_common/hints.h b/ydb/library/yql/parser/lexer_common/hints.h
@@ -29,7 +29,8 @@ using TSQLHints = TMap<NYql::TPosition, TVector<TSQLHint>>;
 // For example: SELECT /*+ Name(Value) */ -- Name2(Value2)
 // in this case TSQLHints will consist of single entry with position of SELECT token
 
-bool CollectSqlHints(ILexer& lexer, const TString& query, const TString& queryName, const TString& queryFile, TSQLHints& hints, NYql::TIssues& issues, size_t maxErrors);
+bool CollectSqlHints(ILexer& lexer, const TString& query, const TString& queryName,
+    const TString& queryFile, TSQLHints& hints, NYql::TIssues& issues, size_t maxErrors, bool utf8Aware);
 
 }
 
diff --git a/ydb/library/yql/parser/lexer_common/parse_hints_impl.cpp b/ydb/library/yql/parser/lexer_common/parse_hints_impl.cpp
@@ -15,13 +15,13 @@ enum EParseState {
     WANT_WS,
 };
 
-TVector<TSQLHint> ParseSqlHints(NYql::TPosition commentPos, const TStringBuf& comment) {
+TVector<TSQLHint> ParseSqlHints(NYql::TPosition commentPos, const TStringBuf& comment, bool utf8Aware) {
     TVector<TSQLHint> result;
     if (!comment.StartsWith("/*+") && !comment.StartsWith("--+")) {
         return result;
     }
     TSQLHint hint;
-    NYql::TTextWalker commentWalker(commentPos);
+    NYql::TTextWalker commentWalker(commentPos, utf8Aware);
     const size_t len = comment.size();
     EParseState state = EParseState::INITIAL;
     for (size_t i = 0; i < len; ++i) {
diff --git a/ydb/library/yql/parser/lexer_common/parse_hints_impl.h b/ydb/library/yql/parser/lexer_common/parse_hints_impl.h
@@ -6,7 +6,7 @@ namespace NSQLTranslation {
 
 namespace NDetail {
 
-TVector<TSQLHint> ParseSqlHints(NYql::TPosition commentPos, const TStringBuf& comment);
+TVector<TSQLHint> ParseSqlHints(NYql::TPosition commentPos, const TStringBuf& comment, bool utf8Aware);
 
 }
 
diff --git a/ydb/library/yql/parser/lexer_common/ut/hints_ut.cpp b/ydb/library/yql/parser/lexer_common/ut/hints_ut.cpp
@@ -15,7 +15,7 @@ TSQLHints CollectHints(const TString& query, bool antlr4Parser) {
     TSQLHints result;
     NYql::TIssues issues;
     size_t maxErrors = 100;
-    UNIT_ASSERT(CollectSqlHints(*lexer, query, "", "", result, issues, maxErrors));
+    UNIT_ASSERT(CollectSqlHints(*lexer, query, "", "", result, issues, maxErrors, false));
     UNIT_ASSERT(issues.Empty());
     return result;
 }
diff --git a/ydb/library/yql/parser/lexer_common/ut/parse_hints_ut.cpp b/ydb/library/yql/parser/lexer_common/ut/parse_hints_ut.cpp
@@ -8,7 +8,7 @@ using namespace NSQLTranslation;
 using namespace NSQLTranslation::NDetail;
 
 void CheckParse(TStringBuf comment, TStringBuf expected) {
-    TString parsed = JoinSeq(",", ParseSqlHints({}, comment));
+    TString parsed = JoinSeq(",", ParseSqlHints({}, comment, false));
     UNIT_ASSERT_NO_DIFF(parsed, expected);
 }
 
diff --git a/ydb/library/yql/public/issue/yql_issue.h b/ydb/library/yql/public/issue/yql_issue.h
@@ -55,7 +55,7 @@ struct TPosition {
 
 class TTextWalker {
 public:
-    TTextWalker(TPosition& position, bool utf8Aware = false)
+    TTextWalker(TPosition& position, bool utf8Aware)
         : Position(position)
         , Utf8Aware(utf8Aware)
         , HaveCr(false)
diff --git a/ydb/library/yql/public/issue/yql_issue_ut.cpp b/ydb/library/yql/public/issue/yql_issue_ut.cpp
@@ -89,7 +89,7 @@ Y_UNIT_TEST_SUITE(TextWalkerTest) {
         TPosition pos;
         pos.Row = 1;
 
-        TTextWalker walker(pos);
+        TTextWalker walker(pos, false);
         walker.Advance(TStringBuf("a\r\taa"));
 
         UNIT_ASSERT_VALUES_EQUAL(pos, TPosition(5, 1));
@@ -101,7 +101,7 @@ Y_UNIT_TEST_SUITE(TextWalkerTest) {
         TPosition pos;
         pos.Row = 1;
 
-        TTextWalker walker(pos);
+        TTextWalker walker(pos, false);
         walker.Advance(TStringBuf("a\raa\r"));
         UNIT_ASSERT_VALUES_EQUAL(pos, TPosition(4, 1));
         walker.Advance('\n');
@@ -115,6 +115,28 @@ Y_UNIT_TEST_SUITE(TextWalkerTest) {
         walker.Advance('a');
         UNIT_ASSERT_VALUES_EQUAL(pos, TPosition(1, 3));
     }
+
+    Y_UNIT_TEST(UnicodeTest) {
+        {
+            TPosition pos;
+            pos.Row = 1;
+
+            TTextWalker walker(pos, false);
+            walker.Advance(TStringBuf("привет"));
+
+            UNIT_ASSERT_VALUES_EQUAL(pos, TPosition(12, 1));
+        }
+
+        {
+            TPosition pos;
+            pos.Row = 1;
+
+            TTextWalker walker(pos, true);
+            walker.Advance(TStringBuf("привет"));
+
+            UNIT_ASSERT_VALUES_EQUAL(pos, TPosition(6, 1));
+        }
+    }
 }
 
 Y_UNIT_TEST_SUITE(ToOneLineStringTest) {
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -5274,8 +5274,8 @@ class TConverter : public IPGParseEvents {
     void ScanRows(const TString& query) {
         QuerySize = query.Size();
         RowStarts.push_back(0);
-        TPosition position(1, 1);
-        TTextWalker walker(position);
+        TPosition position(0, 1);
+        TTextWalker walker(position, true);
         auto prevRow = position.Row;
         for (ui32 i = 0; i < query.Size(); ++i) {
             walker.Advance(query[i]);
diff --git a/ydb/library/yql/sql/v0/node.cpp b/ydb/library/yql/sql/v0/node.cpp
@@ -1771,7 +1771,7 @@ bool TryStringContent(const TString& str, TString& result, ui32& flags, TString&
 
         auto unescapeResult = UnescapeArbitraryAtom(atom, quoteChar, &sout, &readBytes);
         if (unescapeResult != EUnescapeResult::OK) {
-            TTextWalker walker(pos);
+            TTextWalker walker(pos, false);
             walker.Advance(atom.Trunc(readBytes));
             error = UnescapeResultToString(unescapeResult);
             return false;
@@ -1839,7 +1839,7 @@ TString IdContent(TContext& ctx, const TString& s) {
 
     auto unescapeResult = UnescapeArbitraryAtom(atom, endSym, &sout, &readBytes);
     if (unescapeResult != EUnescapeResult::OK) {
-        TTextWalker walker(pos);
+        TTextWalker walker(pos, false);
         walker.Advance(atom.Trunc(readBytes));
         ctx.Error(pos) << "Cannot parse broken identifier: " << UnescapeResultToString(unescapeResult);
         return {};
diff --git a/ydb/library/yql/sql/v1/node.cpp b/ydb/library/yql/sql/v1/node.cpp
@@ -1728,7 +1728,7 @@ TNodePtr IAggregation::WindowTraits(const TNodePtr& type, TContext& ctx) const {
 }
 
 namespace {
-bool UnescapeQuoted(const TString& str, TPosition& pos, char quoteChar, TString& result, TString& error) {
+bool UnescapeQuoted(const TString& str, TPosition& pos, char quoteChar, TString& result, TString& error, bool utf8Aware) {
     result = error = {};
 
     size_t readBytes = 0;
@@ -1739,7 +1739,7 @@ bool UnescapeQuoted(const TString& str, TPosition& pos, char quoteChar, TString&
 
     auto unescapeResult = UnescapeArbitraryAtom(atom, quoteChar, &sout, &readBytes);
     if (unescapeResult != EUnescapeResult::OK) {
-        TTextWalker walker(pos);
+        TTextWalker walker(pos, utf8Aware);
         walker.Advance(atom.Trunc(readBytes));
         error = UnescapeResultToString(unescapeResult);
         return false;
@@ -1834,7 +1834,7 @@ StringContentInternal(TContext& ctx, TPosition pos, const TString& input, EStrin
             result.Content = UnescapeAnsiQuoted(str);
         } else {
             TString error;
-            if (!UnescapeQuoted(str, pos, str[0], result.Content, error)) {
+            if (!UnescapeQuoted(str, pos, str[0], result.Content, error, ctx.Settings.Antlr4Parser)) {
                 ctx.Error(pos) << "Failed to parse string literal: " << error;
                 return {};
             }
@@ -1902,7 +1902,7 @@ TString IdContent(TContext& ctx, const TString& s) {
 
     auto unescapeResult = UnescapeArbitraryAtom(atom, endSym, &sout, &readBytes);
     if (unescapeResult != EUnescapeResult::OK) {
-        TTextWalker walker(pos);
+        TTextWalker walker(pos, ctx.Settings.Antlr4Parser);
         walker.Advance(atom.Trunc(readBytes));
         ctx.Error(pos) << "Cannot parse broken identifier: " << UnescapeResultToString(unescapeResult);
         return {};
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
@@ -95,7 +95,7 @@ NYql::TAstParseResult SqlToYql(const TString& query, const NSQLTranslation::TTra
     NSQLTranslation::TSQLHints hints;
     auto lexer = MakeLexer(settings.AnsiLexer, settings.Antlr4Parser);
     YQL_ENSURE(lexer);
-    if (!CollectSqlHints(*lexer, query, queryName, settings.File, hints, res.Issues, settings.MaxErrors)) {
+    if (!CollectSqlHints(*lexer, query, queryName, settings.File, hints, res.Issues, settings.MaxErrors, settings.Antlr4Parser)) {
         return res;
     }
 
@@ -188,7 +188,7 @@ TVector<NYql::TAstParseResult> SqlToAstStatements(const TString& query, const NS
     NSQLTranslation::TSQLHints hints;
     auto lexer = MakeLexer(settings.AnsiLexer, settings.Antlr4Parser);
     YQL_ENSURE(lexer);
-    if (!CollectSqlHints(*lexer, query, queryName, settings.File, hints, issues, settings.MaxErrors)) {
+    if (!CollectSqlHints(*lexer, query, queryName, settings.File, hints, issues, settings.MaxErrors, settings.Antlr4Parser)) {
         return result;
     }
 
diff --git a/ydb/library/yql/tools/sql2yql/sql2yql.cpp b/ydb/library/yql/tools/sql2yql/sql2yql.cpp
@@ -315,7 +315,8 @@ int BuildAST(int argc, char* argv[]) {
 
                         NSQLTranslation::TSQLHints hints;
                         auto lexer = SqlLexer(query, parseRes.Issues, settings);
-                        if (lexer && CollectSqlHints(*lexer, query, queryFile, settings.File, hints, parseRes.Issues, settings.MaxErrors)) {
+                        if (lexer && CollectSqlHints(*lexer, query, queryFile, settings.File, hints, parseRes.Issues,
+                            settings.MaxErrors, settings.Antlr4Parser)) {
                             parseRes = NSQLTranslation::SqlASTToYql(*ast, hints, settings);
                         }
                    }

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,8 @@ using TSQLHints = TMap<NYql::TPosition, TVector<TSQLHint>>;`
`29`	`29`	`// For example: SELECT /+ Name(Value) / -- Name2(Value2)`
`30`	`30`	`// in this case TSQLHints will consist of single entry with position of SELECT token`
`31`	`31`
`32`		`-bool CollectSqlHints(ILexer& lexer, const TString& query, const TString& queryName, const TString& queryFile, TSQLHints& hints, NYql::TIssues& issues, size_t maxErrors);`
	`32`	`+bool CollectSqlHints(ILexer& lexer, const TString& query, const TString& queryName,`
	`33`	`+ const TString& queryFile, TSQLHints& hints, NYql::TIssues& issues, size_t maxErrors, bool utf8Aware);`
`33`	`34`
`34`	`35`	`}`
`35`	`36`
Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@ namespace NSQLTranslation {`
`6`	`6`
`7`	`7`	`namespace NDetail {`
`8`	`8`
`9`		`-TVector<TSQLHint> ParseSqlHints(NYql::TPosition commentPos, const TStringBuf& comment);`
	`9`	`+TVector<TSQLHint> ParseSqlHints(NYql::TPosition commentPos, const TStringBuf& comment, bool utf8Aware);`
`10`	`10`
`11`	`11`	`}`
`12`	`12`
Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ TSQLHints CollectHints(const TString& query, bool antlr4Parser) {`
`15`	`15`	`TSQLHints result;`
`16`	`16`	`NYql::TIssues issues;`
`17`	`17`	`size_t maxErrors = 100;`
`18`		`- UNIT_ASSERT(CollectSqlHints(*lexer, query, "", "", result, issues, maxErrors));`
	`18`	`+ UNIT_ASSERT(CollectSqlHints(*lexer, query, "", "", result, issues, maxErrors, false));`
`19`	`19`	`UNIT_ASSERT(issues.Empty());`
`20`	`20`	`return result;`
`21`	`21`	`}`
Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@ using namespace NSQLTranslation;`
`8`	`8`	`using namespace NSQLTranslation::NDetail;`
`9`	`9`
`10`	`10`	`void CheckParse(TStringBuf comment, TStringBuf expected) {`
`11`		`- TString parsed = JoinSeq(",", ParseSqlHints({}, comment));`
	`11`	`+ TString parsed = JoinSeq(",", ParseSqlHints({}, comment, false));`
`12`	`12`	`UNIT_ASSERT_NO_DIFF(parsed, expected);`
`13`	`13`	`}`
`14`	`14`