Add parameters support to generic queries. ()

8 files changed, 75 insertions, 31 deletions

diff --git a/ydb/core/kqp/ut/service/kqp_query_service_ut.cpp b/ydb/core/kqp/ut/service/kqp_query_service_ut.cpp
index 2111ebdd288..33f917b03ed 100644
--- a/ydb/core/kqp/ut/service/kqp_query_service_ut.cpp
+++ b/ydb/core/kqp/ut/service/kqp_query_service_ut.cpp
@@ -13,9 +13,13 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
         auto kikimr = DefaultKikimrRunner();
         auto db = kikimr.GetQueryClient();
 
+        auto params = TParamsBuilder()
+            .AddParam("$value").Int64(17).Build()
+            .Build();
+
         auto it = db.StreamExecuteQuery(R"(
-            SELECT 1;
-        )", TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            SELECT $value;
+        )", TTxControl::BeginTx().CommitTx(), params).ExtractValueSync();
         UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
 
         ui64 count = 0;
@@ -39,12 +43,16 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
         auto kikimr = DefaultKikimrRunner();
         auto db = kikimr.GetQueryClient();
 
+        auto params = TParamsBuilder()
+            .AddParam("$value").Int64(17).Build()
+            .Build();
+
         auto result = db.ExecuteQuery(R"(
-            SELECT 1;
-        )", TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+            SELECT $value;
+        )", TTxControl::BeginTx().CommitTx(), params).ExtractValueSync();
         UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
 
-        CompareYson(R"([[1]])", FormatResultSetYson(result.GetResultSet(0)));
+        CompareYson(R"([[17]])", FormatResultSetYson(result.GetResultSet(0)));
     }
 
     Y_UNIT_TEST(StreamExecuteQuery) {
diff --git a/ydb/public/api/protos/draft/ydb_query.proto b/ydb/public/api/protos/draft/ydb_query.proto
index b61e4228dd7..ad2eae0386b 100644
--- a/ydb/public/api/protos/draft/ydb_query.proto
+++ b/ydb/public/api/protos/draft/ydb_query.proto
@@ -146,16 +146,16 @@ enum ExecMode {
     EXEC_MODE_PARSE = 10;
     EXEC_MODE_VALIDATE = 20;
     EXEC_MODE_EXPLAIN = 30;
-    reserved 4;
-    EXEC_MODE_EXECUTE = 3;
+    reserved 40;
+    EXEC_MODE_EXECUTE = 50;
 }
 
 enum StatsMode {
     STATS_MODE_UNSPECIFIED = 0;
-    STATS_MODE_NONE = 1;       // Stats collection is disabled
-    STATS_MODE_BASIC = 2;      // Aggregated stats of reads, updates and deletes per table
-    STATS_MODE_FULL = 3;       // Add execution stats and plan on top of STATS_MODE_BASIC
-    STATS_MODE_PROFILE = 4;    // Detailed execution stats including stats for individual tasks and channels
+    STATS_MODE_NONE = 10;       // Stats collection is disabled
+    STATS_MODE_BASIC = 20;      // Aggregated stats of reads, updates and deletes per table
+    STATS_MODE_FULL = 30;       // Add execution stats and plan on top of STATS_MODE_BASIC
+    STATS_MODE_PROFILE = 40;    // Detailed execution stats including stats for individual tasks and channels
 }
 
 message ExecuteQueryRequest {
@@ -172,9 +172,9 @@ message ExecuteQueryRequest {
     reserved 7; // idempotency_key
     StatsMode stats_mode = 8;
 
-    // For queries with multiple result sets, some of them may be computed in parallel.
+    // For queries with multiple result sets, some of them may be computed concurrently.
     // If true, parts of different results sets may be interleaved in response stream.
-    bool parallel_result_sets = 9;
+    bool concurrent_result_sets = 9;
 }
 
 message ResultSetMeta {
diff --git a/ydb/public/sdk/cpp/client/draft/ydb_query/client.cpp b/ydb/public/sdk/cpp/client/draft/ydb_query/client.cpp
index 34f63848404..9158e3d41fb 100644
--- a/ydb/public/sdk/cpp/client/draft/ydb_query/client.cpp
+++ b/ydb/public/sdk/cpp/client/draft/ydb_query/client.cpp
@@ -22,15 +22,15 @@ public:
     }
 
     TAsyncExecuteQueryIterator StreamExecuteQuery(const TString& query, const TTxControl& txControl,
-        const TExecuteQuerySettings& settings)
+        const TMaybe<TParams>& params, const TExecuteQuerySettings& settings)
     {
-        return TExecQueryImpl::StreamExecuteQuery(Connections_, DbDriverState_, query, txControl, settings);
+        return TExecQueryImpl::StreamExecuteQuery(Connections_, DbDriverState_, query, txControl, params, settings);
     }
 
     TAsyncExecuteQueryResult ExecuteQuery(const TString& query, const TTxControl& txControl,
-        const TExecuteQuerySettings& settings)
+        const TMaybe<TParams>& params, const TExecuteQuerySettings& settings)
     {
-        return TExecQueryImpl::ExecuteQuery(Connections_, DbDriverState_, query, txControl, settings);
+        return TExecQueryImpl::ExecuteQuery(Connections_, DbDriverState_, query, txControl, params, settings);
     }
 
     NThreading::TFuture<TScriptExecutionOperation> ExecuteScript(const TString& script, const TExecuteScriptSettings& settings) {
@@ -129,16 +129,29 @@ TQueryClient::TQueryClient(const TDriver& driver, const TClientSettings& setting
 {
 }
 
-TAsyncExecuteQueryResult TQueryClient::ExecuteQuery(const TString& query,
-    const TTxControl& txControl, const TExecuteQuerySettings& settings)
+TAsyncExecuteQueryResult TQueryClient::ExecuteQuery(const TString& query, const TTxControl& txControl,
+    const TExecuteQuerySettings& settings)
 {
-    return Impl_->ExecuteQuery(query, txControl, settings);
+    return Impl_->ExecuteQuery(query, txControl, {}, settings);
 }
 
-TAsyncExecuteQueryIterator TQueryClient::StreamExecuteQuery(const TString& query,
-    const TTxControl& txControl, const TExecuteQuerySettings& settings)
+TAsyncExecuteQueryResult TQueryClient::ExecuteQuery(const TString& query, const TTxControl& txControl,
+    const TParams& params, const TExecuteQuerySettings& settings)
 {
-    return Impl_->StreamExecuteQuery(query, txControl, settings);
+    return Impl_->ExecuteQuery(query, txControl, params, settings);
+}
+
+
+TAsyncExecuteQueryIterator TQueryClient::StreamExecuteQuery(const TString& query, const TTxControl& txControl,
+    const TExecuteQuerySettings& settings)
+{
+    return Impl_->StreamExecuteQuery(query, txControl, {}, settings);
+}
+
+TAsyncExecuteQueryIterator TQueryClient::StreamExecuteQuery(const TString& query, const TTxControl& txControl,
+    const TParams& params, const TExecuteQuerySettings& settings)
+{
+    return Impl_->StreamExecuteQuery(query, txControl, params, settings);
 }
 
 NThreading::TFuture<TScriptExecutionOperation> TQueryClient::ExecuteScript(const TString& script,
diff --git a/ydb/public/sdk/cpp/client/draft/ydb_query/client.h b/ydb/public/sdk/cpp/client/draft/ydb_query/client.h
index 59daab44697..955a484183c 100644
--- a/ydb/public/sdk/cpp/client/draft/ydb_query/client.h
+++ b/ydb/public/sdk/cpp/client/draft/ydb_query/client.h
@@ -4,6 +4,7 @@
 #include "tx.h"
 
 #include <ydb/public/sdk/cpp/client/ydb_driver/driver.h>
+#include <ydb/public/sdk/cpp/client/ydb_params/params.h>
 
 #include <util/generic/maybe.h>
 #include <util/generic/ptr.h>
@@ -25,9 +26,15 @@ public:
     TAsyncExecuteQueryResult ExecuteQuery(const TString& query, const TTxControl& txControl,
         const TExecuteQuerySettings& settings = TExecuteQuerySettings());
 
+    TAsyncExecuteQueryResult ExecuteQuery(const TString& query, const TTxControl& txControl,
+        const TParams& params, const TExecuteQuerySettings& settings = TExecuteQuerySettings());
+
     TAsyncExecuteQueryIterator StreamExecuteQuery(const TString& query, const TTxControl& txControl,
         const TExecuteQuerySettings& settings = TExecuteQuerySettings());
 
+    TAsyncExecuteQueryIterator StreamExecuteQuery(const TString& query, const TTxControl& txControl,
+        const TParams& params, const TExecuteQuerySettings& settings = TExecuteQuerySettings());
+
     NThreading::TFuture<TScriptExecutionOperation> ExecuteScript(const TString& script,
         const TExecuteScriptSettings& settings = TExecuteScriptSettings());
 
diff --git a/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.cpp b/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.cpp
index 4c93bbf0a8b..0a342a423bb 100644
--- a/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.cpp
+++ b/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.cpp
@@ -167,7 +167,8 @@ struct TExecuteQueryBuffer : public TThrRefBase, TNonCopyable {
 
 TFuture<std::pair<TPlainStatus, TExecuteQueryProcessorPtr>> StreamExecuteQueryImpl(
     const std::shared_ptr<TGRpcConnectionsImpl>& connections, const TDbDriverStatePtr& driverState,
-    const TString& query, const TTxControl& txControl, const TExecuteQuerySettings& settings)
+    const TString& query, const TTxControl& txControl, const ::google::protobuf::Map<TString, Ydb::TypedValue>* params,
+    const TExecuteQuerySettings& settings)
 {
     auto request = MakeRequest<Ydb::Query::ExecuteQueryRequest>();
     request.set_exec_mode(Ydb::Query::EXEC_MODE_EXECUTE);
@@ -182,6 +183,10 @@ TFuture<std::pair<TPlainStatus, TExecuteQueryProcessorPtr>> StreamExecuteQueryIm
         SetTxSettings(*txControl.TxSettings_, requestTxControl->mutable_begin_tx());
     }
 
+    if (params) {
+        *request.mutable_parameters() = *params;
+    }
+
     auto promise = NewPromise<std::pair<TPlainStatus, TExecuteQueryProcessorPtr>>();
 
     connections->StartReadStream<
@@ -203,7 +208,7 @@ TFuture<std::pair<TPlainStatus, TExecuteQueryProcessorPtr>> StreamExecuteQueryIm
 
 TAsyncExecuteQueryIterator TExecQueryImpl::StreamExecuteQuery(const std::shared_ptr<TGRpcConnectionsImpl>& connections,
     const TDbDriverStatePtr& driverState, const TString& query, const TTxControl& txControl,
-    const TExecuteQuerySettings& settings)
+    const TMaybe<TParams>& params, const TExecuteQuerySettings& settings)
 {
     auto promise = NewPromise<TExecuteQueryIterator>();
 
@@ -218,15 +223,20 @@ TAsyncExecuteQueryIterator TExecQueryImpl::StreamExecuteQuery(const std::shared_
         );
     };
 
-    StreamExecuteQueryImpl(connections, driverState, query, txControl, settings).Subscribe(iteratorCallback);
+    auto paramsProto = params
+        ? &params->GetProtoMap()
+        : nullptr;
+
+    StreamExecuteQueryImpl(connections, driverState, query, txControl, paramsProto, settings)
+        .Subscribe(iteratorCallback);
     return promise.GetFuture();
 }
 
 TAsyncExecuteQueryResult TExecQueryImpl::ExecuteQuery(const std::shared_ptr<TGRpcConnectionsImpl>& connections,
     const TDbDriverStatePtr& driverState, const TString& query, const TTxControl& txControl,
-    const TExecuteQuerySettings& settings)
+    const TMaybe<TParams>& params, const TExecuteQuerySettings& settings)
 {
-    return StreamExecuteQuery(connections, driverState, query, txControl, settings)
+    return StreamExecuteQuery(connections, driverState, query, txControl, params, settings)
         .Apply([](TAsyncExecuteQueryIterator itFuture){
             auto it = itFuture.ExtractValue();
 
diff --git a/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.h b/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.h
index ba562dfbc6e..bd7949e588c 100644
--- a/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.h
+++ b/ydb/public/sdk/cpp/client/draft/ydb_query/impl/exec_query.h
@@ -5,6 +5,7 @@
 #include <ydb/public/sdk/cpp/client/draft/ydb_query/query.h>
 #include <ydb/public/sdk/cpp/client/draft/ydb_query/tx.h>
 #include <ydb/public/sdk/cpp/client/impl/ydb_internal/grpc_connections/grpc_connections.h>
+#include <ydb/public/sdk/cpp/client/ydb_params/params.h>
 
 namespace NYdb::NQuery {
 
@@ -12,11 +13,11 @@ class TExecQueryImpl {
 public:
     static TAsyncExecuteQueryIterator StreamExecuteQuery(const std::shared_ptr<TGRpcConnectionsImpl>& connections,
         const TDbDriverStatePtr& driverState, const TString& query, const TTxControl& txControl,
-        const TExecuteQuerySettings& settings);
+        const TMaybe<TParams>& params, const TExecuteQuerySettings& settings);
 
     static TAsyncExecuteQueryResult ExecuteQuery(const std::shared_ptr<TGRpcConnectionsImpl>& connections,
         const TDbDriverStatePtr& driverState, const TString& query, const TTxControl& txControl,
-        const TExecuteQuerySettings& settings);
+        const TMaybe<TParams>& params, const TExecuteQuerySettings& settings);
 };
 
 } // namespace NYdb::NQuery::NImpl
diff --git a/ydb/public/sdk/cpp/client/draft/ydb_query/query.h b/ydb/public/sdk/cpp/client/draft/ydb_query/query.h
index e7f137dd8b4..b75b03053b3 100644
--- a/ydb/public/sdk/cpp/client/draft/ydb_query/query.h
+++ b/ydb/public/sdk/cpp/client/draft/ydb_query/query.h
@@ -25,7 +25,7 @@ enum class EExecMode {
     Parse = 10,
     Validate = 20,
     Explain = 30,
-    Execute = 3,
+    Execute = 50,
 };
 
 enum class EExecStatus {
diff --git a/ydb/public/sdk/cpp/client/ydb_params/params.h b/ydb/public/sdk/cpp/client/ydb_params/params.h
index 6f565b20f59..6e8b204408e 100644
--- a/ydb/public/sdk/cpp/client/ydb_params/params.h
+++ b/ydb/public/sdk/cpp/client/ydb_params/params.h
@@ -26,6 +26,10 @@ namespace NExperimental {
     class TStreamQueryClient;
 }
 
+namespace NQuery {
+    class TExecQueryImpl;
+}
+
 class TParamsBuilder;
 
 class TParams {
@@ -35,6 +39,7 @@ class TParams {
     friend class NTable::TDataQuery;
     friend class NScripting::TScriptingClient;
     friend class NExperimental::TStreamQueryClient;
+    friend class NQuery::TExecQueryImpl;
     friend class NYdb::TProtoAccessor;
 public:
     bool Empty() const;