YQL-14123: EACH: Do not ignore TablePathPrefix

YQL-14123: EACH: Do not ignore TablePathPrefix

ref:ecb7ef0219145840bc26801ccb71ba11c9b054b0

5 files changed, 40 insertions, 9 deletions

diff --git a/ydb/library/yql/core/services/yql_eval_expr.cpp b/ydb/library/yql/core/services/yql_eval_expr.cpp
index 811f22a605..3b8f0f0416 100644
--- a/ydb/library/yql/core/services/yql_eval_expr.cpp
+++ b/ydb/library/yql/core/services/yql_eval_expr.cpp
@@ -12,6 +12,7 @@
 #include <ydb/library/yql/providers/common/schema/expr/yql_expr_schema.h>
 #include <ydb/library/yql/providers/result/expr_nodes/yql_res_expr_nodes.h>
 #include <ydb/library/yql/utils/log/log.h>
+#include <ydb/library/yql/utils/yql_paths.h>
 
 #include <library/cpp/yson/node/node_io.h>
 #include <library/cpp/string_utils/base64/base64.h>
@@ -797,7 +798,15 @@ IGraphTransformer::TStatus EvaluateExpression(const TExprNode::TPtr& input, TExp
 
         if (node->IsCallable("MrTableEach") || node->IsCallable("MrTableEachStrict")) {
             TExprNode::TListType keys;
+            TStringBuf prefix;
+            if (node->TailPtr()->IsAtom()) {
+                prefix = node->TailPtr()->Content();
+            }
             for (const auto& eachKey : node->Children()) {
+                if (eachKey->IsAtom()) {
+                    continue;
+                }
+
                 if (!eachKey->IsCallable("Key")) {
                     ctx.AddError(TIssue(ctx.GetPosition(eachKey->Pos()), TStringBuilder() << "Expected Key"));
                     return nullptr;
@@ -834,7 +843,9 @@ IGraphTransformer::TStatus EvaluateExpression(const TExprNode::TPtr& input, TExp
                         ctx.AddError(TIssue(ctx.GetPosition(node->Pos()), TStringBuilder() << "Expected literal string as table name"));
                         return nullptr;
                     }
-
+                    if (prefix) {
+                        name = ctx.ChangeChild(*name, 0, ctx.NewAtom(node->Pos(), BuildTablePath(prefix, name->Child(0)->Content())));
+                    }
                     keys.push_back(ctx.ReplaceNode(TExprNode::TPtr(eachKey), *list, std::move(name)));
                 }
             }
diff --git a/ydb/library/yql/sql/v1/context.cpp b/ydb/library/yql/sql/v1/context.cpp
index 7beccd30f4..f86f6e1d73 100644
--- a/ydb/library/yql/sql/v1/context.cpp
+++ b/ydb/library/yql/sql/v1/context.cpp
@@ -226,18 +226,27 @@ bool TContext::SetPathPrefix(const TString& value, TMaybe<TString> arg) {
 }
 
 TNodePtr TContext::GetPrefixedPath(const TString& service, const TDeferredAtom& cluster, const TDeferredAtom& path) {
-    auto* clusterPrefix = cluster.GetLiteral() ? ClusterPathPrefixes.FindPtr(*cluster.GetLiteral()) : nullptr;
+    TStringBuf prefixPath = GetPrefixPath(service, cluster);
+    if (prefixPath) {
+        return AddTablePathPrefix(*this, prefixPath, path);
+    }
+    return path.Build();
+}
+
+TStringBuf TContext::GetPrefixPath(const TString& service, const TDeferredAtom& cluster) const {
+    auto* clusterPrefix = cluster.GetLiteral()
+                            ? ClusterPathPrefixes.FindPtr(*cluster.GetLiteral())
+                            : nullptr;
     if (clusterPrefix && !clusterPrefix->empty()) {
-        return AddTablePathPrefix(*this, *clusterPrefix, path);
+        return *clusterPrefix;
     } else {
         auto* providerPrefix = ProviderPathPrefixes.FindPtr(service);
         if (providerPrefix && !providerPrefix->empty()) {
-            return AddTablePathPrefix(*this, *providerPrefix, path);
+            return *providerPrefix;
         } else if (!PathPrefix.empty()) {
-            return AddTablePathPrefix(*this, PathPrefix, path);
+            return PathPrefix;
         }
-
-        return path.Build();
+        return {};
     }
 }
 
diff --git a/ydb/library/yql/sql/v1/context.h b/ydb/library/yql/sql/v1/context.h
index e65d673ea1..bccce4e77b 100644
--- a/ydb/library/yql/sql/v1/context.h
+++ b/ydb/library/yql/sql/v1/context.h
@@ -145,6 +145,7 @@ namespace NSQLTranslationV1 {
         bool SetPathPrefix(const TString& value, TMaybe<TString> arg = TMaybe<TString>());
 
         TNodePtr GetPrefixedPath(const TString& service, const TDeferredAtom& cluster, const TDeferredAtom& path);
+        TStringBuf GetPrefixPath(const TString& service, const TDeferredAtom& cluster) const;
 
         TNodePtr UniversalAlias(const TString& baseName, TNodePtr&& node);
 
@@ -223,6 +224,7 @@ namespace NSQLTranslationV1 {
         bool PragmaAllowDotInAlias = false;
         bool PragmaInferSchema = false;
         bool PragmaAutoCommit = false;
+        bool PragmaUseTablePrefixForEach = false;
         bool SimpleColumns = true;
         bool CoalesceJoinKeysOnQualifiedAll = false;
         bool PragmaDirectRead = false;
diff --git a/ydb/library/yql/sql/v1/query.cpp b/ydb/library/yql/sql/v1/query.cpp
index 23cfe4d344..df41d79a87 100644
--- a/ydb/library/yql/sql/v1/query.cpp
+++ b/ydb/library/yql/sql/v1/query.cpp
@@ -356,13 +356,19 @@ public:
 
                 auto type = Y("ListType", Y("DataType", Q("String")));
                 auto key = Y("Key", Q(Y(Q("table"), Y("EvaluateExpr",
-                    Y("EnsureType", Y("Coalesce", arg.Expr, Y("List", type)), type)))));
+                    Y("EnsureType", Y("Coalesce", arg.Expr,
+                    Y("List", type)), type)))));
                 if (!arg.View.empty()) {
                     key = L(key, Q(Y(Q("view"), Y("String", BuildQuotedAtom(Pos, arg.View)))));
                 }
                 each = L(each, key);
             }
-
+            if (ctx.PragmaUseTablePrefixForEach) {
+                TStringBuf prefixPath = ctx.GetPrefixPath(Service, Cluster);
+                if (prefixPath) {
+                    each = L(each, BuildQuotedAtom(Pos, TString(prefixPath)));
+                }
+            }
             return each;
         }
         else if (func == "folder") {
diff --git a/ydb/library/yql/sql/v1/sql.cpp b/ydb/library/yql/sql/v1/sql.cpp
index 0981adde7a..d9a050170c 100644
--- a/ydb/library/yql/sql/v1/sql.cpp
+++ b/ydb/library/yql/sql/v1/sql.cpp
@@ -9335,6 +9335,9 @@ TNodePtr TSqlQuery::PragmaStatement(const TRule_pragma_stmt& stmt, bool& success
         } else if (normalizedPragma == "autocommit") {
             Ctx.PragmaAutoCommit = true;
             Ctx.IncrementMonCounter("sql_pragma", "AutoCommit");
+        } else if (normalizedPragma == "usetableprefixforeach") {
+            Ctx.PragmaUseTablePrefixForEach = true;
+            Ctx.IncrementMonCounter("sql_pragma", "UseTablePrefixForEach");
         } else if (normalizedPragma == "tablepathprefix") {
             TString value;
             TMaybe<TString> arg;