Add DocumentApiRestricted into SessionActor

3 files changed, 10 insertions, 6 deletions

diff --git a/ydb/core/kqp/kqp_session_actor.cpp b/ydb/core/kqp/kqp_session_actor.cpp
index 3e8c7122791..fb5cd91db15 100644
--- a/ydb/core/kqp/kqp_session_actor.cpp
+++ b/ydb/core/kqp/kqp_session_actor.cpp
@@ -84,6 +84,7 @@ struct TKqpQueryState {
 
     ui64 CurrentTx = 0;
     TString TraceId;
+    bool IsDocumentApiRestricted = false;
 
     TInstant StartTime;
     NYql::TKikimrQueryDeadlines QueryDeadlines;
@@ -332,6 +333,7 @@ public:
         QueryState->Sender = ev->Sender;
         QueryState->ProxyRequestId = proxyRequestId;
         QueryState->TraceId = requestInfo.GetTraceId();
+        QueryState->IsDocumentApiRestricted = IsDocumentApiRestricted(event.GetRequestType());
         QueryState->StartTime = TInstant::Now();
         QueryState->UserToken = event.GetUserToken();
         QueryState->QueryDeadlines = GetQueryDeadlines(queryRequest);
@@ -482,6 +484,10 @@ public:
                 YQL_ENSURE(false);
         }
 
+        if (query) {
+            query->Settings.DocumentApiRestricted = QueryState->IsDocumentApiRestricted;
+        }
+
         auto compileDeadline = QueryState->QueryDeadlines.TimeoutAt;
         if (QueryState->QueryDeadlines.CancelAt) {
             compileDeadline = Min(compileDeadline, QueryState->QueryDeadlines.CancelAt);
diff --git a/ydb/core/kqp/kqp_worker_actor.cpp b/ydb/core/kqp/kqp_worker_actor.cpp
index 2de558f1115..a3fc2966b59 100644
--- a/ydb/core/kqp/kqp_worker_actor.cpp
+++ b/ydb/core/kqp/kqp_worker_actor.cpp
@@ -41,8 +41,6 @@ void FailForcedNewEngineExecutionForTests(bool fail) {
 
 namespace {
 
-constexpr std::string_view DocumentApiRequestType = "_document_api_request"sv;
-
 using TQueryResult = IKqpHost::TQueryResult;
 
 struct TKqpQueryState {
@@ -2041,10 +2039,6 @@ private:
         return Nothing();
     }
 
-    static bool IsDocumentApiRestricted(const TString& requestType) {
-        return requestType != DocumentApiRequestType;
-    }
-
     static TKikimrQueryLimits GetQueryLimits(const TKqpWorkerSettings& settings) {
         const auto& queryLimitsProto = settings.Service.GetQueryLimits();
         const auto& phaseLimitsProto = queryLimitsProto.GetPhaseLimits();
diff --git a/ydb/core/kqp/kqp_worker_common.h b/ydb/core/kqp/kqp_worker_common.h
index 871c8be02a0..c0878d52f92 100644
--- a/ydb/core/kqp/kqp_worker_common.h
+++ b/ydb/core/kqp/kqp_worker_common.h
@@ -127,4 +127,8 @@ void SlowLogQuery(const TActorContext &ctx, const NYql::TKikimrConfiguration* co
 NYql::EKikimrStatsMode GetStatsModeInt(const NKikimrKqp::TQueryRequest& queryRequest);
 NYql::TKikimrQueryLimits GetQueryLimits(const TKqpWorkerSettings& settings);
 
+inline bool IsDocumentApiRestricted(const TString& requestType) {
+    return requestType != "_document_api_request"sv;
+}
+
 }