YQL-19747: Support JOIN ON, WHERE, GROUP/ORDER BY

Before this patch we completed columns only at
SELECT projection. Now we complete it at these
constructions. Also it respects visibility rules.
commit_hash:aa1e6d4900e9b032801ddbf3bcd347750c2939b1

5 files changed, 100 insertions, 10 deletions

diff --git a/yql/essentials/sql/v1/complete/analysis/global/column.cpp b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
index 5cecb40fb13..527871bfdbf 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/column.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
@@ -130,7 +130,12 @@ namespace NSQLComplete {
 
                 TColumnContext imported;
                 for (const TColumnId& qualified : asterisks) {
-                    auto aliased = source.ExtractAliased(qualified.TableAlias);
+                    TMaybe<TStringBuf> alias = qualified.TableAlias;
+                    if (alias->Empty()) {
+                        alias = Nothing();
+                    }
+
+                    auto aliased = source.ExtractAliased(alias);
                     imported = std::move(imported) | std::move(aliased);
                 }
 
@@ -233,16 +238,29 @@ namespace NSQLComplete {
             }
 
             std::any visitSelect_core(SQLv1::Select_coreContext* ctx) override {
-                SQLv1::Join_sourceContext* source = ctx->join_source(0);
-                if (source == nullptr) {
-                    source = ctx->join_source(1);
+                antlr4::ParserRuleContext* source = nullptr;
+                if (IsEnclosingStrict(ctx->expr(0)) ||
+                    IsEnclosingStrict(ctx->group_by_clause()) ||
+                    IsEnclosingStrict(ctx->expr(1)) ||
+                    IsEnclosingStrict(ctx->window_clause()) ||
+                    IsEnclosingStrict(ctx->ext_order_by_clause())) {
+                    source = ctx;
+                } else {
+                    source = ctx->join_source(0);
+                    source = source == nullptr ? ctx->join_source(1) : source;
                 }
+
                 if (source == nullptr) {
                     return {};
                 }
 
                 return TInferenceVisitor().visit(source);
             }
+
+        private:
+            bool IsEnclosingStrict(antlr4::ParserRuleContext* ctx) const {
+                return ctx != nullptr && IsEnclosing(ctx);
+            }
         };
 
     } // namespace
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.cpp b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
index 5bdba0714b3..b69569098c8 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/global.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
@@ -9,6 +9,7 @@
 
 #include <yql/essentials/sql/v1/complete/antlr4/pipeline.h>
 #include <yql/essentials/sql/v1/complete/syntax/ansi.h>
+#include <yql/essentials/sql/v1/complete/text/word.h>
 
 #include <library/cpp/iterator/functools.h>
 
@@ -142,14 +143,11 @@ namespace NSQLComplete {
 
     private:
         bool IsRecoverable(TCompletionInput input) const {
-            static const TStringBuf prev = " ";
-            static const TStringBuf next = " .(";
-
             TStringBuf s = input.Text;
             size_t i = input.CursorPosition;
 
-            return (i < s.size() && prev.Contains(s[i]) || i == s.size()) &&
-                   (i > 0 /*  */ && next.Contains(s[i - 1]));
+            return (i < s.size() && IsWordBoundary(s[i]) || i == s.size()) &&
+                   (i > 0 /*  */ && IsWordBoundary(s[i - 1]));
         }
 
         SQLv1::Sql_queryContext* Parse(TStringBuf input) {
diff --git a/yql/essentials/sql/v1/complete/analysis/global/ya.make b/yql/essentials/sql/v1/complete/analysis/global/ya.make
index 25c83108230..cf2a9c9aada 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/ya.make
+++ b/yql/essentials/sql/v1/complete/analysis/global/ya.make
@@ -16,6 +16,7 @@ SRCS(
 PEERDIR(
     yql/essentials/sql/v1/complete/core
     yql/essentials/sql/v1/complete/syntax
+    yql/essentials/sql/v1/complete/text
     yql/essentials/parser/antlr_ast/gen/v1_antlr4
     yql/essentials/parser/antlr_ast/gen/v1_ansi_antlr4
 )
diff --git a/yql/essentials/sql/v1/complete/analysis/local/parser_call_stack.cpp b/yql/essentials/sql/v1/complete/analysis/local/parser_call_stack.cpp
index 8ac4ff17999..7b30d2a20d7 100644
--- a/yql/essentials/sql/v1/complete/analysis/local/parser_call_stack.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/local/parser_call_stack.cpp
@@ -124,7 +124,10 @@ namespace NSQLComplete {
     }
 
     bool IsLikelyColumnStack(const TParserCallStack& stack) {
-        return Contains({RULE(Result_column)}, stack);
+        return Contains({RULE(Select_core)}, stack) &&
+               (Contains({RULE(Result_column)}, stack) ||
+                Contains({RULE(Expr)}, stack) ||
+                Contains({RULE(Sort_specification)}, stack));
     }
 
     bool IsLikelyBindingStack(const TParserCallStack& stack) {
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
index ca9709cf5e2..ba422168652 100644
--- a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -1210,6 +1210,76 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         }
     }
 
+    Y_UNIT_TEST(ColumnPositions) {
+        auto engine = MakeSqlCompletionEngineUT();
+
+        TVector<TCandidate> expected = {
+            {ColumnName, "Age"},
+            {ColumnName, "Name"},
+        };
+
+        {
+            TString query = "SELECT # FROM example.`/people`";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = "SELECT f(#) FROM example.`/people`";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = "SELECT * FROM example.`/people` WHERE #";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = "SELECT * FROM example.`/people` WHERE f(#)";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = "SELECT * FROM example.`/people` ORDER BY #";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = "SELECT * FROM example.`/people` ORDER BY f(#)";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = "SELECT * FROM example.`/people` GROUP BY #";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = R"(
+                SELECT *
+                FROM example.`/people` AS a
+                JOIN example.`/people` AS b ON a.#
+            )";
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+    }
+
+    Y_UNIT_TEST(ProjectionVisibility) {
+        auto engine = MakeSqlCompletionEngineUT();
+        {
+            TString query = "SELECT Age as a, # FROM example.`/people`";
+
+            TVector<TCandidate> expected = {
+                {ColumnName, "Age"},
+                {ColumnName, "Name"},
+            };
+
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+        {
+            TString query = "SELECT Age as a, b FROM example.`/people` WHERE #";
+
+            TVector<TCandidate> expected = {
+                {ColumnName, "a"},
+                {ColumnName, "b"},
+            };
+
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, query), expected);
+        }
+    }
+
     Y_UNIT_TEST(Typing) {
         const auto queryUtf16 = TUtf16String::FromUtf8(
             "SELECT \n"