KIKIMR-19649: use query cache for query with text and user params

9 files changed, 237 insertions, 9 deletions

diff --git a/ydb/core/kqp/common/simple/settings.h b/ydb/core/kqp/common/simple/settings.h
index cee8a624cb..dac06f4711 100644
--- a/ydb/core/kqp/common/simple/settings.h
+++ b/ydb/core/kqp/common/simple/settings.h
@@ -13,6 +13,7 @@ struct TKqpQuerySettings {
     bool IsInternalCall = false;
     NKikimrKqp::EQueryType QueryType = NKikimrKqp::EQueryType::QUERY_TYPE_UNDEFINED;
     Ydb::Query::Syntax Syntax = Ydb::Query::Syntax::SYNTAX_UNSPECIFIED;
+    bool IsPrepareQuery = false;
 
     explicit TKqpQuerySettings(NKikimrKqp::EQueryType queryType)
         : QueryType(queryType) {}
diff --git a/ydb/core/kqp/compile_service/CMakeLists.darwin-x86_64.txt b/ydb/core/kqp/compile_service/CMakeLists.darwin-x86_64.txt
index 1288e5d843..5065b4c97f 100644
--- a/ydb/core/kqp/compile_service/CMakeLists.darwin-x86_64.txt
+++ b/ydb/core/kqp/compile_service/CMakeLists.darwin-x86_64.txt
@@ -19,6 +19,7 @@ target_link_libraries(core-kqp-compile_service PUBLIC
   kqp-common-simple
   core-kqp-federated_query
   core-kqp-host
+  ydb-core-ydb_convert
 )
 target_sources(core-kqp-compile_service PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
diff --git a/ydb/core/kqp/compile_service/CMakeLists.linux-aarch64.txt b/ydb/core/kqp/compile_service/CMakeLists.linux-aarch64.txt
index 04756092e4..c9b19a6fe0 100644
--- a/ydb/core/kqp/compile_service/CMakeLists.linux-aarch64.txt
+++ b/ydb/core/kqp/compile_service/CMakeLists.linux-aarch64.txt
@@ -20,6 +20,7 @@ target_link_libraries(core-kqp-compile_service PUBLIC
   kqp-common-simple
   core-kqp-federated_query
   core-kqp-host
+  ydb-core-ydb_convert
 )
 target_sources(core-kqp-compile_service PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
diff --git a/ydb/core/kqp/compile_service/CMakeLists.linux-x86_64.txt b/ydb/core/kqp/compile_service/CMakeLists.linux-x86_64.txt
index 04756092e4..c9b19a6fe0 100644
--- a/ydb/core/kqp/compile_service/CMakeLists.linux-x86_64.txt
+++ b/ydb/core/kqp/compile_service/CMakeLists.linux-x86_64.txt
@@ -20,6 +20,7 @@ target_link_libraries(core-kqp-compile_service PUBLIC
   kqp-common-simple
   core-kqp-federated_query
   core-kqp-host
+  ydb-core-ydb_convert
 )
 target_sources(core-kqp-compile_service PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
diff --git a/ydb/core/kqp/compile_service/CMakeLists.windows-x86_64.txt b/ydb/core/kqp/compile_service/CMakeLists.windows-x86_64.txt
index 1288e5d843..5065b4c97f 100644
--- a/ydb/core/kqp/compile_service/CMakeLists.windows-x86_64.txt
+++ b/ydb/core/kqp/compile_service/CMakeLists.windows-x86_64.txt
@@ -19,6 +19,7 @@ target_link_libraries(core-kqp-compile_service PUBLIC
   kqp-common-simple
   core-kqp-federated_query
   core-kqp-host
+  ydb-core-ydb_convert
 )
 target_sources(core-kqp-compile_service PRIVATE
   ${CMAKE_SOURCE_DIR}/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
diff --git a/ydb/core/kqp/compile_service/kqp_compile_service.cpp b/ydb/core/kqp/compile_service/kqp_compile_service.cpp
index a857fdb959..0264dd7e18 100644
--- a/ydb/core/kqp/compile_service/kqp_compile_service.cpp
+++ b/ydb/core/kqp/compile_service/kqp_compile_service.cpp
@@ -7,6 +7,7 @@
 #include <ydb/core/cms/console/configs_dispatcher.h>
 #include <ydb/core/kqp/counters/kqp_counters.h>
 #include <ydb/core/kqp/common/kqp_lwtrace_probes.h>
+#include <ydb/core/ydb_convert/ydb_convert.h>
 #include <ydb/library/aclib/aclib.h>
 
 #include <library/cpp/actors/core/actor_bootstrapped.h>
@@ -35,6 +36,8 @@ public:
         Y_ENSURE(compileResult->Query);
         auto& query = *compileResult->Query;
 
+        YQL_ENSURE(compileResult->PreparedQuery);
+
         auto queryIt = QueryIndex.emplace(query, compileResult->Uid);
         Y_ENSURE(queryIt.second);
 
@@ -675,6 +678,11 @@ private:
                     if (QueryCache.Insert(compileResult)) {
                         Counters->CompileQueryCacheEvicted->Inc();
                     }
+                    if (compileResult->Query && compileResult->Query->Settings.IsPrepareQuery) {
+                        if (InsertPreparingQuery(compileResult, compileRequest.KeepInCache)) {
+                            Counters->CompileQueryCacheEvicted->Inc();
+                        };
+                    }
                 }
 
                 if (ev->Get()->ReplayMessage) {
@@ -741,6 +749,31 @@ private:
     }
 
 private:
+    bool InsertPreparingQuery(const TKqpCompileResult::TConstPtr& compileResult, bool keepInCache) {
+        YQL_ENSURE(compileResult->Query);
+        auto query = *compileResult->Query;
+
+        YQL_ENSURE(compileResult->PreparedQuery);
+        YQL_ENSURE(!query.QueryParameterTypes);
+        if (compileResult->PreparedQuery->GetParameters().empty()) {
+            return false;
+        }
+        auto queryParameterTypes = std::make_shared<std::map<TString, Ydb::Type>>();
+        for (const auto& param : compileResult->PreparedQuery->GetParameters()) {
+            Ydb::Type paramType;
+            ConvertMiniKQLTypeToYdbType(param.GetType(), paramType);
+            queryParameterTypes->insert({param.GetName(), paramType});
+        }
+        query.QueryParameterTypes = queryParameterTypes;
+        if (QueryCache.FindByQuery(query, keepInCache)) {
+            return false;
+        }
+        auto newCompileResult = TKqpCompileResult::Make(CreateGuidAsString(), std::move(query), compileResult->Status, compileResult->Issues, compileResult->MaxReadType);
+        newCompileResult->AllowCache = compileResult->AllowCache;
+        newCompileResult->PreparedQuery = compileResult->PreparedQuery;
+        return QueryCache.Insert(newCompileResult);
+    }
+
     void ProcessQueue(const TActorContext& ctx) {
         auto maxActiveRequests = TableServiceConfig.GetCompileMaxActiveRequests();
 
diff --git a/ydb/core/kqp/compile_service/ya.make b/ydb/core/kqp/compile_service/ya.make
index 53ba78b8be..04f075dd25 100644
--- a/ydb/core/kqp/compile_service/ya.make
+++ b/ydb/core/kqp/compile_service/ya.make
@@ -12,6 +12,7 @@ PEERDIR(
     ydb/core/kqp/common/simple
     ydb/core/kqp/federated_query
     ydb/core/kqp/host
+    ydb/core/ydb_convert
 )
 
 YQL_LAST_ABI_VERSION()
diff --git a/ydb/core/kqp/session_actor/kqp_query_state.cpp b/ydb/core/kqp/session_actor/kqp_query_state.cpp
index 07ecf58f65..dbf5149482 100644
--- a/ydb/core/kqp/session_actor/kqp_query_state.cpp
+++ b/ydb/core/kqp/session_actor/kqp_query_state.cpp
@@ -131,6 +131,7 @@ std::unique_ptr<TEvKqp::TEvCompileRequest> TKqpQueryState::BuildCompileRequest(s
     settings.DocumentApiRestricted = IsDocumentApiRestricted_;
     settings.IsInternalCall = IsInternalCall();
     settings.Syntax = GetSyntax();
+    settings.IsPrepareQuery = GetAction() == NKikimrKqp::QUERY_ACTION_PREPARE;
 
     bool keepInCache = false;
     switch (GetAction()) {
diff --git a/ydb/core/kqp/ut/query/kqp_params_ut.cpp b/ydb/core/kqp/ut/query/kqp_params_ut.cpp
index 0f667481b2..185080d65e 100644
--- a/ydb/core/kqp/ut/query/kqp_params_ut.cpp
+++ b/ydb/core/kqp/ut/query/kqp_params_ut.cpp
@@ -104,9 +104,7 @@ Y_UNIT_TEST_SUITE(KqpParams) {
 
     Y_UNIT_TEST(ImplicitParameterTypes) {
         TKikimrRunner kikimr;
-        if (!kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.GetEnableImplicitQueryParameterTypes()) {
-            return;
-        }
+        kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableImplicitQueryParameterTypes(true);
         auto db = kikimr.GetTableClient();
         auto session = db.CreateSession().GetValueSync().GetSession();
 
@@ -127,6 +125,200 @@ Y_UNIT_TEST_SUITE(KqpParams) {
         UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
     }
 
+    Y_UNIT_TEST(CheckQueryCacheForPreparedQuery) {
+        // All params are declared in the text
+        TKikimrRunner kikimr;
+        kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableImplicitQueryParameterTypes(true);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto query = Q1_(R"(
+            DECLARE $group AS Int32;
+            DECLARE $name AS String;
+
+            SELECT * FROM `/Root/Test` WHERE Group = $group AND Name = $name;
+        )");
+
+        auto prepareResult = session.PrepareDataQuery(query).GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(prepareResult.GetStatus(), EStatus::SUCCESS, prepareResult.GetIssues().ToString());
+
+        NYdb::NTable::TExecDataQuerySettings execSettings;
+        execSettings.KeepInQueryCache(true);
+        execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+        auto params = db.GetParamsBuilder()
+            .AddParam("$name")
+                .String("Sergey")
+                .Build()
+            .AddParam("$group")
+                .Int32(1)
+                .Build()
+            .Build();
+
+        auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, execSettings).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        auto stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+        UNIT_ASSERT_VALUES_EQUAL(stats.compilation().from_cache(), true);
+    }
+
+    Y_UNIT_TEST(CheckQueryCacheForUnpreparedQuery) {
+        // Some params are declared in text, some by user
+        TKikimrRunner kikimr;
+        kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableImplicitQueryParameterTypes(true);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto query = Q1_(R"(
+            DECLARE $group AS Int32;
+
+            SELECT $group, $name;
+        )");
+
+        auto params = db.GetParamsBuilder()
+            .AddParam("$name")
+                .String("Sergey")
+                .Build()
+            .AddParam("$group")
+                .Int32(1)
+                .Build()
+            .AddParam("$phone")
+                .String("80")
+                .Build()
+            .Build();
+
+        NYdb::NTable::TExecDataQuerySettings execSettings;
+        execSettings.KeepInQueryCache(true);
+        execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+        auto firstQueryResult = session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, execSettings).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(firstQueryResult.GetStatus(), EStatus::SUCCESS, firstQueryResult.GetIssues().ToString());
+
+        auto stats = NYdb::TProtoAccessor::GetProto(*firstQueryResult.GetStats());
+        UNIT_ASSERT_VALUES_EQUAL(stats.compilation().from_cache(), false);
+
+        {
+            // The same query with the same params
+            auto params = db.GetParamsBuilder()
+                .AddParam("$name")
+                    .String("Sergey")
+                    .Build()
+                .AddParam("$group")
+                    .Int32(1)
+                    .Build()
+                .AddParam("$phone")
+                    .String("80")
+                    .Build()
+                .Build();
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, execSettings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+            auto stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+            UNIT_ASSERT_VALUES_EQUAL(stats.compilation().from_cache(), true);
+        }
+
+        {
+            // The same query with different type of user param
+            auto params = db.GetParamsBuilder()
+                .AddParam("$name")
+                    .Int64(2)
+                    .Build()
+                .AddParam("$group")
+                    .Int32(1)
+                    .Build()
+                .AddParam("$phone")
+                    .String("80")
+                    .Build()
+                .Build();
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, execSettings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+            auto stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+            UNIT_ASSERT_VALUES_EQUAL(stats.compilation().from_cache(), false);
+        }
+
+        {
+            // The same query with extra param
+            auto params = db.GetParamsBuilder()
+                .AddParam("$name")
+                    .String("Sergey")
+                    .Build()
+                .AddParam("$group")
+                    .Int32(1)
+                    .Build()
+                .AddParam("$phone")
+                    .String("80")
+                    .Build()
+                .AddParam("$age")
+                    .Int32(1)
+                    .Build()
+                .Build();
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, execSettings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+            auto stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+            UNIT_ASSERT_VALUES_EQUAL(stats.compilation().from_cache(), false);
+        }
+
+        {
+            // The same query with less params
+            auto params = db.GetParamsBuilder()
+                .AddParam("$name")
+                    .String("Sergey")
+                    .Build()
+                .AddParam("$group")
+                    .Int32(1)
+                    .Build()
+                .Build();
+
+            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, execSettings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+            auto stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+            UNIT_ASSERT_VALUES_EQUAL(stats.compilation().from_cache(), false);
+        }
+    }
+
+    Y_UNIT_TEST(CheckQueryCacheForExecuteAndPreparedQueries) {
+        // All params are declared in the text
+        TKikimrRunner kikimr;
+        kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableImplicitQueryParameterTypes(true);
+        auto db = kikimr.GetTableClient();
+        auto session = db.CreateSession().GetValueSync().GetSession();
+
+        auto query = Q1_(R"(
+            DECLARE $group AS Int32;
+            DECLARE $name AS String;
+
+            SELECT * FROM `/Root/Test` WHERE Group = $group AND Name = $name;
+        )");
+
+        NYdb::NTable::TExecDataQuerySettings execSettings;
+        execSettings.KeepInQueryCache(true);
+        execSettings.CollectQueryStats(ECollectQueryStatsMode::Basic);
+
+        auto params = db.GetParamsBuilder()
+            .AddParam("$name")
+                .String("Sergey")
+                .Build()
+            .AddParam("$group")
+                .Int32(1)
+                .Build()
+            .Build();
+
+        auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(), params, execSettings).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        auto stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
+        UNIT_ASSERT_VALUES_EQUAL(stats.compilation().from_cache(), false);
+
+        auto prepareResult = session.PrepareDataQuery(query).GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(prepareResult.GetStatus(), EStatus::SUCCESS, prepareResult.GetIssues().ToString());
+    }
+
     Y_UNIT_TEST(ExplicitSameParameterTypesQueryCacheCheck) {
         TKikimrRunner kikimr;
         auto db = kikimr.GetTableClient();
@@ -154,9 +346,7 @@ Y_UNIT_TEST_SUITE(KqpParams) {
 
     Y_UNIT_TEST(ImplicitSameParameterTypesQueryCacheCheck) {
         TKikimrRunner kikimr;
-        if (!kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.GetEnableImplicitQueryParameterTypes()) {
-            return;
-        }
+        kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableImplicitQueryParameterTypes(true);
         auto db = kikimr.GetTableClient();
         auto session = db.CreateSession().GetValueSync().GetSession();
 
@@ -182,9 +372,7 @@ Y_UNIT_TEST_SUITE(KqpParams) {
 
     Y_UNIT_TEST(ImplicitDifferentParameterTypesQueryCacheCheck) {
         TKikimrRunner kikimr;
-        if (!kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.GetEnableImplicitQueryParameterTypes()) {
-            return;
-        }
+        kikimr.GetTestServer().GetRuntime()->GetAppData(0).FeatureFlags.SetEnableImplicitQueryParameterTypes(true);
         auto db = kikimr.GetTableClient();
         auto session = db.CreateSession().GetValueSync().GetSession();