add a whitelist for pooltrees pragma

7 files changed, 67 insertions, 6 deletions

diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp
index a608637a63..6bded487cc 100644
--- a/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp
+++ b/ydb/library/yql/providers/yt/provider/yql_yt_dq_integration.cpp
@@ -10,6 +10,7 @@
 
 #include <ydb/library/yql/providers/common/dq/yql_dq_integration_impl.h>
 #include <ydb/library/yql/providers/common/codec/yql_codec_type_flags.h>
+#include <ydb/library/yql/providers/common/config/yql_dispatch.h>
 #include <ydb/library/yql/providers/dq/common/yql_dq_settings.h>
 #include <ydb/library/yql/providers/dq/expr_nodes/dqs_expr_nodes.h>
 #include <ydb/library/yql/providers/result/expr_nodes/yql_res_expr_nodes.h>
@@ -31,7 +32,8 @@
 
 namespace NYql {
 
-static const THashSet<TStringBuf> UNSUPPORTED_YT_PRAGMAS = {"maxrowweight", "pooltrees", "layerpaths", "operationspec"};
+static const THashSet<TStringBuf> UNSUPPORTED_YT_PRAGMAS = {"maxrowweight",  "layerpaths", "operationspec"};
+static const THashSet<TStringBuf> POOL_TREES_WHITELIST = {"physical",  "cloud", "cloud_default"};
 
 using namespace NNodes;
 
@@ -260,12 +262,24 @@ public:
                 }
             }
 
-            if (node.ChildrenSize() >= 4) {
-                if (node.Child(2)->Content() == "Attr" && UNSUPPORTED_YT_PRAGMAS.contains(node.Child(3)->Content())) {
-                    AddInfo(ctx, TStringBuilder() << "unsupported yt pragma: " << node.Child(3)->Content(), skipIssues);
+            if (node.ChildrenSize() >= 4 && node.Child(2)->Content() == "Attr") {
+                auto pragma = node.Child(3)->Content();
+                if (UNSUPPORTED_YT_PRAGMAS.contains(pragma)) {
+                    AddInfo(ctx, TStringBuilder() << "unsupported yt pragma: " << pragma, skipIssues);
                     State_->OnlyNativeExecution = true;
                     return false;
                 }
+
+                if (pragma == "pooltrees") {
+                    auto pools = NPrivate::GetDefaultParser<TVector<TString>>()(TString{node.Child(4)->Content()});
+                    for (const auto& pool : pools) {
+                        if (!POOL_TREES_WHITELIST.contains(pool)) {
+                            AddInfo(ctx, TStringBuilder() << "unsupported pool tree: " << pool, skipIssues);
+                            State_->OnlyNativeExecution = true;
+                            return false;
+                        }
+                    }
+                }
             }
         }
         return true;
diff --git a/ydb/library/yql/tests/sql/dq_file/part1/canondata/result.json b/ydb/library/yql/tests/sql/dq_file/part1/canondata/result.json
index 0cdb120beb..dad54058c5 100644
--- a/ydb/library/yql/tests/sql/dq_file/part1/canondata/result.json
+++ b/ydb/library/yql/tests/sql/dq_file/part1/canondata/result.json
@@ -1479,7 +1479,7 @@
         {
             "checksum": "6cbf312cad031e66f9b438541821c09b",
             "size": 5391,
-            "uri": "https://{canondata_backend}/1937027/d467c683292423d8a40039f4803bb61721a8cd03/resource.tar.gz#test.test_dq-blacklisted_pragmas1--Analyze_/plan.txt"
+            "uri": "https://{canondata_backend}/1937027/4993073f253c4402c43ec90285bef30ff79516ac/resource.tar.gz#test.test_dq-blacklisted_pragmas1--Analyze_/plan.txt"
         },
         {
             "uri": "file://test.test_dq-blacklisted_pragmas1--Analyze_/extracted"
diff --git a/ydb/library/yql/tests/sql/dq_file/part1/canondata/test.test_dq-blacklisted_pragmas1--Analyze_/extracted b/ydb/library/yql/tests/sql/dq_file/part1/canondata/test.test_dq-blacklisted_pragmas1--Analyze_/extracted
index d1933b71b3..db38395ac9 100644
--- a/ydb/library/yql/tests/sql/dq_file/part1/canondata/test.test_dq-blacklisted_pragmas1--Analyze_/extracted
+++ b/ydb/library/yql/tests/sql/dq_file/part1/canondata/test.test_dq-blacklisted_pragmas1--Analyze_/extracted
@@ -1,3 +1,3 @@
 <tmp_path>/program.sql:<main>: Info: Optimization
 
-    <tmp_path>/program.sql:<main>: Info: DQ cannot execute the query. Cause: unsupported yt pragma: pooltrees
\ No newline at end of file
+    <tmp_path>/program.sql:<main>: Info: DQ cannot execute the query. Cause: unsupported pool tree: test
\ No newline at end of file
diff --git a/ydb/library/yql/tests/sql/dq_file/part2/canondata/result.json b/ydb/library/yql/tests/sql/dq_file/part2/canondata/result.json
index 278ac2b5ba..443728eb79 100644
--- a/ydb/library/yql/tests/sql/dq_file/part2/canondata/result.json
+++ b/ydb/library/yql/tests/sql/dq_file/part2/canondata/result.json
@@ -1419,6 +1419,28 @@
         }
     ],
     "test.test[distinct-distinct_star_inmem-default.txt-Results]": [],
+    "test.test[dq-pool_trees_whitelist--Analyze]": [
+        {
+            "checksum": "14d974e2a4612975604c2c1f36f2bf3e",
+            "size": 4821,
+            "uri": "https://{canondata_backend}/212715/08ef1044277b3dec1aee8253faf70ced39caf2f5/resource.tar.gz#test.test_dq-pool_trees_whitelist--Analyze_/plan.txt"
+        }
+    ],
+    "test.test[dq-pool_trees_whitelist--Debug]": [
+        {
+            "checksum": "b04fb90343168e19fb4ee7ae89ae2bf4",
+            "size": 1898,
+            "uri": "https://{canondata_backend}/1942100/1170a9ef0a48e22329a02a9d3f8e40192481dbe1/resource.tar.gz#test.test_dq-pool_trees_whitelist--Debug_/opt.yql_patched"
+        }
+    ],
+    "test.test[dq-pool_trees_whitelist--Plan]": [
+        {
+            "checksum": "14d974e2a4612975604c2c1f36f2bf3e",
+            "size": 4821,
+            "uri": "https://{canondata_backend}/1942100/1170a9ef0a48e22329a02a9d3f8e40192481dbe1/resource.tar.gz#test.test_dq-pool_trees_whitelist--Plan_/plan.txt"
+        }
+    ],
+    "test.test[dq-pool_trees_whitelist--Results]": [],
     "test.test[epochs-read_modified--Analyze]": [
         {
             "checksum": "63389a1828258ea7b3eb17e61c93c312",
diff --git a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
index 516c4314ee..ec2735de86 100644
--- a/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
+++ b/ydb/library/yql/tests/sql/sql2yql/canondata/result.json
@@ -4584,6 +4584,13 @@
             "uri": "https://{canondata_backend}/1936997/00f46808be87e2ae2d4ac3ac45675b659c5ace45/resource.tar.gz#test_sql2yql.test_dq-mem_limit_/sql.yql"
         }
     ],
+    "test_sql2yql.test[dq-pool_trees_whitelist]": [
+        {
+            "checksum": "5137c665c4b6c390b2ec84c27fefb128",
+            "size": 1529,
+            "uri": "https://{canondata_backend}/1871002/c65ff3a6f73717dfa83b26a43711be000bed4540/resource.tar.gz#test_sql2yql.test_dq-pool_trees_whitelist_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[dq-precompute_asyncfile]": [
         {
             "checksum": "8483231236acc4f5d127668886f5e8b5",
@@ -21727,6 +21734,13 @@
             "uri": "https://{canondata_backend}/1880306/64654158d6bfb1289c66c626a8162239289559d0/resource.tar.gz#test_sql_format.test_dq-mem_limit_/formatted.sql"
         }
     ],
+    "test_sql_format.test[dq-pool_trees_whitelist]": [
+        {
+            "checksum": "20eade00c283958b45d0b952730cff0f",
+            "size": 188,
+            "uri": "https://{canondata_backend}/1871002/c65ff3a6f73717dfa83b26a43711be000bed4540/resource.tar.gz#test_sql_format.test_dq-pool_trees_whitelist_/formatted.sql"
+        }
+    ],
     "test_sql_format.test[dq-precompute_asyncfile]": [
         {
             "checksum": "91c30aca2dc34875c088cf0ccbd57b3b",
diff --git a/ydb/library/yql/tests/sql/suites/dq/pool_trees_whitelist.cfg b/ydb/library/yql/tests/sql/suites/dq/pool_trees_whitelist.cfg
new file mode 100644
index 0000000000..628ba758c4
--- /dev/null
+++ b/ydb/library/yql/tests/sql/suites/dq/pool_trees_whitelist.cfg
@@ -0,0 +1,3 @@
+providers dq
+in Input input.txt
+out Output output.txt
diff --git a/ydb/library/yql/tests/sql/suites/dq/pool_trees_whitelist.sql b/ydb/library/yql/tests/sql/suites/dq/pool_trees_whitelist.sql
new file mode 100644
index 0000000000..cb4b754ff4
--- /dev/null
+++ b/ydb/library/yql/tests/sql/suites/dq/pool_trees_whitelist.sql
@@ -0,0 +1,8 @@
+/* postgres can not */
+/* syntax version 1 */
+use plato;
+pragma OrderedColumns;
+pragma yt.PoolTrees='physical,cloud';
+
+insert into Input
+select key, subkey, value from Input;