YT: Add from expression with alias to query builder

71aa1ef022e2155cccb75c25dbc315714d5afbf9

3 files changed, 24 insertions, 1 deletions

diff --git a/yt/yt/client/query_client/query_builder.cpp b/yt/yt/client/query_client/query_builder.cpp
index e10a0679ef..0c2100e899 100644
--- a/yt/yt/client/query_client/query_builder.cpp
+++ b/yt/yt/client/query_client/query_builder.cpp
@@ -22,6 +22,12 @@ void TQueryBuilder::SetSource(TString source)
     Source_ = std::move(source);
 }
 
+void TQueryBuilder::SetSource(TString source, TString alias)
+{
+    Source_ = std::move(source);
+    SourceAlias_ = std::move(alias);
+}
+
 int TQueryBuilder::AddSelectExpression(TString expression)
 {
     SelectEntries_.push_back(TEntryWithAlias{
@@ -124,7 +130,11 @@ TString TQueryBuilder::Build()
     if (!Source_) {
         THROW_ERROR_EXCEPTION("Source must be specified in query");
     }
-    parts.push_back(Format("FROM [%v]", *Source_));
+    if (!SourceAlias_) {
+        parts.push_back(Format("FROM [%v]", *Source_));
+    } else {
+        parts.push_back(Format("FROM [%v] AS %v", *Source_, *SourceAlias_));
+    }
 
     for (const auto& join : JoinEntries_) {
         TStringBuf joinType = join.Type == ETableJoinType::Inner ? "JOIN" : "LEFT JOIN";
diff --git a/yt/yt/client/query_client/query_builder.h b/yt/yt/client/query_client/query_builder.h
index 95b87d1766..ab509d970c 100644
--- a/yt/yt/client/query_client/query_builder.h
+++ b/yt/yt/client/query_client/query_builder.h
@@ -24,6 +24,7 @@ class TQueryBuilder
 {
 public:
     void SetSource(TString source);
+    void SetSource(TString source, TString alias);
 
     int AddSelectExpression(TString expression);
     int AddSelectExpression(TString expression, TString alias);
@@ -70,6 +71,7 @@ private:
 
 private:
     std::optional<TString> Source_;
+    std::optional<TString> SourceAlias_;
     std::vector<TEntryWithAlias> SelectEntries_;
     std::vector<TString> WhereConjuncts_;
     std::vector<TOrderByEntry> OrderByEntries_;
diff --git a/yt/yt/client/unittests/query_builder_ut.cpp b/yt/yt/client/unittests/query_builder_ut.cpp
index 535dfdf244..2a6313c495 100644
--- a/yt/yt/client/unittests/query_builder_ut.cpp
+++ b/yt/yt/client/unittests/query_builder_ut.cpp
@@ -51,6 +51,17 @@ TEST(TQueryBuilderTest, Simple)
         "LIMIT 43");
 }
 
+TEST(TQueryBuilderTest, SourceAlias)
+{
+    TQueryBuilder b;
+    b.AddSelectExpression("t_alias.x");
+    b.SetSource("//t", "t_alias");
+
+    EXPECT_EQ(b.Build(),
+        "(t_alias.x) "
+        "FROM [//t] AS t_alias");
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace