YQ-3568 fix error sent large get operation (ydb-platform#8035)

GrigoriyPA · Aug 21, 2024 · 8463088 · 8463088
1 parent d3f247f
commit 8463088
Showing 1 changed file with 22 additions and 4 deletions.
diff --git a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.cpp b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.cpp
@@ -19,6 +19,8 @@ namespace NKikimr::NKqp {
 namespace {
 
 class TScriptFinalizerActor : public TActorBootstrapped<TScriptFinalizerActor> {
+    static constexpr size_t MAX_ARTIFACTS_SIZE_BYTES = 40_MB;
+
 public:
     TScriptFinalizerActor(TEvScriptFinalizeRequest::TPtr request,
         const NKikimrConfig::TQueryServiceConfig& queryServiceConfig,
@@ -35,6 +37,14 @@ class TScriptFinalizerActor : public TActorBootstrapped<TScriptFinalizerActor> {
 
     void CompressScriptArtifacts() const {
         auto& description = Request->Get()->Description;
+
+        TString astTruncateDescription;
+        if (size_t planSize = description.QueryPlan.value_or("").size(); description.QueryAst && description.QueryAst->size() + planSize > MAX_ARTIFACTS_SIZE_BYTES) {
+            astTruncateDescription = TStringBuilder() << "Query artifacts size is " << description.QueryAst->size() + planSize << " bytes (plan + ast), that is larger than allowed limit " << MAX_ARTIFACTS_SIZE_BYTES << " bytes, ast was truncated";
+            size_t toRemove = std::min(description.QueryAst->size() + planSize - MAX_ARTIFACTS_SIZE_BYTES, description.QueryAst->size());
+            description.QueryAst = TruncateString(*description.QueryAst, description.QueryAst->size() - toRemove);
+        }
+
         auto ast = description.QueryAst;
         if (Compressor.IsEnabled() && ast) {
             const auto& [astCompressionMethod, astCompressed] = Compressor.Compress(*ast);
@@ -43,12 +53,15 @@ class TScriptFinalizerActor : public TActorBootstrapped<TScriptFinalizerActor> {
         }
 
         if (description.QueryAst && description.QueryAst->size() > NDataShard::NLimits::MaxWriteValueSize) {
-            NYql::TIssue astTruncatedIssue(TStringBuilder() << "Query ast size is " << description.QueryAst->size() << " bytes, that is larger than allowed limit " << NDataShard::NLimits::MaxWriteValueSize << " bytes, ast was truncated");
+            astTruncateDescription = TStringBuilder() << "Query ast size is " << description.QueryAst->size() << " bytes, that is larger than allowed limit " << NDataShard::NLimits::MaxWriteValueSize << " bytes, ast was truncated";
+            description.QueryAst = TruncateString(*ast, NDataShard::NLimits::MaxWriteValueSize - 1_KB);
+            description.QueryAstCompressionMethod = std::nullopt;
+        }
+
+        if (astTruncateDescription) {
+            NYql::TIssue astTruncatedIssue(astTruncateDescription);
             astTruncatedIssue.SetCode(NYql::DEFAULT_ERROR, NYql::TSeverityIds::S_INFO);
             description.Issues.AddIssue(astTruncatedIssue);
-
-            description.QueryAst = ast->substr(0, NDataShard::NLimits::MaxWriteValueSize - 1_KB) + "...\n(TRUNCATED)";
-            description.QueryAstCompressionMethod = std::nullopt;
         }
     }
 
@@ -221,6 +234,11 @@ class TScriptFinalizerActor : public TActorBootstrapped<TScriptFinalizerActor> {
         PassAway();
     }
 
+private:
+    static TString TruncateString(const TString& str, size_t size) {
+        return str.substr(0, std::min(str.size(), size)) + "...\n(TRUNCATED)";
+    }
+
 private:
     const TActorId ReplyActor;
     const TString ExecutionId;