YQL-16041: support ParamRef in parser

3 files changed, 129 insertions, 1 deletions

diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
index 0f8ca060d6b..bd4716e4b54 100644
--- a/ydb/library/yql/sql/pg/pg_sql.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -5,6 +5,7 @@
 #include <ydb/library/yql/parser/pg_wrapper/interface/parser.h>
 #include <ydb/library/yql/parser/pg_wrapper/interface/utils.h>
 #include <ydb/library/yql/parser/pg_wrapper/parser.h>
+#include <ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/catalog/pg_type_d.h>
 #include <ydb/library/yql/parser/pg_catalog/catalog.h>
 #include <ydb/library/yql/providers/common/provider/yql_provider_names.h>
 #include <ydb/library/yql/core/yql_callable_names.h>
@@ -33,6 +34,8 @@ extern "C" {
 #undef SortBy
 }
 
+constexpr auto PREPARED_PARAM_PREFIX =  "$p";
+
 namespace NSQLTranslationPG {
 
 using namespace NYql;
@@ -220,6 +223,11 @@ public:
             Provider = provider;
             break;
         }
+        
+        for (size_t i = 0; i < Settings.PgParameterTypeOids.size(); ++i) {
+            auto paramName = PREPARED_PARAM_PREFIX + ToString(i + 1);
+            ParamNameToTypeOid[paramName] = Settings.PgParameterTypeOids[i];
+        }
     }
 
     void OnResult(const List* raw) {
@@ -250,11 +258,13 @@ public:
             return nullptr;
         }
 
-
         if (Settings.EndOfQueryCommit) {
             Statements.push_back(L(A("let"), A("world"), L(A("CommitAll!"),
                 A("world"))));
         }
+
+        AddVariableDeclarations();
+
         Statements.push_back(L(A("return"), A("world")));
 
 
@@ -1956,6 +1966,14 @@ public:
         return L(A("If"), final.Pred, final.Value, defaultResult);
     }
 
+    TAstNode* ParseParamRefExpr(const ParamRef* value) {
+        const auto varName = PREPARED_PARAM_PREFIX + ToString(value->number);
+        if (!ParamNameToTypeOid.contains(varName)) {
+            ParamNameToTypeOid[varName] = UNKNOWNOID;
+        }
+        return A(varName);
+    }
+
     TAstNode* ParseExpr(const Node* node, const TExprSettings& settings) {
         switch (NodeTag(node)) {
         case T_A_Const: {
@@ -1994,6 +2012,9 @@ public:
         case T_GroupingFunc: {
             return ParseGroupingFunc(CAST_NODE(GroupingFunc, node));
         }
+        case T_ParamRef: {
+            return ParseParamRefExpr(CAST_NODE(ParamRef, node));
+        }
         default:
             NodeNotImplemented(node);
             return nullptr;
@@ -3078,6 +3099,15 @@ public:
 
     }
 
+    void AddVariableDeclarations() {
+      for (const auto &[varName, typeOid] : ParamNameToTypeOid) {
+        const auto &typeName =
+            typeOid != UNKNOWNOID ? NPg::LookupType(typeOid).Name : "text";
+        const auto pgType = L(A("PgType"), QA(typeName));
+        Statements.push_back(L(A("declare"), A(varName), pgType));
+      }
+    }
+
     template <typename T>
     void NodeNotImplementedImpl(const Node* nodeptr) {
         TStringBuilder b;
@@ -3264,6 +3294,8 @@ private:
     ui32 QuerySize;
     TString Provider;
     static const THashMap<TStringBuf, TString> ProviderToInsertModeMap;
+
+    THashMap<TString, Oid> ParamNameToTypeOid;
 };
 
 const THashMap<TStringBuf, TString> TConverter::ProviderToInsertModeMap = {
diff --git a/ydb/library/yql/sql/pg/pg_sql_ut.cpp b/ydb/library/yql/sql/pg/pg_sql_ut.cpp
index b85d4e0ff02..f4ccd414907 100644
--- a/ydb/library/yql/sql/pg/pg_sql_ut.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql_ut.cpp
@@ -3,6 +3,7 @@
 #include <ydb/library/yql/ast/yql_expr.h>
 #include <ydb/library/yql/providers/common/provider/yql_provider_names.h>
 #include <ydb/library/yql/sql/sql.h>
+#include <ydb/library/yql/parser/pg_catalog/catalog.h>
 #include <ydb/library/yql/parser/pg_wrapper/interface/config.h>
 
 #include <library/cpp/testing/unittest/registar.h>
@@ -52,6 +53,18 @@ NYql::TAstParseResult PgSqlToYql(const TString& query, size_t maxErrors = 10, co
     return SqlToYqlWithMode(query, NSQLTranslation::ESqlMode::QUERY, maxErrors, provider, debug);
 }
 
+using TAstNodeVisitFunc = std::function<void(const NYql::TAstNode& root)>;
+
+void VisitAstNodes(const NYql::TAstNode& root, const TAstNodeVisitFunc& visitFunc) {
+    visitFunc(root);
+    if (!root.IsList()) {
+        return;
+    }
+    for (size_t childIdx = 0; childIdx < root.GetChildrenCount(); ++childIdx) {
+        VisitAstNodes(*root.GetChild(childIdx), visitFunc);
+    }
+}
+
 Y_UNIT_TEST_SUITE(PgSqlParsingOnly) {
     Y_UNIT_TEST(InsertStmt) {
         auto res = PgSqlToYql("INSERT INTO plato.Input VALUES (1, 1)");
@@ -223,4 +236,84 @@ Y_UNIT_TEST_SUITE(PgSqlParsingOnly) {
         const auto expectedAst = NYql::ParseAst(program);
         UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
     }
+    
+    TMap<TString, TString> GetParamNameToPgType(const NYql::TAstNode& root) {
+        TMap<TString, TString> actualParamToType;
+
+        VisitAstNodes(root, [&actualParamToType] (const NYql::TAstNode& node) {
+            bool isDeclareNode = 
+                node.IsListOfSize(3) && node.GetChild(0)->IsAtom() 
+                && node.GetChild(0)->GetContent() == "declare";
+            if (isDeclareNode) {
+                const auto varNameNode = node.GetChild(1);
+                UNIT_ASSERT(varNameNode->IsAtom());
+                const auto varName = varNameNode->GetContent();
+
+                const auto varTypeNode = node.GetChild(2);
+                UNIT_ASSERT(varTypeNode->IsListOfSize(2));
+                UNIT_ASSERT(varTypeNode->GetChild(0)->GetContent() == "PgType");
+                actualParamToType[TString(varName)] = varTypeNode->GetChild(1)->ToString();
+            }
+        });
+        
+        return actualParamToType;
+    }
+    
+    Y_UNIT_TEST(ParamRef_IntAndPoint) {
+        TTranslationSettings settings;
+        
+        settings.PgParameterTypeOids = {NYql::NPg::LookupType("int4").TypeId, NYql::NPg::LookupType("point").TypeId};
+        auto res = SqlToYqlWithMode(
+            R"(select $1 as "x", $2 as "y")",
+            NSQLTranslation::ESqlMode::QUERY, 
+            10, 
+            {}, 
+            EDebugOutput::None, 
+            false,
+            settings);
+        TMap<TString, TString> expectedParamToType {
+            {"$p1", "'int4"},
+            {"$p2", "'point"},
+        };
+        UNIT_ASSERT(res.Root);
+        const auto actualParamToTypes = GetParamNameToPgType(*res.Root);
+        UNIT_ASSERT(expectedParamToType.size() == actualParamToTypes.size());
+        UNIT_ASSERT_EQUAL(expectedParamToType, actualParamToTypes);
+    }
+
+    Y_UNIT_TEST(ParamRef_IntUnknownInt) {
+        TTranslationSettings settings;
+        settings.PgParameterTypeOids = {NYql::NPg::LookupType("int4").TypeId, NYql::NPg::LookupType("unknown").TypeId, NYql::NPg::LookupType("int4").TypeId};
+        auto res = SqlToYqlWithMode(
+            R"(select $1 as "x", $2 as "y", $3 as "z")",
+            NSQLTranslation::ESqlMode::QUERY, 
+            10, 
+            {}, 
+            EDebugOutput::None, 
+            false,
+            settings);
+        TMap<TString, TString> expectedParamToType {
+            {"$p1", "'int4"},
+            {"$p2", "'text"},
+            {"$p3", "'int4"},
+        };
+        UNIT_ASSERT(res.Root);
+        const auto actualParamToTypes = GetParamNameToPgType(*res.Root);
+        UNIT_ASSERT(expectedParamToType.size() == actualParamToTypes.size());
+        UNIT_ASSERT_EQUAL(expectedParamToType, actualParamToTypes);
+    }
+
+    Y_UNIT_TEST(ParamRef_NoTypeOids) {
+        TTranslationSettings settings;
+        settings.PgParameterTypeOids = {};
+        auto res = PgSqlToYql(R"(select $1 as "x", $2 as "y", $3 as "z")");
+        TMap<TString, TString> expectedParamToType {
+            {"$p1", "'text"},
+            {"$p2", "'text"},
+            {"$p3", "'text"},
+        };
+        UNIT_ASSERT(res.Root);
+        auto actualParamToTypes = GetParamNameToPgType(*res.Root);
+        UNIT_ASSERT_VALUES_EQUAL(expectedParamToType, actualParamToTypes);
+    }
 }
diff --git a/ydb/library/yql/sql/settings/translation_settings.h b/ydb/library/yql/sql/settings/translation_settings.h
index fb84b74cb21..4d0ff5cdaff 100644
--- a/ydb/library/yql/sql/settings/translation_settings.h
+++ b/ydb/library/yql/sql/settings/translation_settings.h
@@ -3,6 +3,7 @@
 #include <util/generic/hash.h>
 #include <util/generic/hash_set.h>
 #include <util/generic/map.h>
+#include <util/generic/vector.h>
 
 namespace google::protobuf {
     class Arena;
@@ -90,6 +91,8 @@ namespace NSQLTranslation {
         ISqlFeaturePolicy::TPtr DqDefaultAuto;
         bool AssumeYdbOnClusterWithSlash;
         TString DynamicClusterProvider;
+
+        TVector<ui32> PgParameterTypeOids;
     };
 
     bool ParseTranslationSettings(const TString& query, NSQLTranslation::TTranslationSettings& settings, NYql::TIssues& issues);