YQL-19747: Support qualified asterisk

Support `SELECT x.* FROM t AS x`.
commit_hash:64693f65281f385d6c0d3541fb2874b0153aa88c

6 files changed, 145 insertions, 21 deletions

diff --git a/yql/essentials/sql/v1/complete/analysis/global/column.cpp b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
index 00da14c4b2d..9a9fd7f18a2 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/column.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
@@ -11,7 +11,7 @@ namespace NSQLComplete {
 
         // TODO: Extract it to `identifier.cpp` and reuse it also at `use.cpp`
         //       and replace `GetId` at `parse_tree.cpp`.
-        class TIdentifierVisitor: public TSQLv1BaseVisitor {
+        class TIdentifierVisitor: public SQLv1Antlr4BaseVisitor {
         public:
             std::any visitCluster_expr(SQLv1::Cluster_exprContext* ctx) override {
                 if (auto* x = ctx->pure_column_or_named()) {
@@ -27,14 +27,30 @@ namespace NSQLComplete {
                 return {};
             }
 
+            std::any visitUnary_casual_subexpr(SQLv1::Unary_casual_subexprContext* ctx) override {
+                std::any prev;
+                if (auto* x = ctx->id_expr()) {
+                    prev = visit(x);
+                } else if (auto* x = ctx->atom_expr()) {
+                    prev = visit(x);
+                }
+
+                std::any next = visit(ctx->unary_subexpr_suffix());
+                if (!next.has_value()) {
+                    return prev;
+                }
+
+                return {};
+            }
+
             std::any visitTerminal(antlr4::tree::TerminalNode* node) override {
-                TString text = GetText(node);
                 switch (node->getSymbol()->getType()) {
-                    case SQLv1::TOKEN_ID_QUOTED: {
-                        text = Unquoted(std::move(text));
-                    } break;
+                    case SQLv1::TOKEN_ID_QUOTED:
+                        return Unquoted(GetText(node));
+                    case SQLv1::TOKEN_ID_PLAIN:
+                        return GetText(node);
                 }
-                return text;
+                return {};
             }
 
         private:
@@ -101,25 +117,53 @@ namespace NSQLComplete {
             }
 
             std::any visitSelect_core(SQLv1::Select_coreContext* ctx) override {
-                TMaybe<TColumnContext> head = Head(ctx);
-                if (head.Empty() || head->IsAsterisk()) {
-                    return AccumulatingVisit(ctx->join_source());
+                TColumnContext context = AccumulatingVisit(ctx->result_column());
+                auto asterisks = std::ranges::partition(context.Columns, [](const TColumnId& x) {
+                    return x.Name != "*";
+                });
+
+                if (std::ranges::empty(asterisks)) {
+                    return context;
                 }
 
-                return AccumulatingVisit(ctx->result_column());
+                TColumnContext source = AccumulatingVisit(ctx->join_source());
+
+                TColumnContext imported;
+                for (const TColumnId& qualified : asterisks) {
+                    auto aliased = source.ExtractAliased(qualified.TableAlias);
+                    imported = std::move(imported) | std::move(aliased);
+                }
+
+                context.Columns.erase(asterisks.begin(), asterisks.end());
+                imported = std::move(imported).Renamed("");
+                return std::move(context) | std::move(imported);
             }
 
             std::any visitResult_column(SQLv1::Result_columnContext* ctx) override {
-                if (ctx->TOKEN_ASTERISK() != nullptr) {
+                if (ctx->opt_id_prefix() == nullptr && ctx->TOKEN_ASTERISK() != nullptr) {
                     return TColumnContext::Asterisk();
                 }
 
+                if (ctx->opt_id_prefix() != nullptr && ctx->TOKEN_ASTERISK() != nullptr) {
+                    TMaybe<TString> alias = GetId(ctx->opt_id_prefix()->an_id());
+                    if (alias.Empty()) {
+                        return TColumnContext::Asterisk();
+                    }
+
+                    return TColumnContext{
+                        .Columns = {
+                            {.TableAlias = std::move(*alias), .Name = "*"},
+                        },
+                    };
+                }
+
                 TMaybe<TString> column = GetAlias(ctx);
                 if (column.Defined()) {
                     return TColumnContext{
                         .Columns = {
                             {.Name = std::move(*column)},
-                        }};
+                        },
+                    };
                 }
 
                 return {};
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.cpp b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
index 555d5283c14..f11447c1f12 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/global.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
@@ -19,21 +19,38 @@ namespace NSQLComplete {
         return std::tie(lhs.TableAlias, lhs.Name) < std::tie(rhs.TableAlias, rhs.Name);
     }
 
-    TVector<TAliased<TTableId>> TColumnContext::TablesWithAlias(TStringBuf alias) const {
-        if (alias.empty()) {
-            return TVector<TAliased<TTableId>>(Tables.begin(), Tables.end());
+    TColumnContext TColumnContext::ExtractAliased(TMaybe<TStringBuf> alias) {
+        if (alias.Empty()) {
+            return *this;
         }
 
-        auto filtered = NFuncTools::Filter([&](const auto& x) { return x.Alias == alias; }, Tables);
-        return TVector<TAliased<TTableId>>(filtered.begin(), filtered.end());
+        auto aliasedTables = std::ranges::partition(Tables, [&](const auto& table) {
+            return table.Alias != alias;
+        });
+
+        auto aliasedColumns = std::ranges::partition(Columns, [&](const auto& column) {
+            return column.TableAlias != alias;
+        });
+
+        TVector<TAliased<TTableId>> tables(aliasedTables.begin(), aliasedTables.end());
+        TVector<TColumnId> columns(aliasedColumns.begin(), aliasedColumns.end());
+
+        Tables.erase(aliasedTables.begin(), aliasedTables.end());
+        Columns.erase(aliasedColumns.begin(), aliasedColumns.end());
+
+        return {
+            .Tables = std::move(tables),
+            .Columns = std::move(columns),
+        };
     }
 
     bool TColumnContext::IsAsterisk() const {
-        return Columns.size() == 1 && Columns[0].Name == "*";
+        return Columns.size() == 1 &&
+               Columns[0].TableAlias.empty() &&
+               Columns[0].Name == "*";
     }
 
     TColumnContext TColumnContext::Renamed(TStringBuf alias) && {
-        Y_ENSURE(!alias.empty());
         for (TAliased<TTableId>& table : Tables) {
             table.Alias = alias;
         }
@@ -152,6 +169,12 @@ namespace NSQLComplete {
             }
         }
 
+        void DebugPrint(TStringBuf query, antlr4::ParserRuleContext* ctx) {
+            Cerr << "= = = = = = " << Endl;
+            Cerr << query << Endl;
+            Cerr << ctx->toStringTree(&Parser_, true) << Endl;
+        }
+
         antlr4::ANTLRInputStream Chars_;
         G::TLexer Lexer_;
         antlr4::CommonTokenStream Tokens_;
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.h b/yql/essentials/sql/v1/complete/analysis/global/global.h
index 98445540b31..b29d341c1a6 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/global.h
+++ b/yql/essentials/sql/v1/complete/analysis/global/global.h
@@ -29,8 +29,8 @@ namespace NSQLComplete {
         TVector<TAliased<TTableId>> Tables;
         TVector<TColumnId> Columns;
 
-        TVector<TAliased<TTableId>> TablesWithAlias(TStringBuf alias) const;
         bool IsAsterisk() const;
+        TColumnContext ExtractAliased(TMaybe<TStringBuf> alias);
         TColumnContext Renamed(TStringBuf alias) &&;
 
         friend bool operator==(const TColumnContext& lhs, const TColumnContext& rhs) = default;
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp b/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp
index 15aba3ee155..830b3c2a9ee 100644
--- a/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp
+++ b/yql/essentials/sql/v1/complete/analysis/global/global_ut.cpp
@@ -203,6 +203,40 @@ Y_UNIT_TEST_SUITE(GlobalAnalysisTests) {
             };
             UNIT_ASSERT_VALUES_EQUAL(ctx.Column, expected);
         }
+        {
+            TString query = R"(
+                SELECT # FROM (
+                    SELECT x.*, y.name, e
+                    FROM (SELECT a.*, d FROM a AS a JOIN c AS c ON TRUE) AS x
+                    JOIN b AS y
+                )
+            )";
+
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+
+            TColumnContext expected = {
+                .Tables = {
+                    TAliased<TTableId>("", {"", "a"}),
+                },
+                .Columns = {
+                    {.Name = "d"},
+                    {.Name = "e"},
+                },
+            };
+            UNIT_ASSERT_VALUES_EQUAL(ctx.Column, expected);
+        }
+        {
+            TString query = "SELECT # FROM (SELECT 1, *, 2 FROM t)";
+
+            TGlobalContext ctx = global->Analyze(SharpedInput(query), {});
+
+            TColumnContext expected = {
+                .Tables = {
+                    TAliased<TTableId>("", {"", "t"}),
+                },
+            };
+            UNIT_ASSERT_VALUES_EQUAL(ctx.Column, expected);
+        }
     }
 
     Y_UNIT_TEST(Projection) {
diff --git a/yql/essentials/sql/v1/complete/sql_complete.cpp b/yql/essentials/sql/v1/complete/sql_complete.cpp
index 7cb9788be5b..59cbeba1385 100644
--- a/yql/essentials/sql/v1/complete/sql_complete.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete.cpp
@@ -156,8 +156,12 @@ namespace NSQLComplete {
             }
 
             if (context.Column && global.Column) {
+                TMaybe<TStringBuf> table = context.Column->Table;
+                table = !table->empty() ? table : Nothing();
+
                 request.Constraints.Column = TColumnName::TConstraints();
-                request.Constraints.Column->Tables = global.Column->TablesWithAlias(context.Column->Table);
+                request.Constraints.Column->Tables =
+                    TColumnContext(*global.Column).ExtractAliased(table).Tables;
             }
 
             return request;
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
index aa955bc2cc5..ca9709cf5e2 100644
--- a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -1189,6 +1189,25 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
             };
             UNIT_ASSERT_VALUES_EQUAL(CompleteTop(1, engine, query), expected);
         }
+        {
+            TString query = R"(
+                SELECT #
+                FROM (
+                    SELECT epp.*, test
+                    FROM example.`/people` AS epp
+                    JOIN example.`/yql/tutorial` AS eqt ON TRUE
+                    JOIN testing ON TRUE
+                ) AS ep
+            )";
+
+            TVector<TCandidate> expected = {
+                {ColumnName, "ep.test"},
+                {ColumnName, "ep.Age"},
+                {ColumnName, "ep.Name"},
+                {Keyword, "ALL"},
+            };
+            UNIT_ASSERT_VALUES_EQUAL(CompleteTop(4, engine, query), expected);
+        }
     }
 
     Y_UNIT_TEST(Typing) {