From 9aab790974ff2c083a3ff9b76eb274cea5342444 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy <79596613+GrigoriyPA@users.noreply.github.com> Date: Thu, 2 May 2024 13:06:39 +0300 Subject: [PATCH] YQ-3051 add AST compression (#4151) --- .../run/kikimr_services_initializers.cpp | 2 +- .../kqp/common/events/script_executions.h | 1 + .../kqp_finalize_script_actor.cpp | 125 +++++++++++------- .../kqp_finalize_script_actor.h | 2 +- .../kqp_finalize_script_service.cpp | 58 ++++---- .../kqp_finalize_script_service.h | 2 +- .../proxy_service/kqp_script_executions.cpp | 43 +++++- ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp | 23 +++- ydb/core/protos/config.proto | 2 + ydb/core/testlib/test_client.cpp | 2 +- .../kqprun/configuration/app_config.conf | 1 + ydb/tests/tools/kqprun/src/kqp_runner.cpp | 4 + 12 files changed, 174 insertions(+), 91 deletions(-) diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp index 0dae3013e17b..928b1e07ab9f 100644 --- a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp +++ b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp @@ -2073,7 +2073,7 @@ void TKqpServiceInitializer::InitializeServices(NActors::TActorSystemSetup* setu TActorSetupCmd(proxy, TMailboxType::HTSwap, appData->UserPoolId))); // Create finalize script service - auto finalize = NKqp::CreateKqpFinalizeScriptService(Config.GetQueryServiceConfig().GetFinalizeScriptServiceConfig(), Config.GetMetadataProviderConfig(), federatedQuerySetupFactory); + auto finalize = NKqp::CreateKqpFinalizeScriptService(Config.GetQueryServiceConfig(), Config.GetMetadataProviderConfig(), federatedQuerySetupFactory); setup->LocalServices.push_back(std::make_pair( NKqp::MakeKqpFinalizeScriptServiceId(NodeId), TActorSetupCmd(finalize, TMailboxType::HTSwap, appData->UserPoolId))); diff --git a/ydb/core/kqp/common/events/script_executions.h b/ydb/core/kqp/common/events/script_executions.h index 379a46bdaec9..c1173b7ce37e 100644 --- a/ydb/core/kqp/common/events/script_executions.h +++ b/ydb/core/kqp/common/events/script_executions.h @@ -294,6 +294,7 @@ struct TEvScriptFinalizeRequest : public NActors::TEventLocal QueryPlan; std::optional QueryAst; std::optional LeaseGeneration; + std::optional QueryAstCompressionMethod; }; TEvScriptFinalizeRequest(EFinalizationStatus finalizationStatus, const TString& executionId, const TString& database, 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 index 5b2d31a8530e..99ec23e062e6 100644 --- a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.cpp +++ b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.cpp @@ -1,10 +1,13 @@ #include "kqp_finalize_script_actor.h" +#include #include #include #include +#include + #include #include #include @@ -18,21 +21,42 @@ namespace { class TScriptFinalizerActor : public TActorBootstrapped { public: TScriptFinalizerActor(TEvScriptFinalizeRequest::TPtr request, - const NKikimrConfig::TFinalizeScriptServiceConfig& finalizeScriptServiceConfig, + const NKikimrConfig::TQueryServiceConfig& queryServiceConfig, const NKikimrConfig::TMetadataProviderConfig& metadataProviderConfig, const std::optional& federatedQuerySetup) - : ReplyActor_(request->Sender) - , ExecutionId_(request->Get()->Description.ExecutionId) - , Database_(request->Get()->Description.Database) - , FinalizationStatus_(request->Get()->Description.FinalizationStatus) - , Request_(std::move(request)) - , FinalizationTimeout_(TDuration::Seconds(finalizeScriptServiceConfig.GetScriptFinalizationTimeoutSeconds())) - , MaximalSecretsSnapshotWaitTime_(2 * TDuration::Seconds(metadataProviderConfig.GetRefreshPeriodSeconds())) - , FederatedQuerySetup_(federatedQuerySetup) + : ReplyActor(request->Sender) + , ExecutionId(request->Get()->Description.ExecutionId) + , Database(request->Get()->Description.Database) + , FinalizationStatus(request->Get()->Description.FinalizationStatus) + , Request(std::move(request)) + , FinalizationTimeout(TDuration::Seconds(queryServiceConfig.GetFinalizeScriptServiceConfig().GetScriptFinalizationTimeoutSeconds())) + , MaximalSecretsSnapshotWaitTime(2 * TDuration::Seconds(metadataProviderConfig.GetRefreshPeriodSeconds())) + , FederatedQuerySetup(federatedQuerySetup) + , Compressor(queryServiceConfig.GetQueryArtifactsCompressionMethod(), queryServiceConfig.GetQueryArtifactsCompressionMinSize()) {} + void CompressScriptArtifacts() const { + auto& description = Request->Get()->Description; + auto ast = description.QueryAst; + if (Compressor.IsEnabled() && ast) { + const auto& [astCompressionMethod, astCompressed] = Compressor.Compress(*ast); + description.QueryAstCompressionMethod = astCompressionMethod; + description.QueryAst = astCompressed; + } + + 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"); + 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; + } + } + void Bootstrap() { - Register(CreateSaveScriptFinalStatusActor(SelfId(), std::move(Request_))); + CompressScriptArtifacts(); + Register(CreateSaveScriptFinalStatusActor(SelfId(), std::move(Request))); Become(&TScriptFinalizerActor::FetchState); } @@ -46,17 +70,17 @@ class TScriptFinalizerActor : public TActorBootstrapped { return; } - Schedule(FinalizationTimeout_, new TEvents::TEvWakeup()); + Schedule(FinalizationTimeout, new TEvents::TEvWakeup()); Become(&TScriptFinalizerActor::PrepareState); - CustomerSuppliedId_ = ev->Get()->CustomerSuppliedId; - Sinks_ = std::move(ev->Get()->Sinks); - UserToken_ = ev->Get()->UserToken; - SecretNames_ = std::move(ev->Get()->SecretNames); + CustomerSuppliedId = ev->Get()->CustomerSuppliedId; + Sinks = std::move(ev->Get()->Sinks); + UserToken = ev->Get()->UserToken; + SecretNames = std::move(ev->Get()->SecretNames); - if (Sinks_.empty()) { + if (Sinks.empty()) { FinishScriptFinalization(); - } else if (SecretNames_.empty()) { + } else if (SecretNames.empty()) { ComputeScriptExternalEffect(); } else { FetchSecrets(); @@ -75,7 +99,7 @@ class TScriptFinalizerActor : public TActorBootstrapped { } void FetchSecrets() { - RegisterDescribeSecretsActor(SelfId(), UserToken_, SecretNames_, ActorContext().ActorSystem(), MaximalSecretsSnapshotWaitTime_); + RegisterDescribeSecretsActor(SelfId(), UserToken, SecretNames, ActorContext().ActorSystem(), MaximalSecretsSnapshotWaitTime); } void Handle(TEvDescribeSecretsResponse::TPtr& ev) { @@ -90,15 +114,15 @@ class TScriptFinalizerActor : public TActorBootstrapped { void FillSecureParams(const std::vector& secretValues) { std::map secretsMap; for (size_t i = 0; i < secretValues.size(); ++i) { - secretsMap.emplace(SecretNames_[i], secretValues[i]); + secretsMap.emplace(SecretNames[i], secretValues[i]); } - for (const auto& sink : Sinks_) { - auto sinkName = sink.GetSinkName(); + for (const auto& sink : Sinks) { + const auto& sinkName = sink.GetSinkName(); if (sinkName) { const auto& structuredToken = NYql::CreateStructuredTokenParser(sink.GetAuthInfo()).ToBuilder().ReplaceReferences(secretsMap).ToJson(); - SecureParams_.emplace(sinkName, structuredToken); + SecureParams.emplace(sinkName, structuredToken); } } @@ -122,7 +146,7 @@ class TScriptFinalizerActor : public TActorBootstrapped { NYql::NDqProto::TExternalEffect externalEffectS3; externalEffectS3.SetProviderName(TString(NYql::S3ProviderName)); - for (const auto& sink : Sinks_) { + for (const auto& sink : Sinks) { const TString& sinkType = sink.GetType(); if (sinkType == "S3Sink") { @@ -137,20 +161,20 @@ class TScriptFinalizerActor : public TActorBootstrapped { } void RunS3ApplicatorActor(const NYql::NDqProto::TExternalEffect& externalEffect) { - if (!FederatedQuerySetup_) { + if (!FederatedQuerySetup) { FinishScriptFinalization(Ydb::StatusIds::INTERNAL_ERROR, "unable to aplicate s3 external effect, invalid federated query setup"); return; } Register(NYql::NDq::MakeS3ApplicatorActor( SelfId(), - FederatedQuerySetup_->HttpGateway, + FederatedQuerySetup->HttpGateway, CreateGuidAsString(), - CustomerSuppliedId_, + CustomerSuppliedId, std::nullopt, - FinalizationStatus_ == EFinalizationStatus::FS_COMMIT, - THashMap(SecureParams_.begin(), SecureParams_.end()), - FederatedQuerySetup_->CredentialsFactory, + FinalizationStatus == EFinalizationStatus::FS_COMMIT, + THashMap(SecureParams.begin(), SecureParams.end()), + FederatedQuerySetup->CredentialsFactory, externalEffect ).Release()); } @@ -172,7 +196,7 @@ class TScriptFinalizerActor : public TActorBootstrapped { ) void FinishScriptFinalization(std::optional status, NYql::TIssues issues) { - Register(CreateScriptFinalizationFinisherActor(SelfId(), ExecutionId_, Database_, status, std::move(issues))); + Register(CreateScriptFinalizationFinisherActor(SelfId(), ExecutionId, Database, status, std::move(issues))); Become(&TScriptFinalizerActor::FinishState); } @@ -189,38 +213,39 @@ class TScriptFinalizerActor : public TActorBootstrapped { } void Reply(bool operationAlreadyFinalized, Ydb::StatusIds::StatusCode status, NYql::TIssues&& issues) { - Send(ReplyActor_, new TEvScriptExecutionFinished(operationAlreadyFinalized, status, std::move(issues))); - Send(MakeKqpFinalizeScriptServiceId(SelfId().NodeId()), new TEvScriptFinalizeResponse(ExecutionId_)); + Send(ReplyActor, new TEvScriptExecutionFinished(operationAlreadyFinalized, status, std::move(issues))); + Send(MakeKqpFinalizeScriptServiceId(SelfId().NodeId()), new TEvScriptFinalizeResponse(ExecutionId)); PassAway(); } private: - TActorId ReplyActor_; - TString ExecutionId_; - TString Database_; - EFinalizationStatus FinalizationStatus_; - TEvScriptFinalizeRequest::TPtr Request_; - - TDuration FinalizationTimeout_; - TDuration MaximalSecretsSnapshotWaitTime_; - const std::optional& FederatedQuerySetup_; - - TString CustomerSuppliedId_; - std::vector Sinks_; - - TString UserToken_; - std::vector SecretNames_; - std::unordered_map SecureParams_; + const TActorId ReplyActor; + const TString ExecutionId; + const TString Database; + const EFinalizationStatus FinalizationStatus; + TEvScriptFinalizeRequest::TPtr Request; + + const TDuration FinalizationTimeout; + const TDuration MaximalSecretsSnapshotWaitTime; + const std::optional& FederatedQuerySetup; + const NFq::TCompressor Compressor; + + TString CustomerSuppliedId; + std::vector Sinks; + + TString UserToken; + std::vector SecretNames; + std::unordered_map SecureParams; }; } // anonymous namespace IActor* CreateScriptFinalizerActor(TEvScriptFinalizeRequest::TPtr request, - const NKikimrConfig::TFinalizeScriptServiceConfig& finalizeScriptServiceConfig, + const NKikimrConfig::TQueryServiceConfig& queryServiceConfig, const NKikimrConfig::TMetadataProviderConfig& metadataProviderConfig, const std::optional& federatedQuerySetup) { - return new TScriptFinalizerActor(std::move(request), finalizeScriptServiceConfig, metadataProviderConfig, federatedQuerySetup); + return new TScriptFinalizerActor(std::move(request), queryServiceConfig, metadataProviderConfig, federatedQuerySetup); } } // namespace NKikimr::NKqp diff --git a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.h b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.h index b47e55395744..404ff0261ce3 100644 --- a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.h +++ b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_actor.h @@ -7,7 +7,7 @@ namespace NKikimr::NKqp { IActor* CreateScriptFinalizerActor(TEvScriptFinalizeRequest::TPtr request, - const NKikimrConfig::TFinalizeScriptServiceConfig& finalizeScriptServiceConfig, + const NKikimrConfig::TQueryServiceConfig& queryServiceConfig, const NKikimrConfig::TMetadataProviderConfig& metadataProviderConfig, const std::optional& federatedQuerySetup); diff --git a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.cpp b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.cpp index 6c0868e4c42b..605a91a5ceb1 100644 --- a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.cpp +++ b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.cpp @@ -12,16 +12,16 @@ namespace { class TKqpFinalizeScriptService : public TActorBootstrapped { public: - TKqpFinalizeScriptService(const NKikimrConfig::TFinalizeScriptServiceConfig& finalizeScriptServiceConfig, + TKqpFinalizeScriptService(const NKikimrConfig::TQueryServiceConfig& queryServiceConfig, const NKikimrConfig::TMetadataProviderConfig& metadataProviderConfig, IKqpFederatedQuerySetupFactory::TPtr federatedQuerySetupFactory) - : FinalizeScriptServiceConfig_(finalizeScriptServiceConfig) - , MetadataProviderConfig_(metadataProviderConfig) - , FederatedQuerySetupFactory_(federatedQuerySetupFactory) + : QueryServiceConfig(queryServiceConfig) + , MetadataProviderConfig(metadataProviderConfig) + , FederatedQuerySetupFactory(federatedQuerySetupFactory) {} void Bootstrap(const TActorContext &ctx) { - FederatedQuerySetup_ = FederatedQuerySetupFactory_->Make(ctx.ActorSystem()); + FederatedQuerySetup = FederatedQuerySetupFactory->Make(ctx.ActorSystem()); Become(&TKqpFinalizeScriptService::MainState); } @@ -40,10 +40,10 @@ class TKqpFinalizeScriptService : public TActorBootstrappedGet()->Description.ExecutionId; - if (!FinalizationRequestsQueue_.contains(executionId)) { - WaitingFinalizationExecutions_.push(executionId); + if (!FinalizationRequestsQueue.contains(executionId)) { + WaitingFinalizationExecutions.push(executionId); } - FinalizationRequestsQueue_[executionId].emplace_back(std::move(ev)); + FinalizationRequestsQueue[executionId].emplace_back(std::move(ev)); TryStartFinalizeRequest(); } @@ -60,14 +60,14 @@ class TKqpFinalizeScriptService : public TActorBootstrapped= FinalizeScriptServiceConfig_.GetMaxInFlightFinalizationsCount() || WaitingFinalizationExecutions_.empty()) { + if (FinalizationRequestsInFlight >= QueryServiceConfig.GetFinalizeScriptServiceConfig().GetMaxInFlightFinalizationsCount() || WaitingFinalizationExecutions.empty()) { return; } - TString executionId = WaitingFinalizationExecutions_.front(); - WaitingFinalizationExecutions_.pop(); + TString executionId = WaitingFinalizationExecutions.front(); + WaitingFinalizationExecutions.pop(); - auto& queue = FinalizationRequestsQueue_[executionId]; + auto& queue = FinalizationRequestsQueue[executionId]; Y_ENSURE(!queue.empty()); StartFinalizeRequest(std::move(queue.back())); @@ -75,24 +75,24 @@ class TKqpFinalizeScriptService : public TActorBootstrappedGet()->ExecutionId; - if (!FinalizationRequestsQueue_[executionId].empty()) { - WaitingFinalizationExecutions_.push(executionId); + if (!FinalizationRequestsQueue[executionId].empty()) { + WaitingFinalizationExecutions.push(executionId); } else { - FinalizationRequestsQueue_.erase(executionId); + FinalizationRequestsQueue.erase(executionId); } TryStartFinalizeRequest(); } @@ -122,23 +122,23 @@ class TKqpFinalizeScriptService : public TActorBootstrapped FederatedQuerySetup_; + IKqpFederatedQuerySetupFactory::TPtr FederatedQuerySetupFactory; + std::optional FederatedQuerySetup; - ui32 FinalizationRequestsInFlight_ = 0; - std::queue WaitingFinalizationExecutions_; - std::unordered_map> FinalizationRequestsQueue_; + ui32 FinalizationRequestsInFlight = 0; + std::queue WaitingFinalizationExecutions; + std::unordered_map> FinalizationRequestsQueue; }; } // anonymous namespace -IActor* CreateKqpFinalizeScriptService(const NKikimrConfig::TFinalizeScriptServiceConfig& finalizeScriptServiceConfig, +IActor* CreateKqpFinalizeScriptService(const NKikimrConfig::TQueryServiceConfig& queryServiceConfig, const NKikimrConfig::TMetadataProviderConfig& metadataProviderConfig, IKqpFederatedQuerySetupFactory::TPtr federatedQuerySetupFactory) { - return new TKqpFinalizeScriptService(finalizeScriptServiceConfig, metadataProviderConfig, std::move(federatedQuerySetupFactory)); + return new TKqpFinalizeScriptService(queryServiceConfig, metadataProviderConfig, std::move(federatedQuerySetupFactory)); } } // namespace NKikimr::NKqp diff --git a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.h b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.h index 53f0c2ea5f46..998b227f1cb1 100644 --- a/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.h +++ b/ydb/core/kqp/finalize_script_service/kqp_finalize_script_service.h @@ -5,7 +5,7 @@ namespace NKikimr::NKqp { -IActor* CreateKqpFinalizeScriptService(const NKikimrConfig::TFinalizeScriptServiceConfig& finalizeScriptServiceConfig, +IActor* CreateKqpFinalizeScriptService(const NKikimrConfig::TQueryServiceConfig& queryServiceConfig, const NKikimrConfig::TMetadataProviderConfig& metadataProviderConfig, IKqpFederatedQuerySetupFactory::TPtr federatedQuerySetupFactory); diff --git a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp index 54e815a06a49..6283c3d0c58d 100644 --- a/ydb/core/kqp/proxy_service/kqp_script_executions.cpp +++ b/ydb/core/kqp/proxy_service/kqp_script_executions.cpp @@ -1,6 +1,7 @@ #include "kqp_script_executions.h" #include "kqp_script_executions_impl.h" +#include #include #include #include @@ -135,6 +136,8 @@ class TScriptExecutionsTablesCreator : public TActorBootstrapped CurrentUtcTimestamp() OR expire_at IS NULL); @@ -1220,7 +1225,17 @@ class TGetScriptExecutionOperationQueryActor : public TQueryBase { Metadata.mutable_exec_stats()->set_query_plan(*plan); } - const TMaybe ast = result.ColumnParser("ast").GetOptionalUtf8(); + TMaybe ast; + const TMaybe astCompressionMethod = result.ColumnParser("ast_compression_method").GetOptionalUtf8(); + if (astCompressionMethod) { + const TMaybe astCompressed = result.ColumnParser("ast_compressed").GetOptionalString(); + if (astCompressed) { + const NFq::TCompressor compressor(*astCompressionMethod); + ast = compressor.Decompress(*astCompressed); + } + } else { + ast = result.ColumnParser("ast").GetOptionalUtf8(); + } if (ast) { Metadata.mutable_exec_stats()->set_query_ast(*ast); } @@ -2457,7 +2472,9 @@ class TSaveScriptFinalStatusActor : public TQueryBase { DECLARE $issues AS JsonDocument; DECLARE $plan AS JsonDocument; DECLARE $stats AS JsonDocument; - DECLARE $ast AS Text; + DECLARE $ast AS Optional; + DECLARE $ast_compressed AS Optional; + DECLARE $ast_compression_method AS Optional; DECLARE $operation_ttl AS Interval; DECLARE $customer_supplied_id AS Text; DECLARE $user_token AS Text; @@ -2475,6 +2492,8 @@ class TSaveScriptFinalStatusActor : public TQueryBase { end_ts = CurrentUtcTimestamp(), stats = $stats, ast = $ast, + ast_compressed = $ast_compressed, + ast_compression_method = $ast_compression_method, expire_at = IF($operation_ttl > CAST(0 AS Interval), CurrentUtcTimestamp() + $operation_ttl, NULL), customer_supplied_id = IF($applicate_script_external_effect_required, $customer_supplied_id, NULL), user_token = IF($applicate_script_external_effect_required, $user_token, NULL), @@ -2495,6 +2514,16 @@ class TSaveScriptFinalStatusActor : public TQueryBase { serializedStats = NJson::WriteJson(statsJson); } + TMaybe ast; + TMaybe astCompressed; + TMaybe astCompressionMethod; + if (Request.QueryAst && Request.QueryAstCompressionMethod) { + astCompressed = *Request.QueryAst; + astCompressionMethod = *Request.QueryAstCompressionMethod; + } else { + ast = Request.QueryAst.value_or(""); + } + NYdb::TParamsBuilder params; params .AddParam("$database") @@ -2522,7 +2551,13 @@ class TSaveScriptFinalStatusActor : public TQueryBase { .JsonDocument(serializedStats) .Build() .AddParam("$ast") - .Utf8(Request.QueryAst.value_or("")) + .OptionalUtf8(ast) + .Build() + .AddParam("$ast_compressed") + .OptionalString(astCompressed) + .Build() + .AddParam("$ast_compression_method") + .OptionalUtf8(astCompressionMethod) .Build() .AddParam("$operation_ttl") .Interval(static_cast(OperationTtl.MicroSeconds())) diff --git a/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp b/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp index 9b17d3cfede0..5ad21b522c13 100644 --- a/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp +++ b/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp @@ -616,13 +616,16 @@ Y_UNIT_TEST_SUITE(KqpQueryServiceScripts) { } } - auto scriptExecutionOperation = db.ExecuteScript(sql).ExtractValueSync(); + TExecuteScriptSettings settings; + settings.StatsMode(Ydb::Query::StatsMode::STATS_MODE_FULL); + + auto scriptExecutionOperation = db.ExecuteScript(sql, settings).ExtractValueSync(); UNIT_ASSERT_VALUES_EQUAL_C(scriptExecutionOperation.Status().GetStatus(), EStatus::SUCCESS, scriptExecutionOperation.Status().GetIssues().ToString()); - NYdb::NQuery::TScriptExecutionOperation readyOp = WaitScriptExecutionOperation(scriptExecutionOperation.Id(), ydbDriver); - UNIT_ASSERT_EQUAL(readyOp.Metadata().ExecStatus, EExecStatus::Completed); + auto readyOperation = WaitScriptExecutionOperation(scriptExecutionOperation.Id(), ydbDriver); + UNIT_ASSERT_EQUAL(readyOperation.Metadata().ExecStatus, EExecStatus::Completed); - return scriptExecutionOperation; + return readyOperation; } Y_UNIT_TEST(InvalidFetchToken) { @@ -768,6 +771,18 @@ Y_UNIT_TEST_SUITE(KqpQueryServiceScripts) { UNIT_ASSERT_VALUES_EQUAL_C(readyOp.Status().GetStatus(), EStatus::GENERIC_ERROR, readyOp.Status().GetIssues().ToString()); UNIT_ASSERT(HasIssue(readyOp.Status().GetIssues(), NYql::TIssuesIds::KIKIMR_BAD_OPERATION)); } + + Y_UNIT_TEST(TestAstWithCompression) { + NKikimrConfig::TAppConfig appCfg; + appCfg.MutableQueryServiceConfig()->SetQueryArtifactsCompressionMinSize(0); + appCfg.MutableQueryServiceConfig()->SetQueryArtifactsCompressionMethod("zstd_6"); + + auto kikimr = DefaultKikimrRunner({}, appCfg); + auto db = kikimr.GetQueryClient(); + auto scriptExecutionOperation = CreateScriptExecutionOperation(1, db, kikimr.GetDriver()); + + UNIT_ASSERT_STRING_CONTAINS(scriptExecutionOperation.Metadata().ExecStats.query_ast(), "\"idx\" (DataType 'Int32)"); + } } } // namespace NKqp diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto index 5022698dc5c4..b24f920a8654 100644 --- a/ydb/core/protos/config.proto +++ b/ydb/core/protos/config.proto @@ -1026,6 +1026,8 @@ message TQueryServiceConfig { optional uint64 ScriptResultSizeLimit = 4 [default = 0]; // 0 = infinity optional uint64 ScriptResultRowsLimit = 5 [default = 0]; // 0 = infinity repeated string HostnamePatterns = 13; // List of hostname regexps for external data sources; disabled if empty + optional string QueryArtifactsCompressionMethod = 17; + optional uint64 QueryArtifactsCompressionMinSize = 18 [default = 10485760]; // default = 10 MiB optional NYql.TS3GatewayConfig S3 = 6; optional NYql.THttpGatewayConfig HttpGateway = 7; diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp index 16624c6101c0..bf6edaf40473 100644 --- a/ydb/core/testlib/test_client.cpp +++ b/ydb/core/testlib/test_client.cpp @@ -890,7 +890,7 @@ namespace Tests { TActorId kqpProxyServiceId = Runtime->Register(kqpProxyService, nodeIdx); Runtime->RegisterService(NKqp::MakeKqpProxyID(Runtime->GetNodeId(nodeIdx)), kqpProxyServiceId, nodeIdx); - IActor* scriptFinalizeService = NKqp::CreateKqpFinalizeScriptService(Settings->AppConfig->GetQueryServiceConfig().GetFinalizeScriptServiceConfig(), Settings->AppConfig->GetMetadataProviderConfig(), federatedQuerySetupFactory); + IActor* scriptFinalizeService = NKqp::CreateKqpFinalizeScriptService(Settings->AppConfig->GetQueryServiceConfig(), Settings->AppConfig->GetMetadataProviderConfig(), federatedQuerySetupFactory); TActorId scriptFinalizeServiceId = Runtime->Register(scriptFinalizeService, nodeIdx); Runtime->RegisterService(NKqp::MakeKqpFinalizeScriptServiceId(Runtime->GetNodeId(nodeIdx)), scriptFinalizeServiceId, nodeIdx); } diff --git a/ydb/tests/tools/kqprun/configuration/app_config.conf b/ydb/tests/tools/kqprun/configuration/app_config.conf index 7ef43d6a93cb..ab93fea0af0b 100644 --- a/ydb/tests/tools/kqprun/configuration/app_config.conf +++ b/ydb/tests/tools/kqprun/configuration/app_config.conf @@ -24,6 +24,7 @@ LogConfig { QueryServiceConfig { MdbTransformHost: false + QueryArtifactsCompressionMethod: "zstd_6" ScriptResultRowsLimit: 0 ScriptResultSizeLimit: 10485760 diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index 4bf0971935d8..21c31f74ade2 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -153,6 +153,10 @@ class TKqpRunner::TImpl { return false; } + if (!status.Issues.Empty()) { + Cerr << CerrColors_.Red() << "Request finished with issues:" << CerrColors_.Default() << Endl << status.Issues.ToString() << Endl; + } + PrintScriptPlan(ExecutionMeta_.Plan); return true;