KIKIMR-19068 forward pg syntax parameter for ExecuteYqlScript and StreamExecuteYqlScript

14 files changed, 95 insertions, 19 deletions

diff --git a/ydb/core/grpc_services/rpc_execute_yql_script.cpp b/ydb/core/grpc_services/rpc_execute_yql_script.cpp
index 2d51d0780d2..96d9d0a7965 100644
--- a/ydb/core/grpc_services/rpc_execute_yql_script.cpp
+++ b/ydb/core/grpc_services/rpc_execute_yql_script.cpp
@@ -69,7 +69,8 @@ public:
             nullptr, // query_cache_policy
             req->has_operation_params() ? &req->operation_params() : nullptr,
             false, // keep session
-            false // use cancelAfter
+            false, // use cancelAfter
+            req->syntax()
         );
 
         ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release());
diff --git a/ydb/core/grpc_services/rpc_stream_execute_yql_script.cpp b/ydb/core/grpc_services/rpc_stream_execute_yql_script.cpp
index b93232ac307..a7837f41c78 100644
--- a/ydb/core/grpc_services/rpc_stream_execute_yql_script.cpp
+++ b/ydb/core/grpc_services/rpc_stream_execute_yql_script.cpp
@@ -175,7 +175,8 @@ private:
             nullptr, // query_cache_policy
             req->has_operation_params() ? &req->operation_params() : nullptr,
             false, // keep session
-            false // use cancelAfter
+            false, // use cancelAfter
+            req->syntax()
         );
 
         if (!ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release())) {
diff --git a/ydb/core/kqp/host/kqp_host.cpp b/ydb/core/kqp/host/kqp_host.cpp
index 8dde742ecbf..af5c5d4262a 100644
--- a/ydb/core/kqp/host/kqp_host.cpp
+++ b/ydb/core/kqp/host/kqp_host.cpp
@@ -1415,7 +1415,7 @@ private:
         SessionCtx->Query().PreparedQuery.reset();
 
         TMaybe<TSqlVersion> sqlVersion;
-        auto scriptExpr = CompileYqlQuery(script, true, true, ctx, sqlVersion, {});
+        auto scriptExpr = CompileYqlQuery(script, true, true, ctx, sqlVersion, settings.UsePgParser);
         if (!scriptExpr) {
             return nullptr;
         }
@@ -1439,7 +1439,7 @@ private:
         SessionCtx->Query().PreparedQuery.reset();
 
         TMaybe<TSqlVersion> sqlVersion;
-        auto scriptExpr = CompileYqlQuery(script, true, true, ctx, sqlVersion, {});
+        auto scriptExpr = CompileYqlQuery(script, true, true, ctx, sqlVersion, settings.UsePgParser);
         if (!scriptExpr) {
             return nullptr;
         }
diff --git a/ydb/core/kqp/host/kqp_host.h b/ydb/core/kqp/host/kqp_host.h
index 4eaebf1558b..7c35b1b5c1e 100644
--- a/ydb/core/kqp/host/kqp_host.h
+++ b/ydb/core/kqp/host/kqp_host.h
@@ -55,6 +55,8 @@ public:
         NYql::TKikimrQueryDeadlines Deadlines;
         NYql::EKikimrStatsMode StatsMode = NYql::EKikimrStatsMode::None;
         std::shared_ptr<NGRpcService::IRequestCtxMtSafe> RpcCtx;
+        TMaybe<bool> UsePgParser;
+        TMaybe<TSqlVersion> SyntaxVersion;
     };
 
     virtual ~IKqpHost() {}
diff --git a/ydb/core/kqp/session_actor/kqp_worker_actor.cpp b/ydb/core/kqp/session_actor/kqp_worker_actor.cpp
index c10a443d14a..9355f873363 100644
--- a/ydb/core/kqp/session_actor/kqp_worker_actor.cpp
+++ b/ydb/core/kqp/session_actor/kqp_worker_actor.cpp
@@ -568,11 +568,31 @@ private:
         Cleanup(ctx, true);
     }
 
+    IKqpHost::TExecScriptSettings ParseExecScriptSettings() {
+        IKqpHost::TExecScriptSettings execSettings;
+        execSettings.Deadlines = QueryState->QueryDeadlines;
+        execSettings.RpcCtx = QueryState->RequestEv->GetRequestCtx();
+        auto statsMode = GetStatsMode(QueryState->RequestEv.get(), EKikimrStatsMode::Basic);
+        execSettings.StatsMode = statsMode;
+
+        switch (QueryState->RequestEv->GetSyntax()) {
+            case Ydb::Query::Syntax::SYNTAX_YQL_V1:
+                execSettings.UsePgParser = false;
+                execSettings.SyntaxVersion = 1;
+                break;
+            case Ydb::Query::Syntax::SYNTAX_PG:
+                execSettings.UsePgParser = true;
+                break;
+            default:
+                break;
+        }
+
+        return execSettings;
+    }
+
     bool ExecuteQuery(const TString& query, const ::google::protobuf::Map<TProtoStringType, ::Ydb::TypedValue>& parameters,
         NKikimrKqp::EQueryType type, const TActorId& requestActorId)
     {
-        auto statsMode = GetStatsMode(QueryState->RequestEv.get(), EKikimrStatsMode::Basic);
-
         switch (type) {
             case NKikimrKqp::QUERY_TYPE_SQL_DDL: {
                 IKqpHost::TExecSettings execSettings;
@@ -594,19 +614,13 @@ private:
             }
 
             case NKikimrKqp::QUERY_TYPE_SQL_SCRIPT: {
-                IKqpHost::TExecScriptSettings execSettings;
-                execSettings.Deadlines = QueryState->QueryDeadlines;
-                execSettings.RpcCtx = QueryState->RequestEv->GetRequestCtx();
-                execSettings.StatsMode = statsMode;
+                IKqpHost::TExecScriptSettings execSettings = ParseExecScriptSettings();
                 QueryState->AsyncQueryResult = KqpHost->ExecuteYqlScript(query, parameters, execSettings);
                 break;
             }
 
             case NKikimrKqp::QUERY_TYPE_SQL_SCRIPT_STREAMING: {
-                IKqpHost::TExecScriptSettings execSettings;
-                execSettings.Deadlines = QueryState->QueryDeadlines;
-                execSettings.RpcCtx = QueryState->RequestEv->GetRequestCtx();
-                execSettings.StatsMode = statsMode;
+                IKqpHost::TExecScriptSettings execSettings = ParseExecScriptSettings();
                 QueryState->AsyncQueryResult = KqpHost->StreamExecuteYqlScript(query, parameters,
                     requestActorId, execSettings);
                 break;
diff --git a/ydb/core/kqp/ut/yql/CMakeLists.darwin-x86_64.txt b/ydb/core/kqp/ut/yql/CMakeLists.darwin-x86_64.txt
index 39b054c2846..5435e5c7d92 100644
--- a/ydb/core/kqp/ut/yql/CMakeLists.darwin-x86_64.txt
+++ b/ydb/core/kqp/ut/yql/CMakeLists.darwin-x86_64.txt
@@ -21,7 +21,8 @@ target_link_libraries(ydb-core-kqp-ut-yql PUBLIC
   cpp-testing-unittest_main
   ydb-core-kqp
   kqp-ut-common
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
 )
 target_link_options(ydb-core-kqp-ut-yql PRIVATE
   -Wl,-platform_version,macos,11.0,11.0
diff --git a/ydb/core/kqp/ut/yql/CMakeLists.linux-aarch64.txt b/ydb/core/kqp/ut/yql/CMakeLists.linux-aarch64.txt
index 51aac46b986..395fe66328c 100644
--- a/ydb/core/kqp/ut/yql/CMakeLists.linux-aarch64.txt
+++ b/ydb/core/kqp/ut/yql/CMakeLists.linux-aarch64.txt
@@ -21,7 +21,8 @@ target_link_libraries(ydb-core-kqp-ut-yql PUBLIC
   cpp-testing-unittest_main
   ydb-core-kqp
   kqp-ut-common
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
 )
 target_link_options(ydb-core-kqp-ut-yql PRIVATE
   -ldl
diff --git a/ydb/core/kqp/ut/yql/CMakeLists.linux-x86_64.txt b/ydb/core/kqp/ut/yql/CMakeLists.linux-x86_64.txt
index 80e82e8aada..a1884e2913f 100644
--- a/ydb/core/kqp/ut/yql/CMakeLists.linux-x86_64.txt
+++ b/ydb/core/kqp/ut/yql/CMakeLists.linux-x86_64.txt
@@ -22,7 +22,8 @@ target_link_libraries(ydb-core-kqp-ut-yql PUBLIC
   cpp-testing-unittest_main
   ydb-core-kqp
   kqp-ut-common
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
 )
 target_link_options(ydb-core-kqp-ut-yql PRIVATE
   -ldl
diff --git a/ydb/core/kqp/ut/yql/CMakeLists.windows-x86_64.txt b/ydb/core/kqp/ut/yql/CMakeLists.windows-x86_64.txt
index b385857cece..2df2e09b818 100644
--- a/ydb/core/kqp/ut/yql/CMakeLists.windows-x86_64.txt
+++ b/ydb/core/kqp/ut/yql/CMakeLists.windows-x86_64.txt
@@ -21,7 +21,8 @@ target_link_libraries(ydb-core-kqp-ut-yql PUBLIC
   cpp-testing-unittest_main
   ydb-core-kqp
   kqp-ut-common
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
 )
 target_sources(ydb-core-kqp-ut-yql PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/ut/yql/kqp_pragma_ut.cpp
diff --git a/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp b/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp
index 31a7df30bc8..1301e84fbb5 100644
--- a/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp
+++ b/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp
@@ -1101,6 +1101,53 @@ Y_UNIT_TEST_SUITE(KqpScripting) {
             [[1];["Value21"]]
         ])", FormatResultSetYson(result.GetResultSet(0)));
     }
+
+    Y_UNIT_TEST(ExecuteYqlScriptPg) {
+        TKikimrRunner kikimr;
+
+        auto settings = TExecuteYqlRequestSettings()
+            .Syntax(Ydb::Query::SYNTAX_PG);
+
+        TScriptingClient client(kikimr.GetDriver());
+        auto result = client.ExecuteYqlScript(R"(
+            SELECT * FROM (VALUES
+                (1::int8, 'one'),
+                (2::int8, 'two'),
+                (3::int8, 'three')
+            ) AS t;
+        )", settings).GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        CompareYson(R"([
+            ["1";"one"];
+            ["2";"two"];
+            ["3";"three"]
+        ])", FormatResultSetYson(result.GetResultSet(0)));
+    }
+
+    Y_UNIT_TEST(StreamExecuteYqlScriptPg) {
+        TKikimrRunner kikimr;
+        
+        auto settings = TExecuteYqlRequestSettings()
+            .Syntax(Ydb::Query::SYNTAX_PG);
+
+        TScriptingClient client(kikimr.GetDriver());
+
+        auto result = client.StreamExecuteYqlScript(R"(
+            SELECT * FROM (VALUES
+                (1::int8, 'one'),
+                (2::int8, 'two'),
+                (3::int8, 'three')
+            ) AS t;
+        )", settings).GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        CompareYson(R"([[
+            ["1";"one"];
+            ["2";"two"];
+            ["3";"three"]
+        ]])", StreamResultToYson(result));
+    }
 }
 
 } // namespace NKqp
diff --git a/ydb/core/kqp/ut/yql/ya.make b/ydb/core/kqp/ut/yql/ya.make
index 77593d5518d..d6ad9236496 100644
--- a/ydb/core/kqp/ut/yql/ya.make
+++ b/ydb/core/kqp/ut/yql/ya.make
@@ -21,7 +21,8 @@ SRCS(
 PEERDIR(
     ydb/core/kqp
     ydb/core/kqp/ut/common
-    ydb/library/yql/sql/pg_dummy
+    ydb/library/yql/sql/pg
+    ydb/library/yql/parser/pg_wrapper
 )
 
 YQL_LAST_ABI_VERSION()
diff --git a/ydb/public/api/protos/ydb_scripting.proto b/ydb/public/api/protos/ydb_scripting.proto
index 4c1a2954185..894f92e0c40 100644
--- a/ydb/public/api/protos/ydb_scripting.proto
+++ b/ydb/public/api/protos/ydb_scripting.proto
@@ -13,12 +13,14 @@ import "ydb/public/api/protos/ydb_table.proto";
 import "ydb/public/api/protos/ydb_query_stats.proto";
 import "ydb/public/api/protos/ydb_issue_message.proto";
 import "ydb/public/api/protos/ydb_status_codes.proto";
+import "ydb/public/api/protos/ydb_query.proto";
 
 message ExecuteYqlRequest {
     Ydb.Operations.OperationParams operation_params = 1;
     string script = 2;
     map<string, TypedValue> parameters = 3;
     Ydb.Table.QueryStatsCollection.Mode collect_stats = 4;
+    Ydb.Query.Syntax syntax = 5;
 }
 
 message ExecuteYqlResponse {
diff --git a/ydb/public/sdk/cpp/client/draft/ydb_scripting.cpp b/ydb/public/sdk/cpp/client/draft/ydb_scripting.cpp
index 313219f5c4a..9743c22bcd3 100644
--- a/ydb/public/sdk/cpp/client/draft/ydb_scripting.cpp
+++ b/ydb/public/sdk/cpp/client/draft/ydb_scripting.cpp
@@ -157,6 +157,7 @@ public:
         request.set_script(script);
         SetParams(params, &request);
         request.set_collect_stats(GetStatsCollectionMode(settings.CollectQueryStats_));
+        request.set_syntax(settings.Syntax_);
 
         auto promise = NewPromise<TExecuteYqlResult>();
 
@@ -202,6 +203,7 @@ public:
         request.set_script(script);
         SetParams(params, &request);
         request.set_collect_stats(GetStatsCollectionMode(settings.CollectQueryStats_));
+        request.set_syntax(settings.Syntax_);
 
         auto promise = NewPromise<std::pair<TPlainStatus, TYqlScriptProcessorPtr>>();
 
diff --git a/ydb/public/sdk/cpp/client/draft/ydb_scripting.h b/ydb/public/sdk/cpp/client/draft/ydb_scripting.h
index e99f05daa24..324c96597d6 100644
--- a/ydb/public/sdk/cpp/client/draft/ydb_scripting.h
+++ b/ydb/public/sdk/cpp/client/draft/ydb_scripting.h
@@ -2,6 +2,7 @@
 
 #include <ydb/public/sdk/cpp/client/ydb_table/table.h>
 #include <ydb/public/api/protos/ydb_value.pb.h>
+#include <ydb/public/api/protos/ydb_query.pb.h>
 
 namespace NYdb {
 namespace NScripting {
@@ -101,6 +102,7 @@ using TAsyncExplainYqlResult = NThreading::TFuture<TExplainYqlResult>;
 ////////////////////////////////////////////////////////////////////////////////
 
 struct TExecuteYqlRequestSettings : public TOperationRequestSettings<TExecuteYqlRequestSettings> {
+    FLUENT_SETTING_DEFAULT(Ydb::Query::Syntax, Syntax, Ydb::Query::SYNTAX_YQL_V1);
     FLUENT_SETTING_DEFAULT(NTable::ECollectQueryStatsMode, CollectQueryStats, NTable::ECollectQueryStatsMode::None);
 };