YQL-19747 Improve yql_complete tool and add input validation

vityaman · robot-piglet · commit 5af5cb3b7b42 · 2025-04-08T00:17:32.000+03:00
No description --- Pull Request resolved: ytsaurus/ytsaurus#1185 commit_hash:1def5874ff6a9a5b3dcdd0ad285d2e64b16c9306
diff --git a/yql/essentials/parser/pg_wrapper/parser.cpp b/yql/essentials/parser/pg_wrapper/parser.cpp
@@ -2,7 +2,9 @@
 
 #include "arena_ctx.h"
 
+#include <util/charset/utf8.h>
 #include <util/generic/scope.h>
+
 #include <fcntl.h>
 #include <stdint.h>
 
@@ -219,7 +221,7 @@ void PGParse(const TString& input, IPGParseEvents& events) {
                     break;
                 }
 
-                if (!TTextWalker::IsUtf8Intermediate(input[i])) {
+                if (!IsUTF8ContinuationByte(input[i])) {
                     ++codepoints;
                 }
                 walker.Advance(input[i]);
diff --git a/yql/essentials/public/fastcheck/format.cpp b/yql/essentials/public/fastcheck/format.cpp
@@ -5,6 +5,7 @@
 #include <yql/essentials/sql/v1/proto_parser/antlr4/proto_parser.h>
 #include <yql/essentials/sql/v1/proto_parser/antlr4_ansi/proto_parser.h>
 #include <yql/essentials/core/issue/yql_issue.h>
+#include <util/charset/utf8.h>
 #include <util/string/builder.h>
 
 namespace NYql {
@@ -88,20 +89,20 @@ class TFormatRunner : public ICheckRunner {
                     continue;
                 }
 
-                while (i > 0 && TTextWalker::IsUtf8Intermediate(request.Program[i])) {
+                while (i > 0 && IsUTF8ContinuationByte(request.Program[i])) {
                     --i;
                 }
 
                 break;
             }
 
             TString formattedSample = formattedQuery.substr(i, FormatContextLimit);
-            while (!formattedSample.empty() && TTextWalker::IsUtf8Intermediate(formattedQuery.back())) {
+            while (!formattedSample.empty() && IsUTF8ContinuationByte(formattedQuery.back())) {
                 formattedSample.erase(formattedSample.size() - 1);
             }
 
             TString origSample = request.Program.substr(i, FormatContextLimit);
-            while (!origSample.empty() && TTextWalker::IsUtf8Intermediate(origSample.back())) {
+            while (!origSample.empty() && IsUTF8ContinuationByte(origSample.back())) {
                 origSample.erase(origSample.size() - 1);
             }
 
diff --git a/yql/essentials/public/issue/yql_issue.cpp b/yql/essentials/public/issue/yql_issue.cpp
@@ -54,7 +54,7 @@ TTextWalker& TTextWalker::Advance(char c) {
     }
 
     ui32 charDistance = 1;
-    if (Utf8Aware && IsUtf8Intermediate(c)) {
+    if (Utf8Aware && IsUTF8ContinuationByte(c)) {
         charDistance = 0;
     }
 
diff --git a/yql/essentials/public/issue/yql_issue.h b/yql/essentials/public/issue/yql_issue.h
@@ -63,10 +63,6 @@ class TTextWalker {
     {
     }
 
-    static inline bool IsUtf8Intermediate(char c) {
-        return (c & 0xC0) == 0x80;
-    }
-
     template<typename T>
     TTextWalker& Advance(const T& buf) {
         for (char c : buf) {
diff --git a/yql/essentials/sql/v1/complete/sql_complete.cpp b/yql/essentials/sql/v1/complete/sql_complete.cpp
@@ -26,6 +26,15 @@ namespace NSQLComplete {
         }
 
         TCompletion Complete(TCompletionInput input) {
+            if (
+                input.CursorPosition < input.Text.length() &&
+                    IsUTF8ContinuationByte(input.Text.at(input.CursorPosition)) ||
+                input.Text.length() < input.CursorPosition) {
+                ythrow yexception()
+                    << "invalid cursor position " << input.CursorPosition
+                    << " for input size " << input.Text.size();
+            }
+
             auto prefix = input.Text.Head(input.CursorPosition);
             auto completedToken = GetCompletedToken(prefix);
 
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -63,8 +63,8 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         return MakeSqlCompletionEngine(std::move(lexer), std::move(service));
     }
 
-    TVector<TCandidate> Complete(ISqlCompletionEngine::TPtr& engine, TStringBuf prefix) {
-        return engine->Complete({prefix}).Candidates;
+    TVector<TCandidate> Complete(ISqlCompletionEngine::TPtr& engine, TCompletionInput input) {
+        return engine->Complete(input).Candidates;
     }
 
     Y_UNIT_TEST(Beginning) {
@@ -438,17 +438,31 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         };
 
         auto engine = MakeSqlCompletionEngineUT();
-        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, "se"), expected);
-        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, "sE"), expected);
-        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, "Se"), expected);
-        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, "SE"), expected);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"se"}), expected);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"sE"}), expected);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"Se"}), expected);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"SE"}), expected);
     }
 
     Y_UNIT_TEST(InvalidStatementsRecovery) {
         auto engine = MakeSqlCompletionEngineUT();
-        UNIT_ASSERT_GE(Complete(engine, "select select; ").size(), 35);
-        UNIT_ASSERT_GE(Complete(engine, "select select;").size(), 35);
-        UNIT_ASSERT_VALUES_EQUAL_C(Complete(engine, "!;").size(), 0, "Lexer failing");
+        UNIT_ASSERT_GE(Complete(engine, {"select select; "}).size(), 35);
+        UNIT_ASSERT_GE(Complete(engine, {"select select;"}).size(), 35);
+        UNIT_ASSERT_VALUES_EQUAL_C(Complete(engine, {"!;"}).size(), 0, "Lexer failing");
+    }
+
+    Y_UNIT_TEST(InvalidCursorPosition) {
+        auto engine = MakeSqlCompletionEngineUT();
+
+        UNIT_ASSERT_NO_EXCEPTION(Complete(engine, {"", 0}));
+        UNIT_ASSERT_EXCEPTION(Complete(engine, {"", 1}), yexception);
+
+        UNIT_ASSERT_NO_EXCEPTION(Complete(engine, {"s", 0}));
+        UNIT_ASSERT_NO_EXCEPTION(Complete(engine, {"s", 1}));
+
+        UNIT_ASSERT_NO_EXCEPTION(Complete(engine, {"ы", 0}));
+        UNIT_ASSERT_EXCEPTION(Complete(engine, {"ы", 1}), yexception);
+        UNIT_ASSERT_NO_EXCEPTION(Complete(engine, {"ы", 2}));
     }
 
     Y_UNIT_TEST(DefaultNameService) {
diff --git a/yql/essentials/tools/yql_complete/ya.make b/yql/essentials/tools/yql_complete/ya.make
@@ -7,6 +7,7 @@ PEERDIR(
     yql/essentials/sql/v1/complete
     yql/essentials/sql/v1/lexer/antlr4_pure
     yql/essentials/sql/v1/lexer/antlr4_pure_ansi
+    yql/essentials/utils
 )
 
 SRCS(
diff --git a/yql/essentials/tools/yql_complete/yql_complete.cpp b/yql/essentials/tools/yql_complete/yql_complete.cpp
@@ -6,7 +6,11 @@
 #include <yql/essentials/sql/v1/lexer/antlr4_pure/lexer.h>
 #include <yql/essentials/sql/v1/lexer/antlr4_pure_ansi/lexer.h>
 
+#include <yql/essentials/utils/utf8.h>
+
 #include <library/cpp/getopt/last_getopt.h>
+
+#include <util/charset/utf8.h>
 #include <util/stream/file.h>
 
 NSQLComplete::TFrequencyData LoadFrequencyDataFromFile(TString filepath) {
@@ -25,6 +29,11 @@ NSQLComplete::TLexerSupplier MakePureLexerSupplier() {
     };
 }
 
+size_t UTF8PositionToBytes(const TStringBuf text, size_t position) {
+    const TStringBuf substr = SubstrUTF8(text, position, text.length());
+    return substr.begin() - text.begin();
+}
+
 int Run(int argc, char* argv[]) {
     NLastGetopt::TOpts opts = NLastGetopt::TOpts::Default();
 
@@ -60,9 +69,20 @@ int Run(int argc, char* argv[]) {
             std::move(ranking)));
 
     NSQLComplete::TCompletionInput input;
+
     input.Text = queryString;
+    if (!NYql::IsUtf8(input.Text)) {
+        ythrow yexception() << "provided input is not UTF encoded";
+    }
+
     if (pos) {
-        input.CursorPosition = *pos;
+        input.CursorPosition = UTF8PositionToBytes(input.Text, *pos);
+    } else if (Count(input.Text, '#') == 1) {
+        Cerr << "Note: found an only '#', setting the cursor position\n";
+        input.CursorPosition = input.Text.find('#');
+    } else if (Count(input.Text, '#') >= 2) {
+        Cerr << "Note: found multiple '#', defaulting the cursor position\n";
+        input.CursorPosition = queryString.size();
     } else {
         input.CursorPosition = queryString.size();
     }

Original file line number	Diff line number	Diff line change
`@@ -2,7 +2,9 @@`
`2`	`2`
`3`	`3`	`#include "arena_ctx.h"`
`4`	`4`
	`5`	`+#include <util/charset/utf8.h>`
`5`	`6`	`#include <util/generic/scope.h>`
	`7`	`+`
`6`	`8`	`#include <fcntl.h>`
`7`	`9`	`#include <stdint.h>`
`8`	`10`
`@@ -219,7 +221,7 @@ void PGParse(const TString& input, IPGParseEvents& events) {`
`219`	`221`	`break;`
`220`	`222`	`}`
`221`	`223`
`222`		`- if (!TTextWalker::IsUtf8Intermediate(input[i])) {`
	`224`	`+ if (!IsUTF8ContinuationByte(input[i])) {`
`223`	`225`	`++codepoints;`
`224`	`226`	`}`
`225`	`227`	`walker.Advance(input[i]);`
Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@`
`5`	`5`	`#include <yql/essentials/sql/v1/proto_parser/antlr4/proto_parser.h>`
`6`	`6`	`#include <yql/essentials/sql/v1/proto_parser/antlr4_ansi/proto_parser.h>`
`7`	`7`	`#include <yql/essentials/core/issue/yql_issue.h>`
	`8`	`+#include <util/charset/utf8.h>`
`8`	`9`	`#include <util/string/builder.h>`
`9`	`10`
`10`	`11`	`namespace NYql {`
`@@ -88,20 +89,20 @@ class TFormatRunner : public ICheckRunner {`
`88`	`89`	`continue;`
`89`	`90`	`}`
`90`	`91`
`91`		`- while (i > 0 && TTextWalker::IsUtf8Intermediate(request.Program[i])) {`
	`92`	`+ while (i > 0 && IsUTF8ContinuationByte(request.Program[i])) {`
`92`	`93`	`--i;`
`93`	`94`	`}`
`94`	`95`
`95`	`96`	`break;`
`96`	`97`	`}`
`97`	`98`
`98`	`99`	`TString formattedSample = formattedQuery.substr(i, FormatContextLimit);`
`99`		`- while (!formattedSample.empty() && TTextWalker::IsUtf8Intermediate(formattedQuery.back())) {`
	`100`	`+ while (!formattedSample.empty() && IsUTF8ContinuationByte(formattedQuery.back())) {`
`100`	`101`	`formattedSample.erase(formattedSample.size() - 1);`
`101`	`102`	`}`
`102`	`103`
`103`	`104`	`TString origSample = request.Program.substr(i, FormatContextLimit);`
`104`		`- while (!origSample.empty() && TTextWalker::IsUtf8Intermediate(origSample.back())) {`
	`105`	`+ while (!origSample.empty() && IsUTF8ContinuationByte(origSample.back())) {`
`105`	`106`	`origSample.erase(origSample.size() - 1);`
`106`	`107`	`}`
`107`	`108`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ TTextWalker& TTextWalker::Advance(char c) {`
`54`	`54`	`}`
`55`	`55`
`56`	`56`	`ui32 charDistance = 1;`
`57`		`- if (Utf8Aware && IsUtf8Intermediate(c)) {`
	`57`	`+ if (Utf8Aware && IsUTF8ContinuationByte(c)) {`
`58`	`58`	`charDistance = 0;`
`59`	`59`	`}`
`60`	`60`
Original file line number	Diff line number	Diff line change
`@@ -63,10 +63,6 @@ class TTextWalker {`
`63`	`63`	`{`
`64`	`64`	`}`
`65`	`65`
`66`		`- static inline bool IsUtf8Intermediate(char c) {`
`67`		`- return (c & 0xC0) == 0x80;`
`68`		`- }`
`69`		`-`
`70`	`66`	`template<typename T>`
`71`	`67`	`TTextWalker& Advance(const T& buf) {`
`72`	`68`	`for (char c : buf) {`
Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@ PEERDIR(`
`7`	`7`	`yql/essentials/sql/v1/complete`
`8`	`8`	`yql/essentials/sql/v1/lexer/antlr4_pure`
`9`	`9`	`yql/essentials/sql/v1/lexer/antlr4_pure_ansi`
	`10`	`+ yql/essentials/utils`
`10`	`11`	`)`
`11`	`12`
`12`	`13`	`SRCS(`