YQL-16252, YQL-16623: Support search_path

ydb/library/yql/core/pg_settings/guc_settings.cpp

@@ -0,0 +1,29 @@
+#include "guc_settings.h"
+
+void TGUCSettings::Setup(const std::unordered_map<std::string, std::string>& runtimeSettings) {
+    RollbackSettings_ = runtimeSettings;
+    RollBack();
+}
+
+std::optional<std::string> TGUCSettings::Get(const std::string& key) const {
+    auto it = Settings_.find(key);
+    if (it == Settings_.end()) {
+        return std::nullopt;
+    }
+    return it->second;
+}
+
+void TGUCSettings::Set(const std::string& key, const std::string& val, bool isLocal) {
+    Settings_[key] = val;
+    if (!isLocal) {
+        SessionSettings_[key] = val;
+    }
+}
+
+void TGUCSettings::Commit() {
+    RollbackSettings_ = SessionSettings_;
+}
+
+void TGUCSettings::RollBack() {
+    Settings_ = SessionSettings_ = RollbackSettings_;
+}

ydb/library/yql/core/pg_settings/guc_settings.h

@@ -0,0 +1,20 @@
+#pragma once
+#include <unordered_map>
+#include <vector>
+#include <string>
+#include <optional>
+#include <memory>
+
+class TGUCSettings {
+public:
+    using TPtr = std::shared_ptr<TGUCSettings>;
+    void Setup(const std::unordered_map<std::string, std::string>& runtimeSettings);
+    std::optional<std::string> Get(const std::string&) const;
+    void Set(const std::string&, const std::string&, bool isLocal = false);
+    void Commit();
+    void RollBack();
+private:
+    std::unordered_map<std::string, std::string> Settings_;
+    std::unordered_map<std::string, std::string> RollbackSettings_;
+    std::unordered_map<std::string, std::string> SessionSettings_;
+};

ydb/library/yql/core/pg_settings/ya.make

@@ -0,0 +1,7 @@
+LIBRARY()
+
+SRCS(
+    guc_settings.cpp
+)
+
+END()

ydb/library/yql/sql/pg/pg_sql.cpp

@@ -632,6 +632,54 @@ class TConverter : public IPGParseEvents {
         return QL(QA("values"), QVL(valNames.data(), valNames.size()), VL(valueRows));
     }
 
+    TAstNode* ParseSetConfig(const FuncCall* value) {
+        auto length = ListLength(value->args);
+        if (length != 3) {
+            AddError(TStringBuilder() << "Expected 3 arguments, but got: " << length);
+            return;
+        }
+        auto loc = value->location;
+        VariableSetStmt config;
+        auto arg0 = ListNodeNth(value->args, 0);
+        auto arg1 = ListNodeNth(value->args, 1);
+        auto arg2 = ListNodeNth(value->args, 2);
+        if (NodeTag(arg2) != T_TypeCast) {
+            AddError(TStringBuilder() << "Expected type cast node as is_local arg, but got node with tag");
+            return;
+        }
+        auto isLocalCast = CAST_NODE(TypeCast, arg2)->arg;
+        if (NodeTag(isLocalCast) != T_A_Const) {
+            AddError(TStringBuilder() << "Expected a_const in cast, but got something wrong: " << NodeTag(isLocalCast));
+            return nullptr;
+        }
+        auto isLocalConst = CAST_NODE(A_Const, isLocalCast);
+        if (NodeTag(isLocalConst->val) != T_String) {
+            AddError(TStringBuilder() << "Expected string in const, but got something wrong: " << NodeTag(isLocalCast));
+            return nullptr;
+        }
+        auto rawVal = TString(StrVal(isLocalConst->val));
+        if (rawVal != "t" && rawVal != "f") {
+            AddError(TStringBuilder() << "Expected t/f, but got " << rawVal);
+            return;
+        }
+        config.is_local = rawVal == "t";
+
+        if (NodeTag(arg0) != T_A_Const || NodeTag(arg1) != T_A_Const) {
+            AddError(TStringBuilder() << "Expected const with string, but got something else: " << NodeTag(arg0));
+            return nullptr;
+        }
+
+        auto name = CAST_NODE(A_Const, arg0)->val;
+        auto val = CAST_NODE(A_Const, arg1)->val;
+        if (NodeTag(name) != T_String || NodeTag(val) != T_String) {
+            AddError(TStringBuilder() << "Expected string const as name arg, but got something else: " << NodeTag(name));
+            return;
+        }
+        config.name = (char*)StrVal(name);
+        config.args = list_make1((void*)(&val));
+        return ParseVariableSetStmt(&config, true);
+    }
+
     using TTraverseSelectStack = TStack<std::pair<const SelectStmt*, bool>>;
     using TTraverseNodeStack = TStack<std::pair<const Node*, bool>>;
 
@@ -974,6 +1022,34 @@ class TConverter : public IPGParseEvents {
                 res.emplace_back(CreatePgStarResultItem());
                 i++;
             }
+            bool maybeSelectWithJustSetConfig = !inner && !sort && windowItems.empty() && !having && !groupBy && !whereFilter && !x->distinctClause  && ListLength(x->targetList) == 1;
+            if (maybeSelectWithJustSetConfig) {
+                auto node = ListNodeNth(x->targetList, 0);
+                if (NodeTag(node) != T_ResTarget) {
+                    NodeNotImplemented(x, node);
+                    return nullptr;
+                }
+                auto r = CAST_NODE(ResTarget, node);
+                if (!r->val) {
+                    AddError("SelectStmt: expected val");
+                    return nullptr;
+                }
+                auto call = r->val;
+                if (NodeTag(call) == T_FuncCall) {
+                    auto fn = CAST_NODE(FuncCall, call);
+                    if (ListLength(fn->funcname) == 1) {
+                        auto nameNode = ListNodeNth(fn->funcname, 0);
+                        if (NodeTag(nameNode) != T_String) {
+                            AddError("Function name must be string");
+                            return nullptr;
+                        }
+                        auto name = to_lower(TString(StrVal(ListNodeNth(fn->funcname, 0))));
+                        if (name == "set_config") {
+                            return ParseSetConfig(fn);
+                        }
+                    }
+                }
+            }
             for (int targetIndex = 0; targetIndex < ListLength(x->targetList); ++targetIndex) {
                 auto node = ListNodeNth(x->targetList, targetIndex);
                 if (NodeTag(node) != T_ResTarget) {
@@ -1993,13 +2069,38 @@ class TConverter : public IPGParseEvents {
     }
 
     [[nodiscard]]
-    TAstNode* ParseVariableSetStmt(const VariableSetStmt* value) {
+    TAstNode* ParseVariableSetStmt(const VariableSetStmt* value, bool isSetConfig = false) {
         if (value->kind != VAR_SET_VALUE) {
             AddError(TStringBuilder() << "VariableSetStmt, not supported kind: " << (int)value->kind);
             return nullptr;
         }
 
         auto name = to_lower(TString(value->name));
+        if (isSetConfig) {
+            if (ListLength(value->args) != 1) {
+                AddError(TStringBuilder() << "VariableSetStmt, expected 1 arg, but got: " << ListLength(value->args));
+                return nullptr;
+            }
+            auto val = ListNodeNth(value->args, 0);
+            if (NodeTag(val) != T_String) {
+                AddError(TStringBuilder() << "VariableSetStmt, expected string literal for " << value->name << " option");
+                return nullptr;
+            }
+            TString rawStr = TString(StrVal(val));
+            if (name != "search_path") {
+                AddError(TStringBuilder() << "VariableSetStmt, set_config doesn't support that option:" << name);
+                return nullptr;
+            }
+            if (rawStr != "pg_catalog" && rawStr != "public") {
+                AddError(TStringBuilder() << "VariableSetStmt, search path supports only public and pg_catalogue, but got :" << rawStr);
+                return nullptr;
+            }
+            if (Settings.GUCSettings) {
+                Settings.GUCSettings->Set(name, rawStr, value->is_local);
+            }
+            return Statements.back();
+        }
+
         if (name == "useblocks" || name == "emitaggapply") {
             if (ListLength(value->args) != 1) {
                 AddError(TStringBuilder() << "VariableSetStmt, expected 1 arg, but got: " << ListLength(value->args));
@@ -2277,10 +2378,12 @@ class TConverter : public IPGParseEvents {
         case TRANS_STMT_COMMIT:
             Statements.push_back(L(A("let"), A("world"), L(A("CommitAll!"),
                 A("world"))));
+            Settings.GUCSettings->Commit();
             return true;
         case TRANS_STMT_ROLLBACK:
             Statements.push_back(L(A("let"), A("world"), L(A("CommitAll!"),
                 A("world"), QL(QL(QA("mode"), QA("rollback"))))));
+            Settings.GUCSettings->RollBack();
             return true;
         default:
             AddError(TStringBuilder() << "TransactionStmt: kind is not supported: " << (int)value->kind);
@@ -2413,6 +2516,20 @@ class TConverter : public IPGParseEvents {
         return true;
     }
 
+    TString ResolveCluster(const TStringBuf schemaname) {
+        if (schemaname == "public") {
+            return "";
+        }
+        if (schemaname == "" && Settings.GUCSettings) {
+            auto search_path = Settings.GUCSettings->Get("search_path");
+            if (!search_path || *search_path == "public") {
+                return Settings.DefaultCluster;
+            }
+            return TString(*search_path);
+        }
+        return TString(schemaname);
+    }
+
     TAstNode* BuildClusterSinkOrSourceExpression(
         bool isSink, const TStringBuf schemaname) {
       const auto p = Settings.ClusterMapping.FindPtr(schemaname);
@@ -2443,7 +2560,7 @@ class TConverter : public IPGParseEvents {
         return {};
       }
 
-      const auto cluster = !schemaname.Empty() && schemaname != "public" ? schemaname : Settings.DefaultCluster;
+      const auto cluster = ResolveCluster(schemaname);
       const auto sinkOrSource = BuildClusterSinkOrSourceExpression(isSink, cluster);
       const auto key = BuildTableKeyExpression(relname, isScheme);
       return {sinkOrSource, key};
@@ -2470,7 +2587,7 @@ class TConverter : public IPGParseEvents {
             return {};
         }
 
-        const auto cluster = !schemaname.Empty() && schemaname != "public" ? schemaname : Settings.DefaultCluster;
+        const auto cluster = ResolveCluster(schemaname);
         const auto sinkOrSource = BuildClusterSinkOrSourceExpression(true, cluster);
         const auto key = BuildPgObjectExpression(objectName, pgObjectType);
         return {sinkOrSource, key};
@@ -3262,6 +3379,7 @@ class TConverter : public IPGParseEvents {
         }
 
         auto name = names.back();
+
         const bool isAggregateFunc = NYql::NPg::HasAggregation(name);
         const bool hasReturnSet = NYql::NPg::HasReturnSetProc(name);
 

ydb/library/yql/sql/pg/pg_sql_ut.cpp

@@ -456,4 +456,121 @@ SELECT COUNT(*) FROM public.t;");
         auto issue = *(res.Issues.begin());
         UNIT_ASSERT(issue.GetMessage().Contains("VariableSetStmt, not supported BlockEngine option value: foo"));
     }
+
+    Y_UNIT_TEST(SetConfig_SearchPath) {
+        TTranslationSettings settings;
+        settings.GUCSettings = std::make_shared<TGUCSettings>();
+        settings.ClusterMapping["pg_catalog"] = NYql::PgProviderName;
+        settings.DefaultCluster = "";
+
+        auto res = SqlToYqlWithMode(
+            R"(select set_config('search_path', 'pg_catalog', false);)",
+            NSQLTranslation::ESqlMode::QUERY,
+            10,
+            {},
+            EDebugOutput::ToCerr,
+            false,
+            settings);
+        UNIT_ASSERT_C(res.IsOk(), res.Issues.ToString());
+        UNIT_ASSERT(res.Root);
+
+        res = SqlToYqlWithMode(
+            R"(select oid,
+typinput::int4 as typinput,
+typname,
+typnamespace,
+typtype
+from pg_type)",
+            NSQLTranslation::ESqlMode::QUERY,
+            10,
+            {},
+            EDebugOutput::None,
+            false,
+            settings);
+        UNIT_ASSERT(res.IsOk());
+        UNIT_ASSERT(res.Root);
+
+        res = SqlToYqlWithMode(
+            R"(select oid,
+typinput::int4 as typinput,
+typname,
+typnamespace,
+typtype
+from pg_catalog.pg_type)",
+            NSQLTranslation::ESqlMode::QUERY,
+            10,
+            {},
+            EDebugOutput::None,
+            false,
+            settings);
+        UNIT_ASSERT(res.IsOk());
+        UNIT_ASSERT(res.Root);
+
+        res = SqlToYqlWithMode(
+            R"(select set_config('search_path', 'public', false);)",
+            NSQLTranslation::ESqlMode::QUERY,
+            10,
+            {},
+            EDebugOutput::None,
+            false,
+            settings);
+        UNIT_ASSERT(res.IsOk());
+        UNIT_ASSERT(res.Root);
+
+        res = SqlToYqlWithMode(
+            R"(select set_config('search_path', 'yql', false);)",
+            NSQLTranslation::ESqlMode::QUERY,
+            10,
+            {},
+            EDebugOutput::None,
+            false,
+            settings);
+        UNIT_ASSERT(!res.IsOk());
+        UNIT_ASSERT(!res.Root);
+
+        res = SqlToYqlWithMode(
+            R"(select set_config('search_path', 'pg_catalog', false);)",
+            NSQLTranslation::ESqlMode::QUERY,
+            10,
+            {},
+            EDebugOutput::None,
+            false,
+            settings);
+        UNIT_ASSERT(res.IsOk());
+        UNIT_ASSERT(res.Root);
+
+        res = SqlToYqlWithMode(
+            R"(rollback;)",
+            NSQLTranslation::ESqlMode::QUERY,
+            10,
+            {},
+            EDebugOutput::None,
+            false,
+            settings);
+        UNIT_ASSERT(res.IsOk());
+        UNIT_ASSERT(res.Root);
+
+        google::protobuf::Arena arena;
+        const auto service = TString(NYql::YtProviderName);
+        settings.ClusterMapping["hahn"] = NYql::YtProviderName;
+        settings.ClusterMapping["mon"] = NYql::SolomonProviderName;
+        settings.MaxErrors = 10;
+        settings.Mode = NSQLTranslation::ESqlMode::QUERY;
+        settings.Arena = &arena;
+        settings.AnsiLexer = false;
+        settings.SyntaxVersion = 1;
+        settings.PgParser = true;
+
+        res = SqlToYql(
+            R"(select oid,
+typinput::int4 as typinput,
+typname,
+typnamespace,
+typtype
+from pg_type)",
+            settings);
+        UNIT_ASSERT(res.Issues.ToString().Contains("Unknown cluster:"));
+        UNIT_ASSERT(!res.IsOk());
+        UNIT_ASSERT(!res.Root);
+    }
 }

ydb/library/yql/sql/settings/translation_settings.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <ydb/library/yql/core/pg_settings/guc_settings.h>
+
 #include <util/generic/hash.h>
 #include <util/generic/hash_set.h>
 #include <util/generic/map.h>
@@ -110,6 +112,8 @@ namespace NSQLTranslation {
         bool AutoParametrizeEnabled = false;
         bool AutoParametrizeValuesStmt = false;
         THashSet<TString> AutoParametrizeExprDisabledScopes = {};
+
+        TGUCSettings::TPtr GUCSettings = std::make_shared<TGUCSettings>();
     };
 
     bool ParseTranslationSettings(const TString& query, NSQLTranslation::TTranslationSettings& settings, NYql::TIssues& issues);

ydb/library/yql/sql/settings/ya.make

@@ -9,6 +9,7 @@ PEERDIR(
     library/cpp/deprecated/split
     library/cpp/json
     ydb/library/yql/core/issue
+    ydb/library/yql/core/pg_settings
     ydb/library/yql/core/issue/protos
     ydb/library/yql/utils
 )