Support views and TablePathPrefix in partions/partition_list

commit_hash:2e7c16f7cf69959623ab20e6a44f27be38afe70b

1 files changed, 33 insertions, 17 deletions

diff --git a/yql/essentials/sql/v1/query.cpp b/yql/essentials/sql/v1/query.cpp
index 7a3ceef7ebc..4a79786d309 100644
--- a/yql/essentials/sql/v1/query.cpp
+++ b/yql/essentials/sql/v1/query.cpp
@@ -848,7 +848,15 @@ public:
                 ctx.Error(Pos_) << "Expecting expression as argument";
                 return nullptr;
             }
-            return Y(func.EndsWith("strict") ? "MrPartitionListStrict" : "MrPartitionList", Y("EvaluateExpr", arg.Expr));
+
+            auto partitionList = Y(func.EndsWith("strict") ? "MrPartitionListStrict" : "MrPartitionList", Y("EvaluateExpr", arg.Expr));
+            if (ctx.PragmaUseTablePrefixForEach) {
+                TStringBuf prefixPath = ctx.GetPrefixPath(Service_, Cluster_);
+                if (prefixPath) {
+                    partitionList = L(partitionList, BuildQuotedAtom(Pos_, TString(prefixPath)));
+                }
+            }
+            return partitionList;
         } else if (func == "partitions" || func == "partitionsstrict") {
             auto requiredLangVer = MakeLangVersion(2025, 4);
             if (!IsBackwardCompatibleFeatureAvailable(ctx.Settings.LangVer, requiredLangVer, ctx.Settings.BackportMode)) {
@@ -858,37 +866,45 @@ public:
                 return nullptr;
             }
 
-            if (ctx.DiscoveryMode) {
-                ctx.Error(Pos_, TIssuesIds::YQL_NOT_ALLOWED_IN_DISCOVERY) << "PARTITIONS is not allowed in Discovery mode";
+            const size_t minArgs = 2;
+            const size_t maxArgs = 3;
+            if (Args_.size() < minArgs || Args_.size() > maxArgs) {
+                ctx.Error(Pos_) << Func_ << " requires from " << minArgs << " to " << maxArgs << " arguments, but got: " << Args_.size();
                 return nullptr;
             }
 
-            if (Args_.size() != 2) {
-                ctx.Error(Pos_) << "PARTITIONS requires 2 arguments, but got: " << Args_.size();
+            if (ctx.DiscoveryMode) {
+                ctx.Error(Pos_, TIssuesIds::YQL_NOT_ALLOWED_IN_DISCOVERY) << "PARTITIONS is not allowed in Discovery mode";
                 return nullptr;
             }
 
-            if (Args_[0].HasAt || Args_[1].HasAt) {
-                ctx.Error(Pos_) << "Temporary tables are not supported here";
-                return nullptr;
-            }
+            for (auto& arg : Args_) {
+                if (arg.HasAt) {
+                    ctx.Error(Pos_) << "Temporary tables are not supported here";
+                    return nullptr;
+                }
 
-            if (!Args_[1].View.empty()) {
-                ctx.Error(Pos_) << "VIEW should be used only in first argument";
-                return nullptr;
+                if (!arg.View.empty()) {
+                    ctx.Error(Pos_) << "Use the last argument of " << Func_ << " to specify a VIEW." << Endl
+                                    << "Possible arguments are: prefix, pattern, view.";
+                    return nullptr;
+                }
+
+                ExtractTableName(ctx, arg);
             }
 
-            ExtractTableName(ctx, Args_[0]);
-            ExtractTableName(ctx, Args_[1]);
             auto path = ctx.GetPrefixedPath(Service_, Cluster_, Args_[0].Id);
             if (!path) {
                 return nullptr;
             }
             TNodePtr key = Y("Key", Q(Y(Q("table"), Y("String", path))));
-            key = AddView(key, Args_[0].View);
-            if (!ValidateView(GetPos(), ctx, Service_, Args_[0].View)) {
-                return nullptr;
+            if (Args_.size() == maxArgs) {
+                auto& lastArg = Args_.back();
+                if (!lastArg.Id.Empty()) {
+                    key = L(key, Q(Y(Q("view"), Y("String", lastArg.Id.Build()))));
+                }
             }
+
             TDeferredAtom pattern = Args_[1].Id;
             return Y(func.EndsWith("strict") ? "MrPartitionsStrict" : "MrPartitions", key, pattern.Build());
         }