YQL-19747: Support table completion at CONCAT

vityaman · robot-piglet · commit 64d942a7a113 · 2025-06-11T14:25:30.000+03:00
--- - Related to `YQL-19747` - Related to vityaman#62 --- Pull Request resolved: ytsaurus/ytsaurus#1323 commit_hash:ebab6c3290ba984174c85bba35eeb066b53af5aa
diff --git a/yql/essentials/sql/v1/complete/analysis/global/function.cpp b/yql/essentials/sql/v1/complete/analysis/global/function.cpp
@@ -0,0 +1,51 @@
+#include "function.h"
+
+#include "narrowing_visitor.h"
+
+namespace NSQLComplete {
+
+    namespace {
+
+        class TVisitor: public TSQLv1NarrowingVisitor {
+        public:
+            TVisitor(antlr4::TokenStream* tokens, size_t cursorPosition)
+                : TSQLv1NarrowingVisitor(tokens, cursorPosition)
+            {
+            }
+
+            std::any visit(antlr4::tree::ParseTree* tree) override {
+                if (IsEnclosing(tree)) {
+                    return TSQLv1NarrowingVisitor::visit(tree);
+                }
+                return {};
+            }
+
+            std::any visitTable_ref(SQLv1::Table_refContext* ctx) override {
+                auto* function = ctx->an_id_expr();
+                auto* lparen = ctx->TOKEN_LPAREN();
+                if (function == nullptr || lparen == nullptr) {
+                    return {};
+                }
+
+                if (CursorPosition() <= TextInterval(lparen).b) {
+                    return {};
+                }
+
+                return function->getText();
+            }
+        };
+
+    } // namespace
+
+    TMaybe<TString> EnclosingFunction(
+        SQLv1::Sql_queryContext* ctx,
+        antlr4::TokenStream* tokens,
+        size_t cursorPosition) {
+        std::any result = TVisitor(tokens, cursorPosition).visit(ctx);
+        if (!result.has_value()) {
+            return Nothing();
+        }
+        return std::any_cast<std::string>(result);
+    }
+
+} // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/function.h b/yql/essentials/sql/v1/complete/analysis/global/function.h
@@ -0,0 +1,15 @@
+#pragma once
+
+#include "parse_tree.h"
+
+#include <util/generic/maybe.h>
+#include <util/generic/string.h>
+
+namespace NSQLComplete {
+
+    TMaybe<TString> EnclosingFunction(
+        SQLv1::Sql_queryContext* ctx,
+        antlr4::TokenStream* tokens,
+        size_t cursorPosition);
+
+} // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.cpp b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
@@ -1,5 +1,6 @@
 #include "global.h"
 
+#include "function.h"
 #include "named_node.h"
 #include "parse_tree.h"
 #include "use.h"
@@ -52,6 +53,7 @@ namespace NSQLComplete {
             // TODO(YQL-19747): Add ~ParseContext(Tokens, ParseTree, CursorPosition)
             ctx.Use = FindUseStatement(sqlQuery, &Tokens_, input.CursorPosition, env);
             ctx.Names = CollectNamedNodes(sqlQuery, &Tokens_, input.CursorPosition);
+            ctx.EnclosingFunction = EnclosingFunction(sqlQuery, &Tokens_, input.CursorPosition);
 
             return ctx;
         }
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.h b/yql/essentials/sql/v1/complete/analysis/global/global.h
@@ -18,6 +18,7 @@ namespace NSQLComplete {
     struct TGlobalContext {
         TMaybe<TUseContext> Use;
         TVector<TString> Names;
+        TMaybe<TString> EnclosingFunction;
     };
 
     class IGlobalAnalysis {
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp b/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp
@@ -69,4 +69,33 @@ Y_UNIT_TEST_SUITE(GlobalAnalysisTests) {
         UNIT_ASSERT_VALUES_EQUAL(ctx.Names, expected);
     }
 
+    Y_UNIT_TEST(EnclosingFunctionName) {
+        IGlobalAnalysis::TPtr global = MakeGlobalAnalysis();
+        {
+            TString query = "SELECT * FROM Concat(#)";
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+            UNIT_ASSERT_VALUES_EQUAL(ctx.EnclosingFunction, "Concat");
+        }
+        {
+            TString query = "SELECT * FROM Concat(a, #)";
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+            UNIT_ASSERT_VALUES_EQUAL(ctx.EnclosingFunction, "Concat");
+        }
+        {
+            TString query = "SELECT * FROM Concat(a#)";
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+            UNIT_ASSERT_VALUES_EQUAL(ctx.EnclosingFunction, "Concat");
+        }
+        {
+            TString query = "SELECT * FROM Concat(#";
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+            UNIT_ASSERT_VALUES_EQUAL(ctx.EnclosingFunction, Nothing());
+        }
+        {
+            TString query = "SELECT * FROM (#)";
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+            UNIT_ASSERT_VALUES_EQUAL(ctx.EnclosingFunction, Nothing());
+        }
+    }
+
 } // Y_UNIT_TEST_SUITE(GlobalAnalysisTests)
diff --git a/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.cpp b/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.cpp
@@ -12,10 +12,21 @@ namespace NSQLComplete {
         return TextInterval(node).a < static_cast<ssize_t>(CursorPosition_);
     }
 
+    std::any TSQLv1NarrowingVisitor::aggregateResult(std::any aggregate, std::any nextResult) {
+        if (nextResult.has_value()) {
+            return nextResult;
+        }
+        return aggregate;
+    }
+
     bool TSQLv1NarrowingVisitor::IsEnclosing(antlr4::tree::ParseTree* tree) const {
         return TextInterval(tree).properlyContains(CursorInterval());
     }
 
+    ssize_t TSQLv1NarrowingVisitor::CursorPosition() const {
+        return CursorPosition_;
+    }
+
     antlr4::misc::Interval TSQLv1NarrowingVisitor::TextInterval(antlr4::tree::ParseTree* tree) const {
         auto tokens = tree->getSourceInterval();
         if (tokens.b == -1) {
@@ -27,7 +38,8 @@ namespace NSQLComplete {
     }
 
     antlr4::misc::Interval TSQLv1NarrowingVisitor::CursorInterval() const {
-        return antlr4::misc::Interval(CursorPosition_, CursorPosition_);
+        auto cursor = CursorPosition();
+        return antlr4::misc::Interval(cursor, cursor);
     }
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.h b/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.h
@@ -9,13 +9,15 @@ namespace NSQLComplete {
         TSQLv1NarrowingVisitor(antlr4::TokenStream* tokens, size_t cursorPosition);
 
     protected:
-        bool IsEnclosing(antlr4::tree::ParseTree* tree) const;
         bool shouldVisitNextChild(antlr4::tree::ParseTree* node, const std::any& /*currentResult*/) override;
+        std::any aggregateResult(std::any aggregate, std::any nextResult) override;
 
-    private:
+        bool IsEnclosing(antlr4::tree::ParseTree* tree) const;
+        ssize_t CursorPosition() const;
         antlr4::misc::Interval TextInterval(antlr4::tree::ParseTree* tree) const;
         antlr4::misc::Interval CursorInterval() const;
 
+    private:
         antlr4::TokenStream* Tokens_;
         size_t CursorPosition_;
     };
diff --git a/yql/essentials/sql/v1/complete/analysis/global/use.cpp b/yql/essentials/sql/v1/complete/analysis/global/use.cpp
@@ -54,13 +54,6 @@ namespace NSQLComplete {
                 };
             }
 
-            std::any aggregateResult(std::any aggregate, std::any nextResult) override {
-                if (nextResult.has_value()) {
-                    return nextResult;
-                }
-                return aggregate;
-            }
-
         private:
             TMaybe<TString> GetId(SQLv1::Pure_column_or_namedContext* ctx) const {
                 if (auto* x = ctx->bind_parameter()) {
diff --git a/yql/essentials/sql/v1/complete/analysis/global/ya.make b/yql/essentials/sql/v1/complete/analysis/global/ya.make
@@ -2,6 +2,7 @@ LIBRARY()
 
 SRCS(
     evaluate.cpp
+    function.cpp
     global.cpp
     named_node.cpp
     narrowing_visitor.cpp
diff --git a/yql/essentials/sql/v1/complete/sql_complete.cpp b/yql/essentials/sql/v1/complete/sql_complete.cpp
@@ -79,7 +79,7 @@ namespace NSQLComplete {
 
         TNameRequest NameRequestFrom(
             TCompletionInput input,
-            const TLocalSyntaxContext& context,
+            const TLocalSyntaxContext& context, // TODO(YQL-19747): rename to `local`
             const TGlobalContext& global) const {
             TNameRequest request = {
                 .Prefix = TString(GetCompletedToken(input, context.EditRange).Content),
@@ -134,6 +134,14 @@ namespace NSQLComplete {
                 request.Constraints.Cluster = std::move(constraints);
             }
 
+            if (auto name = global.EnclosingFunction.Transform(NormalizeName);
+                name && name == "concat") {
+                auto& object = request.Constraints.Object;
+                object = object.Defined() ? object : TObjectNameConstraints();
+                object->Kinds.emplace(EObjectKind::Folder);
+                object->Kinds.emplace(EObjectKind::Table);
+            }
+
             return request;
         }
 
@@ -196,14 +204,14 @@ namespace NSQLComplete {
 
                 if constexpr (std::is_base_of_v<TFolderName, T>) {
                     name.Indentifier.append('/');
-                    if (!context.Object->IsQuoted) {
+                    if (!context.Object || !context.Object->IsQuoted) {
                         name.Indentifier.prepend('`');
                     }
                     return {ECandidateKind::FolderName, std::move(name.Indentifier)};
                 }
 
                 if constexpr (std::is_base_of_v<TTableName, T>) {
-                    if (!context.Object->IsQuoted) {
+                    if (!context.Object || !context.Object->IsQuoted) {
                         name.Indentifier.prepend('`');
                     }
                     return {ECandidateKind::TableName, std::move(name.Indentifier)};
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -1026,6 +1026,20 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         }
     }
 
+    Y_UNIT_TEST(TableAsFunctionArgument) {
+        auto engine = MakeSqlCompletionEngineUT();
+
+        UNIT_ASSERT_VALUES_EQUAL(
+            CompleteTop(1, engine, "SELECT * FROM Concat(#)").at(0).Kind, FolderName);
+        UNIT_ASSERT_VALUES_EQUAL(
+            CompleteTop(1, engine, "SELECT * FROM CONCAT(#)").at(0).Kind, FolderName);
+        UNIT_ASSERT_VALUES_EQUAL(
+            CompleteTop(1, engine, "SELECT * FROM CONCAT(a, #)").at(0).Kind, FolderName);
+
+        UNIT_ASSERT_VALUES_UNEQUAL(
+            CompleteTop(1, engine, "SELECT Max(#)").at(0).Kind, FolderName);
+    }
+
     Y_UNIT_TEST(Typing) {
         const auto queryUtf16 = TUtf16String::FromUtf8(
             "SELECT \n"
