From 65176c723756fe2df6598ed547b57eac27b3eb22 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Mon, 13 Jan 2025 13:52:33 +0500 Subject: [PATCH 01/18] KqpRun improved multi query mode and storage settings (#13177) --- ydb/core/driver_lib/run/config_helpers.cpp | 31 ++ ydb/core/driver_lib/run/config_helpers.h | 7 + .../run/kikimr_services_initializers.cpp | 22 +- ydb/core/testlib/basics/helpers.h | 2 + ydb/core/testlib/basics/storage.h | 4 +- ydb/core/testlib/test_client.cpp | 32 +- ydb/core/testlib/test_client.h | 2 + ydb/tests/tools/kqprun/.gitignore | 2 + ydb/tests/tools/kqprun/flame_graph.sh | 13 +- ydb/tests/tools/kqprun/kqprun.cpp | 313 ++++++++++++++---- ydb/tests/tools/kqprun/src/actors.cpp | 13 +- ydb/tests/tools/kqprun/src/actors.h | 4 +- ydb/tests/tools/kqprun/src/common.cpp | 29 ++ ydb/tests/tools/kqprun/src/common.h | 35 +- ydb/tests/tools/kqprun/src/kqp_runner.cpp | 77 +++-- .../tools/kqprun/src/proto/storage_meta.proto | 7 + ydb/tests/tools/kqprun/src/proto/ya.make | 9 + ydb/tests/tools/kqprun/src/ya.make | 5 + ydb/tests/tools/kqprun/src/ydb_setup.cpp | 83 +++-- ydb/tests/tools/kqprun/ya.make | 6 + 20 files changed, 519 insertions(+), 177 deletions(-) create mode 100644 ydb/tests/tools/kqprun/src/common.cpp create mode 100644 ydb/tests/tools/kqprun/src/proto/storage_meta.proto create mode 100644 ydb/tests/tools/kqprun/src/proto/ya.make diff --git a/ydb/core/driver_lib/run/config_helpers.cpp b/ydb/core/driver_lib/run/config_helpers.cpp index 2d0b834eb713..66a37eee15f1 100644 --- a/ydb/core/driver_lib/run/config_helpers.cpp +++ b/ydb/core/driver_lib/run/config_helpers.cpp @@ -1,5 +1,9 @@ #include "config_helpers.h" +#include +#include +#include + #include @@ -114,4 +118,31 @@ NActors::TSchedulerConfig CreateSchedulerConfig(const NKikimrConfig::TActorSyste } // namespace NActorSystemConfigHelpers +namespace NKikimrConfigHelpers { + +NMemory::TResourceBrokerConfig CreateMemoryControllerResourceBrokerConfig(const NKikimrConfig::TAppConfig& config) { + NMemory::TResourceBrokerConfig resourceBrokerSelfConfig; // for backward compatibility + auto mergeResourceBrokerConfigs = [&](const NKikimrResourceBroker::TResourceBrokerConfig& resourceBrokerConfig) { + if (resourceBrokerConfig.HasResourceLimit() && resourceBrokerConfig.GetResourceLimit().HasMemory()) { + resourceBrokerSelfConfig.LimitBytes = resourceBrokerConfig.GetResourceLimit().GetMemory(); + } + for (const auto& queue : resourceBrokerConfig.GetQueues()) { + if (queue.GetName() == NLocalDb::KqpResourceManagerQueue) { + if (queue.HasLimit() && queue.GetLimit().HasMemory()) { + resourceBrokerSelfConfig.QueryExecutionLimitBytes = queue.GetLimit().GetMemory(); + } + } + } + }; + if (config.HasBootstrapConfig() && config.GetBootstrapConfig().HasResourceBroker()) { + mergeResourceBrokerConfigs(config.GetBootstrapConfig().GetResourceBroker()); + } + if (config.HasResourceBrokerConfig()) { + mergeResourceBrokerConfigs(config.GetResourceBrokerConfig()); + } + return resourceBrokerSelfConfig; +} + +} // namespace NKikimrConfigHelpers + } // namespace NKikimr diff --git a/ydb/core/driver_lib/run/config_helpers.h b/ydb/core/driver_lib/run/config_helpers.h index d38e416b6f5b..39a0c6cdb053 100644 --- a/ydb/core/driver_lib/run/config_helpers.h +++ b/ydb/core/driver_lib/run/config_helpers.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -15,4 +16,10 @@ NActors::TSchedulerConfig CreateSchedulerConfig(const NKikimrConfig::TActorSyste } // namespace NActorSystemConfigHelpers +namespace NKikimrConfigHelpers { + +NMemory::TResourceBrokerConfig CreateMemoryControllerResourceBrokerConfig(const NKikimrConfig::TAppConfig& config); + +} // namespace NKikimrConfigHelpers + } // namespace NKikimr diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp index 3f073ab52be3..e3d4ef977249 100644 --- a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp +++ b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp @@ -2049,28 +2049,8 @@ void TMemoryControllerInitializer::InitializeServices( NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) { - NMemory::TResourceBrokerConfig resourceBrokerSelfConfig; // for backward compatibility - auto mergeResourceBrokerConfigs = [&](const NKikimrResourceBroker::TResourceBrokerConfig& resourceBrokerConfig) { - if (resourceBrokerConfig.HasResourceLimit() && resourceBrokerConfig.GetResourceLimit().HasMemory()) { - resourceBrokerSelfConfig.LimitBytes = resourceBrokerConfig.GetResourceLimit().GetMemory(); - } - for (const auto& queue : resourceBrokerConfig.GetQueues()) { - if (queue.GetName() == NLocalDb::KqpResourceManagerQueue) { - if (queue.HasLimit() && queue.GetLimit().HasMemory()) { - resourceBrokerSelfConfig.QueryExecutionLimitBytes = queue.GetLimit().GetMemory(); - } - } - } - }; - if (Config.HasBootstrapConfig() && Config.GetBootstrapConfig().HasResourceBroker()) { - mergeResourceBrokerConfigs(Config.GetBootstrapConfig().GetResourceBroker()); - } - if (Config.HasResourceBrokerConfig()) { - mergeResourceBrokerConfigs(Config.GetResourceBrokerConfig()); - } - auto* actor = NMemory::CreateMemoryController(TDuration::Seconds(1), ProcessMemoryInfoProvider, - Config.GetMemoryControllerConfig(), resourceBrokerSelfConfig, + Config.GetMemoryControllerConfig(), NKikimrConfigHelpers::CreateMemoryControllerResourceBrokerConfig(Config), appData->Counters); setup->LocalServices.emplace_back( NMemory::MakeMemoryControllerId(0), diff --git a/ydb/core/testlib/basics/helpers.h b/ydb/core/testlib/basics/helpers.h index e6be7bca98ef..24126aca2cc4 100644 --- a/ydb/core/testlib/basics/helpers.h +++ b/ydb/core/testlib/basics/helpers.h @@ -18,6 +18,8 @@ namespace NFake { ui64 SectorSize = 0; ui64 ChunkSize = 0; ui64 DiskSize = 0; + bool FormatDisk = true; + TString DiskPath; }; struct INode { diff --git a/ydb/core/testlib/basics/storage.h b/ydb/core/testlib/basics/storage.h index 65b58a076fb1..e1aac88071d2 100644 --- a/ydb/core/testlib/basics/storage.h +++ b/ydb/core/testlib/basics/storage.h @@ -54,7 +54,7 @@ namespace NKikimr { static ui64 keySalt = 0; ui64 salt = ++keySalt; - TString baseDir = Runtime.GetTempDir(); + TString baseDir = conf.DiskPath ? conf.DiskPath : Runtime.GetTempDir(); if (Conf.UseDisk) { MakeDirIfNotExist(baseDir.c_str()); @@ -62,7 +62,7 @@ namespace NKikimr { PDiskPath = TStringBuilder() << baseDir << "pdisk_1.dat"; - if (!Mock) { + if (!Mock && conf.FormatDisk) { FormatPDisk(PDiskPath, Conf.DiskSize, Conf.SectorSize, Conf.ChunkSize, PDiskGuid, 0x123 + salt, 0x456 + salt, 0x789 + salt, mainKey, diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp index e9113ad903d5..1c6a8d91908e 100644 --- a/ydb/core/testlib/test_client.cpp +++ b/ydb/core/testlib/test_client.cpp @@ -93,6 +93,7 @@ #include #include #include +#include #include #include #include @@ -610,6 +611,7 @@ namespace Tests { NKikimrBlobStorage::TDefineBox boxConfig; boxConfig.SetBoxId(Settings->BOX_ID); + boxConfig.SetItemConfigGeneration(Settings->StorageGeneration); ui32 nodeId = Runtime->GetNodeId(0); Y_ABORT_UNLESS(nodesInfo->Nodes[0].NodeId == nodeId); @@ -617,11 +619,13 @@ namespace Tests { NKikimrBlobStorage::TDefineHostConfig hostConfig; hostConfig.SetHostConfigId(nodeId); + hostConfig.SetItemConfigGeneration(Settings->StorageGeneration); TString path; if (Settings->UseSectorMap) { path ="SectorMap:test-client[:2000]"; } else { - path = TStringBuilder() << Runtime->GetTempDir() << "pdisk_1.dat"; + TString diskPath = Settings->CustomDiskParams.DiskPath; + path = TStringBuilder() << (diskPath ? diskPath : Runtime->GetTempDir()) << "pdisk_1.dat"; } hostConfig.AddDrive()->SetPath(path); if (Settings->Verbose) { @@ -637,7 +641,9 @@ namespace Tests { for (const auto& [poolKind, storagePool] : Settings->StoragePoolTypes) { if (storagePool.GetNumGroups() > 0) { - bsConfigureRequest->Record.MutableRequest()->AddCommand()->MutableDefineStoragePool()->CopyFrom(storagePool); + auto* command = bsConfigureRequest->Record.MutableRequest()->AddCommand()->MutableDefineStoragePool(); + command->CopyFrom(storagePool); + command->SetItemConfigGeneration(Settings->StorageGeneration); } } @@ -1071,6 +1077,28 @@ namespace Tests { } } + { + if (Settings->NeedStatsCollectors) { + TString filePathPrefix; + if (Settings->AppConfig->HasMonitoringConfig()) { + filePathPrefix = Settings->AppConfig->GetMonitoringConfig().GetMemAllocDumpPathPrefix(); + } + + const TIntrusivePtr processMemoryInfoProvider(MakeIntrusive()); + + IActor* monitorActor = CreateMemProfMonitor(TDuration::Seconds(1), processMemoryInfoProvider, + Runtime->GetAppData(nodeIdx).Counters, filePathPrefix); + const TActorId monitorActorId = Runtime->Register(monitorActor, nodeIdx, Runtime->GetAppData(nodeIdx).BatchPoolId); + Runtime->RegisterService(MakeMemProfMonitorID(Runtime->GetNodeId(nodeIdx)), monitorActorId, nodeIdx); + + IActor* controllerActor = NMemory::CreateMemoryController(TDuration::Seconds(1), processMemoryInfoProvider, + Settings->AppConfig->GetMemoryControllerConfig(), NKikimrConfigHelpers::CreateMemoryControllerResourceBrokerConfig(*Settings->AppConfig), + Runtime->GetAppData(nodeIdx).Counters); + const TActorId controllerActorId = Runtime->Register(controllerActor, nodeIdx, Runtime->GetAppData(nodeIdx).BatchPoolId); + Runtime->RegisterService(NMemory::MakeMemoryControllerId(0), controllerActorId, nodeIdx); + } + } + { IActor* kesusService = NKesus::CreateKesusProxyService(); TActorId kesusServiceId = Runtime->Register(kesusService, nodeIdx, userPoolId); diff --git a/ydb/core/testlib/test_client.h b/ydb/core/testlib/test_client.h index ce556d86c323..5b82a0d5e44b 100644 --- a/ydb/core/testlib/test_client.h +++ b/ydb/core/testlib/test_client.h @@ -122,6 +122,7 @@ namespace Tests { TString DomainName = TestDomainName; ui32 NodeCount = 1; ui32 DynamicNodeCount = 0; + ui64 StorageGeneration = 0; NFake::TStorage CustomDiskParams; TControls Controls; TAppPrepare::TFnReg FrFactory = &DefaultFrFactory; @@ -177,6 +178,7 @@ namespace Tests { TServerSettings& SetDomainName(const TString& value); TServerSettings& SetNodeCount(ui32 value) { NodeCount = value; return *this; } TServerSettings& SetDynamicNodeCount(ui32 value) { DynamicNodeCount = value; return *this; } + TServerSettings& SetStorageGeneration(ui64 value) { StorageGeneration = value; return *this; } TServerSettings& SetCustomDiskParams(const NFake::TStorage& value) { CustomDiskParams = value; return *this; } TServerSettings& SetControls(const TControls& value) { Controls = value; return *this; } TServerSettings& SetFrFactory(const TAppPrepare::TFnReg& value) { FrFactory = value; return *this; } diff --git a/ydb/tests/tools/kqprun/.gitignore b/ydb/tests/tools/kqprun/.gitignore index a4578942a676..86fc65c21b31 100644 --- a/ydb/tests/tools/kqprun/.gitignore +++ b/ydb/tests/tools/kqprun/.gitignore @@ -1,6 +1,8 @@ +storage sync_dir example udfs + *.log *.json *.sql diff --git a/ydb/tests/tools/kqprun/flame_graph.sh b/ydb/tests/tools/kqprun/flame_graph.sh index ead8de30683a..61749c81b21e 100755 --- a/ydb/tests/tools/kqprun/flame_graph.sh +++ b/ydb/tests/tools/kqprun/flame_graph.sh @@ -1,11 +1,12 @@ #!/usr/bin/env bash -# For svg graph download https://github.com/brendangregg/FlameGraph -# and run `FlameGraph/stackcollapse-perf.pl profdata.txt | FlameGraph/flamegraph.pl > profdata.svg` +set -eux -pid=$(pgrep -u $USER kqprun) +kqprun_pid=$(pgrep -u $USER kqprun) -echo "Target process id: ${pid}" - -sudo perf record -F 50 --call-graph dwarf -g --proc-map-timeout=10000 --pid $pid -v -o profdata -- sleep 30 +sudo perf record -F 50 --call-graph dwarf -g --proc-map-timeout=10000 --pid $kqprun_pid -v -o profdata -- sleep ${1:-'30'} sudo perf script -i profdata > profdata.txt + +flame_graph_tool="../../../../contrib/tools/flame-graph/" + +${flame_graph_tool}/stackcollapse-perf.pl profdata.txt | ${flame_graph_tool}/flamegraph.pl > profdata.svg diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index 69ddc226f7f9..8c8807899ef4 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -1,7 +1,5 @@ #include "src/kqp_runner.h" -#include - #include #include @@ -12,16 +10,27 @@ #include #include +#include #include #include + #include +#include + #include #include #include -#include + +#ifdef PROFILE_MEMORY_ALLOCATIONS +#include +#endif +namespace NKqpRun { + +namespace { + struct TExecutionOptions { enum class EExecutionCase { GenericScript, @@ -77,7 +86,7 @@ struct TExecutionOptions { return GetValue(index, ScriptQueryActions, NKikimrKqp::EQueryAction::QUERY_ACTION_EXECUTE); } - NKqpRun::TRequestOptions GetSchemeQueryOptions() const { + TRequestOptions GetSchemeQueryOptions() const { TString sql = SchemeQuery; if (UseTemplates) { ReplaceYqlTokenTemplate(sql); @@ -94,7 +103,7 @@ struct TExecutionOptions { }; } - NKqpRun::TRequestOptions GetScriptQueryOptions(size_t index, size_t queryId, TInstant startTime) const { + TRequestOptions GetScriptQueryOptions(size_t index, size_t queryId, TInstant startTime) const { Y_ABORT_UNLESS(index < ScriptQueries.size()); TString sql = ScriptQueries[index]; @@ -110,27 +119,29 @@ struct TExecutionOptions { .PoolId = GetValue(index, PoolIds, TString()), .UserSID = GetValue(index, UserSIDs, TString(BUILTIN_ACL_ROOT)), .Database = GetValue(index, Databases, TString()), - .Timeout = GetValue(index, Timeouts, TDuration::Zero()) + .Timeout = GetValue(index, Timeouts, TDuration::Zero()), + .QueryId = queryId }; } - void Validate(const NKqpRun::TRunnerOptions& runnerOptions) const { + void Validate(const TRunnerOptions& runnerOptions) const { if (!SchemeQuery && ScriptQueries.empty() && !runnerOptions.YdbSettings.MonitoringEnabled && !runnerOptions.YdbSettings.GrpcEnabled) { ythrow yexception() << "Nothing to execute and is not running as daemon"; } - ValidateOptionsSizes(); + ValidateOptionsSizes(runnerOptions); ValidateSchemeQueryOptions(runnerOptions); ValidateScriptExecutionOptions(runnerOptions); ValidateAsyncOptions(runnerOptions.YdbSettings.AsyncQueriesSettings); - ValidateTraceOpt(runnerOptions.TraceOptType); + ValidateTraceOpt(runnerOptions); + ValidateStorageSettings(runnerOptions.YdbSettings); } private: - void ValidateOptionsSizes() const { + void ValidateOptionsSizes(const TRunnerOptions& runnerOptions) const { const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName) { if (checkSize > numberQueries) { - ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of queries is " << numberQueries; + ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of script queries is " << numberQueries; } }; @@ -141,9 +152,13 @@ struct TExecutionOptions { checker(PoolIds.size(), "pool ids"); checker(UserSIDs.size(), "user SIDs"); checker(Timeouts.size(), "timeouts"); + checker(runnerOptions.ScriptQueryAstOutputs.size(), "ast output files"); + checker(runnerOptions.ScriptQueryPlanOutputs.size(), "plan output files"); + checker(runnerOptions.ScriptQueryTimelineFiles.size(), "timeline files"); + checker(runnerOptions.InProgressStatisticsOutputFiles.size(), "statistics files"); } - void ValidateSchemeQueryOptions(const NKqpRun::TRunnerOptions& runnerOptions) const { + void ValidateSchemeQueryOptions(const TRunnerOptions& runnerOptions) const { if (SchemeQuery) { return; } @@ -152,7 +167,7 @@ struct TExecutionOptions { } } - void ValidateScriptExecutionOptions(const NKqpRun::TRunnerOptions& runnerOptions) const { + void ValidateScriptExecutionOptions(const TRunnerOptions& runnerOptions) const { if (runnerOptions.YdbSettings.SameSession && HasExecutionCase(EExecutionCase::AsyncQuery)) { ythrow yexception() << "Same session can not be used with async quries"; } @@ -175,7 +190,7 @@ struct TExecutionOptions { if (ResultsRowsLimit) { ythrow yexception() << "Result rows limit can not be used without script queries"; } - if (runnerOptions.InProgressStatisticsOutputFile) { + if (!runnerOptions.InProgressStatisticsOutputFiles.empty()) { ythrow yexception() << "Script statistics can not be used without script queries"; } @@ -183,10 +198,10 @@ struct TExecutionOptions { if (HasExecutionCase(EExecutionCase::YqlScript)) { return; } - if (runnerOptions.ScriptQueryAstOutput) { + if (!runnerOptions.ScriptQueryAstOutputs.empty()) { ythrow yexception() << "Script query AST output can not be used without script/yql queries"; } - if (runnerOptions.ScriptQueryPlanOutput) { + if (!runnerOptions.ScriptQueryPlanOutputs.empty()) { ythrow yexception() << "Script query plan output can not be used without script/yql queries"; } if (runnerOptions.YdbSettings.SameSession) { @@ -194,7 +209,7 @@ struct TExecutionOptions { } } - void ValidateAsyncOptions(const NKqpRun::TAsyncQueriesSettings& asyncQueriesSettings) const { + void ValidateAsyncOptions(const TAsyncQueriesSettings& asyncQueriesSettings) const { if (asyncQueriesSettings.InFlightLimit && !HasExecutionCase(EExecutionCase::AsyncQuery)) { ythrow yexception() << "In flight limit can not be used without async queries"; } @@ -205,42 +220,62 @@ struct TExecutionOptions { } } - void ValidateTraceOpt(NKqpRun::TRunnerOptions::ETraceOptType traceOptType) const { - switch (traceOptType) { - case NKqpRun::TRunnerOptions::ETraceOptType::Scheme: { + void ValidateTraceOpt(const TRunnerOptions& runnerOptions) const { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + switch (runnerOptions.TraceOptType) { + case TRunnerOptions::ETraceOptType::Scheme: { if (!SchemeQuery) { ythrow yexception() << "Trace opt type scheme cannot be used without scheme query"; } break; } - case NKqpRun::TRunnerOptions::ETraceOptType::Script: { + case TRunnerOptions::ETraceOptType::Script: { if (ScriptQueries.empty()) { ythrow yexception() << "Trace opt type script cannot be used without script queries"; } } - case NKqpRun::TRunnerOptions::ETraceOptType::All: { + case TRunnerOptions::ETraceOptType::All: { if (!SchemeQuery && ScriptQueries.empty()) { ythrow yexception() << "Trace opt type all cannot be used without any queries"; } } - case NKqpRun::TRunnerOptions::ETraceOptType::Disabled: { + case TRunnerOptions::ETraceOptType::Disabled: { break; } } + + if (const auto traceOptId = runnerOptions.TraceOptScriptId) { + if (runnerOptions.TraceOptType != TRunnerOptions::ETraceOptType::Script) { + ythrow yexception() << "Trace opt id allowed only for trace opt type script (used " << runnerOptions.TraceOptType << ")"; + } + + const ui64 scriptNumber = ScriptQueries.size() * LoopCount; + if (*traceOptId >= scriptNumber) { + ythrow yexception() << "Invalid trace opt id " << *traceOptId << ", it should be less than number of script queries " << scriptNumber; + } + if (scriptNumber == 1) { + Cout << colors.Red() << "Warning: trace opt id is not necessary for single script mode" << Endl; + } + } } -private: - template - static TValue GetValue(size_t index, const std::vector& values, TValue defaultValue) { - if (values.empty()) { - return defaultValue; + static void ValidateStorageSettings(const TYdbSetupSettings& ydbSettings) { + if (ydbSettings.DisableDiskMock) { + if (ydbSettings.NodeCount + ydbSettings.SharedTenants.size() + ydbSettings.DedicatedTenants.size() > 1) { + ythrow yexception() << "Disable disk mock cannot be used for multi node clusters (already disabled)"; + } else if (ydbSettings.PDisksPath) { + ythrow yexception() << "Disable disk mock cannot be used with real PDisks (already disabled)"; + } + } + if (ydbSettings.FormatStorage && !ydbSettings.PDisksPath) { + ythrow yexception() << "Cannot format storage without real PDisks, please use --storage-path"; } - return values[std::min(index, values.size() - 1)]; } +private: static void ReplaceYqlTokenTemplate(TString& sql) { - const TString variableName = TStringBuilder() << "${" << NKqpRun::YQL_TOKEN_VARIABLE << "}"; - if (const TString& yqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE)) { + const TString variableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}"; + if (const TString& yqlToken = GetEnv(YQL_TOKEN_VARIABLE)) { SubstGlobal(sql, variableName, yqlToken); } else if (sql.Contains(variableName)) { ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable\n"; @@ -249,7 +284,7 @@ struct TExecutionOptions { }; -void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TExecutionOptions& executionOptions, NKqpRun::TKqpRunner& runner) { +void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TExecutionOptions& executionOptions, TKqpRunner& runner) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); switch (executionOptions.GetExecutionCase(index)) { @@ -292,7 +327,7 @@ void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TE } -void RunArgumentQueries(const TExecutionOptions& executionOptions, NKqpRun::TKqpRunner& runner) { +void RunArgumentQueries(const TExecutionOptions& executionOptions, TKqpRunner& runner) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); if (executionOptions.SchemeQuery) { @@ -354,11 +389,11 @@ void RunAsDaemon() { } -void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunnerOptions& runnerOptions) { +void RunScript(const TExecutionOptions& executionOptions, const TRunnerOptions& runnerOptions) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization of kqp runner..." << colors.Default() << Endl; - NKqpRun::TKqpRunner runner(runnerOptions); + TKqpRunner runner(runnerOptions); try { RunArgumentQueries(executionOptions, runner); @@ -403,11 +438,13 @@ TIntrusivePtr CreateFunctionRegistr class TMain : public TMainClassArgs { - inline static const TString YqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE); + inline static const TString YqlToken = GetEnv(YQL_TOKEN_VARIABLE); inline static std::vector> FileHolders; + inline static IOutputStream* ProfileAllocationsOutput = nullptr; + inline static NColorizer::TColors CoutColors = NColorizer::AutoColors(Cout); TExecutionOptions ExecutionOptions; - NKqpRun::TRunnerOptions RunnerOptions; + TRunnerOptions RunnerOptions; THashMap TablesMapping; TVector UdfsPaths; @@ -415,6 +452,9 @@ class TMain : public TMainClassArgs { bool ExcludeLinkedUdfs = false; bool EmulateYt = false; + std::optional DefaultLogPriority; + std::unordered_map LogPriorities; + static TString LoadFile(const TString& file) { return TFileInput(file).ReadAll(); } @@ -450,10 +490,29 @@ class TMain : public TMainClassArgs { return choices; } + bool Contains(const TString& choice) const { + return ChoicesMap.contains(choice); + } + private: const std::map ChoicesMap; }; +#ifdef PROFILE_MEMORY_ALLOCATIONS +public: + static void FinishProfileMemoryAllocations() { + if (ProfileAllocationsOutput) { + NAllocProfiler::StopAllocationSampling(*ProfileAllocationsOutput); + } else { + TString output; + TStringOutput stream(output); + NAllocProfiler::StopAllocationSampling(stream); + + Cout << CoutColors.Red() << "Warning: profile memory allocations output is not specified, please use flag `--profile-output` for writing profile info (dump size " << NKikimr::NBlobDepot::FormatByteSize(output.size()) << ")" << CoutColors.Default() << Endl; + } + } +#endif + protected: void RegisterOptions(NLastGetopt::TOpts& options) override { options.SetTitle("KqpRun -- tool to execute queries by using kikimr provider (instead of dq provider in DQrun tool)"); @@ -467,11 +526,13 @@ class TMain : public TMainClassArgs { .Handler1([this](const NLastGetopt::TOptsParser* option) { ExecutionOptions.SchemeQuery = LoadFile(option->CurVal()); }); + options.AddLongOption('p', "script-query", "Script query to execute (typically DML query)") .RequiredArgument("file") .Handler1([this](const NLastGetopt::TOptsParser* option) { ExecutionOptions.ScriptQueries.emplace_back(LoadFile(option->CurVal())); }); + options.AddLongOption("templates", "Enable templates for -s and -p queries, such as ${YQL_TOKEN} and ${QUERY_ID}") .NoArgument() .SetFlag(&ExecutionOptions.UseTemplates); @@ -483,12 +544,11 @@ class TMain : public TMainClassArgs { TStringBuf filePath; TStringBuf(option->CurVal()).Split('@', tableName, filePath); if (tableName.empty() || filePath.empty()) { - ythrow yexception() << "Incorrect table mapping, expected form table@file, e.g. yt.Root/plato.Input@input.txt"; + ythrow yexception() << "Incorrect table mapping, expected form table@file, e. g. yt.Root/plato.Input@input.txt"; } - if (TablesMapping.contains(tableName)) { - ythrow yexception() << "Got duplicate table name: " << tableName; + if (!TablesMapping.emplace(tableName, filePath).second) { + ythrow yexception() << "Got duplicated table name: " << tableName; } - TablesMapping[tableName] = filePath; }); options.AddLongOption('c', "app-config", "File with app config (TAppConfig for ydb tenant)") @@ -507,9 +567,11 @@ class TMain : public TMainClassArgs { options.AddLongOption('u', "udf", "Load shared library with UDF by given path") .RequiredArgument("file") .EmplaceTo(&UdfsPaths); + options.AddLongOption("udfs-dir", "Load all shared libraries with UDFs found in given directory") .RequiredArgument("directory") .StoreResult(&UdfsDirectory); + options.AddLongOption("exclude-linked-udfs", "Exclude linked udfs when same udf passed from -u or --udfs-dir") .NoArgument() .SetFlag(&ExcludeLinkedUdfs); @@ -524,13 +586,50 @@ class TMain : public TMainClassArgs { std::remove(file.c_str()); } }); - TChoices traceOpt({ - {"all", NKqpRun::TRunnerOptions::ETraceOptType::All}, - {"scheme", NKqpRun::TRunnerOptions::ETraceOptType::Scheme}, - {"script", NKqpRun::TRunnerOptions::ETraceOptType::Script}, - {"disabled", NKqpRun::TRunnerOptions::ETraceOptType::Disabled} + + TChoices logPriority({ + {"emerg", NActors::NLog::EPriority::PRI_EMERG}, + {"alert", NActors::NLog::EPriority::PRI_ALERT}, + {"crit", NActors::NLog::EPriority::PRI_CRIT}, + {"error", NActors::NLog::EPriority::PRI_ERROR}, + {"warn", NActors::NLog::EPriority::PRI_WARN}, + {"notice", NActors::NLog::EPriority::PRI_NOTICE}, + {"info", NActors::NLog::EPriority::PRI_INFO}, + {"debug", NActors::NLog::EPriority::PRI_DEBUG}, + {"trace", NActors::NLog::EPriority::PRI_TRACE}, + }); + options.AddLongOption("log-default", "Default log priority") + .RequiredArgument("priority") + .Choices(logPriority.GetChoices()) + .StoreMappedResultT(&DefaultLogPriority, logPriority); + + options.AddLongOption("log", "Component log priority in format = (e. g. KQP_YQL=trace)") + .RequiredArgument("component priority") + .Handler1([this, logPriority](const NLastGetopt::TOptsParser* option) { + TStringBuf component; + TStringBuf priority; + TStringBuf(option->CurVal()).Split('=', component, priority); + if (component.empty() || priority.empty()) { + ythrow yexception() << "Incorrect log setting, expected form component=priority, e. g. KQP_YQL=trace"; + } + + if (!logPriority.Contains(TString(priority))) { + ythrow yexception() << "Incorrect log priority: " << priority; + } + + const auto service = GetLogService(TString(component)); + if (!LogPriorities.emplace(service, logPriority(TString(priority))).second) { + ythrow yexception() << "Got duplicated log service name: " << component; + } + }); + + TChoices traceOpt({ + {"all", TRunnerOptions::ETraceOptType::All}, + {"scheme", TRunnerOptions::ETraceOptType::Scheme}, + {"script", TRunnerOptions::ETraceOptType::Script}, + {"disabled", TRunnerOptions::ETraceOptType::Disabled} }); - options.AddLongOption('T', "trace-opt", "print AST in the begin of each transformation") + options.AddLongOption('T', "trace-opt", "Print AST in the begin of each transformation") .RequiredArgument("trace-opt-query") .DefaultValue("disabled") .Choices(traceOpt.GetChoices()) @@ -539,6 +638,11 @@ class TMain : public TMainClassArgs { RunnerOptions.YdbSettings.TraceOptEnabled = traceOptType != NKqpRun::TRunnerOptions::ETraceOptType::Disabled; return traceOptType; }); + + options.AddLongOption('I', "trace-opt-index", "Index of -p query to use --trace-opt, starts from zero") + .RequiredArgument("uint") + .StoreResult(&RunnerOptions.TraceOptScriptId); + options.AddLongOption("trace-id", "Trace id for -p queries") .RequiredArgument("id") .EmplaceTo(&ExecutionOptions.TraceIds); @@ -547,14 +651,16 @@ class TMain : public TMainClassArgs { .RequiredArgument("file") .DefaultValue("-") .StoreMappedResultT(&RunnerOptions.ResultOutput, &GetDefaultOutput); + options.AddLongOption('L', "result-rows-limit", "Rows limit for script execution results") .RequiredArgument("uint") .DefaultValue(0) .StoreResult(&ExecutionOptions.ResultsRowsLimit); - TChoices resultFormat({ - {"rows", NKqpRun::TRunnerOptions::EResultOutputFormat::RowsJson}, - {"full-json", NKqpRun::TRunnerOptions::EResultOutputFormat::FullJson}, - {"full-proto", NKqpRun::TRunnerOptions::EResultOutputFormat::FullProto} + + TChoices resultFormat({ + {"rows", TRunnerOptions::EResultOutputFormat::RowsJson}, + {"full-json", TRunnerOptions::EResultOutputFormat::FullJson}, + {"full-proto", TRunnerOptions::EResultOutputFormat::FullProto} }); options.AddLongOption('R', "result-format", "Script query result format") .RequiredArgument("result-format") @@ -568,19 +674,26 @@ class TMain : public TMainClassArgs { options.AddLongOption("script-ast-file", "File with script query ast (use '-' to write in stdout)") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.ScriptQueryAstOutput, &GetDefaultOutput); + .Handler1([this](const NLastGetopt::TOptsParser* option) { + RunnerOptions.ScriptQueryAstOutputs.emplace_back(GetDefaultOutput(TString(option->CurValOrDef()))); + }); options.AddLongOption("script-plan-file", "File with script query plan (use '-' to write in stdout)") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.ScriptQueryPlanOutput, &GetDefaultOutput); + .Handler1([this](const NLastGetopt::TOptsParser* option) { + RunnerOptions.ScriptQueryPlanOutputs.emplace_back(GetDefaultOutput(TString(option->CurValOrDef()))); + }); + options.AddLongOption("script-statistics", "File with script inprogress statistics") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.InProgressStatisticsOutputFile, [](const TString& file) { + .Handler1([this](const NLastGetopt::TOptsParser* option) { + const TString file(option->CurValOrDef()); if (file == "-") { ythrow yexception() << "Script in progress statistics cannot be printed to stdout, please specify file name"; } - return file; + RunnerOptions.InProgressStatisticsOutputFiles.emplace_back(file); }); + TChoices planFormat({ {"pretty", NYdb::NConsoleClient::EDataFormat::Pretty}, {"table", NYdb::NConsoleClient::EDataFormat::PrettyTable}, @@ -594,13 +707,18 @@ class TMain : public TMainClassArgs { options.AddLongOption("script-timeline-file", "File with script query timline in svg format") .RequiredArgument("file") - .StoreMappedResultT(&RunnerOptions.ScriptQueryTimelineFile, [](const TString& file) { + .Handler1([this](const NLastGetopt::TOptsParser* option) { + const TString file(option->CurValOrDef()); if (file == "-") { ythrow yexception() << "Script timline cannot be printed to stdout, please specify file name"; } - return file; + RunnerOptions.ScriptQueryTimelineFiles.emplace_back(file); }); + options.AddLongOption("profile-output", "File with profile memory allocations output (use '-' to write in stdout)") + .RequiredArgument("file") + .StoreMappedResultT(&ProfileAllocationsOutput, &GetDefaultOutput); + // Pipeline settings TChoices executionCase({ @@ -616,13 +734,20 @@ class TMain : public TMainClassArgs { TString choice(option->CurValOrDef()); ExecutionOptions.ExecutionCases.emplace_back(executionCase(choice)); }); + options.AddLongOption("inflight-limit", "In flight limit for async queries (use 0 for unlimited)") .RequiredArgument("uint") .DefaultValue(0) .StoreResult(&RunnerOptions.YdbSettings.AsyncQueriesSettings.InFlightLimit); - TChoices verbose({ - {"each-query", NKqpRun::TAsyncQueriesSettings::EVerbose::EachQuery}, - {"final", NKqpRun::TAsyncQueriesSettings::EVerbose::Final} + + options.AddLongOption("verbose", "Common verbose level (max level 2)") + .RequiredArgument("uint") + .DefaultValue(1) + .StoreResult(&RunnerOptions.YdbSettings.VerboseLevel); + + TChoices verbose({ + {"each-query", TAsyncQueriesSettings::EVerbose::EachQuery}, + {"final", TAsyncQueriesSettings::EVerbose::Final} }); options.AddLongOption("async-verbose", "Verbose type for async queries") .RequiredArgument("type") @@ -660,10 +785,12 @@ class TMain : public TMainClassArgs { .RequiredArgument("uint") .DefaultValue(ExecutionOptions.LoopCount) .StoreResult(&ExecutionOptions.LoopCount); + options.AddLongOption("loop-delay", "Delay in milliseconds between loop steps") .RequiredArgument("uint") .DefaultValue(0) .StoreMappedResultT(&ExecutionOptions.LoopDelay, &TDuration::MilliSeconds); + options.AddLongOption("continue-after-fail", "Don't not stop requests execution after fails") .NoArgument() .SetFlag(&ExecutionOptions.ContinueAfterFail); @@ -742,9 +869,13 @@ class TMain : public TMainClassArgs { return static_cast(diskSize) << 30; }); - options.AddLongOption("real-pdisks", "Use real PDisks instead of in memory PDisks (also disable disk mock)") + options.AddLongOption("storage-path", "Use real PDisks by specified path instead of in memory PDisks (also disable disk mock), use '-' to use temp directory") + .RequiredArgument("directory") + .StoreResult(&RunnerOptions.YdbSettings.PDisksPath); + + options.AddLongOption("format-storage", "Clear storage if it exists on --storage-path") .NoArgument() - .SetFlag(&RunnerOptions.YdbSettings.UseRealPDisks); + .SetFlag(&RunnerOptions.YdbSettings.FormatStorage); options.AddLongOption("disable-disk-mock", "Disable disk mock on single node cluster") .NoArgument() @@ -767,18 +898,21 @@ class TMain : public TMainClassArgs { int DoRun(NLastGetopt::TOptsParseResult&&) override { ExecutionOptions.Validate(RunnerOptions); - if (RunnerOptions.YdbSettings.DisableDiskMock && RunnerOptions.YdbSettings.NodeCount + RunnerOptions.YdbSettings.SharedTenants.size() + RunnerOptions.YdbSettings.DedicatedTenants.size() > 1) { - ythrow yexception() << "Disable disk mock cannot be used for multi node clusters"; - } - RunnerOptions.YdbSettings.YqlToken = YqlToken; RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(UdfsDirectory, UdfsPaths, ExcludeLinkedUdfs).Get(); + + auto& appConfig = RunnerOptions.YdbSettings.AppConfig; if (ExecutionOptions.ResultsRowsLimit) { - RunnerOptions.YdbSettings.AppConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(ExecutionOptions.ResultsRowsLimit); + appConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(ExecutionOptions.ResultsRowsLimit); + } + + if (DefaultLogPriority) { + appConfig.MutableLogConfig()->SetDefaultLevel(*DefaultLogPriority); } + ModifyLogPriorities(LogPriorities, *appConfig.MutableLogConfig()); if (EmulateYt) { - const auto& fileStorageConfig = RunnerOptions.YdbSettings.AppConfig.GetQueryServiceConfig().GetFileStorage(); + const auto& fileStorageConfig = appConfig.GetQueryServiceConfig().GetFileStorage(); auto fileStorage = WithAsync(CreateFileStorage(fileStorageConfig, {MakeYtDownloader(fileStorageConfig)})); auto ytFileServices = NYql::NFile::TYtFileServices::Make(RunnerOptions.YdbSettings.FunctionRegistry.Get(), TablesMapping, fileStorage); RunnerOptions.YdbSettings.YtGateway = NYql::CreateYtFileGateway(ytFileServices); @@ -787,7 +921,26 @@ class TMain : public TMainClassArgs { ythrow yexception() << "Tables mapping is not supported without emulate YT mode"; } +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.YdbSettings.VerboseLevel >= 1) { + Cout << CoutColors.Cyan() << "Starting profile memory allocations" << CoutColors.Default() << Endl; + } + NAllocProfiler::StartAllocationSampling(true); +#else + if (ProfileAllocationsOutput) { + ythrow yexception() << "Profile memory allocations disabled, please rebuild kqprun with flag `-D PROFILE_MEMORY_ALLOCATIONS`"; + } +#endif + RunScript(ExecutionOptions, RunnerOptions); + +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.YdbSettings.VerboseLevel >= 1) { + Cout << CoutColors.Cyan() << "Finishing profile memory allocations" << CoutColors.Default() << Endl; + } + FinishProfileMemoryAllocations(); +#endif + return 0; } }; @@ -814,13 +967,31 @@ void SegmentationFaultHandler(int) { abort(); } +#ifdef PROFILE_MEMORY_ALLOCATIONS +void InterruptHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cout << colors.Red() << "Execution interrupted, finishing profile memory allocations..." << colors.Default() << Endl; + TMain::FinishProfileMemoryAllocations(); + + abort(); +} +#endif + +} // anonymous namespace + +} // namespace NKqpRun int main(int argc, const char* argv[]) { - std::set_terminate(KqprunTerminateHandler); - signal(SIGSEGV, &SegmentationFaultHandler); + std::set_terminate(NKqpRun::KqprunTerminateHandler); + signal(SIGSEGV, &NKqpRun::SegmentationFaultHandler); + +#ifdef PROFILE_MEMORY_ALLOCATIONS + signal(SIGINT, &NKqpRun::InterruptHandler); +#endif try { - TMain().Run(argc, argv); + NKqpRun::TMain().Run(argc, argv); } catch (...) { NColorizer::TColors colors = NColorizer::AutoColors(Cerr); diff --git a/ydb/tests/tools/kqprun/src/actors.cpp b/ydb/tests/tools/kqprun/src/actors.cpp index 66b28b109388..ce9dda414fa9 100644 --- a/ydb/tests/tools/kqprun/src/actors.cpp +++ b/ydb/tests/tools/kqprun/src/actors.cpp @@ -14,6 +14,7 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped promise, TProgressCallback progressCallback) : TargetNode_(request.TargetNode) + , QueryId_(request.QueryId) , Request_(std::move(request.Event)) , Promise_(promise) , ResultRowsLimit_(std::numeric_limits::max()) @@ -83,12 +84,14 @@ class TRunScriptActorMock : public NActors::TActorBootstrappedGet()->Record); + ProgressCallback_(QueryId_, ev->Get()->Record); } } private: - ui32 TargetNode_ = 0; + const ui32 TargetNode_ = 0; + const size_t QueryId_ = 0; + std::unique_ptr Request_; NThreading::TPromise Promise_; ui64 ResultRowsLimit_; @@ -296,6 +299,7 @@ class TSessionHolderActor : public NActors::TActorBootstrapped openPromise, NThreading::TPromise closePromise) : TargetNode_(request.TargetNode) , TraceId_(request.Event->Record.GetTraceId()) + , VerboseLevel_(request.VerboseLevel) , Request_(std::move(request.Event)) , OpenPromise_(openPromise) , ClosePromise_(closePromise) @@ -314,7 +318,9 @@ class TSessionHolderActor : public NActors::TActorBootstrapped= 1) { + Cout << CoutColors_.Cyan() << "Created new session on node " << TargetNode_ << " with id " << SessionId_ << "\n"; + } PingSession(); } @@ -390,6 +396,7 @@ class TSessionHolderActor : public NActors::TActorBootstrapped Request_; diff --git a/ydb/tests/tools/kqprun/src/actors.h b/ydb/tests/tools/kqprun/src/actors.h index cd477239ddb4..abe6be68060f 100644 --- a/ydb/tests/tools/kqprun/src/actors.h +++ b/ydb/tests/tools/kqprun/src/actors.h @@ -18,11 +18,13 @@ struct TQueryRequest { ui32 TargetNode; ui64 ResultRowsLimit; ui64 ResultSizeLimit; + size_t QueryId; }; struct TCreateSessionRequest { std::unique_ptr Event; ui32 TargetNode; + ui8 VerboseLevel; }; struct TEvPrivate { @@ -75,7 +77,7 @@ struct TEvPrivate { }; }; -using TProgressCallback = std::function; +using TProgressCallback = std::function; NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPromise promise, TProgressCallback progressCallback); diff --git a/ydb/tests/tools/kqprun/src/common.cpp b/ydb/tests/tools/kqprun/src/common.cpp new file mode 100644 index 000000000000..0241ca6a932c --- /dev/null +++ b/ydb/tests/tools/kqprun/src/common.cpp @@ -0,0 +1,29 @@ +#include "common.h" + + +namespace NKqpRun { + +NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName) { + NKikimrServices::EServiceKikimr service; + if (!NKikimrServices::EServiceKikimr_Parse(serviceName, &service)) { + ythrow yexception() << "Invalid kikimr service name " << serviceName; + } + return service; +} + +void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig) { + for (auto& entry : *logConfig.MutableEntry()) { + const auto it = logPriorities.find(GetLogService(entry.GetComponent())); + if (it != logPriorities.end()) { + entry.SetLevel(it->second); + logPriorities.erase(it); + } + } + for (const auto& [service, priority] : logPriorities) { + auto* entry = logConfig.AddEntry(); + entry->SetComponent(NKikimrServices::EServiceKikimr_Name(service)); + entry->SetLevel(priority); + } +} + +} // namespace NKqpRun diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index 0c55b7bc04de..81de1525889e 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -1,14 +1,18 @@ #pragma once #include -#include #include +#include +#include + +#include +#include + #include #include -#include -#include +#include namespace NKqpRun { @@ -35,7 +39,8 @@ struct TYdbSetupSettings { bool SameSession = false; bool DisableDiskMock = false; - bool UseRealPDisks = false; + bool FormatStorage = false; + std::optional PDisksPath; ui64 DiskSize = 32_GB; bool MonitoringEnabled = false; @@ -46,6 +51,7 @@ struct TYdbSetupSettings { bool TraceOptEnabled = false; TString LogOutputFile; + ui8 VerboseLevel = 1; TString YqlToken; TIntrusivePtr FunctionRegistry; @@ -72,14 +78,15 @@ struct TRunnerOptions { IOutputStream* ResultOutput = nullptr; IOutputStream* SchemeQueryAstOutput = nullptr; - IOutputStream* ScriptQueryAstOutput = nullptr; - IOutputStream* ScriptQueryPlanOutput = nullptr; - TString ScriptQueryTimelineFile; - TString InProgressStatisticsOutputFile; + std::vector ScriptQueryAstOutputs; + std::vector ScriptQueryPlanOutputs; + std::vector ScriptQueryTimelineFiles; + std::vector InProgressStatisticsOutputFiles; EResultOutputFormat ResultOutputFormat = EResultOutputFormat::RowsJson; NYdb::NConsoleClient::EDataFormat PlanOutputFormat = NYdb::NConsoleClient::EDataFormat::Default; ETraceOptType TraceOptType = ETraceOptType::Disabled; + std::optional TraceOptScriptId; TDuration ScriptCancelAfter; @@ -95,6 +102,18 @@ struct TRequestOptions { TString UserSID; TString Database; TDuration Timeout; + size_t QueryId = 0; }; +template +TValue GetValue(size_t index, const std::vector& values, TValue defaultValue) { + if (values.empty()) { + return defaultValue; + } + return values[std::min(index, values.size() - 1)]; +} + +NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName); +void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig); + } // namespace NKqpRun diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index e880554d2c34..f926538af6d3 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -120,7 +120,7 @@ class TKqpRunner::TImpl { } bool ExecuteScript(const TRequestOptions& script) { - StartScriptTraceOpt(); + StartScriptTraceOpt(script.QueryId); TRequestResult status = YdbSetup_.ScriptRequest(script, ExecutionOperation_); @@ -132,11 +132,11 @@ class TKqpRunner::TImpl { ExecutionMeta_ = TExecutionMeta(); ExecutionMeta_.Database = script.Database; - return WaitScriptExecutionOperation(); + return WaitScriptExecutionOperation(script.QueryId); } bool ExecuteQuery(const TRequestOptions& query, EQueryType queryType) { - StartScriptTraceOpt(); + StartScriptTraceOpt(query.QueryId); StartTime_ = TInstant::Now(); TString queryTypeStr; @@ -164,9 +164,9 @@ class TKqpRunner::TImpl { meta.Plan = ExecutionMeta_.Plan; } - PrintScriptAst(meta.Ast); - PrintScriptProgress(meta.Plan); - PrintScriptPlan(meta.Plan); + PrintScriptAst(query.QueryId, meta.Ast); + PrintScriptProgress(query.QueryId, meta.Plan); + PrintScriptPlan(query.QueryId, meta.Plan); PrintScriptFinish(meta, queryTypeStr); if (!status.IsSuccess()) { @@ -224,7 +224,7 @@ class TKqpRunner::TImpl { if (Options_.ResultOutput) { Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Writing script query results..." << CoutColors_.Default() << Endl; for (size_t i = 0; i < ResultSets_.size(); ++i) { - if (ResultSets_.size() > 1) { + if (ResultSets_.size() > 1 && Options_.YdbSettings.VerboseLevel >= 1) { *Options_.ResultOutput << CoutColors_.Cyan() << "Result set " << i + 1 << ":" << CoutColors_.Default() << Endl; } PrintScriptResult(ResultSets_[i]); @@ -233,7 +233,7 @@ class TKqpRunner::TImpl { } private: - bool WaitScriptExecutionOperation() { + bool WaitScriptExecutionOperation(ui64 queryId) { StartTime_ = TInstant::Now(); TDuration getOperationPeriod = TDuration::Seconds(1); @@ -244,7 +244,7 @@ class TKqpRunner::TImpl { TRequestResult status; while (true) { status = YdbSetup_.GetScriptExecutionOperationRequest(ExecutionMeta_.Database, ExecutionOperation_, ExecutionMeta_); - PrintScriptProgress(ExecutionMeta_.Plan); + PrintScriptProgress(queryId, ExecutionMeta_.Plan); if (ExecutionMeta_.Ready) { break; @@ -267,9 +267,11 @@ class TKqpRunner::TImpl { Sleep(getOperationPeriod); } - PrintScriptAst(ExecutionMeta_.Ast); - PrintScriptProgress(ExecutionMeta_.Plan); - PrintScriptPlan(ExecutionMeta_.Plan); + TYdbSetup::StopTraceOpt(); + + PrintScriptAst(queryId, ExecutionMeta_.Ast); + PrintScriptProgress(queryId, ExecutionMeta_.Plan); + PrintScriptPlan(queryId, ExecutionMeta_.Plan); PrintScriptFinish(ExecutionMeta_, "Script"); if (!status.IsSuccess() || ExecutionMeta_.ExecutionStatus != NYdb::NQuery::EExecStatus::Completed) { @@ -290,23 +292,33 @@ class TKqpRunner::TImpl { } } - void StartScriptTraceOpt() const { - if (Options_.TraceOptType == TRunnerOptions::ETraceOptType::All || Options_.TraceOptType == TRunnerOptions::ETraceOptType::Script) { + void StartScriptTraceOpt(size_t queryId) const { + bool startTraceOpt = Options_.TraceOptType == TRunnerOptions::ETraceOptType::All; + + if (Options_.TraceOptType == TRunnerOptions::ETraceOptType::Script) { + startTraceOpt |= !Options_.TraceOptScriptId || *Options_.TraceOptScriptId == queryId; + } + + if (startTraceOpt) { YdbSetup_.StartTraceOpt(); } } void PrintSchemeQueryAst(const TString& ast) const { if (Options_.SchemeQueryAstOutput) { - Cout << CoutColors_.Cyan() << "Writing scheme query ast" << CoutColors_.Default() << Endl; + if (Options_.YdbSettings.VerboseLevel >= 1) { + Cout << CoutColors_.Cyan() << "Writing scheme query ast" << CoutColors_.Default() << Endl; + } Options_.SchemeQueryAstOutput->Write(ast); } } - void PrintScriptAst(const TString& ast) const { - if (Options_.ScriptQueryAstOutput) { - Cout << CoutColors_.Cyan() << "Writing script query ast" << CoutColors_.Default() << Endl; - Options_.ScriptQueryAstOutput->Write(ast); + void PrintScriptAst(size_t queryId, const TString& ast) const { + if (const auto output = GetValue(queryId, Options_.ScriptQueryAstOutputs, nullptr)) { + if (Options_.YdbSettings.VerboseLevel >= 1) { + Cout << CoutColors_.Cyan() << "Writing script query ast" << CoutColors_.Default() << Endl; + } + output->Write(ast); } } @@ -325,16 +337,18 @@ class TKqpRunner::TImpl { printer.Print(plan); } - void PrintScriptPlan(const TString& plan) const { - if (Options_.ScriptQueryPlanOutput) { - Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; - PrintPlan(plan, Options_.ScriptQueryPlanOutput); + void PrintScriptPlan(size_t queryId, const TString& plan) const { + if (const auto output = GetValue(queryId, Options_.ScriptQueryPlanOutputs, nullptr)) { + if (Options_.YdbSettings.VerboseLevel >= 1) { + Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; + } + PrintPlan(plan, output); } } - void PrintScriptProgress(const TString& plan) const { - if (Options_.InProgressStatisticsOutputFile) { - TFileOutput outputStream(Options_.InProgressStatisticsOutputFile); + void PrintScriptProgress(size_t queryId, const TString& plan) const { + if (const auto& output = GetValue(queryId, Options_.InProgressStatisticsOutputFiles, {})) { + TFileOutput outputStream(output); outputStream << TInstant::Now().ToIsoStringLocal() << " Script in progress statistics" << Endl; auto convertedPlan = plan; @@ -361,8 +375,8 @@ class TKqpRunner::TImpl { outputStream.Finish(); } - if (Options_.ScriptQueryTimelineFile) { - TFileOutput outputStream(Options_.ScriptQueryTimelineFile); + if (const auto& output = GetValue(queryId, Options_.ScriptQueryTimelineFiles, {})) { + TFileOutput outputStream(output); TPlanVisualizer planVisualizer; planVisualizer.LoadPlans(plan); @@ -373,10 +387,10 @@ class TKqpRunner::TImpl { } TProgressCallback GetProgressCallback() { - return [this](const NKikimrKqp::TEvExecuterProgress& executerProgress) mutable { + return [this](ui64 queryId, const NKikimrKqp::TEvExecuterProgress& executerProgress) mutable { const TString& plan = executerProgress.GetQueryPlan(); ExecutionMeta_.Plan = plan; - PrintScriptProgress(plan); + PrintScriptProgress(queryId, plan); }; } @@ -411,6 +425,9 @@ class TKqpRunner::TImpl { } void PrintScriptFinish(const TQueryMeta& meta, const TString& queryType) const { + if (Options_.YdbSettings.VerboseLevel < 1) { + return; + } Cout << CoutColors_.Cyan() << queryType << " request finished."; if (meta.TotalDuration) { Cout << " Total duration: " << meta.TotalDuration; diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto new file mode 100644 index 000000000000..fabc54f2ba64 --- /dev/null +++ b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto @@ -0,0 +1,7 @@ +syntax = "proto3"; + +package NKqpRun; + +message TStorageMeta { + uint64 StorageGeneration = 1; +} diff --git a/ydb/tests/tools/kqprun/src/proto/ya.make b/ydb/tests/tools/kqprun/src/proto/ya.make new file mode 100644 index 000000000000..8cbcbe198594 --- /dev/null +++ b/ydb/tests/tools/kqprun/src/proto/ya.make @@ -0,0 +1,9 @@ +PROTO_LIBRARY() + +ONLY_TAGS(CPP_PROTO) + +SRCS( + storage_meta.proto +) + +END() diff --git a/ydb/tests/tools/kqprun/src/ya.make b/ydb/tests/tools/kqprun/src/ya.make index 8b74f61204d0..59097668037a 100644 --- a/ydb/tests/tools/kqprun/src/ya.make +++ b/ydb/tests/tools/kqprun/src/ya.make @@ -2,14 +2,19 @@ LIBRARY() SRCS( actors.cpp + common.cpp kqp_runner.cpp ydb_setup.cpp ) PEERDIR( ydb/core/testlib + + ydb/tests/tools/kqprun/src/proto ) +GENERATE_ENUM_SERIALIZATION(common.h) + YQL_LAST_ABI_VERSION() END() diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 97ab3a0b9b5f..856ab445ed7d 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -9,6 +9,9 @@ #include #include + +#include + #include @@ -65,7 +68,7 @@ class TStaticSecuredCredentialsFactory : public NYql::ISecuredServiceAccountCred class TSessionState { public: - explicit TSessionState(NActors::TTestActorRuntime* runtime, ui32 targetNodeIndex, const TString& database, const TString& traceId) + explicit TSessionState(NActors::TTestActorRuntime* runtime, ui32 targetNodeIndex, const TString& database, const TString& traceId, ui8 verboseLevel) : Runtime_(runtime) , TargetNodeIndex_(targetNodeIndex) { @@ -78,7 +81,8 @@ class TSessionState { auto closePromise = NThreading::NewPromise(); SessionHolderActor_ = Runtime_->Register(CreateSessionHolderActor(TCreateSessionRequest{ .Event = std::move(event), - .TargetNode = Runtime_->GetNodeId(targetNodeIndex) + .TargetNode = Runtime_->GetNodeId(targetNodeIndex), + .VerboseLevel = verboseLevel }, openPromise, closePromise)); SessionId_ = openPromise.GetFuture().GetValueSync(); @@ -149,13 +153,8 @@ class TYdbSetup::TImpl { } } - for (auto setting : Settings_.AppConfig.GetLogConfig().get_arr_entry()) { - NKikimrServices::EServiceKikimr service; - if (!NKikimrServices::EServiceKikimr_Parse(setting.GetComponent(), &service)) { - ythrow yexception() << "Invalid kikimr service name " << setting.GetComponent(); - } - - runtime.SetLogPriority(service, NActors::NLog::EPriority(setting.GetLevel())); + for (const auto& setting : Settings_.AppConfig.GetLogConfig().get_arr_entry()) { + runtime.SetLogPriority(GetLogService(setting.GetComponent()), NActors::NLog::EPriority(setting.GetLevel())); } runtime.SetLogBackendFactory([this]() { return CreateLogBackend(); }); @@ -177,15 +176,46 @@ class TYdbSetup::TImpl { } void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const { + TString diskPath; + if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") { + diskPath = TStringBuilder() << *Settings_.PDisksPath << "/"; + } + + bool formatDisk = true; + NKqpRun::TStorageMeta storageMeta; + if (diskPath) { + TFsPath storageMetaPath(diskPath); + storageMetaPath.MkDirs(); + + storageMetaPath = storageMetaPath.Child("kqprun_storage_meta.conf"); + if (storageMetaPath.Exists() && !Settings_.FormatStorage) { + if (!google::protobuf::TextFormat::ParseFromString(TFileInput(storageMetaPath.GetPath()).ReadAll(), &storageMeta)) { + ythrow yexception() << "Storage meta is corrupted, please use --format-storage"; + } + storageMeta.SetStorageGeneration(storageMeta.GetStorageGeneration() + 1); + formatDisk = false; + } + + TString storageMetaStr; + google::protobuf::TextFormat::PrintToString(storageMeta, &storageMetaStr); + + TFileOutput storageMetaOutput(storageMetaPath.GetPath()); + storageMetaOutput.Write(storageMetaStr); + storageMetaOutput.Finish(); + } + const NKikimr::NFake::TStorage storage = { - .UseDisk = Settings_.UseRealPDisks, + .UseDisk = !!Settings_.PDisksPath, .SectorSize = NKikimr::TTestStorageFactory::SECTOR_SIZE, - .ChunkSize = Settings_.UseRealPDisks ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE, - .DiskSize = Settings_.DiskSize + .ChunkSize = Settings_.PDisksPath ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE, + .DiskSize = Settings_.DiskSize, + .FormatDisk = formatDisk, + .DiskPath = diskPath }; - serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.UseRealPDisks); + serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.PDisksPath); serverSettings.SetCustomDiskParams(storage); + serverSettings.SetStorageGeneration(storageMeta.GetStorageGeneration()); } NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) { @@ -210,7 +240,7 @@ class TYdbSetup::TImpl { serverSettings.SetYtGateway(Settings_.YtGateway); serverSettings.S3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); serverSettings.SetInitializeFederatedQuerySetupFactory(true); - serverSettings.SetVerbose(false); + serverSettings.SetVerbose(Settings_.VerboseLevel >= 2); SetLoggerSettings(serverSettings); SetFunctionRegistry(serverSettings); @@ -315,21 +345,7 @@ class TYdbSetup::TImpl { return; } - bool found = false; - for (auto& entry : *Settings_.AppConfig.MutableLogConfig()->MutableEntry()) { - if (entry.GetComponent() == "KQP_YQL") { - entry.SetLevel(NActors::NLog::PRI_TRACE); - found = true; - break; - } - } - - if (!found) { - auto entry = Settings_.AppConfig.MutableLogConfig()->AddEntry(); - entry->SetComponent("KQP_YQL"); - entry->SetLevel(NActors::NLog::PRI_TRACE); - } - + ModifyLogPriorities({{NKikimrServices::EServiceKikimr::KQP_YQL, NActors::NLog::PRI_TRACE}}, *Settings_.AppConfig.MutableLogConfig()); NYql::NLog::InitLogger(NActors::CreateNullBackend()); } @@ -355,13 +371,13 @@ class TYdbSetup::TImpl { InitializeServer(grpcPort); WaitResourcesPublishing(); - if (Settings_.MonitoringEnabled) { + if (Settings_.MonitoringEnabled && Settings_.VerboseLevel >= 1) { for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount; ++nodeIndex) { Cout << CoutColors_.Cyan() << "Monitoring port" << (Settings_.NodeCount > 1 ? TStringBuilder() << " for node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; } } - if (Settings_.GrpcEnabled) { + if (Settings_.GrpcEnabled && Settings_.VerboseLevel >= 1) { Cout << CoutColors_.Cyan() << "Domain gRPC port: " << CoutColors_.Default() << grpcPort << Endl; } } @@ -516,7 +532,7 @@ class TYdbSetup::TImpl { if (Settings_.SameSession) { if (!SessionState_) { - SessionState_ = TSessionState(GetRuntime(), targetNodeIndex, database, query.TraceId); + SessionState_ = TSessionState(GetRuntime(), targetNodeIndex, database, query.TraceId, Settings_.VerboseLevel); } request->SetSessionId(SessionState_->GetSessionId()); } @@ -535,7 +551,8 @@ class TYdbSetup::TImpl { .Event = std::move(event), .TargetNode = GetRuntime()->GetNodeId(targetNodeIndex), .ResultRowsLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultRowsLimit(), - .ResultSizeLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultSizeLimit() + .ResultSizeLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultSizeLimit(), + .QueryId = query.QueryId }; } diff --git a/ydb/tests/tools/kqprun/ya.make b/ydb/tests/tools/kqprun/ya.make index 1ed5ff21cdd1..05b119d8e099 100644 --- a/ydb/tests/tools/kqprun/ya.make +++ b/ydb/tests/tools/kqprun/ya.make @@ -1,5 +1,11 @@ PROGRAM(kqprun) +IF (PROFILE_MEMORY_ALLOCATIONS) + MESSAGE("Enabled profile memory allocations") + ALLOCATOR(LF_DBG) + CFLAGS(-D PROFILE_MEMORY_ALLOCATIONS) +ENDIF() + SRCS( kqprun.cpp ) From 2f9ecc6c0121fd394a74d29fe8f5f18b64458494 Mon Sep 17 00:00:00 2001 From: Hor911 Date: Mon, 27 Jan 2025 14:46:09 +0300 Subject: [PATCH 02/18] FPE Handler (#13868) --- ydb/tests/tools/kqprun/kqprun.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index 8c8807899ef4..1d5e29034ed3 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -967,6 +967,16 @@ void SegmentationFaultHandler(int) { abort(); } +void FloatingPointExceptionHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << "======= floating point exception call stack ========" << colors.Default() << Endl; + FormatBackTrace(&Cerr); + Cerr << colors.Red() << "==============================================" << colors.Default() << Endl; + + abort(); +} + #ifdef PROFILE_MEMORY_ALLOCATIONS void InterruptHandler(int) { NColorizer::TColors colors = NColorizer::AutoColors(Cerr); @@ -985,6 +995,7 @@ void InterruptHandler(int) { int main(int argc, const char* argv[]) { std::set_terminate(NKqpRun::KqprunTerminateHandler); signal(SIGSEGV, &NKqpRun::SegmentationFaultHandler); + signal(SIGFPE, &NKqpRun::FloatingPointExceptionHandler); #ifdef PROFILE_MEMORY_ALLOCATIONS signal(SIGINT, &NKqpRun::InterruptHandler); From f583efed299f85d35a663575dfc87499f0fa36b2 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Wed, 29 Jan 2025 20:11:47 +0500 Subject: [PATCH 03/18] YQ-4046 KqpRun improved templates (#13951) --- .../kqprun/configuration/app_config.conf | 2 + ydb/tests/tools/kqprun/kqprun.cpp | 70 ++++++++++++---- ydb/tests/tools/kqprun/src/actors.cpp | 81 ++++++++++++++++--- ydb/tests/tools/kqprun/src/actors.h | 9 ++- ydb/tests/tools/kqprun/src/common.h | 3 +- .../tools/kqprun/src/proto/storage_meta.proto | 1 + ydb/tests/tools/kqprun/src/ydb_setup.cpp | 22 ++++- 7 files changed, 154 insertions(+), 34 deletions(-) diff --git a/ydb/tests/tools/kqprun/configuration/app_config.conf b/ydb/tests/tools/kqprun/configuration/app_config.conf index 749a24bf9520..ad27a75fc61a 100644 --- a/ydb/tests/tools/kqprun/configuration/app_config.conf +++ b/ydb/tests/tools/kqprun/configuration/app_config.conf @@ -46,6 +46,7 @@ ActorSystemConfig { ColumnShardConfig { DisabledOnSchemeShard: false + WritingInFlightRequestBytesLimit: 104857600 } FeatureFlags { @@ -53,6 +54,7 @@ FeatureFlags { EnableScriptExecutionOperations: true EnableExternalSourceSchemaInference: true EnableTempTables: true + EnableReplaceIfExistsForExternalEntities: true } KQPConfig { diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index 1d5e29034ed3..ceb4a9114866 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -88,9 +88,6 @@ struct TExecutionOptions { TRequestOptions GetSchemeQueryOptions() const { TString sql = SchemeQuery; - if (UseTemplates) { - ReplaceYqlTokenTemplate(sql); - } return { .Query = sql, @@ -108,7 +105,6 @@ struct TExecutionOptions { TString sql = ScriptQueries[index]; if (UseTemplates) { - ReplaceYqlTokenTemplate(sql); SubstGlobal(sql, "${QUERY_ID}", ToString(queryId)); } @@ -271,16 +267,6 @@ struct TExecutionOptions { ythrow yexception() << "Cannot format storage without real PDisks, please use --storage-path"; } } - -private: - static void ReplaceYqlTokenTemplate(TString& sql) { - const TString variableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}"; - if (const TString& yqlToken = GetEnv(YQL_TOKEN_VARIABLE)) { - SubstGlobal(sql, variableName, yqlToken); - } else if (sql.Contains(variableName)) { - ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable\n"; - } - } }; @@ -446,6 +432,7 @@ class TMain : public TMainClassArgs { TExecutionOptions ExecutionOptions; TRunnerOptions RunnerOptions; + std::unordered_map Templates; THashMap TablesMapping; TVector UdfsPaths; TString UdfsDirectory; @@ -537,6 +524,31 @@ class TMain : public TMainClassArgs { .NoArgument() .SetFlag(&ExecutionOptions.UseTemplates); + options.AddLongOption("var-template", "Add template from environment variables or file for -s and -p queries (use variable@file for files)") + .RequiredArgument("variable") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TStringBuf variable; + TStringBuf filePath; + TStringBuf(option->CurVal()).Split('@', variable, filePath); + if (variable.empty()) { + ythrow yexception() << "Variable name should not be empty"; + } + + TString value; + if (!filePath.empty()) { + value = LoadFile(TString(filePath)); + } else { + value = GetEnv(TString(variable)); + if (!value) { + ythrow yexception() << "Invalid env template, can not find value for variable '" << variable << "'"; + } + } + + if (!Templates.emplace(variable, value).second) { + ythrow yexception() << "Got duplicated template variable name '" << variable << "'"; + } + }); + options.AddLongOption('t', "table", "File with input table (can be used by YT with -E flag), table@file") .RequiredArgument("table@file") .Handler1([this](const NLastGetopt::TOptsParser* option) { @@ -845,6 +857,11 @@ class TMain : public TMainClassArgs { .NoArgument() .SetFlag(&EmulateYt); + options.AddLongOption('H', "health-check", "Level of health check before start (max level 2)") + .RequiredArgument("uint") + .DefaultValue(1) + .StoreResult(&RunnerOptions.YdbSettings.HealthCheckLevel); + options.AddLongOption("domain", "Test cluster domain name") .RequiredArgument("name") .DefaultValue(RunnerOptions.YdbSettings.DomainName) @@ -862,9 +879,8 @@ class TMain : public TMainClassArgs { .RequiredArgument("path") .InsertTo(&RunnerOptions.YdbSettings.ServerlessTenants); - options.AddLongOption("storage-size", "Domain storage size in gigabytes") + options.AddLongOption("storage-size", "Domain storage size in gigabytes (32 GiB by default)") .RequiredArgument("uint") - .DefaultValue(32) .StoreMappedResultT(&RunnerOptions.YdbSettings.DiskSize, [](ui32 diskSize) { return static_cast(diskSize) << 30; }); @@ -898,6 +914,11 @@ class TMain : public TMainClassArgs { int DoRun(NLastGetopt::TOptsParseResult&&) override { ExecutionOptions.Validate(RunnerOptions); + ReplaceTemplates(ExecutionOptions.SchemeQuery); + for (auto& sql : ExecutionOptions.ScriptQueries) { + ReplaceTemplates(sql); + } + RunnerOptions.YdbSettings.YqlToken = YqlToken; RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(UdfsDirectory, UdfsPaths, ExcludeLinkedUdfs).Get(); @@ -943,6 +964,21 @@ class TMain : public TMainClassArgs { return 0; } + +private: + void ReplaceTemplates(TString& sql) const { + for (const auto& [variable, value] : Templates) { + SubstGlobal(sql, TStringBuilder() << "${" << variable <<"}", value); + } + if (ExecutionOptions.UseTemplates) { + const TString tokenVariableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}"; + if (const TString& yqlToken = GetEnv(YQL_TOKEN_VARIABLE)) { + SubstGlobal(sql, tokenVariableName, yqlToken); + } else if (sql.Contains(tokenVariableName)) { + ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable"; + } + } + } }; @@ -972,7 +1008,7 @@ void FloatingPointExceptionHandler(int) { Cerr << colors.Red() << "======= floating point exception call stack ========" << colors.Default() << Endl; FormatBackTrace(&Cerr); - Cerr << colors.Red() << "==============================================" << colors.Default() << Endl; + Cerr << colors.Red() << "====================================================" << colors.Default() << Endl; abort(); } diff --git a/ydb/tests/tools/kqprun/src/actors.cpp b/ydb/tests/tools/kqprun/src/actors.cpp index ce9dda414fa9..c3556a9923c9 100644 --- a/ydb/tests/tools/kqprun/src/actors.cpp +++ b/ydb/tests/tools/kqprun/src/actors.cpp @@ -232,40 +232,78 @@ class TResourcesWaiterActor : public NActors::TActorBootstrapped promise, i32 expectedNodeCount) - : ExpectedNodeCount_(expectedNodeCount) + TResourcesWaiterActor(NThreading::TPromise promise, const TWaitResourcesSettings& settings) + : Settings_(settings) , Promise_(promise) {} void Bootstrap() { - Become(&TResourcesWaiterActor::StateFunc); + if (Settings_.HealthCheckLevel < 1) { + Finish(); + return; + } + + Become(&TResourcesWaiterActor::StateWaitNodeCont); CheckResourcesPublish(); } - void Handle(NActors::TEvents::TEvWakeup::TPtr&) { + void HandleWaitNodeCountWakeup() { CheckResourcesPublish(); } void Handle(TEvPrivate::TEvResourcesInfo::TPtr& ev) { - if (ev->Get()->NodeCount == ExpectedNodeCount_) { - Promise_.SetValue(); - PassAway(); + const auto nodeCont = ev->Get()->NodeCount; + if (nodeCont == Settings_.ExpectedNodeCount) { + if (Settings_.HealthCheckLevel < 2) { + Finish(); + } else { + Become(&TResourcesWaiterActor::StateWaitScript); + StartScriptQuery(); + } return; } + if (Settings_.VerboseLevel >= 2) { + Cout << CoutColors_.Cyan() << "Retry invalid node count, got " << nodeCont << ", expected " << Settings_.ExpectedNodeCount << CoutColors_.Default() << Endl; + } Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); } - STRICT_STFUNC(StateFunc, - hFunc(NActors::TEvents::TEvWakeup, Handle); + void HandleWaitScriptWakeup() { + StartScriptQuery(); + } + + void Handle(NKikimr::NKqp::TEvKqp::TEvScriptResponse::TPtr& ev) { + const auto status = ev->Get()->Status; + if (status == Ydb::StatusIds::SUCCESS) { + Finish(); + return; + } + + if (Settings_.VerboseLevel >= 2) { + Cout << CoutColors_.Cyan() << "Retry script creation fail with status " << status << ", reason:\n" << CoutColors_.Default() << ev->Get()->Issues.ToString() << Endl; + } + Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); + } + + STRICT_STFUNC(StateWaitNodeCont, + sFunc(NActors::TEvents::TEvWakeup, HandleWaitNodeCountWakeup); hFunc(TEvPrivate::TEvResourcesInfo, Handle); ) + STRICT_STFUNC(StateWaitScript, + sFunc(NActors::TEvents::TEvWakeup, HandleWaitScriptWakeup); + hFunc(NKikimr::NKqp::TEvKqp::TEvScriptResponse, Handle); + ) + private: void CheckResourcesPublish() { GetResourceManager(); if (!ResourceManager_) { + if (Settings_.VerboseLevel >= 2) { + Cout << CoutColors_.Cyan() << "Retry uninitialized resource manager" << CoutColors_.Default() << Endl; + } Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); return; } @@ -287,8 +325,27 @@ class TResourcesWaiterActor : public NActors::TActorBootstrapped(); + event->Record.SetUserToken(NACLib::TUserToken("", BUILTIN_ACL_ROOT, {}).SerializeAsString()); + + auto request = event->Record.MutableRequest(); + request->SetQuery("SELECT 42"); + request->SetType(NKikimrKqp::QUERY_TYPE_SQL_GENERIC_SCRIPT); + request->SetAction(NKikimrKqp::EQueryAction::QUERY_ACTION_EXECUTE); + request->SetDatabase(Settings_.Database); + + Send(NKikimr::NKqp::MakeKqpProxyID(SelfId().NodeId()), event.Release()); + } + + void Finish() { + Promise_.SetValue(); + PassAway(); + } + private: - const i32 ExpectedNodeCount_; + const TWaitResourcesSettings Settings_; + const NColorizer::TColors CoutColors_ = NColorizer::AutoColors(Cout); NThreading::TPromise Promise_; std::shared_ptr ResourceManager_; @@ -415,8 +472,8 @@ NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settin return new TAsyncQueryRunnerActor(settings); } -NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, i32 expectedNodeCount) { - return new TResourcesWaiterActor(promise, expectedNodeCount); +NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, const TWaitResourcesSettings& settings) { + return new TResourcesWaiterActor(promise, settings); } NActors::IActor* CreateSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise openPromise, NThreading::TPromise closePromise) { diff --git a/ydb/tests/tools/kqprun/src/actors.h b/ydb/tests/tools/kqprun/src/actors.h index abe6be68060f..5f73fed3ec2d 100644 --- a/ydb/tests/tools/kqprun/src/actors.h +++ b/ydb/tests/tools/kqprun/src/actors.h @@ -27,6 +27,13 @@ struct TCreateSessionRequest { ui8 VerboseLevel; }; +struct TWaitResourcesSettings { + i32 ExpectedNodeCount; + ui8 HealthCheckLevel; + ui8 VerboseLevel; + TString Database; +}; + struct TEvPrivate { enum EEv : ui32 { EvStartAsyncQuery = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), @@ -83,7 +90,7 @@ NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPr NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settings); -NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, i32 expectedNodeCount); +NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise promise, const TWaitResourcesSettings& settings); NActors::IActor* CreateSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise openPromise, NThreading::TPromise closePromise); diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index 81de1525889e..c97984ec6b78 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -36,12 +36,13 @@ struct TYdbSetupSettings { std::unordered_set SharedTenants; std::unordered_set ServerlessTenants; TDuration InitializationTimeout = TDuration::Seconds(10); + ui8 HealthCheckLevel = 1; bool SameSession = false; bool DisableDiskMock = false; bool FormatStorage = false; std::optional PDisksPath; - ui64 DiskSize = 32_GB; + std::optional DiskSize; bool MonitoringEnabled = false; ui16 MonitoringPortOffset = 0; diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto index fabc54f2ba64..98401b8200d0 100644 --- a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto +++ b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto @@ -4,4 +4,5 @@ package NKqpRun; message TStorageMeta { uint64 StorageGeneration = 1; + uint64 StorageSize = 2; } diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 856ab445ed7d..88216f6c7a5c 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -178,7 +178,10 @@ class TYdbSetup::TImpl { void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const { TString diskPath; if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") { - diskPath = TStringBuilder() << *Settings_.PDisksPath << "/"; + if (Settings_.PDisksPath->empty()) { + ythrow yexception() << "Storage directory path should not be empty"; + } + diskPath = TStringBuilder() << Settings_.PDisksPath << (Settings_.PDisksPath->back() != '/' ? "/" : ""); } bool formatDisk = true; @@ -196,6 +199,13 @@ class TYdbSetup::TImpl { formatDisk = false; } + if (Settings_.DiskSize && storageMeta.GetStorageSize() != *Settings_.DiskSize) { + if (!formatDisk) { + ythrow yexception() << "Cannot change disk size without formatting storage, please use --format-storage"; + } + storageMeta.SetStorageSize(*Settings_.DiskSize); + } + TString storageMetaStr; google::protobuf::TextFormat::PrintToString(storageMeta, &storageMetaStr); @@ -208,7 +218,7 @@ class TYdbSetup::TImpl { .UseDisk = !!Settings_.PDisksPath, .SectorSize = NKikimr::TTestStorageFactory::SECTOR_SIZE, .ChunkSize = Settings_.PDisksPath ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE, - .DiskSize = Settings_.DiskSize, + .DiskSize = Settings_.DiskSize ? *Settings_.DiskSize : 32_GB, .FormatDisk = formatDisk, .DiskPath = diskPath }; @@ -351,7 +361,13 @@ class TYdbSetup::TImpl { void WaitResourcesPublishing() const { auto promise = NThreading::NewPromise(); - GetRuntime()->Register(CreateResourcesWaiterActor(promise, Settings_.NodeCount), 0, GetRuntime()->GetAppData().SystemPoolId); + const TWaitResourcesSettings settings = { + .ExpectedNodeCount = static_cast(Settings_.NodeCount), + .HealthCheckLevel = Settings_.HealthCheckLevel, + .VerboseLevel = Settings_.VerboseLevel, + .Database = NKikimr::CanonizePath(Settings_.DomainName) + }; + GetRuntime()->Register(CreateResourcesWaiterActor(promise, settings), 0, GetRuntime()->GetAppData().SystemPoolId); try { promise.GetFuture().GetValue(Settings_.InitializationTimeout); From e4edadd79b3573847f022163da310856d5f16cf0 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Fri, 31 Jan 2025 16:59:48 +0500 Subject: [PATCH 04/18] YQ-4046 KqpRun fixed storage settings (#14065) --- ydb/tests/tools/kqprun/kqprun.cpp | 23 ++-- ydb/tests/tools/kqprun/src/actors.cpp | 130 +++++++++++++--------- ydb/tests/tools/kqprun/src/actors.h | 7 +- ydb/tests/tools/kqprun/src/common.h | 22 +++- ydb/tests/tools/kqprun/src/kqp_runner.cpp | 26 ++++- ydb/tests/tools/kqprun/src/ydb_setup.cpp | 43 ++++--- 6 files changed, 167 insertions(+), 84 deletions(-) diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index ceb4a9114866..ec8ce224a64c 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -752,10 +752,12 @@ class TMain : public TMainClassArgs { .DefaultValue(0) .StoreResult(&RunnerOptions.YdbSettings.AsyncQueriesSettings.InFlightLimit); - options.AddLongOption("verbose", "Common verbose level (max level 2)") + options.AddLongOption("verbose", TStringBuilder() << "Common verbose level (max level " << static_cast(TYdbSetupSettings::EVerbose::Max) - 1 << ")") .RequiredArgument("uint") - .DefaultValue(1) - .StoreResult(&RunnerOptions.YdbSettings.VerboseLevel); + .DefaultValue(static_cast(TYdbSetupSettings::EVerbose::Info)) + .StoreMappedResultT(&RunnerOptions.YdbSettings.VerboseLevel, [](ui8 value) { + return static_cast(std::min(value, static_cast(TYdbSetupSettings::EVerbose::Max))); + }); TChoices verbose({ {"each-query", TAsyncQueriesSettings::EVerbose::EachQuery}, @@ -857,10 +859,17 @@ class TMain : public TMainClassArgs { .NoArgument() .SetFlag(&EmulateYt); - options.AddLongOption('H', "health-check", "Level of health check before start (max level 2)") + options.AddLongOption('H', "health-check", TStringBuilder() << "Level of health check before start (max level " << static_cast(TYdbSetupSettings::EHealthCheck::Max) - 1 << ")") + .RequiredArgument("uint") + .DefaultValue(static_cast(TYdbSetupSettings::EHealthCheck::NodesCount)) + .StoreMappedResultT(&RunnerOptions.YdbSettings.HealthCheckLevel, [](ui8 value) { + return static_cast(std::min(value, static_cast(TYdbSetupSettings::EHealthCheck::Max))); + }); + + options.AddLongOption("health-check-timeout", "Health check timeout in seconds") .RequiredArgument("uint") - .DefaultValue(1) - .StoreResult(&RunnerOptions.YdbSettings.HealthCheckLevel); + .DefaultValue(10) + .StoreMappedResultT(&RunnerOptions.YdbSettings.HealthCheckTimeout, &TDuration::Seconds); options.AddLongOption("domain", "Test cluster domain name") .RequiredArgument("name") @@ -879,7 +888,7 @@ class TMain : public TMainClassArgs { .RequiredArgument("path") .InsertTo(&RunnerOptions.YdbSettings.ServerlessTenants); - options.AddLongOption("storage-size", "Domain storage size in gigabytes (32 GiB by default)") + options.AddLongOption("storage-size", TStringBuilder() << "Domain storage size in gigabytes (" << NKikimr::NBlobDepot::FormatByteSize(DEFAULT_STORAGE_SIZE) << " by default)") .RequiredArgument("uint") .StoreMappedResultT(&RunnerOptions.YdbSettings.DiskSize, [](ui32 diskSize) { return static_cast(diskSize) << 30; diff --git a/ydb/tests/tools/kqprun/src/actors.cpp b/ydb/tests/tools/kqprun/src/actors.cpp index c3556a9923c9..c548e0025660 100644 --- a/ydb/tests/tools/kqprun/src/actors.cpp +++ b/ydb/tests/tools/kqprun/src/actors.cpp @@ -83,14 +83,19 @@ class TRunScriptActorMock : public NActors::TActorBootstrappedGet()->Record); + try { + if (ProgressCallback_) { + ProgressCallback_(QueryId_, ev->Get()->Record); + } + } catch (...) { + Cerr << CerrColors_.Red() << "Got unexpected exception during progress callback: " << CurrentExceptionMessage() << CerrColors_.Default() << Endl; } } private: const ui32 TargetNode_ = 0; const size_t QueryId_ = 0; + const NColorizer::TColors CerrColors_ = NColorizer::AutoColors(Cerr); std::unique_ptr Request_; NThreading::TPromise Promise_; @@ -229,48 +234,61 @@ class TAsyncQueryRunnerActor : public NActors::TActor { }; class TResourcesWaiterActor : public NActors::TActorBootstrapped { + using IRetryPolicy = IRetryPolicy; + using EVerbose = TYdbSetupSettings::EVerbose; + using EHealthCheck = TYdbSetupSettings::EHealthCheck; + static constexpr TDuration REFRESH_PERIOD = TDuration::MilliSeconds(10); public: TResourcesWaiterActor(NThreading::TPromise promise, const TWaitResourcesSettings& settings) : Settings_(settings) + , RetryPolicy_(IRetryPolicy::GetExponentialBackoffPolicy( + &TResourcesWaiterActor::Retryable, REFRESH_PERIOD, + TDuration::MilliSeconds(100), TDuration::Seconds(1), + std::numeric_limits::max(), std::max(2 * REFRESH_PERIOD, Settings_.HealthCheckTimeout) + )) , Promise_(promise) {} void Bootstrap() { - if (Settings_.HealthCheckLevel < 1) { + Become(&TResourcesWaiterActor::StateFunc); + + HealthCheckStage_ = EHealthCheck::NodesCount; + DoHealthCheck(); + } + + void DoHealthCheck() { + if (Settings_.HealthCheckLevel < HealthCheckStage_) { Finish(); return; } - Become(&TResourcesWaiterActor::StateWaitNodeCont); - CheckResourcesPublish(); - } + switch (HealthCheckStage_) { + case TYdbSetupSettings::EHealthCheck::NodesCount: + CheckResourcesPublish(); + break; - void HandleWaitNodeCountWakeup() { - CheckResourcesPublish(); - } + case TYdbSetupSettings::EHealthCheck::ScriptRequest: + StartScriptQuery(); + break; - void Handle(TEvPrivate::TEvResourcesInfo::TPtr& ev) { - const auto nodeCont = ev->Get()->NodeCount; - if (nodeCont == Settings_.ExpectedNodeCount) { - if (Settings_.HealthCheckLevel < 2) { + case TYdbSetupSettings::EHealthCheck::None: + case TYdbSetupSettings::EHealthCheck::Max: Finish(); - } else { - Become(&TResourcesWaiterActor::StateWaitScript); - StartScriptQuery(); - } - return; + break; } + } - if (Settings_.VerboseLevel >= 2) { - Cout << CoutColors_.Cyan() << "Retry invalid node count, got " << nodeCont << ", expected " << Settings_.ExpectedNodeCount << CoutColors_.Default() << Endl; + void Handle(TEvPrivate::TEvResourcesInfo::TPtr& ev) { + const auto nodeCount = ev->Get()->NodeCount; + if (nodeCount == Settings_.ExpectedNodeCount) { + HealthCheckStage_ = EHealthCheck::ScriptRequest; + DoHealthCheck(); + return; } - Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); - } - void HandleWaitScriptWakeup() { - StartScriptQuery(); + Retry(TStringBuilder() << "invalid node count, got " << nodeCount << ", expected " << Settings_.ExpectedNodeCount, true); } void Handle(NKikimr::NKqp::TEvKqp::TEvScriptResponse::TPtr& ev) { @@ -280,45 +298,26 @@ class TResourcesWaiterActor : public NActors::TActorBootstrapped= 2) { - Cout << CoutColors_.Cyan() << "Retry script creation fail with status " << status << ", reason:\n" << CoutColors_.Default() << ev->Get()->Issues.ToString() << Endl; - } - Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); + Retry(TStringBuilder() << "script creation fail with status " << status << ", reason:\n" << CoutColors_.Default() << ev->Get()->Issues.ToString(), true); } - STRICT_STFUNC(StateWaitNodeCont, - sFunc(NActors::TEvents::TEvWakeup, HandleWaitNodeCountWakeup); + STRICT_STFUNC(StateFunc, + sFunc(NActors::TEvents::TEvWakeup, DoHealthCheck); hFunc(TEvPrivate::TEvResourcesInfo, Handle); - ) - - STRICT_STFUNC(StateWaitScript, - sFunc(NActors::TEvents::TEvWakeup, HandleWaitScriptWakeup); hFunc(NKikimr::NKqp::TEvKqp::TEvScriptResponse, Handle); ) private: void CheckResourcesPublish() { - GetResourceManager(); - if (!ResourceManager_) { - if (Settings_.VerboseLevel >= 2) { - Cout << CoutColors_.Cyan() << "Retry uninitialized resource manager" << CoutColors_.Default() << Endl; - } - Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup()); - return; + ResourceManager_ = NKikimr::NKqp::TryGetKqpResourceManager(SelfId().NodeId()); } - UpdateResourcesInfo(); - } - - void GetResourceManager() { - if (ResourceManager_) { + if (!ResourceManager_) { + Retry("uninitialized resource manager", true); return; } - ResourceManager_ = NKikimr::NKqp::TryGetKqpResourceManager(SelfId().NodeId()); - } - void UpdateResourcesInfo() const { ResourceManager_->RequestClusterResourcesInfo( [selfId = SelfId(), actorContext = ActorContext()](TVector&& resources) { actorContext.Send(selfId, new TEvPrivate::TEvResourcesInfo(resources.size())); @@ -338,20 +337,49 @@ class TResourcesWaiterActor : public NActors::TActorBootstrappedCreateRetryState(); + } + + if (auto delay = RetryState_->GetNextRetryDelay(shortRetry)) { + if (Settings_.VerboseLevel >= EVerbose::InitLogs) { + Cout << CoutColors_.Cyan() << "Retry in " << *delay << " " << message << CoutColors_.Default() << Endl; + } + Schedule(*delay, new NActors::TEvents::TEvWakeup()); + } else { + Fail(TStringBuilder() << "Health check timeout " << Settings_.HealthCheckTimeout << " exceeded, use --health-check-timeout for increasing it or check out health check logs by using --verbose " << static_cast(EVerbose::InitLogs)); + } + } + void Finish() { Promise_.SetValue(); PassAway(); } + void Fail(const TString& error) { + Promise_.SetException(error); + PassAway(); + } + + static ERetryErrorClass Retryable(bool shortRetry) { + return shortRetry ? ERetryErrorClass::ShortRetry : ERetryErrorClass::LongRetry; + } + private: const TWaitResourcesSettings Settings_; const NColorizer::TColors CoutColors_ = NColorizer::AutoColors(Cout); + const IRetryPolicy::TPtr RetryPolicy_; + IRetryPolicy::IRetryState::TPtr RetryState_ = nullptr; NThreading::TPromise Promise_; + EHealthCheck HealthCheckStage_ = EHealthCheck::None; std::shared_ptr ResourceManager_; }; class TSessionHolderActor : public NActors::TActorBootstrapped { + using EVerbose = TYdbSetupSettings::EVerbose; + public: TSessionHolderActor(TCreateSessionRequest request, NThreading::TPromise openPromise, NThreading::TPromise closePromise) : TargetNode_(request.TargetNode) @@ -375,7 +403,7 @@ class TSessionHolderActor : public NActors::TActorBootstrapped= 1) { + if (VerboseLevel_ >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Created new session on node " << TargetNode_ << " with id " << SessionId_ << "\n"; } @@ -453,7 +481,7 @@ class TSessionHolderActor : public NActors::TActorBootstrapped Request_; diff --git a/ydb/tests/tools/kqprun/src/actors.h b/ydb/tests/tools/kqprun/src/actors.h index 5f73fed3ec2d..410794491b90 100644 --- a/ydb/tests/tools/kqprun/src/actors.h +++ b/ydb/tests/tools/kqprun/src/actors.h @@ -24,13 +24,14 @@ struct TQueryRequest { struct TCreateSessionRequest { std::unique_ptr Event; ui32 TargetNode; - ui8 VerboseLevel; + TYdbSetupSettings::EVerbose VerboseLevel; }; struct TWaitResourcesSettings { i32 ExpectedNodeCount; - ui8 HealthCheckLevel; - ui8 VerboseLevel; + TYdbSetupSettings::EHealthCheck HealthCheckLevel; + TDuration HealthCheckTimeout; + TYdbSetupSettings::EVerbose VerboseLevel; TString Database; }; diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index c97984ec6b78..e85378c8667d 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -18,6 +18,7 @@ namespace NKqpRun { constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN"; +constexpr ui64 DEFAULT_STORAGE_SIZE = 32_GB; struct TAsyncQueriesSettings { enum class EVerbose { @@ -30,13 +31,28 @@ struct TAsyncQueriesSettings { }; struct TYdbSetupSettings { + enum class EVerbose { + None, + Info, + QueriesText, + InitLogs, + Max + }; + + enum class EHealthCheck { + None, + NodesCount, + ScriptRequest, + Max + }; + ui32 NodeCount = 1; TString DomainName = "Root"; std::unordered_set DedicatedTenants; std::unordered_set SharedTenants; std::unordered_set ServerlessTenants; - TDuration InitializationTimeout = TDuration::Seconds(10); - ui8 HealthCheckLevel = 1; + TDuration HealthCheckTimeout = TDuration::Seconds(10); + EHealthCheck HealthCheckLevel = EHealthCheck::NodesCount; bool SameSession = false; bool DisableDiskMock = false; @@ -52,7 +68,7 @@ struct TYdbSetupSettings { bool TraceOptEnabled = false; TString LogOutputFile; - ui8 VerboseLevel = 1; + EVerbose VerboseLevel = EVerbose::Info; TString YqlToken; TIntrusivePtr FunctionRegistry; diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index f926538af6d3..ef98165db8ca 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -87,6 +87,8 @@ void PrintStatistics(const TString& fullStat, const THashMap& flat //// TKqpRunner::TImpl class TKqpRunner::TImpl { + using EVerbose = TYdbSetupSettings::EVerbose; + public: enum class EQueryType { ScriptQuery, @@ -96,6 +98,7 @@ class TKqpRunner::TImpl { explicit TImpl(const TRunnerOptions& options) : Options_(options) + , VerboseLevel_(Options_.YdbSettings.VerboseLevel) , YdbSetup_(options.YdbSettings) , StatProcessor_(NFq::CreateStatProcessor("stat_full")) , CerrColors_(NColorizer::AutoColors(Cerr)) @@ -105,6 +108,10 @@ class TKqpRunner::TImpl { bool ExecuteSchemeQuery(const TRequestOptions& query) const { StartSchemeTraceOpt(); + if (VerboseLevel_ >= EVerbose::QueriesText) { + Cout << CoutColors_.Cyan() << "Starting scheme request:\n" << CoutColors_.Default() << query.Query << Endl; + } + TSchemeMeta meta; TRequestResult status = YdbSetup_.SchemeQueryRequest(query, meta); TYdbSetup::StopTraceOpt(); @@ -122,6 +129,10 @@ class TKqpRunner::TImpl { bool ExecuteScript(const TRequestOptions& script) { StartScriptTraceOpt(script.QueryId); + if (VerboseLevel_ >= EVerbose::QueriesText) { + Cout << CoutColors_.Cyan() << "Starting script request:\n" << CoutColors_.Default() << script.Query << Endl; + } + TRequestResult status = YdbSetup_.ScriptRequest(script, ExecutionOperation_); if (!status.IsSuccess()) { @@ -139,6 +150,10 @@ class TKqpRunner::TImpl { StartScriptTraceOpt(query.QueryId); StartTime_ = TInstant::Now(); + if (VerboseLevel_ >= EVerbose::QueriesText) { + Cout << CoutColors_.Cyan() << "Starting query request:\n" << CoutColors_.Default() << query.Query << Endl; + } + TString queryTypeStr; TQueryMeta meta; TRequestResult status; @@ -224,7 +239,7 @@ class TKqpRunner::TImpl { if (Options_.ResultOutput) { Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Writing script query results..." << CoutColors_.Default() << Endl; for (size_t i = 0; i < ResultSets_.size(); ++i) { - if (ResultSets_.size() > 1 && Options_.YdbSettings.VerboseLevel >= 1) { + if (ResultSets_.size() > 1 && Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { *Options_.ResultOutput << CoutColors_.Cyan() << "Result set " << i + 1 << ":" << CoutColors_.Default() << Endl; } PrintScriptResult(ResultSets_[i]); @@ -306,7 +321,7 @@ class TKqpRunner::TImpl { void PrintSchemeQueryAst(const TString& ast) const { if (Options_.SchemeQueryAstOutput) { - if (Options_.YdbSettings.VerboseLevel >= 1) { + if (Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Writing scheme query ast" << CoutColors_.Default() << Endl; } Options_.SchemeQueryAstOutput->Write(ast); @@ -315,7 +330,7 @@ class TKqpRunner::TImpl { void PrintScriptAst(size_t queryId, const TString& ast) const { if (const auto output = GetValue(queryId, Options_.ScriptQueryAstOutputs, nullptr)) { - if (Options_.YdbSettings.VerboseLevel >= 1) { + if (Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Writing script query ast" << CoutColors_.Default() << Endl; } output->Write(ast); @@ -339,7 +354,7 @@ class TKqpRunner::TImpl { void PrintScriptPlan(size_t queryId, const TString& plan) const { if (const auto output = GetValue(queryId, Options_.ScriptQueryPlanOutputs, nullptr)) { - if (Options_.YdbSettings.VerboseLevel >= 1) { + if (Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; } PrintPlan(plan, output); @@ -425,7 +440,7 @@ class TKqpRunner::TImpl { } void PrintScriptFinish(const TQueryMeta& meta, const TString& queryType) const { - if (Options_.YdbSettings.VerboseLevel < 1) { + if (Options_.YdbSettings.VerboseLevel < EVerbose::Info) { return; } Cout << CoutColors_.Cyan() << queryType << " request finished."; @@ -439,6 +454,7 @@ class TKqpRunner::TImpl { private: TRunnerOptions Options_; + EVerbose VerboseLevel_; TYdbSetup YdbSetup_; std::unique_ptr StatProcessor_; diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 88216f6c7a5c..0b4fe93eb3cb 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -2,9 +2,9 @@ #include +#include #include #include - #include #include @@ -68,7 +68,7 @@ class TStaticSecuredCredentialsFactory : public NYql::ISecuredServiceAccountCred class TSessionState { public: - explicit TSessionState(NActors::TTestActorRuntime* runtime, ui32 targetNodeIndex, const TString& database, const TString& traceId, ui8 verboseLevel) + explicit TSessionState(NActors::TTestActorRuntime* runtime, ui32 targetNodeIndex, const TString& database, const TString& traceId, TYdbSetupSettings::EVerbose verboseLevel) : Runtime_(runtime) , TargetNodeIndex_(targetNodeIndex) { @@ -134,6 +134,8 @@ void FillQueryMeta(TQueryMeta& meta, const NKikimrKqp::TQueryResponse& response) //// TYdbSetup::TImpl class TYdbSetup::TImpl { + using EVerbose = TYdbSetupSettings::EVerbose; + private: TAutoPtr CreateLogBackend() const { if (Settings_.LogOutputFile) { @@ -176,21 +178,26 @@ class TYdbSetup::TImpl { } void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const { - TString diskPath; + TFsPath diskPath; if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") { - if (Settings_.PDisksPath->empty()) { + diskPath = *Settings_.PDisksPath; + if (!diskPath) { ythrow yexception() << "Storage directory path should not be empty"; } - diskPath = TStringBuilder() << Settings_.PDisksPath << (Settings_.PDisksPath->back() != '/' ? "/" : ""); + if (diskPath.IsRelative()) { + diskPath = TFsPath::Cwd() / diskPath; + } + diskPath.Fix(); + diskPath.MkDir(); + if (Settings_.VerboseLevel >= EVerbose::InitLogs) { + Cout << CoutColors_.Cyan() << "Setup storage by path: " << diskPath.GetPath() << CoutColors_.Default() << Endl; + } } bool formatDisk = true; NKqpRun::TStorageMeta storageMeta; if (diskPath) { - TFsPath storageMetaPath(diskPath); - storageMetaPath.MkDirs(); - - storageMetaPath = storageMetaPath.Child("kqprun_storage_meta.conf"); + const auto storageMetaPath = TFsPath(diskPath).Child("kqprun_storage_meta.conf"); if (storageMetaPath.Exists() && !Settings_.FormatStorage) { if (!google::protobuf::TextFormat::ParseFromString(TFileInput(storageMetaPath.GetPath()).ReadAll(), &storageMeta)) { ythrow yexception() << "Storage meta is corrupted, please use --format-storage"; @@ -201,9 +208,11 @@ class TYdbSetup::TImpl { if (Settings_.DiskSize && storageMeta.GetStorageSize() != *Settings_.DiskSize) { if (!formatDisk) { - ythrow yexception() << "Cannot change disk size without formatting storage, please use --format-storage"; + ythrow yexception() << "Cannot change disk size without formatting storage, current disk size " << NKikimr::NBlobDepot::FormatByteSize(storageMeta.GetStorageSize()) << ", please use --format-storage"; } storageMeta.SetStorageSize(*Settings_.DiskSize); + } else if (!storageMeta.GetStorageSize()) { + storageMeta.SetStorageSize(DEFAULT_STORAGE_SIZE); } TString storageMetaStr; @@ -214,13 +223,16 @@ class TYdbSetup::TImpl { storageMetaOutput.Finish(); } + TString storagePath = diskPath.GetPath(); + SlashFolderLocal(storagePath); + const NKikimr::NFake::TStorage storage = { .UseDisk = !!Settings_.PDisksPath, .SectorSize = NKikimr::TTestStorageFactory::SECTOR_SIZE, .ChunkSize = Settings_.PDisksPath ? NKikimr::TTestStorageFactory::CHUNK_SIZE : NKikimr::TTestStorageFactory::MEM_CHUNK_SIZE, .DiskSize = Settings_.DiskSize ? *Settings_.DiskSize : 32_GB, .FormatDisk = formatDisk, - .DiskPath = diskPath + .DiskPath = storagePath }; serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.PDisksPath); @@ -250,7 +262,7 @@ class TYdbSetup::TImpl { serverSettings.SetYtGateway(Settings_.YtGateway); serverSettings.S3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); serverSettings.SetInitializeFederatedQuerySetupFactory(true); - serverSettings.SetVerbose(Settings_.VerboseLevel >= 2); + serverSettings.SetVerbose(Settings_.VerboseLevel >= EVerbose::InitLogs); SetLoggerSettings(serverSettings); SetFunctionRegistry(serverSettings); @@ -364,13 +376,14 @@ class TYdbSetup::TImpl { const TWaitResourcesSettings settings = { .ExpectedNodeCount = static_cast(Settings_.NodeCount), .HealthCheckLevel = Settings_.HealthCheckLevel, + .HealthCheckTimeout = Settings_.HealthCheckTimeout, .VerboseLevel = Settings_.VerboseLevel, .Database = NKikimr::CanonizePath(Settings_.DomainName) }; GetRuntime()->Register(CreateResourcesWaiterActor(promise, settings), 0, GetRuntime()->GetAppData().SystemPoolId); try { - promise.GetFuture().GetValue(Settings_.InitializationTimeout); + promise.GetFuture().GetValue(2 * Settings_.HealthCheckTimeout); } catch (...) { ythrow yexception() << "Failed to initialize all resources: " << CurrentExceptionMessage(); } @@ -387,13 +400,13 @@ class TYdbSetup::TImpl { InitializeServer(grpcPort); WaitResourcesPublishing(); - if (Settings_.MonitoringEnabled && Settings_.VerboseLevel >= 1) { + if (Settings_.MonitoringEnabled && Settings_.VerboseLevel >= EVerbose::Info) { for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount; ++nodeIndex) { Cout << CoutColors_.Cyan() << "Monitoring port" << (Settings_.NodeCount > 1 ? TStringBuilder() << " for node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; } } - if (Settings_.GrpcEnabled && Settings_.VerboseLevel >= 1) { + if (Settings_.GrpcEnabled && Settings_.VerboseLevel >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Domain gRPC port: " << CoutColors_.Default() << grpcPort << Endl; } } From a37f865340ea74b97ef50f6231fd1e998854b854 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Fri, 7 Feb 2025 13:43:11 +0500 Subject: [PATCH 05/18] YQ-4092 KqpRun improve tenants flags (#14279) --- ydb/tests/tools/kqprun/kqprun.cpp | 53 ++++- ydb/tests/tools/kqprun/src/common.h | 6 +- .../tools/kqprun/src/proto/storage_meta.proto | 14 ++ ydb/tests/tools/kqprun/src/ydb_setup.cpp | 193 ++++++++++++------ 4 files changed, 192 insertions(+), 74 deletions(-) diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index ec8ce224a64c..1ddbd4adc78d 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -95,7 +95,7 @@ struct TExecutionOptions { .TraceId = DefaultTraceId, .PoolId = "", .UserSID = BUILTIN_ACL_ROOT, - .Database = "", + .Database = GetValue(0, Databases, TString()), .Timeout = TDuration::Zero() }; } @@ -135,15 +135,15 @@ struct TExecutionOptions { private: void ValidateOptionsSizes(const TRunnerOptions& runnerOptions) const { - const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName) { - if (checkSize > numberQueries) { + const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName, bool useInSchemeQuery = false) { + if (checkSize > std::max(numberQueries, static_cast(useInSchemeQuery ? 1 : 0))) { ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of script queries is " << numberQueries; } }; checker(ExecutionCases.size(), "execution cases"); checker(ScriptQueryActions.size(), "script query actions"); - checker(Databases.size(), "databases"); + checker(Databases.size(), "databases", true); checker(TraceIds.size(), "trace ids"); checker(PoolIds.size(), "pool ids"); checker(UserSIDs.size(), "user SIDs"); @@ -257,7 +257,7 @@ struct TExecutionOptions { static void ValidateStorageSettings(const TYdbSetupSettings& ydbSettings) { if (ydbSettings.DisableDiskMock) { - if (ydbSettings.NodeCount + ydbSettings.SharedTenants.size() + ydbSettings.DedicatedTenants.size() > 1) { + if (ydbSettings.NodeCount + ydbSettings.Tenants.size() > 1) { ythrow yexception() << "Disable disk mock cannot be used for multi node clusters (already disabled)"; } else if (ydbSettings.PDisksPath) { ythrow yexception() << "Disable disk mock cannot be used with real PDisks (already disabled)"; @@ -876,17 +876,50 @@ class TMain : public TMainClassArgs { .DefaultValue(RunnerOptions.YdbSettings.DomainName) .StoreResult(&RunnerOptions.YdbSettings.DomainName); - options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain") + const auto addTenant = [this](const TString& type, TStorageMeta::TTenant::EType protoType, const NLastGetopt::TOptsParser* option) { + TStringBuf tenant; + TStringBuf nodesCountStr; + TStringBuf(option->CurVal()).Split(':', tenant, nodesCountStr); + if (tenant.empty()) { + ythrow yexception() << type << " tenant name should not be empty"; + } + + TStorageMeta::TTenant tenantInfo; + tenantInfo.SetType(protoType); + tenantInfo.SetNodesCount(nodesCountStr ? FromString(nodesCountStr) : 1); + if (tenantInfo.GetNodesCount() == 0) { + ythrow yexception() << type << " tenant should have at least one node"; + } + + if (!RunnerOptions.YdbSettings.Tenants.emplace(tenant, tenantInfo).second) { + ythrow yexception() << "Got duplicated tenant name: " << tenant; + } + }; + options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain (for node count use dedicated-name:node-count)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.DedicatedTenants); + .Handler1(std::bind(addTenant, "Dedicated", TStorageMeta::TTenant::DEDICATED, std::placeholders::_1)); - options.AddLongOption("shared", "Shared tenant path, relative inside domain") + options.AddLongOption("shared", "Shared tenant path, relative inside domain (for node count use dedicated-name:node-count)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.SharedTenants); + .Handler1(std::bind(addTenant, "Shared", TStorageMeta::TTenant::SHARED, std::placeholders::_1)); options.AddLongOption("serverless", "Serverless tenant path, relative inside domain (use string serverless-name@shared-name to specify shared database)") .RequiredArgument("path") - .InsertTo(&RunnerOptions.YdbSettings.ServerlessTenants); + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TStringBuf serverless; + TStringBuf shared; + TStringBuf(option->CurVal()).Split('@', serverless, shared); + if (serverless.empty()) { + ythrow yexception() << "Serverless tenant name should not be empty"; + } + + TStorageMeta::TTenant tenantInfo; + tenantInfo.SetType(TStorageMeta::TTenant::SERVERLESS); + tenantInfo.SetSharedTenant(TString(shared)); + if (!RunnerOptions.YdbSettings.Tenants.emplace(serverless, tenantInfo).second) { + ythrow yexception() << "Got duplicated tenant name: " << serverless; + } + }); options.AddLongOption("storage-size", TStringBuilder() << "Domain storage size in gigabytes (" << NKikimr::NBlobDepot::FormatByteSize(DEFAULT_STORAGE_SIZE) << " by default)") .RequiredArgument("uint") diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index e85378c8667d..4f4f4ac36913 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -9,6 +9,8 @@ #include #include +#include + #include #include @@ -48,9 +50,7 @@ struct TYdbSetupSettings { ui32 NodeCount = 1; TString DomainName = "Root"; - std::unordered_set DedicatedTenants; - std::unordered_set SharedTenants; - std::unordered_set ServerlessTenants; + std::map Tenants; TDuration HealthCheckTimeout = TDuration::Seconds(10); EHealthCheck HealthCheckLevel = EHealthCheck::NodesCount; bool SameSession = false; diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto index 98401b8200d0..ee4a7c4162a2 100644 --- a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto +++ b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto @@ -3,6 +3,20 @@ syntax = "proto3"; package NKqpRun; message TStorageMeta { + message TTenant { + enum EType { + DEDICATED = 0; + SHARED = 1; + SERVERLESS = 2; + } + + EType Type = 1; + uint32 NodesCount = 2; + string SharedTenant = 3; // Only for serverless tenants + } + uint64 StorageGeneration = 1; uint64 StorageSize = 2; + string DomainName = 3; + map Tenants = 4; } diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 0b4fe93eb3cb..ea0b0ce1c0a8 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -177,7 +177,7 @@ class TYdbSetup::TImpl { serverSettings.SetFrFactory(functionRegistryFactory); } - void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) const { + void SetStorageSettings(NKikimr::Tests::TServerSettings& serverSettings) { TFsPath diskPath; if (Settings_.PDisksPath && *Settings_.PDisksPath != "-") { diskPath = *Settings_.PDisksPath; @@ -195,32 +195,33 @@ class TYdbSetup::TImpl { } bool formatDisk = true; - NKqpRun::TStorageMeta storageMeta; if (diskPath) { - const auto storageMetaPath = TFsPath(diskPath).Child("kqprun_storage_meta.conf"); - if (storageMetaPath.Exists() && !Settings_.FormatStorage) { - if (!google::protobuf::TextFormat::ParseFromString(TFileInput(storageMetaPath.GetPath()).ReadAll(), &storageMeta)) { + StorageMetaPath_ = TFsPath(diskPath).Child("kqprun_storage_meta.conf"); + if (StorageMetaPath_.Exists() && !Settings_.FormatStorage) { + if (!google::protobuf::TextFormat::ParseFromString(TFileInput(StorageMetaPath_.GetPath()).ReadAll(), &StorageMeta_)) { ythrow yexception() << "Storage meta is corrupted, please use --format-storage"; } - storageMeta.SetStorageGeneration(storageMeta.GetStorageGeneration() + 1); + StorageMeta_.SetStorageGeneration(StorageMeta_.GetStorageGeneration() + 1); formatDisk = false; } - if (Settings_.DiskSize && storageMeta.GetStorageSize() != *Settings_.DiskSize) { + if (Settings_.DiskSize && StorageMeta_.GetStorageSize() != *Settings_.DiskSize) { if (!formatDisk) { - ythrow yexception() << "Cannot change disk size without formatting storage, current disk size " << NKikimr::NBlobDepot::FormatByteSize(storageMeta.GetStorageSize()) << ", please use --format-storage"; + ythrow yexception() << "Cannot change disk size without formatting storage, current disk size " << NKikimr::NBlobDepot::FormatByteSize(StorageMeta_.GetStorageSize()) << ", please use --format-storage"; } - storageMeta.SetStorageSize(*Settings_.DiskSize); - } else if (!storageMeta.GetStorageSize()) { - storageMeta.SetStorageSize(DEFAULT_STORAGE_SIZE); + StorageMeta_.SetStorageSize(*Settings_.DiskSize); + } else if (!StorageMeta_.GetStorageSize()) { + StorageMeta_.SetStorageSize(DEFAULT_STORAGE_SIZE); } - TString storageMetaStr; - google::protobuf::TextFormat::PrintToString(storageMeta, &storageMetaStr); + const TString& domainName = NKikimr::CanonizePath(Settings_.DomainName); + if (!StorageMeta_.GetDomainName()) { + StorageMeta_.SetDomainName(domainName); + } else if (StorageMeta_.GetDomainName() != domainName) { + ythrow yexception() << "Cannot change domain name without formatting storage, current name " << StorageMeta_.GetDomainName() << ", please use --format-storage"; + } - TFileOutput storageMetaOutput(storageMetaPath.GetPath()); - storageMetaOutput.Write(storageMetaStr); - storageMetaOutput.Finish(); + UpdateStorageMeta(); } TString storagePath = diskPath.GetPath(); @@ -237,7 +238,7 @@ class TYdbSetup::TImpl { serverSettings.SetEnableMockOnSingleNode(!Settings_.DisableDiskMock && !Settings_.PDisksPath); serverSettings.SetCustomDiskParams(storage); - serverSettings.SetStorageGeneration(storageMeta.GetStorageGeneration()); + serverSettings.SetStorageGeneration(StorageMeta_.GetStorageGeneration()); } NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) { @@ -278,28 +279,54 @@ class TYdbSetup::TImpl { serverSettings.SetGrpcPort(grpcPort); } - if (!Settings_.SharedTenants.empty() || !Settings_.DedicatedTenants.empty()) { - serverSettings.SetDynamicNodeCount(Settings_.SharedTenants.size() + Settings_.DedicatedTenants.size()); - for (const TString& dedicatedTenant : Settings_.DedicatedTenants) { - serverSettings.AddStoragePoolType(dedicatedTenant); - } - for (const auto& sharedTenant : Settings_.SharedTenants) { - serverSettings.AddStoragePoolType(sharedTenant); + for (const auto& [tenantPath, tenantInfo] : StorageMeta_.GetTenants()) { + Settings_.Tenants.emplace(tenantPath, tenantInfo); + } + + ui32 dynNodesCount = 0; + for (const auto& [tenantPath, tenantInfo] : Settings_.Tenants) { + if (tenantInfo.GetType() != TStorageMeta::TTenant::SERVERLESS) { + serverSettings.AddStoragePoolType(tenantPath); + dynNodesCount += tenantInfo.GetNodesCount(); } } + serverSettings.SetDynamicNodeCount(dynNodesCount); return serverSettings; } - void CreateTenant(Ydb::Cms::CreateDatabaseRequest&& request, const TString& type) const { - const auto path = request.path(); - Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating " << type << " tenant " << path << "..." << CoutColors_.Default() << Endl; - Tenants_->CreateTenant(std::move(request)); + void CreateTenant(Ydb::Cms::CreateDatabaseRequest&& request, const TString& relativePath, const TString& type, TStorageMeta::TTenant tenantInfo) { + const auto absolutePath = request.path(); + const auto [it, inserted] = StorageMeta_.MutableTenants()->emplace(relativePath, tenantInfo); + if (inserted) { + if (Settings_.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating " << type << " tenant " << absolutePath << "..." << CoutColors_.Default() << Endl; + } + Tenants_->CreateTenant(std::move(request), tenantInfo.GetNodesCount()); + UpdateStorageMeta(); + } else { + if (it->second.GetType() != tenantInfo.GetType()) { + ythrow yexception() << "Can not change tenant " << absolutePath << " type without formatting storage, current type " << TStorageMeta::TTenant::EType_Name(it->second.GetType()) << ", please use --format-storage"; + } + if (it->second.GetSharedTenant() != tenantInfo.GetSharedTenant()) { + ythrow yexception() << "Can not change tenant " << absolutePath << " shared resources without formatting storage from '" << it->second.GetSharedTenant() << "', please use --format-storage"; + } + if (it->second.GetNodesCount() != tenantInfo.GetNodesCount()) { + it->second.SetNodesCount(tenantInfo.GetNodesCount()); + UpdateStorageMeta(); + } + if (Settings_.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Starting " << type << " tenant " << absolutePath << "..." << CoutColors_.Default() << Endl; + } + if (!request.has_serverless_resources()) { + Tenants_->Run(absolutePath, tenantInfo.GetNodesCount()); + } + } if (Settings_.MonitoringEnabled) { - ui32 nodeIndex = GetNodeIndexForDatabase(path); + ui32 nodeIndex = GetNodeIndexForDatabase(absolutePath); NActors::TActorId edgeActor = GetRuntime()->AllocateEdgeActor(nodeIndex); - GetRuntime()->Register(NKikimr::CreateBoardPublishActor(NKikimr::MakeEndpointsBoardPath(path), "", edgeActor, 0, true), nodeIndex, GetRuntime()->GetAppData(nodeIndex).UserPoolId); + GetRuntime()->Register(NKikimr::CreateBoardPublishActor(NKikimr::MakeEndpointsBoardPath(absolutePath), "", edgeActor, 0, true), nodeIndex, GetRuntime()->GetAppData(nodeIndex).UserPoolId); } } @@ -309,39 +336,50 @@ class TYdbSetup::TImpl { } void CreateTenants() { - for (const TString& dedicatedTenant : Settings_.DedicatedTenants) { + std::set sharedTenants; + std::map serverlessTenants; + for (const auto& [tenantPath, tenantInfo] : Settings_.Tenants) { Ydb::Cms::CreateDatabaseRequest request; - request.set_path(GetTenantPath(dedicatedTenant)); - AddTenantStoragePool(request.mutable_resources()->add_storage_units(), dedicatedTenant); - CreateTenant(std::move(request), "dedicated"); + request.set_path(GetTenantPath(tenantPath)); + + switch (tenantInfo.GetType()) { + case TStorageMeta::TTenant::DEDICATED: + AddTenantStoragePool(request.mutable_resources()->add_storage_units(), tenantPath); + CreateTenant(std::move(request), tenantPath, "dedicated", tenantInfo); + break; + + case TStorageMeta::TTenant::SHARED: + sharedTenants.emplace(tenantPath); + AddTenantStoragePool(request.mutable_shared_resources()->add_storage_units(), tenantPath); + CreateTenant(std::move(request), tenantPath, "shared", tenantInfo); + break; + + case TStorageMeta::TTenant::SERVERLESS: + serverlessTenants.emplace(tenantPath, tenantInfo); + break; + + default: + ythrow yexception() << "Unexpected tenant type: " << TStorageMeta::TTenant::EType_Name(tenantInfo.GetType()); + break; + } } - for (const TString& sharedTenant : Settings_.SharedTenants) { - Ydb::Cms::CreateDatabaseRequest request; - request.set_path(GetTenantPath(sharedTenant)); - AddTenantStoragePool(request.mutable_shared_resources()->add_storage_units(), sharedTenant); - CreateTenant(std::move(request), "shared"); - } + for (auto [tenantPath, tenantInfo] : serverlessTenants) { + if (!tenantInfo.GetSharedTenant()) { + if (sharedTenants.empty()) { + ythrow yexception() << "Can not create serverless tenant, there is no shared tenants, please use `--shared `"; + } + if (sharedTenants.size() > 1) { + ythrow yexception() << "Can not create serverless tenant, there is more than one shared tenant, please use `--serverless " << tenantPath << "@`"; + } + tenantInfo.SetSharedTenant(*sharedTenants.begin()); + } - ServerlessToShared_.reserve(Settings_.ServerlessTenants.size()); - for (const TString& serverlessTenant : Settings_.ServerlessTenants) { Ydb::Cms::CreateDatabaseRequest request; - if (serverlessTenant.Contains('@')) { - TStringBuf serverless; - TStringBuf shared; - TStringBuf(serverlessTenant).Split('@', serverless, shared); - - request.set_path(GetTenantPath(TString(serverless))); - request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(TString(shared))); - } else if (!Settings_.SharedTenants.empty()) { - request.set_path(GetTenantPath(serverlessTenant)); - request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(*Settings_.SharedTenants.begin())); - } else { - ythrow yexception() << "Can not create serverless tenant " << serverlessTenant << ", there is no shared tenants"; - } + request.set_path(GetTenantPath(tenantPath)); + request.mutable_serverless_resources()->set_shared_database_path(GetTenantPath(tenantInfo.GetSharedTenant())); ServerlessToShared_[request.path()] = request.serverless_resources().shared_database_path(); - - CreateTenant(std::move(request), "serverless"); + CreateTenant(std::move(request), tenantPath, "serverless", tenantInfo); } } @@ -402,8 +440,18 @@ class TYdbSetup::TImpl { if (Settings_.MonitoringEnabled && Settings_.VerboseLevel >= EVerbose::Info) { for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount; ++nodeIndex) { - Cout << CoutColors_.Cyan() << "Monitoring port" << (Settings_.NodeCount > 1 ? TStringBuilder() << " for node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; + Cout << CoutColors_.Cyan() << "Monitoring port" << (Server_->StaticNodes() + Server_->DynamicNodes() > 1 ? TStringBuilder() << " for static node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl; } + const auto printTenantNodes = [this](const std::pair& tenantInfo) { + if (tenantInfo.second.GetType() == TStorageMeta::TTenant::SERVERLESS) { + return; + } + const auto& nodes = Tenants_->List(GetTenantPath(tenantInfo.first)); + for (auto it = nodes.rbegin(); it != nodes.rend(); ++it) { + Cout << CoutColors_.Cyan() << "Monitoring port for dynamic node " << *it + 1 << " [" << tenantInfo.first << "]: " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(*it) << Endl; + } + }; + std::for_each(Settings_.Tenants.rbegin(), Settings_.Tenants.rend(), std::bind(printTenantNodes, std::placeholders::_1)); } if (Settings_.GrpcEnabled && Settings_.VerboseLevel >= EVerbose::Info) { @@ -590,12 +638,12 @@ class TYdbSetup::TImpl { } TString GetDatabasePath(const TString& database) const { - return NKikimr::CanonizePath(database ? database : Settings_.DomainName); + return NKikimr::CanonizePath(database ? database : GetDefaultDatabase()); } ui32 GetNodeIndexForDatabase(const TString& path) const { - auto canonizedPath = NKikimr::CanonizePath(path); - if (canonizedPath.empty() || canonizedPath == NKikimr::CanonizePath(Settings_.DomainName)) { + auto canonizedPath = NKikimr::CanonizePath(path ? path : GetDefaultDatabase()); + if (canonizedPath == NKikimr::CanonizePath(Settings_.DomainName)) { return RandomNumber(Settings_.NodeCount); } @@ -610,6 +658,27 @@ class TYdbSetup::TImpl { ythrow yexception() << "Unknown tenant '" << canonizedPath << "'"; } + TString GetDefaultDatabase() const { + if (StorageMeta_.TenantsSize() > 1) { + ythrow yexception() << "Can not choose default database, there is more than one tenants, please use `-D `"; + } + if (StorageMeta_.TenantsSize() == 1) { + return StorageMeta_.GetTenants().begin()->first; + } + return Settings_.DomainName; + } + + void UpdateStorageMeta() const { + if (StorageMetaPath_) { + TString storageMetaStr; + google::protobuf::TextFormat::PrintToString(StorageMeta_, &storageMetaStr); + + TFileOutput storageMetaOutput(StorageMetaPath_.GetPath()); + storageMetaOutput.Write(storageMetaStr); + storageMetaOutput.Finish(); + } + } + private: TYdbSetupSettings Settings_; NColorizer::TColors CoutColors_; @@ -622,6 +691,8 @@ class TYdbSetup::TImpl { std::unordered_map ServerlessToShared_; std::optional AsyncQueryRunnerActorId_; std::optional SessionState_; + TFsPath StorageMetaPath_; + NKqpRun::TStorageMeta StorageMeta_; }; From 26b0dc5a8b2fb7983966c14b7e61eae619c3c426 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Sun, 9 Feb 2025 19:11:13 +0500 Subject: [PATCH 06/18] YQ-3561 first version of FQ run tool (#14259) --- ydb/core/testlib/test_client.cpp | 8 +- ydb/core/testlib/test_client.h | 2 + ydb/tests/tools/fqrun/.gitignore | 4 + .../tools/fqrun/configuration/fq_config.conf | 141 +++++++++ ydb/tests/tools/fqrun/fqprun.cpp | 166 +++++++++++ ydb/tests/tools/fqrun/src/common.h | 29 ++ ydb/tests/tools/fqrun/src/fq_runner.cpp | 136 +++++++++ ydb/tests/tools/fqrun/src/fq_runner.h | 24 ++ ydb/tests/tools/fqrun/src/fq_setup.cpp | 261 +++++++++++++++++ ydb/tests/tools/fqrun/src/fq_setup.h | 33 +++ ydb/tests/tools/fqrun/src/ya.make | 25 ++ ydb/tests/tools/fqrun/ya.make | 18 ++ ydb/tests/tools/kqprun/kqprun.cpp | 187 +----------- ydb/tests/tools/kqprun/runlib/application.cpp | 113 ++++++++ ydb/tests/tools/kqprun/runlib/application.h | 30 ++ ydb/tests/tools/kqprun/runlib/settings.h | 25 ++ ydb/tests/tools/kqprun/runlib/utils.cpp | 274 ++++++++++++++++++ ydb/tests/tools/kqprun/runlib/utils.h | 95 ++++++ ydb/tests/tools/kqprun/runlib/ya.make | 28 ++ ydb/tests/tools/kqprun/src/common.cpp | 29 -- ydb/tests/tools/kqprun/src/common.h | 24 +- ydb/tests/tools/kqprun/src/kqp_runner.cpp | 163 +---------- ydb/tests/tools/kqprun/src/ya.make | 2 +- ydb/tests/tools/kqprun/src/ydb_setup.cpp | 40 +-- ydb/tests/tools/kqprun/src/ydb_setup.h | 19 +- ydb/tests/tools/kqprun/ya.make | 1 + ydb/tests/tools/ya.make | 1 + 27 files changed, 1441 insertions(+), 437 deletions(-) create mode 100644 ydb/tests/tools/fqrun/.gitignore create mode 100644 ydb/tests/tools/fqrun/configuration/fq_config.conf create mode 100644 ydb/tests/tools/fqrun/fqprun.cpp create mode 100644 ydb/tests/tools/fqrun/src/common.h create mode 100644 ydb/tests/tools/fqrun/src/fq_runner.cpp create mode 100644 ydb/tests/tools/fqrun/src/fq_runner.h create mode 100644 ydb/tests/tools/fqrun/src/fq_setup.cpp create mode 100644 ydb/tests/tools/fqrun/src/fq_setup.h create mode 100644 ydb/tests/tools/fqrun/src/ya.make create mode 100644 ydb/tests/tools/fqrun/ya.make create mode 100644 ydb/tests/tools/kqprun/runlib/application.cpp create mode 100644 ydb/tests/tools/kqprun/runlib/application.h create mode 100644 ydb/tests/tools/kqprun/runlib/settings.h create mode 100644 ydb/tests/tools/kqprun/runlib/utils.cpp create mode 100644 ydb/tests/tools/kqprun/runlib/utils.h create mode 100644 ydb/tests/tools/kqprun/runlib/ya.make delete mode 100644 ydb/tests/tools/kqprun/src/common.cpp diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp index 1c6a8d91908e..fc49d26427ea 100644 --- a/ydb/core/testlib/test_client.cpp +++ b/ydb/core/testlib/test_client.cpp @@ -440,7 +440,7 @@ namespace Tests { GRpcServer->AddService(new NGRpcService::TGRpcMonitoringService(system, counters, grpcRequestProxies[0], true)); GRpcServer->AddService(new NGRpcService::TGRpcYdbQueryService(system, counters, grpcRequestProxies, true, 1)); GRpcServer->AddService(new NGRpcService::TGRpcYdbTabletService(system, counters, grpcRequestProxies, true, 1)); - if (Settings->EnableYq) { + if (Settings->EnableYq || Settings->EnableYqGrpc) { GRpcServer->AddService(new NGRpcService::TGRpcFederatedQueryService(system, counters, grpcRequestProxies[0])); GRpcServer->AddService(new NGRpcService::TGRpcFqPrivateTaskService(system, counters, grpcRequestProxies[0])); } @@ -1099,6 +1099,12 @@ namespace Tests { } } + { + auto statActor = NStat::CreateStatService(); + const TActorId statActorId = Runtime->Register(statActor.Release(), nodeIdx, Runtime->GetAppData(nodeIdx).UserPoolId); + Runtime->RegisterService(NStat::MakeStatServiceID(Runtime->GetNodeId(nodeIdx)), statActorId, nodeIdx); + } + { IActor* kesusService = NKesus::CreateKesusProxyService(); TActorId kesusServiceId = Runtime->Register(kesusService, nodeIdx, userPoolId); diff --git a/ydb/core/testlib/test_client.h b/ydb/core/testlib/test_client.h index 5b82a0d5e44b..157ae243065b 100644 --- a/ydb/core/testlib/test_client.h +++ b/ydb/core/testlib/test_client.h @@ -140,6 +140,7 @@ namespace Tests { bool UseRealThreads = true; bool EnableKqpSpilling = false; bool EnableYq = false; + bool EnableYqGrpc = false; TDuration KeepSnapshotTimeout = TDuration::Zero(); ui64 ChangesQueueItemsLimit = 0; ui64 ChangesQueueBytesLimit = 0; @@ -205,6 +206,7 @@ namespace Tests { TServerSettings& SetEnableDbCounters(bool value) { FeatureFlags.SetEnableDbCounters(value); return *this; } TServerSettings& SetEnablePersistentQueryStats(bool value) { FeatureFlags.SetEnablePersistentQueryStats(value); return *this; } TServerSettings& SetEnableYq(bool value) { EnableYq = value; return *this; } + TServerSettings& SetEnableYqGrpc(bool value) { EnableYqGrpc = value; return *this; } TServerSettings& SetKeepSnapshotTimeout(TDuration value) { KeepSnapshotTimeout = value; return *this; } TServerSettings& SetChangesQueueItemsLimit(ui64 value) { ChangesQueueItemsLimit = value; return *this; } TServerSettings& SetChangesQueueBytesLimit(ui64 value) { ChangesQueueBytesLimit = value; return *this; } diff --git a/ydb/tests/tools/fqrun/.gitignore b/ydb/tests/tools/fqrun/.gitignore new file mode 100644 index 000000000000..51aaf6608d57 --- /dev/null +++ b/ydb/tests/tools/fqrun/.gitignore @@ -0,0 +1,4 @@ +sync_dir + +*.log +*.sql diff --git a/ydb/tests/tools/fqrun/configuration/fq_config.conf b/ydb/tests/tools/fqrun/configuration/fq_config.conf new file mode 100644 index 000000000000..8a980fa5daec --- /dev/null +++ b/ydb/tests/tools/fqrun/configuration/fq_config.conf @@ -0,0 +1,141 @@ +Enabled: true +EnableDynamicNameservice: true +EnableTaskCounters: true + +CheckpointCoordinator { + Enabled: true + + Storage { + TablePrefix: "yq/checkpoints" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } +} + +Common { + MdbGateway: "https://mdb.api.cloud.yandex.net:443" + ObjectStorageEndpoint: "https://storage-internal.cloud.yandex.net" + IdsPrefix: "kr" + QueryArtifactsCompressionMethod: "zstd_6" + MonitoringEndpoint: "monitoring.api.cloud.yandex.net" + KeepInternalErrors: true + UseNativeProtocolForClickHouse: true + DisableSslForGenericDataSources: true + ShowQueryTimeline: true +} + +ControlPlaneProxy { + Enabled: true +} + +ControlPlaneStorage { + Enabled: true + StatsMode: STATS_MODE_PROFILE + + AvailableConnection: "OBJECT_STORAGE" + AvailableConnection: "DATA_STREAMS" + AvailableConnection: "MONITORING" + AvailableConnection: "POSTGRESQL_CLUSTER" + AvailableConnection: "CLICKHOUSE_CLUSTER" + AvailableConnection: "YDB_DATABASE" + AvailableConnection: "GREENPLUM_CLUSTER" + AvailableConnection: "MYSQL_CLUSTER" + + AvailableStreamingConnection: "OBJECT_STORAGE" + AvailableStreamingConnection: "DATA_STREAMS" + AvailableStreamingConnection: "MONITORING" + AvailableStreamingConnection: "YDB_DATABASE" + + AvailableBinding: "OBJECT_STORAGE" + AvailableBinding: "DATA_STREAMS" + + Storage { + TablePrefix: "yq/control_plane" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } +} + +DbPool { + Enabled: true + + Storage { + TablePrefix: "yq/db_pool" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } +} + +NodesManager { + Enabled: true +} + +PendingFetcher { + Enabled: true +} + +PrivateApi { + Enabled: true +} + +PrivateProxy { + Enabled: true +} + +QuotasManager { + Enabled: true +} + +RateLimiter { + Enabled: true + ControlPlaneEnabled: true + DataPlaneEnabled: true + + Database { + TablePrefix: "yq/rate_limiter" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } + + Limiters { + CoordinationNodePath: "limiter_alpha" + } +} + +ResourceManager { + Enabled: true +} + +RowDispatcher { + Enabled: true + SendStatusPeriodSec: 10 + TimeoutBeforeStartSessionSec: 10 + + CompileService { + ParallelCompilationLimit: 20 + } + + Coordinator { + CoordinationNodePath: "yq/row_dispatcher" + + Database { + TablePrefix: "yq/row_dispatcher" + ClientTimeoutSec: 70 + OperationTimeoutSec: 60 + CancelAfterSec: 60 + } + } + + JsonParser { + BatchSizeBytes: 1048576 + BatchCreationTimeoutMs: 1000 + } +} + +TestConnection { + Enabled: true +} diff --git a/ydb/tests/tools/fqrun/fqprun.cpp b/ydb/tests/tools/fqrun/fqprun.cpp new file mode 100644 index 000000000000..277aa8d2ed90 --- /dev/null +++ b/ydb/tests/tools/fqrun/fqprun.cpp @@ -0,0 +1,166 @@ +#include +#include + +#include + +#include +#include +#include + +using namespace NKikimrRun; + +namespace NFqRun { + +namespace { + +struct TExecutionOptions { + TString Query; + + bool HasResults() const { + return !Query.empty(); + } + + TRequestOptions GetQueryOptions() const { + return { + .Query = Query + }; + } + + void Validate(const TRunnerOptions& runnerOptions) const { + if (!Query && !runnerOptions.FqSettings.MonitoringEnabled && !runnerOptions.FqSettings.GrpcEnabled) { + ythrow yexception() << "Nothing to execute and is not running as daemon"; + } + } +}; + +void RunArgumentQueries(const TExecutionOptions& executionOptions, TFqRunner& runner) { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + + if (executionOptions.Query) { + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing query..." << colors.Default() << Endl; + if (!runner.ExecuteStreamQuery(executionOptions.GetQueryOptions())) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Query execution failed"; + } + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Fetching query results..." << colors.Default() << Endl; + if (!runner.FetchQueryResults()) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Fetch query results failed"; + } + } + + if (executionOptions.HasResults()) { + try { + runner.PrintQueryResults(); + } catch (...) { + ythrow yexception() << "Failed to print script results, reason:\n" << CurrentExceptionMessage(); + } + } +} + +void RunAsDaemon() { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization finished" << colors.Default() << Endl; + while (true) { + Sleep(TDuration::Seconds(1)); + } +} + +void RunScript(const TExecutionOptions& executionOptions, const TRunnerOptions& runnerOptions) { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization of fq runner..." << colors.Default() << Endl; + TFqRunner runner(runnerOptions); + + try { + RunArgumentQueries(executionOptions, runner); + } catch (const yexception& exception) { + if (runnerOptions.FqSettings.MonitoringEnabled) { + Cerr << colors.Red() << CurrentExceptionMessage() << colors.Default() << Endl; + } else { + throw exception; + } + } + + if (runnerOptions.FqSettings.MonitoringEnabled || runnerOptions.FqSettings.GrpcEnabled) { + RunAsDaemon(); + } + + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Finalization of fq runner..." << colors.Default() << Endl; +} + +class TMain : public TMainBase { +protected: + void RegisterOptions(NLastGetopt::TOpts& options) override { + options.SetTitle("FqRun -- tool to execute stream queries through FQ proxy"); + options.AddHelpOption('h'); + options.SetFreeArgsNum(0); + + // Inputs + + options.AddLongOption('p', "query", "Query to execute") + .RequiredArgument("file") + .StoreMappedResult(&ExecutionOptions.Query, &LoadFile); + + options.AddLongOption("fq-cfg", "File with FQ config (NFq::NConfig::TConfig for FQ proxy)") + .RequiredArgument("file") + .DefaultValue("./configuration/fq_config.conf") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(TString(option->CurValOrDef())), &RunnerOptions.FqSettings.FqConfig)) { + ythrow yexception() << "Bad format of FQ configuration"; + } + }); + + // Outputs + + options.AddLongOption("result-file", "File with query results (use '-' to write in stdout)") + .RequiredArgument("file") + .DefaultValue("-") + .StoreMappedResultT(&RunnerOptions.ResultOutput, &GetDefaultOutput); + + TChoices resultFormat({ + {"rows", EResultOutputFormat::RowsJson}, + {"full-json", EResultOutputFormat::FullJson}, + {"full-proto", EResultOutputFormat::FullProto} + }); + options.AddLongOption('R', "result-format", "Query result format") + .RequiredArgument("result-format") + .DefaultValue("rows") + .Choices(resultFormat.GetChoices()) + .StoreMappedResultT(&RunnerOptions.ResultOutputFormat, resultFormat); + + RegisterKikimrOptions(options, RunnerOptions.FqSettings); + } + + int DoRun(NLastGetopt::TOptsParseResult&&) override { + ExecutionOptions.Validate(RunnerOptions); + + auto& logConfig = RunnerOptions.FqSettings.LogConfig; + logConfig.SetDefaultLevel(NActors::NLog::EPriority::PRI_CRIT); + FillLogConfig(logConfig); + + RunScript(ExecutionOptions, RunnerOptions); + + return 0; + } + +private: + TExecutionOptions ExecutionOptions; + TRunnerOptions RunnerOptions; +}; + +} // anonymous namespace + +} // namespace NFqRun + +int main(int argc, const char* argv[]) { + SetupSignalActions(); + + try { + NFqRun::TMain().Run(argc, argv); + } catch (...) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << CurrentExceptionMessage() << colors.Default() << Endl; + return 1; + } +} diff --git a/ydb/tests/tools/fqrun/src/common.h b/ydb/tests/tools/fqrun/src/common.h new file mode 100644 index 000000000000..5af4b019dbc3 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/common.h @@ -0,0 +1,29 @@ +#pragma once + +#include + +#include +#include +#include + +namespace NFqRun { + +constexpr i64 MAX_RESULT_SET_ROWS = 1000; + +struct TFqSetupSettings : public NKikimrRun::TServerSettings { + NFq::NConfig::TConfig FqConfig; + NKikimrConfig::TLogConfig LogConfig; +}; + +struct TRunnerOptions { + IOutputStream* ResultOutput = nullptr; + NKikimrRun::EResultOutputFormat ResultOutputFormat = NKikimrRun::EResultOutputFormat::RowsJson; + + TFqSetupSettings FqSettings; +}; + +struct TRequestOptions { + TString Query; +}; + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_runner.cpp b/ydb/tests/tools/fqrun/src/fq_runner.cpp new file mode 100644 index 000000000000..6fe51cfb50f9 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_runner.cpp @@ -0,0 +1,136 @@ +#include "fq_runner.h" +#include "fq_setup.h" + +#include + +using namespace NKikimrRun; + +namespace NFqRun { + +class TFqRunner::TImpl { + static constexpr TDuration REFRESH_PERIOD = TDuration::Seconds(1); + +public: + explicit TImpl(const TRunnerOptions& options) + : Options(options) + , FqSetup(options.FqSettings) + , CerrColors(NColorizer::AutoColors(Cerr)) + , CoutColors(NColorizer::AutoColors(Cout)) + {} + + bool ExecuteStreamQuery(const TRequestOptions& query) { + const TRequestResult status = FqSetup.StreamRequest(query, StreamQueryId); + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to start stream request execution, reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + + return WaitStreamQuery(); + } + + bool FetchQueryResults() { + ResultSets.clear(); + ResultSets.resize(ExecutionMeta.ResultSetSizes.size()); + for (i32 resultSetId = 0; resultSetId < static_cast(ExecutionMeta.ResultSetSizes.size()); ++resultSetId) { + const auto rowsCount = ExecutionMeta.ResultSetSizes[resultSetId]; + if (rowsCount > MAX_RESULT_SET_ROWS) { + Cerr << CerrColors.Red() << "Result set with id " << resultSetId << " have " << rowsCount << " rows, it is larger than allowed limit " << MAX_RESULT_SET_ROWS << ", results will be truncated" << CerrColors.Default() << Endl; + } + + const TRequestResult status = FqSetup.FetchQueryResults(StreamQueryId, resultSetId, ResultSets[resultSetId]); + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to fetch result set with id " << resultSetId << ", reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + } + + return true; + } + + void PrintQueryResults() { + if (Options.ResultOutput) { + Cout << CoutColors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Writing query results..." << CoutColors.Default() << Endl; + for (size_t i = 0; i < ResultSets.size(); ++i) { + if (ResultSets.size() > 1) { + *Options.ResultOutput << CoutColors.Cyan() << "Result set " << i + 1 << ":" << CoutColors.Default() << Endl; + } + PrintResultSet(Options.ResultOutputFormat, *Options.ResultOutput, ResultSets[i]); + } + } + } + +private: + static bool IsFinalStatus(FederatedQuery::QueryMeta::ComputeStatus status) { + using EStatus = FederatedQuery::QueryMeta; + return IsIn({EStatus::FAILED, EStatus::COMPLETED, EStatus::ABORTED_BY_USER, EStatus::ABORTED_BY_SYSTEM}, status); + } + + bool WaitStreamQuery() { + StartTime = TInstant::Now(); + + while (true) { + const TRequestResult status = FqSetup.DescribeQuery(StreamQueryId, ExecutionMeta); + + if (IsFinalStatus(ExecutionMeta.Status)) { + break; + } + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to describe query, reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + + Sleep(REFRESH_PERIOD); + } + + Cout << CoutColors.Cyan() << "Query finished. Duration: " << TInstant::Now() - StartTime << CoutColors.Default() << Endl; + + if (ExecutionMeta.Status != FederatedQuery::QueryMeta::COMPLETED) { + Cerr << CerrColors.Red() << "Failed to execute query, invalid final status " << FederatedQuery::QueryMeta::ComputeStatus_Name(ExecutionMeta.Status) << ", issues:" << CerrColors.Default() << Endl << ExecutionMeta.Issues.ToString() << Endl; + if (ExecutionMeta.TransientIssues) { + Cerr << CerrColors.Red() << "Transient issues:" << CerrColors.Default() << Endl << ExecutionMeta.TransientIssues.ToString() << Endl; + } + return false; + } + + if (ExecutionMeta.Issues) { + Cerr << CerrColors.Red() << "Query finished with issues:" << CerrColors.Default() << Endl << ExecutionMeta.Issues.ToString() << Endl; + } + + if (ExecutionMeta.TransientIssues) { + Cerr << CerrColors.Red() << "Query finished with transient issues:" << CerrColors.Default() << Endl << ExecutionMeta.TransientIssues.ToString() << Endl; + } + + return true; + } + +private: + const TRunnerOptions Options; + const TFqSetup FqSetup; + const NColorizer::TColors CerrColors; + const NColorizer::TColors CoutColors; + + TString StreamQueryId; + TInstant StartTime; + TExecutionMeta ExecutionMeta; + std::vector ResultSets; +}; + +TFqRunner::TFqRunner(const TRunnerOptions& options) + : Impl(new TImpl(options)) +{} + +bool TFqRunner::ExecuteStreamQuery(const TRequestOptions& query) const { + return Impl->ExecuteStreamQuery(query); +} + +bool TFqRunner::FetchQueryResults() const { + return Impl->FetchQueryResults(); +} + +void TFqRunner::PrintQueryResults() const { + Impl->PrintQueryResults(); +} + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_runner.h b/ydb/tests/tools/fqrun/src/fq_runner.h new file mode 100644 index 000000000000..7b803648dcf9 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_runner.h @@ -0,0 +1,24 @@ +#pragma once + +#include "common.h" + +#include + +namespace NFqRun { + +class TFqRunner { +public: + explicit TFqRunner(const TRunnerOptions& options); + + bool ExecuteStreamQuery(const TRequestOptions& query) const; + + bool FetchQueryResults() const; + + void PrintQueryResults() const; + +private: + class TImpl; + std::shared_ptr Impl; +}; + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_setup.cpp b/ydb/tests/tools/fqrun/src/fq_setup.cpp new file mode 100644 index 000000000000..6b16b4f5af15 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_setup.cpp @@ -0,0 +1,261 @@ +#include "fq_setup.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +using namespace NKikimrRun; + +namespace NFqRun { + +namespace { + +Ydb::StatusIds::StatusCode GetStatus(const NYql::TIssues& issues) { + return issues ? Ydb::StatusIds::BAD_REQUEST : Ydb::StatusIds::SUCCESS; +} + +} // anonymous namespace + +class TFqSetup::TImpl { +private: + TAutoPtr CreateLogBackend() const { + if (Settings.LogOutputFile) { + return NActors::CreateFileBackend(Settings.LogOutputFile); + } else { + return NActors::CreateStderrBackend(); + } + } + + void SetLoggerSettings(NKikimr::Tests::TServerSettings& serverSettings) const { + auto loggerInitializer = [this](NActors::TTestActorRuntime& runtime) { + InitLogSettings(Settings.LogConfig, runtime); + runtime.SetLogBackendFactory([this]() { return CreateLogBackend(); }); + }; + + serverSettings.SetLoggerInitializer(loggerInitializer); + } + + NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) { + NKikimr::Tests::TServerSettings serverSettings(PortManager.GetPort()); + + serverSettings.SetDomainName(Settings.DomainName); + serverSettings.SetVerbose(false); + + NKikimrConfig::TAppConfig config; + *config.MutableLogConfig() = Settings.LogConfig; + serverSettings.SetAppConfig(config); + + SetLoggerSettings(serverSettings); + + if (Settings.MonitoringEnabled) { + serverSettings.InitKikimrRunConfig(); + serverSettings.SetMonitoringPortOffset(Settings.MonitoringPortOffset, true); + serverSettings.SetNeedStatsCollectors(true); + } + + serverSettings.SetGrpcPort(grpcPort); + serverSettings.SetEnableYqGrpc(true); + + return serverSettings; + } + + void InitializeServer(ui32 grpcPort) { + const auto& serverSettings = GetServerSettings(grpcPort); + + Server = MakeIntrusive(serverSettings); + Server->GetRuntime()->SetDispatchTimeout(TDuration::Max()); + + Server->EnableGRpc(NYdbGrpc::TServerOptions() + .SetHost("localhost") + .SetPort(grpcPort) + .SetLogger(NYdbGrpc::CreateActorSystemLogger(*GetRuntime()->GetActorSystem(0), NKikimrServices::GRPC_SERVER)) + .SetGRpcShutdownDeadline(TDuration::Zero()) + ); + + Client = std::make_unique(serverSettings); + Client->InitRootScheme(); + } + + NFq::NConfig::TConfig GetFqProxyConfig(ui32 grpcPort, ui32 httpPort) const { + auto fqConfig = Settings.FqConfig; + + fqConfig.MutableControlPlaneStorage()->AddSuperUsers(BUILTIN_ACL_ROOT); + fqConfig.MutablePrivateProxy()->AddGrantedUsers(BUILTIN_ACL_ROOT); + + const TString endpoint = TStringBuilder() << "localhost:" << grpcPort; + const TString database = NKikimr::CanonizePath(Settings.DomainName); + const auto fillStorageConfig = [endpoint, database](NFq::NConfig::TYdbStorageConfig* config) { + config->SetEndpoint(endpoint); + config->SetDatabase(database); + }; + fillStorageConfig(fqConfig.MutableControlPlaneStorage()->MutableStorage()); + fillStorageConfig(fqConfig.MutableDbPool()->MutableStorage()); + fillStorageConfig(fqConfig.MutableCheckpointCoordinator()->MutableStorage()); + fillStorageConfig(fqConfig.MutableRateLimiter()->MutableDatabase()); + fillStorageConfig(fqConfig.MutableRowDispatcher()->MutableCoordinator()->MutableDatabase()); + + auto* privateApiConfig = fqConfig.MutablePrivateApi(); + privateApiConfig->SetTaskServiceEndpoint(endpoint); + privateApiConfig->SetTaskServiceDatabase(database); + + auto* nodesMenagerConfig = fqConfig.MutableNodesManager(); + nodesMenagerConfig->SetPort(grpcPort); + nodesMenagerConfig->SetHost("localhost"); + + fqConfig.MutableCommon()->SetYdbMvpCloudEndpoint(TStringBuilder() << "http://localhost:" << httpPort << "/yql-mock/abc"); + + return fqConfig; + } + + void InitializeFqProxy(ui32 grpcPort) { + const ui32 httpPort = PortManager.GetPort(); + const auto& fqConfig = GetFqProxyConfig(grpcPort, httpPort); + const auto counters = GetRuntime()->GetAppData().Counters->GetSubgroup("counters", "yq"); + YqSharedResources = NFq::CreateYqSharedResources(fqConfig, NKikimr::CreateYdbCredentialsProviderFactory, counters); + + const auto actorRegistrator = [runtime = GetRuntime()](NActors::TActorId serviceActorId, NActors::IActor* actor) { + auto actorId = runtime->Register(actor, 0, runtime->GetAppData().UserPoolId); + runtime->RegisterService(serviceActorId, actorId); + }; + + const auto folderServiceFactory = [](auto& config) { + return NKikimr::NFolderService::CreateMockFolderServiceAdapterActor(config, ""); + }; + + NFq::Init( + fqConfig, GetRuntime()->GetNodeId(), actorRegistrator, &GetRuntime()->GetAppData(), + Settings.DomainName, nullptr, YqSharedResources, folderServiceFactory, 0, {} + ); + + NFq::InitTest(GetRuntime(), httpPort, grpcPort, YqSharedResources); + } + +public: + explicit TImpl(const TFqSetupSettings& settings) + : Settings(settings) + { + const ui32 grpcPort = Settings.GrpcPort ? Settings.GrpcPort : PortManager.GetPort(); + InitializeServer(grpcPort); + InitializeFqProxy(grpcPort); + + if (Settings.MonitoringEnabled) { + Cout << CoutColors.Cyan() << "Monitoring port: " << CoutColors.Default() << GetRuntime()->GetMonPort() << Endl; + } + + if (Settings.GrpcEnabled) { + Cout << CoutColors.Cyan() << "Domain gRPC port: " << CoutColors.Default() << grpcPort << Endl; + } + } + + ~TImpl() { + if (YqSharedResources) { + YqSharedResources->Stop(); + } + } + + NFq::TEvControlPlaneProxy::TEvCreateQueryResponse::TPtr StreamRequest(const TRequestOptions& query) const { + FederatedQuery::CreateQueryRequest request; + request.set_execute_mode(FederatedQuery::ExecuteMode::RUN); + + auto& content = *request.mutable_content(); + content.set_type(FederatedQuery::QueryContent::STREAMING); + content.set_text(query.Query); + content.mutable_acl()->set_visibility(::FederatedQuery::Acl::SCOPE); + + return RunControlPlaneProxyRequest(request); + } + + NFq::TEvControlPlaneProxy::TEvDescribeQueryResponse::TPtr DescribeQuery(const TString& queryId) const { + FederatedQuery::DescribeQueryRequest request; + request.set_query_id(queryId); + + return RunControlPlaneProxyRequest(request); + } + + NFq::TEvControlPlaneProxy::TEvGetResultDataResponse::TPtr FetchQueryResults(const TString& queryId, i32 resultSetId) const { + FederatedQuery::GetResultDataRequest request; + request.set_query_id(queryId); + request.set_result_set_index(resultSetId); + request.set_limit(MAX_RESULT_SET_ROWS); + + return RunControlPlaneProxyRequest(request); + } + +private: + NActors::TTestActorRuntime* GetRuntime() const { + return Server->GetRuntime(); + } + + template + typename TResponse::TPtr RunControlPlaneProxyRequest(const TProto& request) const { + auto event = std::make_unique("yandexcloud://kqprun", request, BUILTIN_ACL_ROOT, BUILTIN_ACL_ROOT, TVector{}); + return RunControlPlaneProxyRequest(std::move(event)); + } + + template + typename TResponse::TPtr RunControlPlaneProxyRequest(std::unique_ptr event) const { + NActors::TActorId edgeActor = GetRuntime()->AllocateEdgeActor(); + NActors::TActorId controlPlaneProxy = NFq::ControlPlaneProxyActorId(); + + GetRuntime()->Send(controlPlaneProxy, edgeActor, event.release()); + + return GetRuntime()->GrabEdgeEvent(edgeActor); + } + +private: + const TFqSetupSettings Settings; + const NColorizer::TColors CoutColors; + + NKikimr::Tests::TServer::TPtr Server; + std::unique_ptr Client; + NFq::IYqSharedResources::TPtr YqSharedResources; + TPortManager PortManager; +}; + +TFqSetup::TFqSetup(const TFqSetupSettings& settings) + : Impl(new TImpl(settings)) +{} + +TRequestResult TFqSetup::StreamRequest(const TRequestOptions& query, TString& queryId) const { + const auto response = Impl->StreamRequest(query); + + queryId = response->Get()->Result.query_id(); + + const auto& issues = response->Get()->Issues; + return TRequestResult(GetStatus(issues), issues); +} + +TRequestResult TFqSetup::DescribeQuery(const TString& queryId, TExecutionMeta& meta) const { + const auto response = Impl->DescribeQuery(queryId); + + const auto& result = response->Get()->Result.query(); + meta.Status = result.meta().status(); + NYql::IssuesFromMessage(result.issue(), meta.Issues); + NYql::IssuesFromMessage(result.transient_issue(), meta.TransientIssues); + + meta.ResultSetSizes.clear(); + for (const auto& resultMeta : result.result_set_meta()) { + meta.ResultSetSizes.emplace_back(resultMeta.rows_count()); + } + + const auto& issues = response->Get()->Issues; + return TRequestResult(GetStatus(issues), issues); +} + +TRequestResult TFqSetup::FetchQueryResults(const TString& queryId, i32 resultSetId, Ydb::ResultSet& resultSet) const { + const auto response = Impl->FetchQueryResults(queryId, resultSetId); + + resultSet = response->Get()->Result.result_set(); + + const auto& issues = response->Get()->Issues; + return TRequestResult(GetStatus(issues), issues); +} + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_setup.h b/ydb/tests/tools/fqrun/src/fq_setup.h new file mode 100644 index 000000000000..c78fad43e2dd --- /dev/null +++ b/ydb/tests/tools/fqrun/src/fq_setup.h @@ -0,0 +1,33 @@ +#pragma once + +#include "common.h" + +#include + +namespace NFqRun { + +struct TExecutionMeta { + FederatedQuery::QueryMeta::ComputeStatus Status; + NYql::TIssues Issues; + NYql::TIssues TransientIssues; + std::vector ResultSetSizes; +}; + +class TFqSetup { + using TRequestResult = NKikimrRun::TRequestResult; + +public: + explicit TFqSetup(const TFqSetupSettings& settings); + + TRequestResult StreamRequest(const TRequestOptions& query, TString& queryId) const; + + TRequestResult DescribeQuery(const TString& queryId, TExecutionMeta& meta) const; + + TRequestResult FetchQueryResults(const TString& queryId, i32 resultSetId, Ydb::ResultSet& resultSet) const; + +private: + class TImpl; + std::shared_ptr Impl; +}; + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/ya.make b/ydb/tests/tools/fqrun/src/ya.make new file mode 100644 index 000000000000..5173c30f4be2 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/ya.make @@ -0,0 +1,25 @@ +LIBRARY() + +SRCS( + fq_runner.cpp + fq_setup.cpp +) + +PEERDIR( + library/cpp/colorizer + library/cpp/testing/unittest + util + ydb/core/fq/libs/config/protos + ydb/core/fq/libs/control_plane_proxy/events + ydb/core/fq/libs/init + ydb/core/fq/libs/mock + ydb/core/testlib + ydb/library/folder_service/mock + ydb/library/grpc/server/actors + ydb/library/security + ydb/tests/tools/kqprun/runlib +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/tests/tools/fqrun/ya.make b/ydb/tests/tools/fqrun/ya.make new file mode 100644 index 000000000000..fa7212f0f8bc --- /dev/null +++ b/ydb/tests/tools/fqrun/ya.make @@ -0,0 +1,18 @@ +PROGRAM(fqprun) + +SRCS( + fqprun.cpp +) + +PEERDIR( + library/cpp/colorizer + util + ydb/tests/tools/fqrun/src + ydb/tests/tools/kqprun/runlib + yql/essentials/parser/pg_wrapper + yql/essentials/sql/pg +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index 1ddbd4adc78d..4505610ce417 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -1,5 +1,3 @@ -#include "src/kqp_runner.h" - #include #include @@ -9,11 +7,12 @@ #include #include -#include #include - #include #include +#include +#include +#include #include #include @@ -26,6 +25,7 @@ #include #endif +using namespace NKikimrRun; namespace NKqpRun { @@ -423,9 +423,8 @@ TIntrusivePtr CreateFunctionRegistr } -class TMain : public TMainClassArgs { +class TMain : public TMainBase { inline static const TString YqlToken = GetEnv(YQL_TOKEN_VARIABLE); - inline static std::vector> FileHolders; inline static IOutputStream* ProfileAllocationsOutput = nullptr; inline static NColorizer::TColors CoutColors = NColorizer::AutoColors(Cout); @@ -439,52 +438,6 @@ class TMain : public TMainClassArgs { bool ExcludeLinkedUdfs = false; bool EmulateYt = false; - std::optional DefaultLogPriority; - std::unordered_map LogPriorities; - - static TString LoadFile(const TString& file) { - return TFileInput(file).ReadAll(); - } - - static IOutputStream* GetDefaultOutput(const TString& file) { - if (file == "-") { - return &Cout; - } - if (file) { - FileHolders.emplace_back(new TFileOutput(file)); - return FileHolders.back().get(); - } - return nullptr; - } - - template - class TChoices { - public: - explicit TChoices(std::map choicesMap) - : ChoicesMap(std::move(choicesMap)) - {} - - TResult operator()(const TString& choice) const { - return ChoicesMap.at(choice); - } - - TVector GetChoices() const { - TVector choices; - choices.reserve(ChoicesMap.size()); - for (const auto& [choice, _] : ChoicesMap) { - choices.emplace_back(choice); - } - return choices; - } - - bool Contains(const TString& choice) const { - return ChoicesMap.contains(choice); - } - - private: - const std::map ChoicesMap; - }; - #ifdef PROFILE_MEMORY_ALLOCATIONS public: static void FinishProfileMemoryAllocations() { @@ -590,51 +543,6 @@ class TMain : public TMainClassArgs { // Outputs - options.AddLongOption("log-file", "File with execution logs (writes in stderr if empty)") - .RequiredArgument("file") - .StoreResult(&RunnerOptions.YdbSettings.LogOutputFile) - .Handler1([](const NLastGetopt::TOptsParser* option) { - if (const TString& file = option->CurVal()) { - std::remove(file.c_str()); - } - }); - - TChoices logPriority({ - {"emerg", NActors::NLog::EPriority::PRI_EMERG}, - {"alert", NActors::NLog::EPriority::PRI_ALERT}, - {"crit", NActors::NLog::EPriority::PRI_CRIT}, - {"error", NActors::NLog::EPriority::PRI_ERROR}, - {"warn", NActors::NLog::EPriority::PRI_WARN}, - {"notice", NActors::NLog::EPriority::PRI_NOTICE}, - {"info", NActors::NLog::EPriority::PRI_INFO}, - {"debug", NActors::NLog::EPriority::PRI_DEBUG}, - {"trace", NActors::NLog::EPriority::PRI_TRACE}, - }); - options.AddLongOption("log-default", "Default log priority") - .RequiredArgument("priority") - .Choices(logPriority.GetChoices()) - .StoreMappedResultT(&DefaultLogPriority, logPriority); - - options.AddLongOption("log", "Component log priority in format = (e. g. KQP_YQL=trace)") - .RequiredArgument("component priority") - .Handler1([this, logPriority](const NLastGetopt::TOptsParser* option) { - TStringBuf component; - TStringBuf priority; - TStringBuf(option->CurVal()).Split('=', component, priority); - if (component.empty() || priority.empty()) { - ythrow yexception() << "Incorrect log setting, expected form component=priority, e. g. KQP_YQL=trace"; - } - - if (!logPriority.Contains(TString(priority))) { - ythrow yexception() << "Incorrect log priority: " << priority; - } - - const auto service = GetLogService(TString(component)); - if (!LogPriorities.emplace(service, logPriority(TString(priority))).second) { - ythrow yexception() << "Got duplicated log service name: " << component; - } - }); - TChoices traceOpt({ {"all", TRunnerOptions::ETraceOptType::All}, {"scheme", TRunnerOptions::ETraceOptType::Scheme}, @@ -669,10 +577,10 @@ class TMain : public TMainClassArgs { .DefaultValue(0) .StoreResult(&ExecutionOptions.ResultsRowsLimit); - TChoices resultFormat({ - {"rows", TRunnerOptions::EResultOutputFormat::RowsJson}, - {"full-json", TRunnerOptions::EResultOutputFormat::FullJson}, - {"full-proto", TRunnerOptions::EResultOutputFormat::FullProto} + TChoices resultFormat({ + {"rows", EResultOutputFormat::RowsJson}, + {"full-json", EResultOutputFormat::FullJson}, + {"full-proto", EResultOutputFormat::FullProto} }); options.AddLongOption('R', "result-format", "Script query result format") .RequiredArgument("result-format") @@ -837,24 +745,6 @@ class TMain : public TMainClassArgs { return nodeCount; }); - options.AddLongOption('M', "monitoring", "Embedded UI port (use 0 to start on random free port), if used kqprun will be run as daemon") - .RequiredArgument("uint") - .Handler1([this](const NLastGetopt::TOptsParser* option) { - if (const TString& port = option->CurVal()) { - RunnerOptions.YdbSettings.MonitoringEnabled = true; - RunnerOptions.YdbSettings.MonitoringPortOffset = FromString(port); - } - }); - - options.AddLongOption('G', "grpc", "gRPC port (use 0 to start on random free port), if used kqprun will be run as daemon") - .RequiredArgument("uint") - .Handler1([this](const NLastGetopt::TOptsParser* option) { - if (const TString& port = option->CurVal()) { - RunnerOptions.YdbSettings.GrpcEnabled = true; - RunnerOptions.YdbSettings.GrpcPort = FromString(port); - } - }); - options.AddLongOption('E', "emulate-yt", "Emulate YT tables (use file gateway instead of native gateway)") .NoArgument() .SetFlag(&EmulateYt); @@ -871,11 +761,6 @@ class TMain : public TMainClassArgs { .DefaultValue(10) .StoreMappedResultT(&RunnerOptions.YdbSettings.HealthCheckTimeout, &TDuration::Seconds); - options.AddLongOption("domain", "Test cluster domain name") - .RequiredArgument("name") - .DefaultValue(RunnerOptions.YdbSettings.DomainName) - .StoreResult(&RunnerOptions.YdbSettings.DomainName); - const auto addTenant = [this](const TString& type, TStorageMeta::TTenant::EType protoType, const NLastGetopt::TOptsParser* option) { TStringBuf tenant; TStringBuf nodesCountStr; @@ -939,18 +824,7 @@ class TMain : public TMainClassArgs { .NoArgument() .SetFlag(&RunnerOptions.YdbSettings.DisableDiskMock); - TChoices> backtrace({ - {"heavy", &NKikimr::EnableYDBBacktraceFormat}, - {"light", []() { SetFormatBackTraceFn(FormatBackTrace); }} - }); - options.AddLongOption("backtrace", "Default backtrace format function") - .RequiredArgument("backtrace-type") - .DefaultValue("heavy") - .Choices(backtrace.GetChoices()) - .Handler1([backtrace](const NLastGetopt::TOptsParser* option) { - TString choice(option->CurValOrDef()); - backtrace(choice)(); - }); + RegisterKikimrOptions(options, RunnerOptions.YdbSettings); } int DoRun(NLastGetopt::TOptsParseResult&&) override { @@ -969,10 +843,7 @@ class TMain : public TMainClassArgs { appConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(ExecutionOptions.ResultsRowsLimit); } - if (DefaultLogPriority) { - appConfig.MutableLogConfig()->SetDefaultLevel(*DefaultLogPriority); - } - ModifyLogPriorities(LogPriorities, *appConfig.MutableLogConfig()); + FillLogConfig(*appConfig.MutableLogConfig()); if (EmulateYt) { const auto& fileStorageConfig = appConfig.GetQueryServiceConfig().GetFileStorage(); @@ -1023,38 +894,6 @@ class TMain : public TMainClassArgs { } }; - -void KqprunTerminateHandler() { - NColorizer::TColors colors = NColorizer::AutoColors(Cerr); - - Cerr << colors.Red() << "======= terminate() call stack ========" << colors.Default() << Endl; - FormatBackTrace(&Cerr); - Cerr << colors.Red() << "=======================================" << colors.Default() << Endl; - - abort(); -} - - -void SegmentationFaultHandler(int) { - NColorizer::TColors colors = NColorizer::AutoColors(Cerr); - - Cerr << colors.Red() << "======= segmentation fault call stack ========" << colors.Default() << Endl; - FormatBackTrace(&Cerr); - Cerr << colors.Red() << "==============================================" << colors.Default() << Endl; - - abort(); -} - -void FloatingPointExceptionHandler(int) { - NColorizer::TColors colors = NColorizer::AutoColors(Cerr); - - Cerr << colors.Red() << "======= floating point exception call stack ========" << colors.Default() << Endl; - FormatBackTrace(&Cerr); - Cerr << colors.Red() << "====================================================" << colors.Default() << Endl; - - abort(); -} - #ifdef PROFILE_MEMORY_ALLOCATIONS void InterruptHandler(int) { NColorizer::TColors colors = NColorizer::AutoColors(Cerr); @@ -1071,9 +910,7 @@ void InterruptHandler(int) { } // namespace NKqpRun int main(int argc, const char* argv[]) { - std::set_terminate(NKqpRun::KqprunTerminateHandler); - signal(SIGSEGV, &NKqpRun::SegmentationFaultHandler); - signal(SIGFPE, &NKqpRun::FloatingPointExceptionHandler); + SetupSignalActions(); #ifdef PROFILE_MEMORY_ALLOCATIONS signal(SIGINT, &NKqpRun::InterruptHandler); diff --git a/ydb/tests/tools/kqprun/runlib/application.cpp b/ydb/tests/tools/kqprun/runlib/application.cpp new file mode 100644 index 000000000000..9a35d584f19e --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/application.cpp @@ -0,0 +1,113 @@ +#include "application.h" +#include "utils.h" + +#include + +#include + +#include + +namespace NKikimrRun { + +void TMainBase::RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettings& settings) { + options.AddLongOption("log-file", "File with execution logs (writes in stderr if empty)") + .RequiredArgument("file") + .StoreResult(&settings.LogOutputFile) + .Handler1([](const NLastGetopt::TOptsParser* option) { + if (const TString& file = option->CurVal()) { + std::remove(file.c_str()); + } + }); + + TChoices logPriority({ + {"emerg", NActors::NLog::EPriority::PRI_EMERG}, + {"alert", NActors::NLog::EPriority::PRI_ALERT}, + {"crit", NActors::NLog::EPriority::PRI_CRIT}, + {"error", NActors::NLog::EPriority::PRI_ERROR}, + {"warn", NActors::NLog::EPriority::PRI_WARN}, + {"notice", NActors::NLog::EPriority::PRI_NOTICE}, + {"info", NActors::NLog::EPriority::PRI_INFO}, + {"debug", NActors::NLog::EPriority::PRI_DEBUG}, + {"trace", NActors::NLog::EPriority::PRI_TRACE}, + }); + options.AddLongOption("log-default", "Default log priority") + .RequiredArgument("priority") + .Choices(logPriority.GetChoices()) + .StoreMappedResultT(&DefaultLogPriority, logPriority); + + options.AddLongOption("log", "Component log priority in format = (e. g. KQP_YQL=trace)") + .RequiredArgument("component priority") + .Handler1([this, logPriority](const NLastGetopt::TOptsParser* option) { + TStringBuf component; + TStringBuf priority; + TStringBuf(option->CurVal()).Split('=', component, priority); + if (component.empty() || priority.empty()) { + ythrow yexception() << "Incorrect log setting, expected form component=priority, e. g. KQP_YQL=trace"; + } + + if (!logPriority.Contains(TString(priority))) { + ythrow yexception() << "Incorrect log priority: " << priority; + } + + const auto service = GetLogService(TString(component)); + if (!LogPriorities.emplace(service, logPriority(TString(priority))).second) { + ythrow yexception() << "Got duplicated log service name: " << component; + } + }); + + options.AddLongOption('M', "monitoring", "Embedded UI port (use 0 to start on random free port), if used will be run as daemon") + .RequiredArgument("uint") + .Handler1([&settings](const NLastGetopt::TOptsParser* option) { + if (const TString& port = option->CurVal()) { + settings.MonitoringEnabled = true; + settings.MonitoringPortOffset = FromString(port); + } + }); + + options.AddLongOption('G', "grpc", "gRPC port (use 0 to start on random free port), if used will be run as daemon") + .RequiredArgument("uint") + .Handler1([&settings](const NLastGetopt::TOptsParser* option) { + if (const TString& port = option->CurVal()) { + settings.GrpcEnabled = true; + settings.GrpcPort = FromString(port); + } + }); + + options.AddLongOption("domain", "Test cluster domain name") + .RequiredArgument("name") + .DefaultValue(settings.DomainName) + .StoreResult(&settings.DomainName); + + TChoices> backtrace({ + {"heavy", &NKikimr::EnableYDBBacktraceFormat}, + {"light", []() { SetFormatBackTraceFn(FormatBackTrace); }} + }); + options.AddLongOption("backtrace", "Default backtrace format function") + .RequiredArgument("backtrace-type") + .DefaultValue("heavy") + .Choices(backtrace.GetChoices()) + .Handler1([backtrace](const NLastGetopt::TOptsParser* option) { + TString choice(option->CurValOrDef()); + backtrace(choice)(); + }); +} + +void TMainBase::FillLogConfig(NKikimrConfig::TLogConfig& config) const { + if (DefaultLogPriority) { + config.SetDefaultLevel(*DefaultLogPriority); + } + ModifyLogPriorities(LogPriorities, config); +} + +IOutputStream* TMainBase::GetDefaultOutput(const TString& file) { + if (file == "-") { + return &Cout; + } + if (file) { + FileHolders.emplace_back(new TFileOutput(file)); + return FileHolders.back().get(); + } + return nullptr; +} + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/application.h b/ydb/tests/tools/kqprun/runlib/application.h new file mode 100644 index 000000000000..6004d30fcb7d --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/application.h @@ -0,0 +1,30 @@ +#pragma once + +#include "settings.h" + +#include + +#include + +#include +#include +#include + +namespace NKikimrRun { + +class TMainBase : public TMainClassArgs { +protected: + void RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettings& settings); + + void FillLogConfig(NKikimrConfig::TLogConfig& config) const; + + static IOutputStream* GetDefaultOutput(const TString& file); + +private: + inline static std::vector> FileHolders; + + std::optional DefaultLogPriority; + std::unordered_map LogPriorities; +}; + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/settings.h b/ydb/tests/tools/kqprun/runlib/settings.h new file mode 100644 index 000000000000..c77caf1876f2 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/settings.h @@ -0,0 +1,25 @@ +#pragma once + +#include + +namespace NKikimrRun { + +struct TServerSettings { + TString DomainName = "Root"; + + bool MonitoringEnabled = false; + ui16 MonitoringPortOffset = 0; + + bool GrpcEnabled = false; + ui16 GrpcPort = 0; + + TString LogOutputFile; +}; + +enum class EResultOutputFormat { + RowsJson, // Rows in json format + FullJson, // Columns, rows and types in json format + FullProto, // Columns, rows and types in proto string format +}; + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/utils.cpp b/ydb/tests/tools/kqprun/runlib/utils.cpp new file mode 100644 index 000000000000..208cf4daf045 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/utils.cpp @@ -0,0 +1,274 @@ +#include "utils.h" + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include + +#include + +namespace NKikimrRun { + +namespace { + +void TerminateHandler() { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << "======= terminate() call stack ========" << colors.Default() << Endl; + FormatBackTrace(&Cerr); + Cerr << colors.Red() << "=======================================" << colors.Default() << Endl; + + abort(); +} + + +void SegmentationFaultHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << "======= segmentation fault call stack ========" << colors.Default() << Endl; + FormatBackTrace(&Cerr); + Cerr << colors.Red() << "==============================================" << colors.Default() << Endl; + + abort(); +} + +void FloatingPointExceptionHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cerr << colors.Red() << "======= floating point exception call stack ========" << colors.Default() << Endl; + FormatBackTrace(&Cerr); + Cerr << colors.Red() << "====================================================" << colors.Default() << Endl; + + abort(); +} + +} // nonymous namespace + + +TRequestResult::TRequestResult() + : Status(Ydb::StatusIds::STATUS_CODE_UNSPECIFIED) +{} + +TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues) + : Status(status) + , Issues(issues) +{} + +TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues) + : Status(status) +{ + NYql::IssuesFromMessage(issues, Issues); +} + +bool TRequestResult::IsSuccess() const { + return Status == Ydb::StatusIds::SUCCESS; +} + +TString TRequestResult::ToString() const { + return TStringBuilder() << "Request finished with status: " << Status << "\nIssues:\n" << Issues.ToString() << "\n"; +} + +TStatsPrinter::TStatsPrinter(NYdb::NConsoleClient::EDataFormat planFormat) + : PlanFormat(planFormat) + , StatProcessor(NFq::CreateStatProcessor("stat_full")) +{} + +void TStatsPrinter::PrintPlan(const TString& plan, IOutputStream& output) const { + if (!plan) { + return; + } + + NJson::TJsonValue planJson; + NJson::ReadJsonTree(plan, &planJson, true); + if (!planJson.GetMapSafe().contains("meta")) { + return; + } + + NYdb::NConsoleClient::TQueryPlanPrinter printer(PlanFormat, true, output); + printer.Print(plan); +} + +void TStatsPrinter::PrintInProgressStatistics(const TString& plan, IOutputStream& output) const { + output << TInstant::Now().ToIsoStringLocal() << " Script in progress statistics" << Endl; + + auto convertedPlan = plan; + try { + convertedPlan = StatProcessor->ConvertPlan(plan); + } catch (const NJson::TJsonException& ex) { + output << "Error plan conversion: " << ex.what() << Endl; + return; + } + + try { + double cpuUsage = 0.0; + auto fullStat = StatProcessor->GetQueryStat(convertedPlan, cpuUsage, nullptr); + auto flatStat = StatProcessor->GetFlatStat(convertedPlan); + auto publicStat = StatProcessor->GetPublicStat(fullStat); + + output << "\nCPU usage: " << cpuUsage << Endl; + PrintStatistics(fullStat, flatStat, publicStat, output); + } catch (const NJson::TJsonException& ex) { + output << "Error stat conversion: " << ex.what() << Endl; + return; + } + + output << "\nPlan visualization:" << Endl; + PrintPlan(convertedPlan, output); +} + +void TStatsPrinter::PrintTimeline(const TString& plan, IOutputStream& output) { + TPlanVisualizer planVisualizer; + planVisualizer.LoadPlans(plan); + output.Write(planVisualizer.PrintSvg()); +} + +void TStatsPrinter::PrintStatistics(const TString& fullStat, const THashMap& flatStat, const NFq::TPublicStat& publicStat, IOutputStream& output) { + output << "\nFlat statistics:" << Endl; + for (const auto& [propery, value] : flatStat) { + TString valueString = ToString(value); + if (propery.Contains("Bytes")) { + valueString = NKikimr::NBlobDepot::FormatByteSize(value); + } else if (propery.Contains("TimeUs")) { + valueString = NFq::FormatDurationUs(value); + } else if (propery.Contains("TimeMs")) { + valueString = NFq::FormatDurationMs(value); + } else { + valueString = FormatNumber(value); + } + output << propery << " = " << valueString << Endl; + } + + output << "\nPublic statistics:" << Endl; + if (auto memoryUsageBytes = publicStat.MemoryUsageBytes) { + output << "MemoryUsage = " << NKikimr::NBlobDepot::FormatByteSize(*memoryUsageBytes) << Endl; + } + if (auto cpuUsageUs = publicStat.CpuUsageUs) { + output << "CpuUsage = " << NFq::FormatDurationUs(*cpuUsageUs) << Endl; + } + if (auto inputBytes = publicStat.InputBytes) { + output << "InputSize = " << NKikimr::NBlobDepot::FormatByteSize(*inputBytes) << Endl; + } + if (auto outputBytes = publicStat.OutputBytes) { + output << "OutputSize = " << NKikimr::NBlobDepot::FormatByteSize(*outputBytes) << Endl; + } + if (auto sourceInputRecords = publicStat.SourceInputRecords) { + output << "SourceInputRecords = " << FormatNumber(*sourceInputRecords) << Endl; + } + if (auto sinkOutputRecords = publicStat.SinkOutputRecords) { + output << "SinkOutputRecords = " << FormatNumber(*sinkOutputRecords) << Endl; + } + if (auto runningTasks = publicStat.RunningTasks) { + output << "RunningTasks = " << FormatNumber(*runningTasks) << Endl; + } + + output << "\nFull statistics:" << Endl; + NJson::TJsonValue statsJson; + NJson::ReadJsonTree(fullStat, &statsJson); + NJson::WriteJson(&output, &statsJson, true, true, true); + output << Endl; +} + +TString TStatsPrinter::FormatNumber(i64 number) { + struct TSeparator : public std::numpunct { + char do_thousands_sep() const final { + return '.'; + } + + std::string do_grouping() const final { + return "\03"; + } + }; + + std::ostringstream stream; + stream.imbue(std::locale(stream.getloc(), new TSeparator())); + stream << number; + return stream.str(); +} + +TString LoadFile(const TString& file) { + return TFileInput(file).ReadAll(); +} + +NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName) { + NKikimrServices::EServiceKikimr service; + if (!NKikimrServices::EServiceKikimr_Parse(serviceName, &service)) { + ythrow yexception() << "Invalid kikimr service name " << serviceName; + } + return service; +} + +void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig) { + for (auto& entry : *logConfig.MutableEntry()) { + const auto it = logPriorities.find(GetLogService(entry.GetComponent())); + if (it != logPriorities.end()) { + entry.SetLevel(it->second); + logPriorities.erase(it); + } + } + for (const auto& [service, priority] : logPriorities) { + auto* entry = logConfig.AddEntry(); + entry->SetComponent(NKikimrServices::EServiceKikimr_Name(service)); + entry->SetLevel(priority); + } +} + +void InitLogSettings(const NKikimrConfig::TLogConfig& logConfig, NActors::TTestActorRuntimeBase& runtime) { + if (logConfig.HasDefaultLevel()) { + auto priority = NActors::NLog::EPriority(logConfig.GetDefaultLevel()); + auto descriptor = NKikimrServices::EServiceKikimr_descriptor(); + for (int i = 0; i < descriptor->value_count(); ++i) { + runtime.SetLogPriority(static_cast(descriptor->value(i)->number()), priority); + } + } + + for (const auto& setting : logConfig.get_arr_entry()) { + runtime.SetLogPriority(GetLogService(setting.GetComponent()), NActors::NLog::EPriority(setting.GetLevel())); + } +} + +void SetupSignalActions() { + std::set_terminate(&TerminateHandler); + signal(SIGSEGV, &SegmentationFaultHandler); + signal(SIGFPE, &FloatingPointExceptionHandler); +} + +void PrintResultSet(EResultOutputFormat format, IOutputStream& output, const Ydb::ResultSet& resultSet) { + switch (format) { + case EResultOutputFormat::RowsJson: { + NYdb::TResultSet result(resultSet); + NYdb::TResultSetParser parser(result); + while (parser.TryNextRow()) { + NJsonWriter::TBuf writer(NJsonWriter::HEM_UNSAFE, &output); + writer.SetWriteNanAsString(true); + NYdb::FormatResultRowJson(parser, result.GetColumnsMeta(), writer, NYdb::EBinaryStringEncoding::Unicode); + output << Endl; + } + break; + } + + case EResultOutputFormat::FullJson: { + resultSet.PrintJSON(output); + output << Endl; + break; + } + + case EResultOutputFormat::FullProto: { + TString resultSetString; + google::protobuf::TextFormat::Printer printer; + printer.SetSingleLineMode(false); + printer.SetUseUtf8StringEscaping(true); + printer.PrintToString(resultSet, &resultSetString); + output << resultSetString; + break; + } + } +} + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/utils.h b/ydb/tests/tools/kqprun/runlib/utils.h new file mode 100644 index 000000000000..9e4d69c14501 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/utils.h @@ -0,0 +1,95 @@ +#pragma once + +#include "settings.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace NKikimrRun { + +struct TRequestResult { + Ydb::StatusIds::StatusCode Status; + NYql::TIssues Issues; + + TRequestResult(); + + TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues); + + TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues); + + bool IsSuccess() const; + + TString ToString() const; +}; + +template +class TChoices { +public: + explicit TChoices(std::map choicesMap) + : ChoicesMap(std::move(choicesMap)) + {} + + TResult operator()(const TString& choice) const { + return ChoicesMap.at(choice); + } + + TVector GetChoices() const { + TVector choices; + choices.reserve(ChoicesMap.size()); + for (const auto& [choice, _] : ChoicesMap) { + choices.emplace_back(choice); + } + return choices; + } + + bool Contains(const TString& choice) const { + return ChoicesMap.contains(choice); + } + +private: + const std::map ChoicesMap; +}; + +class TStatsPrinter { +public: + explicit TStatsPrinter(NYdb::NConsoleClient::EDataFormat planFormat); + + void PrintPlan(const TString& plan, IOutputStream& output) const; + + void PrintInProgressStatistics(const TString& plan, IOutputStream& output) const; + + static void PrintTimeline(const TString& plan, IOutputStream& output); + + static void PrintStatistics(const TString& fullStat, const THashMap& flatStat, const NFq::TPublicStat& publicStat, IOutputStream& output); + + // Function adds thousands separators + // 123456789 -> 123.456.789 + static TString FormatNumber(i64 number); + +private: + const NYdb::NConsoleClient::EDataFormat PlanFormat; + const std::unique_ptr StatProcessor; +}; + +TString LoadFile(const TString& file); + +NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName); + +void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig); + +void InitLogSettings(const NKikimrConfig::TLogConfig& logConfig, NActors::TTestActorRuntimeBase& runtime); + +void SetupSignalActions(); + +void PrintResultSet(EResultOutputFormat format, IOutputStream& output, const Ydb::ResultSet& resultSet); + +} // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/ya.make b/ydb/tests/tools/kqprun/runlib/ya.make new file mode 100644 index 000000000000..352a3a423105 --- /dev/null +++ b/ydb/tests/tools/kqprun/runlib/ya.make @@ -0,0 +1,28 @@ +LIBRARY() + +SRCS( + application.cpp + utils.cpp +) + +PEERDIR( + library/cpp/colorizer + library/cpp/getopt + library/cpp/json + util + ydb/core/base + ydb/core/blob_depot + ydb/core/fq/libs/compute/common + ydb/core/protos + ydb/library/actors/core + ydb/library/actors/testlib + ydb/library/services + ydb/public/api/protos + ydb/public/lib/json_value + ydb/public/lib/ydb_cli/common + yql/essentials/public/issue +) + +YQL_LAST_ABI_VERSION() + +END() diff --git a/ydb/tests/tools/kqprun/src/common.cpp b/ydb/tests/tools/kqprun/src/common.cpp deleted file mode 100644 index 0241ca6a932c..000000000000 --- a/ydb/tests/tools/kqprun/src/common.cpp +++ /dev/null @@ -1,29 +0,0 @@ -#include "common.h" - - -namespace NKqpRun { - -NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName) { - NKikimrServices::EServiceKikimr service; - if (!NKikimrServices::EServiceKikimr_Parse(serviceName, &service)) { - ythrow yexception() << "Invalid kikimr service name " << serviceName; - } - return service; -} - -void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig) { - for (auto& entry : *logConfig.MutableEntry()) { - const auto it = logPriorities.find(GetLogService(entry.GetComponent())); - if (it != logPriorities.end()) { - entry.SetLevel(it->second); - logPriorities.erase(it); - } - } - for (const auto& [service, priority] : logPriorities) { - auto* entry = logConfig.AddEntry(); - entry->SetComponent(NKikimrServices::EServiceKikimr_Name(service)); - entry->SetLevel(priority); - } -} - -} // namespace NKqpRun diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h index 4f4f4ac36913..427dc6c03947 100644 --- a/ydb/tests/tools/kqprun/src/common.h +++ b/ydb/tests/tools/kqprun/src/common.h @@ -2,12 +2,11 @@ #include #include - #include #include - #include #include +#include #include @@ -32,7 +31,7 @@ struct TAsyncQueriesSettings { EVerbose Verbose = EVerbose::EachQuery; }; -struct TYdbSetupSettings { +struct TYdbSetupSettings : public NKikimrRun::TServerSettings { enum class EVerbose { None, Info, @@ -49,7 +48,6 @@ struct TYdbSetupSettings { }; ui32 NodeCount = 1; - TString DomainName = "Root"; std::map Tenants; TDuration HealthCheckTimeout = TDuration::Seconds(10); EHealthCheck HealthCheckLevel = EHealthCheck::NodesCount; @@ -60,14 +58,7 @@ struct TYdbSetupSettings { std::optional PDisksPath; std::optional DiskSize; - bool MonitoringEnabled = false; - ui16 MonitoringPortOffset = 0; - - bool GrpcEnabled = false; - ui16 GrpcPort = 0; - bool TraceOptEnabled = false; - TString LogOutputFile; EVerbose VerboseLevel = EVerbose::Info; TString YqlToken; @@ -87,12 +78,6 @@ struct TRunnerOptions { All, }; - enum class EResultOutputFormat { - RowsJson, // Rows in json format - FullJson, // Columns, rows and types in json format - FullProto, // Columns, rows and types in proto string format - }; - IOutputStream* ResultOutput = nullptr; IOutputStream* SchemeQueryAstOutput = nullptr; std::vector ScriptQueryAstOutputs; @@ -100,7 +85,7 @@ struct TRunnerOptions { std::vector ScriptQueryTimelineFiles; std::vector InProgressStatisticsOutputFiles; - EResultOutputFormat ResultOutputFormat = EResultOutputFormat::RowsJson; + NKikimrRun::EResultOutputFormat ResultOutputFormat = NKikimrRun::EResultOutputFormat::RowsJson; NYdb::NConsoleClient::EDataFormat PlanOutputFormat = NYdb::NConsoleClient::EDataFormat::Default; ETraceOptType TraceOptType = ETraceOptType::Disabled; std::optional TraceOptScriptId; @@ -130,7 +115,4 @@ TValue GetValue(size_t index, const std::vector& values, TValue defaultV return values[std::min(index, values.size() - 1)]; } -NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName); -void ModifyLogPriorities(std::unordered_map logPriorities, NKikimrConfig::TLogConfig& logConfig); - } // namespace NKqpRun diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index ef98165db8ca..f2378965461b 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -2,88 +2,11 @@ #include "ydb_setup.h" #include -#include - -#include -#include - -#include -#include -#include +using namespace NKikimrRun; namespace NKqpRun { -namespace { - -// Function adds thousands separators -// 123456789 -> 123.456.789 -TString FormatNumber(i64 number) { - struct TSeparator : public std::numpunct { - char do_thousands_sep() const final { - return '.'; - } - - std::string do_grouping() const final { - return "\03"; - } - }; - - std::ostringstream stream; - stream.imbue(std::locale(stream.getloc(), new TSeparator())); - stream << number; - return stream.str(); -} - -void PrintStatistics(const TString& fullStat, const THashMap& flatStat, const NFq::TPublicStat& publicStat, IOutputStream& output) { - output << "\nFlat statistics:" << Endl; - for (const auto& [propery, value] : flatStat) { - TString valueString = ToString(value); - if (propery.Contains("Bytes")) { - valueString = NKikimr::NBlobDepot::FormatByteSize(value); - } else if (propery.Contains("TimeUs")) { - valueString = NFq::FormatDurationUs(value); - } else if (propery.Contains("TimeMs")) { - valueString = NFq::FormatDurationMs(value); - } else { - valueString = FormatNumber(value); - } - output << propery << " = " << valueString << Endl; - } - - output << "\nPublic statistics:" << Endl; - if (auto memoryUsageBytes = publicStat.MemoryUsageBytes) { - output << "MemoryUsage = " << NKikimr::NBlobDepot::FormatByteSize(*memoryUsageBytes) << Endl; - } - if (auto cpuUsageUs = publicStat.CpuUsageUs) { - output << "CpuUsage = " << NFq::FormatDurationUs(*cpuUsageUs) << Endl; - } - if (auto inputBytes = publicStat.InputBytes) { - output << "InputSize = " << NKikimr::NBlobDepot::FormatByteSize(*inputBytes) << Endl; - } - if (auto outputBytes = publicStat.OutputBytes) { - output << "OutputSize = " << NKikimr::NBlobDepot::FormatByteSize(*outputBytes) << Endl; - } - if (auto sourceInputRecords = publicStat.SourceInputRecords) { - output << "SourceInputRecords = " << FormatNumber(*sourceInputRecords) << Endl; - } - if (auto sinkOutputRecords = publicStat.SinkOutputRecords) { - output << "SinkOutputRecords = " << FormatNumber(*sinkOutputRecords) << Endl; - } - if (auto runningTasks = publicStat.RunningTasks) { - output << "RunningTasks = " << FormatNumber(*runningTasks) << Endl; - } - - output << "\nFull statistics:" << Endl; - NJson::TJsonValue statsJson; - NJson::ReadJsonTree(fullStat, &statsJson); - NJson::WriteJson(&output, &statsJson, true, true, true); - output << Endl; -} - -} // anonymous namespace - - //// TKqpRunner::TImpl class TKqpRunner::TImpl { @@ -100,7 +23,7 @@ class TKqpRunner::TImpl { : Options_(options) , VerboseLevel_(Options_.YdbSettings.VerboseLevel) , YdbSetup_(options.YdbSettings) - , StatProcessor_(NFq::CreateStatProcessor("stat_full")) + , StatsPrinter_(Options_.PlanOutputFormat) , CerrColors_(NColorizer::AutoColors(Cerr)) , CoutColors_(NColorizer::AutoColors(Cout)) {} @@ -242,7 +165,7 @@ class TKqpRunner::TImpl { if (ResultSets_.size() > 1 && Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { *Options_.ResultOutput << CoutColors_.Cyan() << "Result set " << i + 1 << ":" << CoutColors_.Default() << Endl; } - PrintScriptResult(ResultSets_[i]); + PrintResultSet(Options_.ResultOutputFormat, *Options_.ResultOutput, ResultSets_[i]); } } } @@ -337,66 +260,24 @@ class TKqpRunner::TImpl { } } - void PrintPlan(const TString& plan, IOutputStream* output) const { - if (!plan) { - return; - } - - NJson::TJsonValue planJson; - NJson::ReadJsonTree(plan, &planJson, true); - if (!planJson.GetMapSafe().contains("meta")) { - return; - } - - NYdb::NConsoleClient::TQueryPlanPrinter printer(Options_.PlanOutputFormat, true, *output); - printer.Print(plan); - } - void PrintScriptPlan(size_t queryId, const TString& plan) const { if (const auto output = GetValue(queryId, Options_.ScriptQueryPlanOutputs, nullptr)) { if (Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; } - PrintPlan(plan, output); + StatsPrinter_.PrintPlan(plan, *output); } } void PrintScriptProgress(size_t queryId, const TString& plan) const { if (const auto& output = GetValue(queryId, Options_.InProgressStatisticsOutputFiles, {})) { TFileOutput outputStream(output); - outputStream << TInstant::Now().ToIsoStringLocal() << " Script in progress statistics" << Endl; - - auto convertedPlan = plan; - try { - convertedPlan = StatProcessor_->ConvertPlan(plan); - } catch (const NJson::TJsonException& ex) { - outputStream << "Error plan conversion: " << ex.what() << Endl; - } - - try { - double cpuUsage = 0.0; - auto fullStat = StatProcessor_->GetQueryStat(convertedPlan, cpuUsage, nullptr); - auto flatStat = StatProcessor_->GetFlatStat(convertedPlan); - auto publicStat = StatProcessor_->GetPublicStat(fullStat); - - outputStream << "\nCPU usage: " << cpuUsage << Endl; - PrintStatistics(fullStat, flatStat, publicStat, outputStream); - } catch (const NJson::TJsonException& ex) { - outputStream << "Error stat conversion: " << ex.what() << Endl; - } - - outputStream << "\nPlan visualization:" << Endl; - PrintPlan(convertedPlan, &outputStream); - + StatsPrinter_.PrintInProgressStatistics(plan, outputStream); outputStream.Finish(); } if (const auto& output = GetValue(queryId, Options_.ScriptQueryTimelineFiles, {})) { TFileOutput outputStream(output); - - TPlanVisualizer planVisualizer; - planVisualizer.LoadPlans(plan); - outputStream.Write(planVisualizer.PrintSvg()); - + StatsPrinter_.PrintTimeline(plan, outputStream); outputStream.Finish(); } } @@ -409,36 +290,6 @@ class TKqpRunner::TImpl { }; } - void PrintScriptResult(const Ydb::ResultSet& resultSet) const { - switch (Options_.ResultOutputFormat) { - case TRunnerOptions::EResultOutputFormat::RowsJson: { - NYdb::TResultSet result(resultSet); - NYdb::TResultSetParser parser(result); - while (parser.TryNextRow()) { - NJsonWriter::TBuf writer(NJsonWriter::HEM_UNSAFE, Options_.ResultOutput); - writer.SetWriteNanAsString(true); - NYdb::FormatResultRowJson(parser, result.GetColumnsMeta(), writer, NYdb::EBinaryStringEncoding::Unicode); - *Options_.ResultOutput << Endl; - } - break; - } - - case TRunnerOptions::EResultOutputFormat::FullJson: - resultSet.PrintJSON(*Options_.ResultOutput); - *Options_.ResultOutput << Endl; - break; - - case TRunnerOptions::EResultOutputFormat::FullProto: - TString resultSetString; - google::protobuf::TextFormat::Printer printer; - printer.SetSingleLineMode(false); - printer.SetUseUtf8StringEscaping(true); - printer.PrintToString(resultSet, &resultSetString); - *Options_.ResultOutput << resultSetString; - break; - } - } - void PrintScriptFinish(const TQueryMeta& meta, const TString& queryType) const { if (Options_.YdbSettings.VerboseLevel < EVerbose::Info) { return; @@ -457,7 +308,7 @@ class TKqpRunner::TImpl { EVerbose VerboseLevel_; TYdbSetup YdbSetup_; - std::unique_ptr StatProcessor_; + TStatsPrinter StatsPrinter_; NColorizer::TColors CerrColors_; NColorizer::TColors CoutColors_; diff --git a/ydb/tests/tools/kqprun/src/ya.make b/ydb/tests/tools/kqprun/src/ya.make index 59097668037a..ed5f5a619bd5 100644 --- a/ydb/tests/tools/kqprun/src/ya.make +++ b/ydb/tests/tools/kqprun/src/ya.make @@ -2,7 +2,6 @@ LIBRARY() SRCS( actors.cpp - common.cpp kqp_runner.cpp ydb_setup.cpp ) @@ -10,6 +9,7 @@ SRCS( PEERDIR( ydb/core/testlib + ydb/tests/tools/kqprun/runlib ydb/tests/tools/kqprun/src/proto ) diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index ea0b0ce1c0a8..8ff68914a656 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -14,6 +14,7 @@ #include +using namespace NKikimrRun; namespace NKqpRun { @@ -147,18 +148,7 @@ class TYdbSetup::TImpl { void SetLoggerSettings(NKikimr::Tests::TServerSettings& serverSettings) const { auto loggerInitializer = [this](NActors::TTestActorRuntime& runtime) { - if (Settings_.AppConfig.GetLogConfig().HasDefaultLevel()) { - auto priority = NActors::NLog::EPriority(Settings_.AppConfig.GetLogConfig().GetDefaultLevel()); - auto descriptor = NKikimrServices::EServiceKikimr_descriptor(); - for (int i = 0; i < descriptor->value_count(); ++i) { - runtime.SetLogPriority(static_cast(descriptor->value(i)->number()), priority); - } - } - - for (const auto& setting : Settings_.AppConfig.GetLogConfig().get_arr_entry()) { - runtime.SetLogPriority(GetLogService(setting.GetComponent()), NActors::NLog::EPriority(setting.GetLevel())); - } - + InitLogSettings(Settings_.AppConfig.GetLogConfig(), runtime); runtime.SetLogBackendFactory([this]() { return CreateLogBackend(); }); }; @@ -696,32 +686,6 @@ class TYdbSetup::TImpl { }; -//// TRequestResult - -TRequestResult::TRequestResult() - : Status(Ydb::StatusIds::STATUS_CODE_UNSPECIFIED) -{} - -TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues) - : Status(status) - , Issues(issues) -{} - -TRequestResult::TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues) - : Status(status) -{ - NYql::IssuesFromMessage(issues, Issues); -} - -bool TRequestResult::IsSuccess() const { - return Status == Ydb::StatusIds::SUCCESS; -} - -TString TRequestResult::ToString() const { - return TStringBuilder() << "Request finished with status: " << Status << "\nIssues:\n" << Issues.ToString() << "\n"; -} - - //// TYdbSetup TYdbSetup::TYdbSetup(const TYdbSetupSettings& settings) diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.h b/ydb/tests/tools/kqprun/src/ydb_setup.h index 393873b511bf..4edd849b9d26 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.h +++ b/ydb/tests/tools/kqprun/src/ydb_setup.h @@ -4,6 +4,7 @@ #include "actors.h" #include +#include namespace NKqpRun { @@ -30,23 +31,9 @@ struct TExecutionMeta : public TQueryMeta { }; -struct TRequestResult { - Ydb::StatusIds::StatusCode Status; - NYql::TIssues Issues; - - TRequestResult(); - - TRequestResult(Ydb::StatusIds::StatusCode status, const NYql::TIssues& issues); - - TRequestResult(Ydb::StatusIds::StatusCode status, const google::protobuf::RepeatedPtrField& issues); - - bool IsSuccess() const; - - TString ToString() const; -}; - - class TYdbSetup { + using TRequestResult = NKikimrRun::TRequestResult; + public: explicit TYdbSetup(const TYdbSetupSettings& settings); diff --git a/ydb/tests/tools/kqprun/ya.make b/ydb/tests/tools/kqprun/ya.make index 05b119d8e099..81689d3bd5a3 100644 --- a/ydb/tests/tools/kqprun/ya.make +++ b/ydb/tests/tools/kqprun/ya.make @@ -17,6 +17,7 @@ PEERDIR( yt/yql/providers/yt/gateway/file yql/essentials/sql/pg + ydb/tests/tools/kqprun/runlib ydb/tests/tools/kqprun/src ) diff --git a/ydb/tests/tools/ya.make b/ydb/tests/tools/ya.make index b7586673bba3..bcae74c00dc2 100644 --- a/ydb/tests/tools/ya.make +++ b/ydb/tests/tools/ya.make @@ -2,6 +2,7 @@ RECURSE( canondata_sync datastreams_helpers fq_runner + fqrun idx_test kqprun mdb_mock From 5e5de154ac837a9fa541ea39ce3368536de15fc7 Mon Sep 17 00:00:00 2001 From: Dmitry Kardymon Date: Thu, 19 Dec 2024 14:42:37 +0300 Subject: [PATCH 07/18] YQ-3721 PQ sink: write to file (#10501) Co-authored-by: Fiodar Miron Co-authored-by: yumkam --- ydb/core/fq/libs/init/init.cpp | 5 +- .../pq/async_io/dq_pq_write_actor.cpp | 28 ++- .../providers/pq/async_io/dq_pq_write_actor.h | 4 +- .../dummy/yql_pq_file_topic_client.cpp | 182 +++++++++++++++++- .../providers/pq/provider/ut/yql_pq_ut.cpp | 2 +- ydb/library/yql/tools/dq/worker_node/main.cpp | 5 +- ydb/library/yql/tools/dqrun/dqrun.cpp | 2 +- ydb/library/yql/tools/mrrun/mrrun.cpp | 5 +- ydb/tests/fq/pq_async_io/ut_helpers.cpp | 9 + 9 files changed, 219 insertions(+), 23 deletions(-) diff --git a/ydb/core/fq/libs/init/init.cpp b/ydb/core/fq/libs/init/init.cpp index 23c139aea2b6..edca3f733212 100644 --- a/ydb/core/fq/libs/init/init.cpp +++ b/ydb/core/fq/libs/init/init.cpp @@ -225,7 +225,8 @@ void Init( std::make_shared(protoConfig.GetGateways().GetPq()), appData->FunctionRegistry ); - RegisterDqPqReadActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, NYql::CreatePqNativeGateway(std::move(pqServices)), + auto pqGateway = NYql::CreatePqNativeGateway(std::move(pqServices)); + RegisterDqPqReadActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, pqGateway, yqCounters->GetSubgroup("subsystem", "DqSourceTracker"), protoConfig.GetCommon().GetPqReconnectPeriod()); s3ActorsFactory->RegisterS3ReadActorFactory(*asyncIoFactory, credentialsFactory, httpGateway, s3HttpRetryPolicy, readActorFactoryCfg, @@ -234,7 +235,7 @@ void Init( httpGateway, s3HttpRetryPolicy); RegisterGenericProviderFactories(*asyncIoFactory, credentialsFactory, connectorClient); - RegisterDqPqWriteActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, yqCounters->GetSubgroup("subsystem", "DqSinkTracker")); + RegisterDqPqWriteActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, pqGateway, yqCounters->GetSubgroup("subsystem", "DqSinkTracker")); RegisterDQSolomonWriteActorFactory(*asyncIoFactory, credentialsFactory); } diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp index 62e105cb0978..21ef29fcc48a 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp @@ -142,7 +142,8 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu std::shared_ptr credentialsProviderFactory, IDqComputeActorAsyncOutput::ICallbacks* callbacks, const ::NMonitoring::TDynamicCounterPtr& counters, - i64 freeSpace) + i64 freeSpace, + const IPqGateway::TPtr& pqGateway) : TActor(&TDqPqWriteActor::StateFunc) , OutputIndex(outputIndex) , TxId(txId) @@ -153,7 +154,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu , Callbacks(callbacks) , LogPrefix(TStringBuilder() << "SelfId: " << this->SelfId() << ", TxId: " << TxId << ", TaskId: " << taskId << ", PQ sink. ") , FreeSpace(freeSpace) - , TopicClient(Driver, GetTopicClientSettings()) + , PqGateway(pqGateway) { EgressStats.Level = statsLevel; } @@ -303,6 +304,13 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu : NYdb::NTopic::ECodec::GZIP); } + ITopicClient& GetTopicClient() { + if (!TopicClient) { + TopicClient = PqGateway->GetTopicClient(Driver, GetTopicClientSettings()); + } + return *TopicClient; + } + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() { return NYdb::NTopic::TTopicClientSettings() .Database(SinkParams.GetDatabase()) @@ -317,7 +325,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu void CreateSessionIfNotExists() { if (!WriteSession) { - WriteSession = TopicClient.CreateWriteSession(GetWriteSessionSettings()); + WriteSession = GetTopicClient().CreateWriteSession(GetWriteSessionSettings()); SubscribeOnNextEvent(); } } @@ -475,7 +483,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu i64 FreeSpace = 0; bool Finished = false; - NYdb::NTopic::TTopicClient TopicClient; + ITopicClient::TPtr TopicClient; std::shared_ptr WriteSession; TString SourceId; ui64 NextSeqNo = 1; @@ -486,6 +494,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu std::queue Buffer; std::queue WaitingAcks; // Size of items which are waiting for acks (used to update free space) std::queue> DeferredCheckpoints; + IPqGateway::TPtr PqGateway; }; std::pair CreateDqPqWriteActor( @@ -499,6 +508,7 @@ std::pair CreateDqPqWriteActor( ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, IDqComputeActorAsyncOutput::ICallbacks* callbacks, const ::NMonitoring::TDynamicCounterPtr& counters, + IPqGateway::TPtr pqGateway, i64 freeSpace) { const TString& tokenName = settings.GetToken().GetName(); @@ -515,13 +525,14 @@ std::pair CreateDqPqWriteActor( CreateCredentialsProviderFactoryForStructuredToken(credentialsFactory, token, addBearerToToken), callbacks, counters, - freeSpace); + freeSpace, + pqGateway); return {actor, actor}; } -void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const ::NMonitoring::TDynamicCounterPtr& counters) { +void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const IPqGateway::TPtr& pqGateway, const ::NMonitoring::TDynamicCounterPtr& counters) { factory.RegisterSink("PqSink", - [driver = std::move(driver), credentialsFactory = std::move(credentialsFactory), counters]( + [driver = std::move(driver), credentialsFactory = std::move(credentialsFactory), counters, pqGateway]( NPq::NProto::TDqPqTopicSink&& settings, IDqAsyncIoFactory::TSinkArguments&& args) { @@ -536,7 +547,8 @@ void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver dri driver, credentialsFactory, args.Callback, - counters + counters, + pqGateway ); }); } diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h index 3af8c6e63ec8..39f388a79dcc 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -31,8 +32,9 @@ std::pair CreateDqPqWriteActor( ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, IDqComputeActorAsyncOutput::ICallbacks* callbacks, const ::NMonitoring::TDynamicCounterPtr& counters, + IPqGateway::TPtr pqGateway, i64 freeSpace = DqPqDefaultFreeSpace); -void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const ::NMonitoring::TDynamicCounterPtr& counters = MakeIntrusive<::NMonitoring::TDynamicCounters>()); +void RegisterDqPqWriteActorFactory(TDqAsyncIoFactory& factory, NYdb::TDriver driver, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, const IPqGateway::TPtr& pqGateway, const ::NMonitoring::TDynamicCounterPtr& counters = MakeIntrusive<::NMonitoring::TDynamicCounters>()); } // namespace NYql::NDq diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp index 9cb64a242f65..abbaf50e9d04 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp @@ -36,8 +36,8 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); Y_UNUSED(maxByteSize); TVector res; - for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() <= maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { - res.push_back(*event); + for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { + res.push_back(std::move(*event)); } return res; } @@ -59,7 +59,6 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); bool Close(TDuration timeout = TDuration::Max()) override { Y_UNUSED(timeout); - // TOOD send TSessionClosedEvent EventsQ_.Stop(); Pool_.Stop(); @@ -133,7 +132,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); Sleep(FILE_POLL_PERIOD); } } - + TFile File_; TBlockingEQueue EventsQ_ {4_MB}; NYdb::NTopic::TPartitionSession::TPtr Session_; @@ -146,6 +145,172 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); ui64 SeqNo_ = 0; }; +class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb::NTopic::TContinuationTokenIssuer { +public: + TFileTopicWriteSession(TFile file): + File_(std::move(file)), FileWriter_([this] () { + PushToFile(); + }), Counters_() + { + Pool_.Start(1); + EventsQ_.Push(NYdb::NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); + } + + NThreading::TFuture WaitEvent() override { + return NThreading::Async([this] () { + EventsQ_.BlockUntilEvent(); + return NThreading::MakeFuture(); + }, Pool_); + } + + TMaybe GetEvent(bool block) override { + return EventsQ_.Pop(block); + } + + TVector GetEvents(bool block, TMaybe maxEventsCount) override { + TVector res; + for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { + res.push_back(std::move(*event)); + } + return res; + } + + NThreading::TFuture GetInitSeqNo() override { + return NThreading::MakeFuture(SeqNo_); + } + + void Write(NYdb::NTopic::TContinuationToken&&, NYdb::NTopic::TWriteMessage&& message, + NYdb::NTable::TTransaction* tx) override { + Y_UNUSED(tx); + + auto size = message.Data.size(); + EventsMsgQ_.Push(TOwningWriteMessage(std::move(message)), size); + } + + void Write(NYdb::NTopic::TContinuationToken&& token, TStringBuf data, TMaybe seqNo, + TMaybe createTimestamp) override { + NYdb::NTopic::TWriteMessage message(data); + if (seqNo.Defined()) { + message.SeqNo(*seqNo); + } + if (createTimestamp.Defined()) { + message.CreateTimestamp(*createTimestamp); + } + + Write(std::move(token), std::move(message), nullptr); + } + + // Ignores codec in message and always writes raw for debugging purposes + void WriteEncoded(NYdb::NTopic::TContinuationToken&& token, NYdb::NTopic::TWriteMessage&& params, + NYdb::NTable::TTransaction* tx) override { + Y_UNUSED(tx); + + NYdb::NTopic::TWriteMessage message(params.Data); + + if (params.CreateTimestamp_.Defined()) { + message.CreateTimestamp(*params.CreateTimestamp_); + } + if (params.SeqNo_) { + message.SeqNo(*params.SeqNo_); + } + message.MessageMeta(params.MessageMeta_); + + Write(std::move(token), std::move(message), nullptr); + } + + // Ignores codec in message and always writes raw for debugging purposes + void WriteEncoded(NYdb::NTopic::TContinuationToken&& token, TStringBuf data, NYdb::NTopic::ECodec codec, ui32 originalSize, + TMaybe seqNo, TMaybe createTimestamp) override { + Y_UNUSED(codec); + Y_UNUSED(originalSize); + + NYdb::NTopic::TWriteMessage message(data); + if (seqNo.Defined()) { + message.SeqNo(*seqNo); + } + if (createTimestamp.Defined()) { + message.CreateTimestamp(*createTimestamp); + } + + Write(std::move(token), std::move(message), nullptr); + } + + bool Close(TDuration timeout = TDuration::Max()) override { + Y_UNUSED(timeout); + EventsQ_.Stop(); + EventsMsgQ_.Stop(); + Pool_.Stop(); + + if (FileWriter_.joinable()) { + FileWriter_.join(); + } + return true; + } + + NYdb::NTopic::TWriterCounters::TPtr GetCounters() override { + return Counters_; + } + + ~TFileTopicWriteSession() override { + EventsQ_.Stop(); + EventsMsgQ_.Stop(); + Pool_.Stop(); + if (FileWriter_.joinable()) { + FileWriter_.join(); + } + } + +private: + void PushToFile() { + TFileOutput fo(File_); + ui64 offset = 0; // FIXME dummy + ui64 partitionId = 0; // FIXME dummy + while (auto maybeMsg = EventsMsgQ_.Pop(true)) { + NYdb::NTopic::TWriteSessionEvent::TAcksEvent acks; + do { + auto& [content, msg] = *maybeMsg; + NYdb::NTopic::TWriteSessionEvent::TWriteAck ack; + if (msg.SeqNo_.Defined()) { // FIXME should be auto generated otherwise + ack.SeqNo = *msg.SeqNo_; + } + ack.State = NYdb::NTopic::TWriteSessionEvent::TWriteAck::EES_WRITTEN; + ack.Details.ConstructInPlace(offset, partitionId); + acks.Acks.emplace_back(std::move(ack)); + offset += content.size() + 1; + fo.Write(content); + fo.Write('\n'); + } while ((maybeMsg = EventsMsgQ_.Pop(false))); + fo.Flush(); + EventsQ_.Push(std::move(acks), 1 + acks.Acks.size()); + EventsQ_.Push(NYdb::NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}, 1); + if (EventsQ_.IsStopped()) { + break; + } + } + } + + TFile File_; + + // We acquire ownership of messages immediately + // TODO: remove extra message copying to and from queue + struct TOwningWriteMessage { + TString content; + NYdb::NTopic::TWriteMessage msg; + + TOwningWriteMessage(NYdb::NTopic::TWriteMessage&& msg): content(msg.Data), msg(std::move(msg)) { + msg.Data = content; + } + }; + TBlockingEQueue EventsMsgQ_ {4_MB}; + + TBlockingEQueue EventsQ_ {128_KB}; + std::thread FileWriter_; + + TThreadPool Pool_; + NYdb::NTopic::TWriterCounters::TPtr Counters_; + ui64 SeqNo_ = 0; +}; + struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession { TDummyPartitionSession(ui64 sessionId, const TString& topicPath, ui64 partId) { PartitionSessionId = sessionId; @@ -231,8 +396,13 @@ std::shared_ptr TFileTopicClient::Cre } std::shared_ptr TFileTopicClient::CreateWriteSession(const NYdb::NTopic::TWriteSessionSettings& settings) { - Y_UNUSED(settings); - return nullptr; + TString topicPath = settings.Path_; + auto topicsIt = Topics_.find(make_pair("pq", topicPath)); + Y_ENSURE(topicsIt != Topics_.end()); + auto filePath = topicsIt->second.FilePath; + Y_ENSURE(filePath); + + return std::make_shared(TFile(*filePath, EOpenMode::TEnum::RdWr)); } NYdb::TAsyncStatus TFileTopicClient::CommitOffset(const TString& path, ui64 partitionId, const TString& consumerName, ui64 offset, diff --git a/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp b/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp index d6f1bd5a46d2..3cd2c1538d4d 100644 --- a/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp +++ b/ydb/library/yql/providers/pq/provider/ut/yql_pq_ut.cpp @@ -45,7 +45,7 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory(const NYdb::TDriver& driver, c auto factory = MakeIntrusive(); RegisterDqPqReadActorFactory(*factory, driver, nullptr, pqGateway); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); return factory; } diff --git a/ydb/library/yql/tools/dq/worker_node/main.cpp b/ydb/library/yql/tools/dq/worker_node/main.cpp index f43e31656864..9a349b0f493c 100644 --- a/ydb/library/yql/tools/dq/worker_node/main.cpp +++ b/ydb/library/yql/tools/dq/worker_node/main.cpp @@ -115,11 +115,12 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory(const NYdb::TDriver& driver, I std::make_shared(), nullptr ); - RegisterDqPqReadActorFactory(*factory, driver, nullptr, CreatePqNativeGateway(std::move(pqServices))); + auto pqGateway = CreatePqNativeGateway(std::move(pqServices)); + RegisterDqPqReadActorFactory(*factory, driver, nullptr, pqGateway); RegisterYdbReadActorFactory(*factory, driver, nullptr); RegisterClickHouseReadActorFactory(*factory, nullptr, httpGateway); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); auto retryPolicy = GetHTTPDefaultRetryPolicy(); diff --git a/ydb/library/yql/tools/dqrun/dqrun.cpp b/ydb/library/yql/tools/dqrun/dqrun.cpp index 64af597ec8b2..0b109a3af273 100644 --- a/ydb/library/yql/tools/dqrun/dqrun.cpp +++ b/ydb/library/yql/tools/dqrun/dqrun.cpp @@ -311,7 +311,7 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory( RegisterDQSolomonReadActorFactory(*factory, nullptr); RegisterClickHouseReadActorFactory(*factory, nullptr, httpGateway); RegisterGenericProviderFactories(*factory, credentialsFactory, genericClient); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); s3ActorsFactory->RegisterS3WriteActorFactory(*factory, nullptr, httpGateway, GetHTTPDefaultRetryPolicy()); diff --git a/ydb/library/yql/tools/mrrun/mrrun.cpp b/ydb/library/yql/tools/mrrun/mrrun.cpp index 8e1a09477269..37fdd07c5802 100644 --- a/ydb/library/yql/tools/mrrun/mrrun.cpp +++ b/ydb/library/yql/tools/mrrun/mrrun.cpp @@ -245,10 +245,11 @@ NDq::IDqAsyncIoFactory::TPtr CreateAsyncIoFactory(const NYdb::TDriver& driver, I std::make_shared(), nullptr ); - RegisterDqPqReadActorFactory(*factory, driver, nullptr, CreatePqNativeGateway(std::move(pqServices))); + auto pqGateway = CreatePqNativeGateway(std::move(pqServices)); + RegisterDqPqReadActorFactory(*factory, driver, nullptr, pqGateway); RegisterYdbReadActorFactory(*factory, driver, nullptr); RegisterClickHouseReadActorFactory(*factory, nullptr, httpGateway); - RegisterDqPqWriteActorFactory(*factory, driver, nullptr); + RegisterDqPqWriteActorFactory(*factory, driver, nullptr, pqGateway); auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory(); auto retryPolicy = GetHTTPDefaultRetryPolicy(); diff --git a/ydb/tests/fq/pq_async_io/ut_helpers.cpp b/ydb/tests/fq/pq_async_io/ut_helpers.cpp index df15829e670d..e011029479ca 100644 --- a/ydb/tests/fq/pq_async_io/ut_helpers.cpp +++ b/ydb/tests/fq/pq_async_io/ut_helpers.cpp @@ -124,6 +124,14 @@ void TPqIoTestFixture::InitAsyncOutput( { const THashMap secureParams; + TPqGatewayServices pqServices( + Driver, + nullptr, + nullptr, + std::make_shared(), + nullptr + ); + CaSetup->Execute([&](TFakeActor& actor) { auto [dqAsyncOutput, dqAsyncOutputAsActor] = CreateDqPqWriteActor( std::move(settings), @@ -136,6 +144,7 @@ void TPqIoTestFixture::InitAsyncOutput( nullptr, &actor.GetAsyncOutputCallbacks(), MakeIntrusive(), + CreatePqNativeGateway(std::move(pqServices)), freeSpace); actor.InitAsyncOutput(dqAsyncOutput, dqAsyncOutputAsActor); From a27f08dbe6a55b246924f402397712d0b84558e1 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Mon, 17 Feb 2025 16:13:37 +0500 Subject: [PATCH 08/18] YQ-3561 FQrun supported connections / bindings (#14552) --- .../libs/actors/clusters_from_connections.cpp | 27 ++- ydb/core/fq/libs/actors/pending_fetcher.cpp | 14 +- ydb/core/fq/libs/actors/proxy.h | 4 +- ydb/core/fq/libs/actors/run_actor.cpp | 2 +- .../libs/compute/common/run_actor_params.cpp | 4 +- .../fq/libs/compute/common/run_actor_params.h | 5 +- ydb/core/fq/libs/init/init.cpp | 10 +- ydb/core/fq/libs/init/init.h | 4 +- .../pq/gateway/dummy/yql_pq_dummy_gateway.cpp | 2 +- ydb/tests/tools/fqrun/.gitignore | 3 + ydb/tests/tools/fqrun/README.md | 35 ++++ .../tools/fqrun/configuration/fq_config.conf | 172 +++++++++++++++++- .../tools/fqrun/{fqprun.cpp => fqrun.cpp} | 101 ++++++++++ ydb/tests/tools/fqrun/src/common.cpp | 11 ++ ydb/tests/tools/fqrun/src/common.h | 18 ++ ydb/tests/tools/fqrun/src/fq_runner.cpp | 85 ++++++++- ydb/tests/tools/fqrun/src/fq_runner.h | 4 + ydb/tests/tools/fqrun/src/fq_setup.cpp | 70 +++++-- ydb/tests/tools/fqrun/src/fq_setup.h | 4 + ydb/tests/tools/fqrun/src/ya.make | 2 + ydb/tests/tools/fqrun/ya.make | 11 +- ydb/tests/tools/kqprun/README.md | 62 +++++++ .../kqprun/configuration/app_config.conf | 2 +- ydb/tests/tools/kqprun/kqprun.cpp | 12 +- ydb/tests/tools/kqprun/runlib/application.h | 2 +- ydb/tests/tools/kqprun/src/kqp_runner.cpp | 8 +- 26 files changed, 614 insertions(+), 60 deletions(-) create mode 100644 ydb/tests/tools/fqrun/README.md rename ydb/tests/tools/fqrun/{fqprun.cpp => fqrun.cpp} (52%) create mode 100644 ydb/tests/tools/fqrun/src/common.cpp create mode 100644 ydb/tests/tools/kqprun/README.md diff --git a/ydb/core/fq/libs/actors/clusters_from_connections.cpp b/ydb/core/fq/libs/actors/clusters_from_connections.cpp index 1249aeafb679..3438e85fd5ca 100644 --- a/ydb/core/fq/libs/actors/clusters_from_connections.cpp +++ b/ydb/core/fq/libs/actors/clusters_from_connections.cpp @@ -88,6 +88,19 @@ std::pair ParseHttpEndpoint(const TString& endpoint) { return std::make_pair(ToString(host), scheme != "http"); } +std::pair ParseGrpcEndpoint(const TString& endpoint) { + TStringBuf scheme; + TStringBuf address; + TStringBuf uri; + NHttp::CrackURL(endpoint, scheme, address, uri); + + TString hostname; + TIpPort port; + NHttp::CrackAddress(TString(address), hostname, port); + + return {hostname, port}; +} + void FillSolomonClusterConfig(NYql::TSolomonClusterConfig& clusterConfig, const TString& name, const TString& authToken, @@ -230,8 +243,18 @@ void AddClustersFromConnections( clusterCfg->SetKind(NYql::EGenericDataSourceKind::YDB); clusterCfg->SetProtocol(NYql::EGenericProtocol::NATIVE); clusterCfg->SetName(connectionName); - clusterCfg->SetDatabaseId(db.database_id()); - clusterCfg->SetUseSsl(!common.GetDisableSslForGenericDataSources()); + if (const auto& databaseId = db.database_id()) { + clusterCfg->SetDatabaseId(databaseId); + clusterCfg->SetUseSsl(!common.GetDisableSslForGenericDataSources()); + } else { + const auto& [host, port] = ParseGrpcEndpoint(db.endpoint()); + + auto& endpoint = *clusterCfg->MutableEndpoint(); + endpoint.set_host(host); + endpoint.set_port(port); + clusterCfg->SetUseSsl(db.secure()); + clusterCfg->SetDatabaseName(db.database()); + } FillClusterAuth(*clusterCfg, db.auth(), authToken, accountIdSignatures); clusters.emplace(connectionName, GenericProviderName); break; diff --git a/ydb/core/fq/libs/actors/pending_fetcher.cpp b/ydb/core/fq/libs/actors/pending_fetcher.cpp index 9358b8afc80a..c9d437a361f7 100644 --- a/ydb/core/fq/libs/actors/pending_fetcher.cpp +++ b/ydb/core/fq/libs/actors/pending_fetcher.cpp @@ -157,7 +157,8 @@ class TPendingFetcher : public NActors::TActorBootstrapped { const ::NMonitoring::TDynamicCounterPtr& clientCounters, const TString& tenantName, NActors::TMon* monitoring, - std::shared_ptr s3ActorsFactory + std::shared_ptr s3ActorsFactory, + NYql::IPqGateway::TPtr defaultPqGateway ) : YqSharedResources(yqSharedResources) , CredentialsProviderFactory(credentialsProviderFactory) @@ -180,6 +181,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { , Monitoring(monitoring) , ComputeConfig(config.GetCompute()) , S3ActorsFactory(std::move(s3ActorsFactory)) + , DefaultPqGateway(std::move(defaultPqGateway)) { Y_ENSURE(GetYqlDefaultModuleResolverWithContext(ModuleResolver)); } @@ -472,7 +474,8 @@ class TPendingFetcher : public NActors::TActorBootstrapped { NProtoInterop::CastFromProto(task.result_ttl()), std::map(task.parameters().begin(), task.parameters().end()), S3ActorsFactory, - ComputeConfig.GetWorkloadManagerConfig(task.scope()) + ComputeConfig.GetWorkloadManagerConfig(task.scope()), + DefaultPqGateway ); auto runActorId = @@ -548,6 +551,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { NActors::TMon* Monitoring; TComputeConfig ComputeConfig; std::shared_ptr S3ActorsFactory; + NYql::IPqGateway::TPtr DefaultPqGateway; }; @@ -567,7 +571,8 @@ NActors::IActor* CreatePendingFetcher( const ::NMonitoring::TDynamicCounterPtr& clientCounters, const TString& tenantName, NActors::TMon* monitoring, - std::shared_ptr s3ActorsFactory) + std::shared_ptr s3ActorsFactory, + NYql::IPqGateway::TPtr defaultPqGateway) { return new TPendingFetcher( yqSharedResources, @@ -585,7 +590,8 @@ NActors::IActor* CreatePendingFetcher( clientCounters, tenantName, monitoring, - std::move(s3ActorsFactory)); + std::move(s3ActorsFactory), + defaultPqGateway); } TActorId MakePendingFetcherId(ui32 nodeId) { diff --git a/ydb/core/fq/libs/actors/proxy.h b/ydb/core/fq/libs/actors/proxy.h index 92cc9dd13faa..9f457859189b 100644 --- a/ydb/core/fq/libs/actors/proxy.h +++ b/ydb/core/fq/libs/actors/proxy.h @@ -16,6 +16,7 @@ #include #include #include +#include #include #include @@ -53,7 +54,8 @@ NActors::IActor* CreatePendingFetcher( const ::NMonitoring::TDynamicCounterPtr& clientCounters, const TString& tenantName, NActors::TMon* monitoring, - std::shared_ptr s3ActorsFactory + std::shared_ptr s3ActorsFactory, + NYql::IPqGateway::TPtr defaultPqGateway ); NActors::IActor* CreateRunActor( diff --git a/ydb/core/fq/libs/actors/run_actor.cpp b/ydb/core/fq/libs/actors/run_actor.cpp index 950f935c03a1..f18b666a4423 100644 --- a/ydb/core/fq/libs/actors/run_actor.cpp +++ b/ydb/core/fq/libs/actors/run_actor.cpp @@ -1977,7 +1977,7 @@ class TRunActor : public NActors::TActorBootstrapped { std::make_shared(gatewaysConfig.GetPq()), Params.FunctionRegistry ); - const auto pqGateway = NYql::CreatePqNativeGateway(pqServices); + const auto pqGateway = Params.DefaultPqGateway ? Params.DefaultPqGateway : NYql::CreatePqNativeGateway(pqServices); dataProvidersInit.push_back(GetPqDataProviderInitializer(pqGateway, false, dbResolver)); } diff --git a/ydb/core/fq/libs/compute/common/run_actor_params.cpp b/ydb/core/fq/libs/compute/common/run_actor_params.cpp index a6a26c0f9cea..82071360cab9 100644 --- a/ydb/core/fq/libs/compute/common/run_actor_params.cpp +++ b/ydb/core/fq/libs/compute/common/run_actor_params.cpp @@ -60,7 +60,8 @@ TRunActorParams::TRunActorParams( TDuration resultTtl, std::map&& queryParameters, std::shared_ptr s3ActorsFactory, - const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager + const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager, + NYql::IPqGateway::TPtr defaultPqGateway ) : YqSharedResources(yqSharedResources) , CredentialsProviderFactory(credentialsProviderFactory) @@ -117,6 +118,7 @@ TRunActorParams::TRunActorParams( , QueryParameters(std::move(queryParameters)) , S3ActorsFactory(std::move(s3ActorsFactory)) , WorkloadManager(workloadManager) + , DefaultPqGateway(defaultPqGateway) { } diff --git a/ydb/core/fq/libs/compute/common/run_actor_params.h b/ydb/core/fq/libs/compute/common/run_actor_params.h index 7ae4eda044ff..1ff1db905cd5 100644 --- a/ydb/core/fq/libs/compute/common/run_actor_params.h +++ b/ydb/core/fq/libs/compute/common/run_actor_params.h @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -79,7 +80,8 @@ struct TRunActorParams { // TODO2 : Change name TDuration resultTtl, std::map&& queryParameters, std::shared_ptr s3ActorsFactory, - const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager + const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager, + NYql::IPqGateway::TPtr defaultPqGateway ); TRunActorParams(const TRunActorParams& params) = default; @@ -145,6 +147,7 @@ struct TRunActorParams { // TODO2 : Change name std::map QueryParameters; std::shared_ptr S3ActorsFactory; ::NFq::NConfig::TWorkloadManagerConfig WorkloadManager; + NYql::IPqGateway::TPtr DefaultPqGateway; }; } /* NFq */ diff --git a/ydb/core/fq/libs/init/init.cpp b/ydb/core/fq/libs/init/init.cpp index edca3f733212..3bdc089b3677 100644 --- a/ydb/core/fq/libs/init/init.cpp +++ b/ydb/core/fq/libs/init/init.cpp @@ -66,7 +66,8 @@ void Init( const IYqSharedResources::TPtr& iyqSharedResources, const std::function& folderServiceFactory, ui32 icPort, - const std::vector& additionalCompNodeFactories + const std::vector& additionalCompNodeFactories, + NYql::IPqGateway::TPtr defaultPqGateway ) { Y_ABORT_UNLESS(iyqSharedResources, "No YQ shared resources created"); @@ -204,7 +205,7 @@ void Init( credentialsFactory, tenant, yqCounters->GetSubgroup("subsystem", "row_dispatcher"), - CreatePqNativeGateway(pqServices), + defaultPqGateway ? defaultPqGateway : CreatePqNativeGateway(pqServices), appData->Mon, appData->Counters); actorRegistrator(NFq::RowDispatcherServiceActorId(), rowDispatcher.release()); @@ -225,7 +226,7 @@ void Init( std::make_shared(protoConfig.GetGateways().GetPq()), appData->FunctionRegistry ); - auto pqGateway = NYql::CreatePqNativeGateway(std::move(pqServices)); + auto pqGateway = defaultPqGateway ? defaultPqGateway : NYql::CreatePqNativeGateway(std::move(pqServices)); RegisterDqPqReadActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, pqGateway, yqCounters->GetSubgroup("subsystem", "DqSourceTracker"), protoConfig.GetCommon().GetPqReconnectPeriod()); @@ -345,7 +346,8 @@ void Init( clientCounters, tenant, appData->Mon, - s3ActorsFactory + s3ActorsFactory, + defaultPqGateway ); actorRegistrator(MakePendingFetcherId(nodeId), fetcher); diff --git a/ydb/core/fq/libs/init/init.h b/ydb/core/fq/libs/init/init.h index 6dcf35afd6d1..517557b04d8d 100644 --- a/ydb/core/fq/libs/init/init.h +++ b/ydb/core/fq/libs/init/init.h @@ -12,6 +12,7 @@ #include #include +#include #include @@ -36,7 +37,8 @@ void Init( const IYqSharedResources::TPtr& yqSharedResources, const std::function& folderServiceFactory, ui32 icPort, - const std::vector& additionalCompNodeFactories + const std::vector& additionalCompNodeFactories, + NYql::IPqGateway::TPtr defaultPqGateway = nullptr ); } // NFq diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp index f1575567e468..59add76e5ca8 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp @@ -11,7 +11,7 @@ namespace NYql { NThreading::TFuture TDummyPqGateway::OpenSession(const TString& sessionId, const TString& username) { with_lock (Mutex) { Y_ENSURE(sessionId); - Y_ENSURE(username); + Y_UNUSED(username); Y_ENSURE(!IsIn(OpenedSessions, sessionId), "Session " << sessionId << " is already opened in pq gateway"); OpenedSessions.insert(sessionId); diff --git a/ydb/tests/tools/fqrun/.gitignore b/ydb/tests/tools/fqrun/.gitignore index 51aaf6608d57..99dcb6bb049f 100644 --- a/ydb/tests/tools/fqrun/.gitignore +++ b/ydb/tests/tools/fqrun/.gitignore @@ -2,3 +2,6 @@ sync_dir *.log *.sql +*.conf +*.parquet +*.json diff --git a/ydb/tests/tools/fqrun/README.md b/ydb/tests/tools/fqrun/README.md new file mode 100644 index 000000000000..cd133168dcd4 --- /dev/null +++ b/ydb/tests/tools/fqrun/README.md @@ -0,0 +1,35 @@ +# FQ run tool + +Tool can be used to execute streaming queries by using FQ proxy infrastructure. + +## Examples + +### Queries + +* Run select 42: + ```(bash) + ./fqrun -s "SELECT 42" + ``` + +### Logs + +* Setup log settings: + ```(bash) + ./fqrun -s "SELECT 42" --log-default=warn --log FQ_RUN_ACTOR=trace --log-file query.log + ``` + +### Cluster + +* Embedded UI: + ```(bash) + ./fqrun -M 32000 + ``` + + Monitoring endpoint: https://localhost:32000 + +* gRPC endpoint: + ```(bash) + ./fqrun -G 32000 + ``` + + Connect with ydb CLI: `ydb -e grpc://localhost:32000 -d /Root` diff --git a/ydb/tests/tools/fqrun/configuration/fq_config.conf b/ydb/tests/tools/fqrun/configuration/fq_config.conf index 8a980fa5daec..576b0d394bdb 100644 --- a/ydb/tests/tools/fqrun/configuration/fq_config.conf +++ b/ydb/tests/tools/fqrun/configuration/fq_config.conf @@ -4,6 +4,12 @@ EnableTaskCounters: true CheckpointCoordinator { Enabled: true + CheckpointingPeriodMillis: 30000 + MaxInflight: 1 + + CheckpointGarbageConfig { + Enabled: true + } Storage { TablePrefix: "yq/checkpoints" @@ -14,15 +20,24 @@ CheckpointCoordinator { } Common { + YdbMvpCloudEndpoint: "https://ydbc.ydb.cloud.yandex.net:8789/ydbc/cloud-prod" MdbGateway: "https://mdb.api.cloud.yandex.net:443" - ObjectStorageEndpoint: "https://storage-internal.cloud.yandex.net" + MdbTransformHost: false + ObjectStorageEndpoint: "https://storage.yandexcloud.net" IdsPrefix: "kr" QueryArtifactsCompressionMethod: "zstd_6" MonitoringEndpoint: "monitoring.api.cloud.yandex.net" KeepInternalErrors: true UseNativeProtocolForClickHouse: true - DisableSslForGenericDataSources: true ShowQueryTimeline: true + MaxTasksPerOperation: 400 + MaxTasksPerStage: 50 + PqReconnectPeriod: "30m" + + YdbDriverConfig { + ClientThreadsNum: 6 + NetworkThreadsNum: 6 + } } ControlPlaneProxy { @@ -32,6 +47,15 @@ ControlPlaneProxy { ControlPlaneStorage { Enabled: true StatsMode: STATS_MODE_PROFILE + DumpRawStatistics: true + TasksBatchSize: 100 + NumTasksProportion: 4 + AnalyticsRetryCounterLimit: 20 + StreamingRetryCounterLimit: 20 + AnalyticsRetryCounterUpdateTime: "1d" + StreamingRetryCounterUpdateTime: "1d" + TaskLeaseTtl: "30s" + DisableCurrentIam: false AvailableConnection: "OBJECT_STORAGE" AvailableConnection: "DATA_STREAMS" @@ -45,7 +69,11 @@ ControlPlaneStorage { AvailableStreamingConnection: "OBJECT_STORAGE" AvailableStreamingConnection: "DATA_STREAMS" AvailableStreamingConnection: "MONITORING" + AvailableStreamingConnection: "POSTGRESQL_CLUSTER" + AvailableStreamingConnection: "CLICKHOUSE_CLUSTER" AvailableStreamingConnection: "YDB_DATABASE" + AvailableStreamingConnection: "GREENPLUM_CLUSTER" + AvailableStreamingConnection: "MYSQL_CLUSTER" AvailableBinding: "OBJECT_STORAGE" AvailableBinding: "DATA_STREAMS" @@ -69,6 +97,114 @@ DbPool { } } +Gateways { + Enabled: true + + Dq { + DefaultSettings { + Name: "HashShuffleTasksRatio" + Value: "1" + } + DefaultSettings { + Name: "UseFinalizeByKey" + Value: "true" + } + } + + Generic { + MdbGateway: "https://mdb.api.cloud.yandex.net:443" + + Connector { + UseSsl: false + + Endpoint { + host: "localhost" + port: 2130 + } + } + + DefaultSettings { + Name: "DateTimeFormat" + Value: "string" + } + } + + HttpGateway { + BuffersSizePerStream: 5000000 + ConnectionTimeoutSeconds: 15 + LowSpeedBytesLimit: 1024 + LowSpeedTimeSeconds: 20 + MaxInFlightCount: 2000 + MaxSimulatenousDownloadsSize: 2000000000 + RequestTimeoutSeconds: 0 + } + + Pq { + ClusterMapping { + Name: "pq" + Endpoint: "localhost:2135" + Database: "local" + ClusterType: CT_DATA_STREAMS + UseSsl: true + SharedReading: true + ReadGroup: "fqrun" + } + } + + Solomon { + DefaultSettings { + Name: "_EnableReading" + Value: "true" + } + } + + S3 { + AllowConcurrentListings: true + AllowLocalFiles: true + FileSizeLimit: 100000000000 + GeneratorPathsLimit: 50000 + ListingCallbackPerThreadQueueSize: 100 + ListingCallbackThreadCount: 1 + MaxDirectoriesAndFilesPerQuery: 500000 + MaxDiscoveryFilesPerQuery: 1000 + MaxFilesPerQuery: 500000 + MaxInflightListsPerQuery: 100 + MinDesiredDirectoriesOfFilesPerQuery: 1000 + RegexpCacheSize: 100 + + FormatSizeLimit { + Name: "parquet" + FileSizeLimit: 52428800 + } + FormatSizeLimit { + Name: "raw" + FileSizeLimit: 52428800 + } + + DefaultSettings { + Name: "AtomicUploadCommit" + Value: "true" + } + DefaultSettings { + Name: "UseBlocksSource" + Value: "true" + } + } + + YqlCore { + Flags { + Name: "_EnableMatchRecognize" + } + Flags { + Name: "_EnableStreamLookupJoin" + } + } +} + +Health { + Enabled: true +} + NodesManager { Enabled: true } @@ -79,6 +215,7 @@ PendingFetcher { PrivateApi { Enabled: true + Loopback: true } PrivateProxy { @@ -87,6 +224,23 @@ PrivateProxy { QuotasManager { Enabled: true + QuotaDescriptions { + SubjectType: "cloud" + MetricName: "yq.cpuPercent.count" + HardLimit: 7500 + DefaultLimit: 3500 + } + QuotaDescriptions { + SubjectType: "cloud" + MetricName: "yq.streamingQueryDurationMinutes.count" + DefaultLimit: 10080 + } + QuotaDescriptions { + SubjectType: "cloud" + MetricName: "yq.analyticsQueryDurationMinutes.count" + HardLimit: 1440 + DefaultLimit: 30 + } } RateLimiter { @@ -106,14 +260,26 @@ RateLimiter { } } +ReadActorsFactoryConfig { + PqReadActorFactoryConfig { + CookieCommitMode: false + } +} + ResourceManager { Enabled: true + MkqlInitialMemoryLimit: 16777216 + MkqlTotalMemoryLimit: 193273528320 + MkqlAllocSize: 16777216 + MkqlTaskHardMemoryLimit: 24696061952 } RowDispatcher { Enabled: true SendStatusPeriodSec: 10 - TimeoutBeforeStartSessionSec: 10 + TimeoutBeforeStartSessionSec: 0 + MaxSessionUsedMemory: 16000000 + WithoutConsumer: false CompileService { ParallelCompilationLimit: 20 diff --git a/ydb/tests/tools/fqrun/fqprun.cpp b/ydb/tests/tools/fqrun/fqrun.cpp similarity index 52% rename from ydb/tests/tools/fqrun/fqprun.cpp rename to ydb/tests/tools/fqrun/fqrun.cpp index 277aa8d2ed90..587bc405b034 100644 --- a/ydb/tests/tools/fqrun/fqprun.cpp +++ b/ydb/tests/tools/fqrun/fqrun.cpp @@ -3,6 +3,7 @@ #include +#include #include #include #include @@ -15,6 +16,8 @@ namespace { struct TExecutionOptions { TString Query; + std::vector Connections; + std::vector Bindings; bool HasResults() const { return !Query.empty(); @@ -36,6 +39,20 @@ struct TExecutionOptions { void RunArgumentQueries(const TExecutionOptions& executionOptions, TFqRunner& runner) { NColorizer::TColors colors = NColorizer::AutoColors(Cout); + if (!executionOptions.Connections.empty()) { + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating connections..." << colors.Default() << Endl; + if (!runner.CreateConnections(executionOptions.Connections)) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Failed to create connections"; + } + } + + if (!executionOptions.Bindings.empty()) { + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating bindings..." << colors.Default() << Endl; + if (!runner.CreateBindings(executionOptions.Bindings)) { + ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Failed to create bindings"; + } + } + if (executionOptions.Query) { Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing query..." << colors.Default() << Endl; if (!runner.ExecuteStreamQuery(executionOptions.GetQueryOptions())) { @@ -89,6 +106,8 @@ void RunScript(const TExecutionOptions& executionOptions, const TRunnerOptions& } class TMain : public TMainBase { + using EVerbose = TFqSetupSettings::EVerbose; + protected: void RegisterOptions(NLastGetopt::TOpts& options) override { options.SetTitle("FqRun -- tool to execute stream queries through FQ proxy"); @@ -101,6 +120,33 @@ class TMain : public TMainBase { .RequiredArgument("file") .StoreMappedResult(&ExecutionOptions.Query, &LoadFile); + options.AddLongOption('s', "sql", "Query SQL text to execute") + .RequiredArgument("str") + .StoreResult(&ExecutionOptions.Query); + options.MutuallyExclusive("query", "sql"); + + options.AddLongOption('c', "connection", "External datasource connection protobuf FederatedQuery::ConnectionContent") + .RequiredArgument("file") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + auto& connection = ExecutionOptions.Connections.emplace_back(); + const TString file(TString(option->CurValOrDef())); + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(file), &connection)) { + ythrow yexception() << "Bad format of FQ connection in file '" << file << "'"; + } + SetupAcl(connection.mutable_acl()); + }); + + options.AddLongOption('b', "binding", "External datasource binding protobuf FederatedQuery::BindingContent") + .RequiredArgument("file") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + auto& binding = ExecutionOptions.Bindings.emplace_back(); + const TString file(TString(option->CurValOrDef())); + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(file), &binding)) { + ythrow yexception() << "Bad format of FQ binding in file '" << file << "'"; + } + SetupAcl(binding.mutable_acl()); + }); + options.AddLongOption("fq-cfg", "File with FQ config (NFq::NConfig::TConfig for FQ proxy)") .RequiredArgument("file") .DefaultValue("./configuration/fq_config.conf") @@ -110,6 +156,24 @@ class TMain : public TMainBase { } }); + options.AddLongOption("emulate-s3", "Enable readings by s3 provider from files, `bucket` value in connection - path to folder with files") + .NoArgument() + .SetFlag(&RunnerOptions.FqSettings.EmulateS3); + + options.AddLongOption("emulate-pq", "Emulate YDS with local file, accepts list of tables to emulate with following format: topic@file (can be used in query from cluster `pq`)") + .RequiredArgument("topic@file") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TStringBuf topicName; + TStringBuf filePath; + TStringBuf(option->CurVal()).Split('@', topicName, filePath); + if (topicName.empty() || filePath.empty()) { + ythrow yexception() << "Incorrect PQ file mapping, expected form topic@file"; + } + if (!PqFilesMapping.emplace(topicName, filePath).second) { + ythrow yexception() << "Got duplicated topic name: " << topicName; + } + }); + // Outputs options.AddLongOption("result-file", "File with query results (use '-' to write in stdout)") @@ -128,24 +192,61 @@ class TMain : public TMainBase { .Choices(resultFormat.GetChoices()) .StoreMappedResultT(&RunnerOptions.ResultOutputFormat, resultFormat); + // Pipeline settings + + options.AddLongOption("verbose", TStringBuilder() << "Common verbose level (max level " << static_cast(EVerbose::Max) - 1 << ")") + .RequiredArgument("uint") + .DefaultValue(static_cast(EVerbose::Info)) + .StoreMappedResultT(&RunnerOptions.FqSettings.VerboseLevel, [](ui8 value) { + return static_cast(std::min(value, static_cast(EVerbose::Max))); + }); + RegisterKikimrOptions(options, RunnerOptions.FqSettings); } int DoRun(NLastGetopt::TOptsParseResult&&) override { ExecutionOptions.Validate(RunnerOptions); + RunnerOptions.FqSettings.YqlToken = GetEnv(YQL_TOKEN_VARIABLE); + + auto& gatewayConfig = *RunnerOptions.FqSettings.FqConfig.mutable_gateways(); + FillTokens(gatewayConfig.mutable_pq()); + FillTokens(gatewayConfig.mutable_s3()); + FillTokens(gatewayConfig.mutable_generic()); + FillTokens(gatewayConfig.mutable_ydb()); + FillTokens(gatewayConfig.mutable_solomon()); + auto& logConfig = RunnerOptions.FqSettings.LogConfig; logConfig.SetDefaultLevel(NActors::NLog::EPriority::PRI_CRIT); FillLogConfig(logConfig); + if (!PqFilesMapping.empty()) { + auto fileGateway = MakeIntrusive(); + for (const auto& [topic, file] : PqFilesMapping) { + fileGateway->AddDummyTopic(NYql::TDummyTopic("pq", TString(topic), TString(file))); + } + RunnerOptions.FqSettings.PqGateway = std::move(fileGateway); + } + RunScript(ExecutionOptions, RunnerOptions); return 0; } +private: + template + void FillTokens(TGatewayConfig* gateway) const { + for (auto& cluster : *gateway->mutable_clustermapping()) { + if (!cluster.GetToken()) { + cluster.SetToken(RunnerOptions.FqSettings.YqlToken); + } + } + } + private: TExecutionOptions ExecutionOptions; TRunnerOptions RunnerOptions; + std::unordered_map PqFilesMapping; }; } // anonymous namespace diff --git a/ydb/tests/tools/fqrun/src/common.cpp b/ydb/tests/tools/fqrun/src/common.cpp new file mode 100644 index 000000000000..7d0fc1804532 --- /dev/null +++ b/ydb/tests/tools/fqrun/src/common.cpp @@ -0,0 +1,11 @@ +#include "common.h" + +namespace NFqRun { + +void SetupAcl(FederatedQuery::Acl* acl) { + if (acl->visibility() == FederatedQuery::Acl::VISIBILITY_UNSPECIFIED) { + acl->set_visibility(FederatedQuery::Acl::SCOPE); + } +} + +} // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/common.h b/ydb/tests/tools/fqrun/src/common.h index 5af4b019dbc3..53dfbe654e93 100644 --- a/ydb/tests/tools/fqrun/src/common.h +++ b/ydb/tests/tools/fqrun/src/common.h @@ -4,13 +4,29 @@ #include #include +#include #include namespace NFqRun { +constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN"; constexpr i64 MAX_RESULT_SET_ROWS = 1000; struct TFqSetupSettings : public NKikimrRun::TServerSettings { + enum class EVerbose { + None, + Info, + QueriesText, + InitLogs, + Max + }; + + bool EmulateS3 = false; + + EVerbose VerboseLevel = EVerbose::Info; + + TString YqlToken; + NYql::IPqGateway::TPtr PqGateway; NFq::NConfig::TConfig FqConfig; NKikimrConfig::TLogConfig LogConfig; }; @@ -26,4 +42,6 @@ struct TRequestOptions { TString Query; }; +void SetupAcl(FederatedQuery::Acl* acl); + } // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_runner.cpp b/ydb/tests/tools/fqrun/src/fq_runner.cpp index 6fe51cfb50f9..33995c17fac8 100644 --- a/ydb/tests/tools/fqrun/src/fq_runner.cpp +++ b/ydb/tests/tools/fqrun/src/fq_runner.cpp @@ -8,17 +8,24 @@ using namespace NKikimrRun; namespace NFqRun { class TFqRunner::TImpl { + using EVerbose = TFqSetupSettings::EVerbose; + static constexpr TDuration REFRESH_PERIOD = TDuration::Seconds(1); public: explicit TImpl(const TRunnerOptions& options) : Options(options) + , VerboseLevel(options.FqSettings.VerboseLevel) , FqSetup(options.FqSettings) , CerrColors(NColorizer::AutoColors(Cerr)) , CoutColors(NColorizer::AutoColors(Cout)) {} bool ExecuteStreamQuery(const TRequestOptions& query) { + if (VerboseLevel >= EVerbose::QueriesText) { + Cout << CoutColors.Cyan() << "Starting stream request:\n" << CoutColors.Default() << query.Query << Endl; + } + const TRequestResult status = FqSetup.StreamRequest(query, StreamQueryId); if (!status.IsSuccess()) { @@ -52,7 +59,7 @@ class TFqRunner::TImpl { if (Options.ResultOutput) { Cout << CoutColors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Writing query results..." << CoutColors.Default() << Endl; for (size_t i = 0; i < ResultSets.size(); ++i) { - if (ResultSets.size() > 1) { + if (ResultSets.size() > 1 && VerboseLevel >= EVerbose::Info) { *Options.ResultOutput << CoutColors.Cyan() << "Result set " << i + 1 << ":" << CoutColors.Default() << Endl; } PrintResultSet(Options.ResultOutputFormat, *Options.ResultOutput, ResultSets[i]); @@ -60,6 +67,53 @@ class TFqRunner::TImpl { } } + bool CreateConnections(const std::vector& connections) { + for (const auto& connection : connections) { + if (VerboseLevel >= EVerbose::QueriesText) { + Cout << CoutColors.Cyan() << "Creating connection:\n" << CoutColors.Default() << Endl << connection.DebugString() << Endl; + } + + TString connectionId; + const TRequestResult status = FqSetup.CreateConnection(connection, connectionId); + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to create connection '" << connection.name() << "', reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + + if (!ConnectionNameToId.emplace(connection.name(), connectionId).second) { + Cerr << CerrColors.Red() << "Got duplicated connection name '" << connection.name() << "'" << CerrColors.Default() << Endl; + return false; + } + } + + return true; + } + + bool CreateBindings(const std::vector& bindings) const { + for (auto binding : bindings) { + if (VerboseLevel >= EVerbose::QueriesText) { + Cout << CoutColors.Cyan() << "Creating binding:\n" << CoutColors.Default() << Endl << binding.DebugString() << Endl; + } + + const auto it = ConnectionNameToId.find(binding.connection_id()); + if (it == ConnectionNameToId.end()) { + Cerr << CerrColors.Red() << "Failed to create binding '" << binding.name() << "', connection with name '" << binding.connection_id() << "' not found" << CerrColors.Default() << Endl; + return false; + } + + binding.set_connection_id(it->second); + const TRequestResult status = FqSetup.CreateBinding(binding); + + if (!status.IsSuccess()) { + Cerr << CerrColors.Red() << "Failed to create binding '" << binding.name() << "', reason:" << CerrColors.Default() << Endl << status.ToString() << Endl; + return false; + } + } + + return true; + } + private: static bool IsFinalStatus(FederatedQuery::QueryMeta::ComputeStatus status) { using EStatus = FederatedQuery::QueryMeta; @@ -70,7 +124,13 @@ class TFqRunner::TImpl { StartTime = TInstant::Now(); while (true) { - const TRequestResult status = FqSetup.DescribeQuery(StreamQueryId, ExecutionMeta); + TExecutionMeta meta; + const TRequestResult status = FqSetup.DescribeQuery(StreamQueryId, meta); + + if (meta.TransientIssues.Size() != ExecutionMeta.TransientIssues.Size() && VerboseLevel >= EVerbose::Info) { + Cerr << CerrColors.Red() << "Query transient issues updated:" << CerrColors.Default() << Endl << meta.TransientIssues.ToString() << Endl; + } + ExecutionMeta = meta; if (IsFinalStatus(ExecutionMeta.Status)) { break; @@ -84,13 +144,12 @@ class TFqRunner::TImpl { Sleep(REFRESH_PERIOD); } - Cout << CoutColors.Cyan() << "Query finished. Duration: " << TInstant::Now() - StartTime << CoutColors.Default() << Endl; + if (VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Query finished. Duration: " << TInstant::Now() - StartTime << CoutColors.Default() << Endl; + } if (ExecutionMeta.Status != FederatedQuery::QueryMeta::COMPLETED) { Cerr << CerrColors.Red() << "Failed to execute query, invalid final status " << FederatedQuery::QueryMeta::ComputeStatus_Name(ExecutionMeta.Status) << ", issues:" << CerrColors.Default() << Endl << ExecutionMeta.Issues.ToString() << Endl; - if (ExecutionMeta.TransientIssues) { - Cerr << CerrColors.Red() << "Transient issues:" << CerrColors.Default() << Endl << ExecutionMeta.TransientIssues.ToString() << Endl; - } return false; } @@ -98,15 +157,12 @@ class TFqRunner::TImpl { Cerr << CerrColors.Red() << "Query finished with issues:" << CerrColors.Default() << Endl << ExecutionMeta.Issues.ToString() << Endl; } - if (ExecutionMeta.TransientIssues) { - Cerr << CerrColors.Red() << "Query finished with transient issues:" << CerrColors.Default() << Endl << ExecutionMeta.TransientIssues.ToString() << Endl; - } - return true; } private: const TRunnerOptions Options; + const EVerbose VerboseLevel; const TFqSetup FqSetup; const NColorizer::TColors CerrColors; const NColorizer::TColors CoutColors; @@ -115,6 +171,7 @@ class TFqRunner::TImpl { TInstant StartTime; TExecutionMeta ExecutionMeta; std::vector ResultSets; + std::unordered_map ConnectionNameToId; }; TFqRunner::TFqRunner(const TRunnerOptions& options) @@ -133,4 +190,12 @@ void TFqRunner::PrintQueryResults() const { Impl->PrintQueryResults(); } +bool TFqRunner::CreateConnections(const std::vector& connections) const { + return Impl->CreateConnections(connections); +} + +bool TFqRunner::CreateBindings(const std::vector& bindings) const { + return Impl->CreateBindings(bindings); +} + } // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_runner.h b/ydb/tests/tools/fqrun/src/fq_runner.h index 7b803648dcf9..b6e85db142bc 100644 --- a/ydb/tests/tools/fqrun/src/fq_runner.h +++ b/ydb/tests/tools/fqrun/src/fq_runner.h @@ -16,6 +16,10 @@ class TFqRunner { void PrintQueryResults() const; + bool CreateConnections(const std::vector& connections) const; + + bool CreateBindings(const std::vector& bindings) const; + private: class TImpl; std::shared_ptr Impl; diff --git a/ydb/tests/tools/fqrun/src/fq_setup.cpp b/ydb/tests/tools/fqrun/src/fq_setup.cpp index 6b16b4f5af15..394bfc5c3e16 100644 --- a/ydb/tests/tools/fqrun/src/fq_setup.cpp +++ b/ydb/tests/tools/fqrun/src/fq_setup.cpp @@ -17,13 +17,15 @@ namespace NFqRun { namespace { -Ydb::StatusIds::StatusCode GetStatus(const NYql::TIssues& issues) { - return issues ? Ydb::StatusIds::BAD_REQUEST : Ydb::StatusIds::SUCCESS; +TRequestResult GetStatus(const NYql::TIssues& issues) { + return TRequestResult(issues ? Ydb::StatusIds::BAD_REQUEST : Ydb::StatusIds::SUCCESS, issues); } } // anonymous namespace class TFqSetup::TImpl { + using EVerbose = TFqSetupSettings::EVerbose; + private: TAutoPtr CreateLogBackend() const { if (Settings.LogOutputFile) { @@ -46,7 +48,7 @@ class TFqSetup::TImpl { NKikimr::Tests::TServerSettings serverSettings(PortManager.GetPort()); serverSettings.SetDomainName(Settings.DomainName); - serverSettings.SetVerbose(false); + serverSettings.SetVerbose(Settings.VerboseLevel >= EVerbose::InitLogs); NKikimrConfig::TAppConfig config; *config.MutableLogConfig() = Settings.LogConfig; @@ -83,7 +85,7 @@ class TFqSetup::TImpl { Client->InitRootScheme(); } - NFq::NConfig::TConfig GetFqProxyConfig(ui32 grpcPort, ui32 httpPort) const { + NFq::NConfig::TConfig GetFqProxyConfig(ui32 grpcPort) const { auto fqConfig = Settings.FqConfig; fqConfig.MutableControlPlaneStorage()->AddSuperUsers(BUILTIN_ACL_ROOT); @@ -109,14 +111,19 @@ class TFqSetup::TImpl { nodesMenagerConfig->SetPort(grpcPort); nodesMenagerConfig->SetHost("localhost"); - fqConfig.MutableCommon()->SetYdbMvpCloudEndpoint(TStringBuilder() << "http://localhost:" << httpPort << "/yql-mock/abc"); + auto* healthConfig = fqConfig.MutableHealth(); + healthConfig->SetPort(grpcPort); + healthConfig->SetDatabase(database); + + if (Settings.EmulateS3) { + fqConfig.MutableCommon()->SetObjectStorageEndpoint("file://"); + } return fqConfig; } void InitializeFqProxy(ui32 grpcPort) { - const ui32 httpPort = PortManager.GetPort(); - const auto& fqConfig = GetFqProxyConfig(grpcPort, httpPort); + const auto& fqConfig = GetFqProxyConfig(grpcPort); const auto counters = GetRuntime()->GetAppData().Counters->GetSubgroup("counters", "yq"); YqSharedResources = NFq::CreateYqSharedResources(fqConfig, NKikimr::CreateYdbCredentialsProviderFactory, counters); @@ -131,10 +138,9 @@ class TFqSetup::TImpl { NFq::Init( fqConfig, GetRuntime()->GetNodeId(), actorRegistrator, &GetRuntime()->GetAppData(), - Settings.DomainName, nullptr, YqSharedResources, folderServiceFactory, 0, {} + Settings.DomainName, nullptr, YqSharedResources, folderServiceFactory, 0, {}, Settings.PqGateway ); - - NFq::InitTest(GetRuntime(), httpPort, grpcPort, YqSharedResources); + YqSharedResources->Init(GetRuntime()->GetActorSystem(0)); } public: @@ -145,11 +151,11 @@ class TFqSetup::TImpl { InitializeServer(grpcPort); InitializeFqProxy(grpcPort); - if (Settings.MonitoringEnabled) { + if (Settings.MonitoringEnabled && Settings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors.Cyan() << "Monitoring port: " << CoutColors.Default() << GetRuntime()->GetMonPort() << Endl; } - if (Settings.GrpcEnabled) { + if (Settings.GrpcEnabled && Settings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors.Cyan() << "Domain gRPC port: " << CoutColors.Default() << grpcPort << Endl; } } @@ -167,7 +173,7 @@ class TFqSetup::TImpl { auto& content = *request.mutable_content(); content.set_type(FederatedQuery::QueryContent::STREAMING); content.set_text(query.Query); - content.mutable_acl()->set_visibility(::FederatedQuery::Acl::SCOPE); + SetupAcl(content.mutable_acl()); return RunControlPlaneProxyRequest(request); } @@ -188,6 +194,20 @@ class TFqSetup::TImpl { return RunControlPlaneProxyRequest(request); } + NFq::TEvControlPlaneProxy::TEvCreateConnectionResponse::TPtr CreateConnection(const FederatedQuery::ConnectionContent& connection) const { + FederatedQuery::CreateConnectionRequest request; + *request.mutable_content() = connection; + + return RunControlPlaneProxyRequest(request); + } + + NFq::TEvControlPlaneProxy::TEvCreateBindingResponse::TPtr CreateBinding(const FederatedQuery::BindingContent& binding) const { + FederatedQuery::CreateBindingRequest request; + *request.mutable_content() = binding; + + return RunControlPlaneProxyRequest(request); + } + private: NActors::TTestActorRuntime* GetRuntime() const { return Server->GetRuntime(); @@ -195,7 +215,7 @@ class TFqSetup::TImpl { template typename TResponse::TPtr RunControlPlaneProxyRequest(const TProto& request) const { - auto event = std::make_unique("yandexcloud://kqprun", request, BUILTIN_ACL_ROOT, BUILTIN_ACL_ROOT, TVector{}); + auto event = std::make_unique("yandexcloud://fqrun", request, BUILTIN_ACL_ROOT, Settings.YqlToken ? Settings.YqlToken : "fqrun", TVector{}); return RunControlPlaneProxyRequest(std::move(event)); } @@ -228,8 +248,7 @@ TRequestResult TFqSetup::StreamRequest(const TRequestOptions& query, TString& qu queryId = response->Get()->Result.query_id(); - const auto& issues = response->Get()->Issues; - return TRequestResult(GetStatus(issues), issues); + return GetStatus(response->Get()->Issues); } TRequestResult TFqSetup::DescribeQuery(const TString& queryId, TExecutionMeta& meta) const { @@ -245,8 +264,7 @@ TRequestResult TFqSetup::DescribeQuery(const TString& queryId, TExecutionMeta& m meta.ResultSetSizes.emplace_back(resultMeta.rows_count()); } - const auto& issues = response->Get()->Issues; - return TRequestResult(GetStatus(issues), issues); + return GetStatus(response->Get()->Issues); } TRequestResult TFqSetup::FetchQueryResults(const TString& queryId, i32 resultSetId, Ydb::ResultSet& resultSet) const { @@ -254,8 +272,20 @@ TRequestResult TFqSetup::FetchQueryResults(const TString& queryId, i32 resultSet resultSet = response->Get()->Result.result_set(); - const auto& issues = response->Get()->Issues; - return TRequestResult(GetStatus(issues), issues); + return GetStatus(response->Get()->Issues); +} + +TRequestResult TFqSetup::CreateConnection(const FederatedQuery::ConnectionContent& connection, TString& connectionId) const { + const auto response = Impl->CreateConnection(connection); + + connectionId = response->Get()->Result.connection_id(); + + return GetStatus(response->Get()->Issues); +} + +TRequestResult TFqSetup::CreateBinding(const FederatedQuery::BindingContent& binding) const { + const auto response = Impl->CreateBinding(binding); + return GetStatus(response->Get()->Issues); } } // namespace NFqRun diff --git a/ydb/tests/tools/fqrun/src/fq_setup.h b/ydb/tests/tools/fqrun/src/fq_setup.h index c78fad43e2dd..2ef7733bf028 100644 --- a/ydb/tests/tools/fqrun/src/fq_setup.h +++ b/ydb/tests/tools/fqrun/src/fq_setup.h @@ -25,6 +25,10 @@ class TFqSetup { TRequestResult FetchQueryResults(const TString& queryId, i32 resultSetId, Ydb::ResultSet& resultSet) const; + TRequestResult CreateConnection(const FederatedQuery::ConnectionContent& connection, TString& connectionId) const; + + TRequestResult CreateBinding(const FederatedQuery::BindingContent& binding) const; + private: class TImpl; std::shared_ptr Impl; diff --git a/ydb/tests/tools/fqrun/src/ya.make b/ydb/tests/tools/fqrun/src/ya.make index 5173c30f4be2..6f44cf16c38d 100644 --- a/ydb/tests/tools/fqrun/src/ya.make +++ b/ydb/tests/tools/fqrun/src/ya.make @@ -1,6 +1,7 @@ LIBRARY() SRCS( + common.cpp fq_runner.cpp fq_setup.cpp ) @@ -17,6 +18,7 @@ PEERDIR( ydb/library/folder_service/mock ydb/library/grpc/server/actors ydb/library/security + ydb/library/yql/providers/pq/provider ydb/tests/tools/kqprun/runlib ) diff --git a/ydb/tests/tools/fqrun/ya.make b/ydb/tests/tools/fqrun/ya.make index fa7212f0f8bc..6553c9748b08 100644 --- a/ydb/tests/tools/fqrun/ya.make +++ b/ydb/tests/tools/fqrun/ya.make @@ -1,18 +1,25 @@ -PROGRAM(fqprun) +PROGRAM(fqrun) SRCS( - fqprun.cpp + fqrun.cpp ) PEERDIR( library/cpp/colorizer + library/cpp/getopt util + ydb/library/yql/providers/pq/gateway/dummy ydb/tests/tools/fqrun/src ydb/tests/tools/kqprun/runlib yql/essentials/parser/pg_wrapper yql/essentials/sql/pg ) +PEERDIR( + yql/essentials/udfs/common/compress_base + yql/essentials/udfs/common/re2 +) + YQL_LAST_ABI_VERSION() END() diff --git a/ydb/tests/tools/kqprun/README.md b/ydb/tests/tools/kqprun/README.md new file mode 100644 index 000000000000..e5d82a018c56 --- /dev/null +++ b/ydb/tests/tools/kqprun/README.md @@ -0,0 +1,62 @@ +# KqpRun tool + +Tool can be used to execute queries by using kikimr provider. + +For profiling memory allocations build kqprun with ya make flag `-D PROFILE_MEMORY_ALLOCATIONS`. + +## Examples + +### Queries + +* Run select 42: + ```(bash) + ./kqprun --sql "SELECT 42" + ``` + +* Queries shooting: + ```(bash) + ./kqprun --sql "SELECT 42" -C async --loop-count 0 --loop-delay 100 --inflight-limit 10 + ``` + +### Logs + +* Setup log settings (`-C query` for clear logs): + ```(bash) + ./kqprun --sql "SELECT 42" -C query --log-default=warn --log KQP_YQL=trace --log-file query.log + ``` + +* Trace opt: + ```(bash) + ./kqprun --sql "SELECT 42" -C query -T script + ``` + +* Runtime statistics: + ```(bash) + ./kqprun --sql "SELECT 42" --script-statistics stats.log --script-timeline-file timeline.svg + ``` + +### Cluster + +* Embedded UI: + ```(bash) + ./kqprun -M 32000 + ``` + + Monitoring endpoint: https://localhost:32000 + +* gRPC endpoint: + ```(bash) + ./kqprun -G 32000 + ``` + + Connect with ydb CLI: `ydb -e grpc://localhost:32000 -d /Root` + +* Static storage: + ```(bash) + ./kqprun -M 32000 --storage-path ./storage --storage-size 32 + ``` + +* Create serverless domain and execute query in this domain: + ```(bash) + ./kqprun -M 32000 --shared my-shared --serverless my-serverless --sql "SELECT 42" -D /Root/my-serverless + ``` diff --git a/ydb/tests/tools/kqprun/configuration/app_config.conf b/ydb/tests/tools/kqprun/configuration/app_config.conf index ad27a75fc61a..ba4ff05f187b 100644 --- a/ydb/tests/tools/kqprun/configuration/app_config.conf +++ b/ydb/tests/tools/kqprun/configuration/app_config.conf @@ -99,7 +99,7 @@ QueryServiceConfig { Endpoint { host: "localhost" - port: 50051 + port: 2130 } } } diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index 4505610ce417..e0bbe707bbea 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -424,6 +424,8 @@ TIntrusivePtr CreateFunctionRegistr class TMain : public TMainBase { + using EVerbose = TYdbSetupSettings::EVerbose; + inline static const TString YqlToken = GetEnv(YQL_TOKEN_VARIABLE); inline static IOutputStream* ProfileAllocationsOutput = nullptr; inline static NColorizer::TColors CoutColors = NColorizer::AutoColors(Cout); @@ -473,6 +475,10 @@ class TMain : public TMainBase { ExecutionOptions.ScriptQueries.emplace_back(LoadFile(option->CurVal())); }); + options.AddLongOption("sql", "Script query SQL text to execute (typically DML query)") + .RequiredArgument("str") + .AppendTo(&ExecutionOptions.ScriptQueries); + options.AddLongOption("templates", "Enable templates for -s and -p queries, such as ${YQL_TOKEN} and ${QUERY_ID}") .NoArgument() .SetFlag(&ExecutionOptions.UseTemplates); @@ -660,11 +666,11 @@ class TMain : public TMainBase { .DefaultValue(0) .StoreResult(&RunnerOptions.YdbSettings.AsyncQueriesSettings.InFlightLimit); - options.AddLongOption("verbose", TStringBuilder() << "Common verbose level (max level " << static_cast(TYdbSetupSettings::EVerbose::Max) - 1 << ")") + options.AddLongOption("verbose", TStringBuilder() << "Common verbose level (max level " << static_cast(EVerbose::Max) - 1 << ")") .RequiredArgument("uint") - .DefaultValue(static_cast(TYdbSetupSettings::EVerbose::Info)) + .DefaultValue(static_cast(EVerbose::Info)) .StoreMappedResultT(&RunnerOptions.YdbSettings.VerboseLevel, [](ui8 value) { - return static_cast(std::min(value, static_cast(TYdbSetupSettings::EVerbose::Max))); + return static_cast(std::min(value, static_cast(EVerbose::Max))); }); TChoices verbose({ diff --git a/ydb/tests/tools/kqprun/runlib/application.h b/ydb/tests/tools/kqprun/runlib/application.h index 6004d30fcb7d..88eb2e289011 100644 --- a/ydb/tests/tools/kqprun/runlib/application.h +++ b/ydb/tests/tools/kqprun/runlib/application.h @@ -2,7 +2,7 @@ #include "settings.h" -#include +#include #include diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index f2378965461b..bbc432be16e4 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -162,7 +162,7 @@ class TKqpRunner::TImpl { if (Options_.ResultOutput) { Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Writing script query results..." << CoutColors_.Default() << Endl; for (size_t i = 0; i < ResultSets_.size(); ++i) { - if (ResultSets_.size() > 1 && Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { + if (ResultSets_.size() > 1 && VerboseLevel_ >= EVerbose::Info) { *Options_.ResultOutput << CoutColors_.Cyan() << "Result set " << i + 1 << ":" << CoutColors_.Default() << Endl; } PrintResultSet(Options_.ResultOutputFormat, *Options_.ResultOutput, ResultSets_[i]); @@ -244,7 +244,7 @@ class TKqpRunner::TImpl { void PrintSchemeQueryAst(const TString& ast) const { if (Options_.SchemeQueryAstOutput) { - if (Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { + if (VerboseLevel_ >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Writing scheme query ast" << CoutColors_.Default() << Endl; } Options_.SchemeQueryAstOutput->Write(ast); @@ -253,7 +253,7 @@ class TKqpRunner::TImpl { void PrintScriptAst(size_t queryId, const TString& ast) const { if (const auto output = GetValue(queryId, Options_.ScriptQueryAstOutputs, nullptr)) { - if (Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { + if (VerboseLevel_ >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Writing script query ast" << CoutColors_.Default() << Endl; } output->Write(ast); @@ -262,7 +262,7 @@ class TKqpRunner::TImpl { void PrintScriptPlan(size_t queryId, const TString& plan) const { if (const auto output = GetValue(queryId, Options_.ScriptQueryPlanOutputs, nullptr)) { - if (Options_.YdbSettings.VerboseLevel >= EVerbose::Info) { + if (VerboseLevel_ >= EVerbose::Info) { Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; } StatsPrinter_.PrintPlan(plan, *output); From 5839c03a56f890fe5f815528a0bdd817f854e227 Mon Sep 17 00:00:00 2001 From: Dmitry Kardymon Date: Thu, 20 Feb 2025 13:05:40 +0300 Subject: [PATCH 09/18] YQ-3697 Add partition count to dqrun (#9837) --- .../pq/gateway/dummy/yql_pq_dummy_gateway.cpp | 4 ++-- .../pq/gateway/dummy/yql_pq_dummy_gateway.h | 11 +++++----- .../dummy/yql_pq_file_topic_client.cpp | 22 +++++++++++++------ ydb/tests/tools/fqrun/fqrun.cpp | 20 +++++++++-------- 4 files changed, 34 insertions(+), 23 deletions(-) diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp index 59add76e5ca8..3253d0f47d19 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp @@ -53,8 +53,8 @@ NThreading::TFuture TDummyPqGateway::ListStreams(const TDummyPqGateway& TDummyPqGateway::AddDummyTopic(const TDummyTopic& topic) { with_lock (Mutex) { Y_ENSURE(topic.Cluster); - Y_ENSURE(topic.Path); - const auto key = std::make_pair(topic.Cluster, topic.Path); + Y_ENSURE(topic.TopicName); + const auto key = std::make_pair(topic.Cluster, topic.TopicName); Y_ENSURE(Topics.emplace(key, topic).second, "Already inserted dummy topic {" << topic.Cluster << ", " << topic.Path << "}"); return *this; } diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h index 84a394531ee0..a39fc4f3882e 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h @@ -9,10 +9,11 @@ namespace NYql { struct TDummyTopic { - TDummyTopic(const TString& cluster, const TString& path, const TMaybe& filePath = {}) + TDummyTopic(const TString& cluster, const TString& topicName, const TMaybe& path = {}, size_t partitionCount = 1) : Cluster(cluster) + , TopicName(topicName) , Path(path) - , FilePath(filePath) + , PartitionsCount(partitionCount) { } @@ -22,9 +23,9 @@ struct TDummyTopic { } TString Cluster; - TString Path; - TMaybe FilePath; - size_t PartitionsCount = 1; + TString TopicName; + TMaybe Path; + size_t PartitionsCount; }; // Dummy Pq gateway for tests. diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp index abbaf50e9d04..3348e3a2b503 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp @@ -5,6 +5,7 @@ #include +#include #include #include "yql_pq_blocking_queue.h" @@ -325,17 +326,24 @@ struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession { std::shared_ptr TFileTopicClient::CreateReadSession(const NYdb::NTopic::TReadSessionSettings& settings) { Y_ENSURE(!settings.Topics_.empty()); - TString topicPath = settings.Topics_.front().Path_; - + const auto& topic = settings.Topics_.front(); + auto topicPath = topic.Path_; + Y_ENSURE(topic.PartitionIds_.size() >= 1); + ui64 partitionId = topic.PartitionIds_.front(); auto topicsIt = Topics_.find(make_pair("pq", topicPath)); Y_ENSURE(topicsIt != Topics_.end()); - auto filePath = topicsIt->second.FilePath; + auto filePath = topicsIt->second.Path; Y_ENSURE(filePath); - + + TFsPath fsPath(*filePath); + if (fsPath.IsDirectory()) { + filePath = TStringBuilder() << *filePath << "/" << ToString(partitionId); + } else if (!fsPath.Exists()) { + filePath = TStringBuilder() << *filePath << "_" << partitionId; + } + // TODO ui64 sessionId = 0; - ui64 partitionId = 0; - return std::make_shared( TFile(*filePath, EOpenMode::TEnum::RdOnly), MakeIntrusive(sessionId, topicPath, partitionId) @@ -399,7 +407,7 @@ std::shared_ptr TFileTopicClient::CreateWriteSessio TString topicPath = settings.Path_; auto topicsIt = Topics_.find(make_pair("pq", topicPath)); Y_ENSURE(topicsIt != Topics_.end()); - auto filePath = topicsIt->second.FilePath; + auto filePath = topicsIt->second.Path; Y_ENSURE(filePath); return std::make_shared(TFile(*filePath, EOpenMode::TEnum::RdWr)); diff --git a/ydb/tests/tools/fqrun/fqrun.cpp b/ydb/tests/tools/fqrun/fqrun.cpp index 587bc405b034..c5c0133fc4de 100644 --- a/ydb/tests/tools/fqrun/fqrun.cpp +++ b/ydb/tests/tools/fqrun/fqrun.cpp @@ -163,13 +163,15 @@ class TMain : public TMainBase { options.AddLongOption("emulate-pq", "Emulate YDS with local file, accepts list of tables to emulate with following format: topic@file (can be used in query from cluster `pq`)") .RequiredArgument("topic@file") .Handler1([this](const NLastGetopt::TOptsParser* option) { - TStringBuf topicName; - TStringBuf filePath; - TStringBuf(option->CurVal()).Split('@', topicName, filePath); - if (topicName.empty() || filePath.empty()) { - ythrow yexception() << "Incorrect PQ file mapping, expected form topic@file"; + TStringBuf topicName, others; + TStringBuf(option->CurVal()).Split('@', topicName, others); + TStringBuf path, partitionCountStr; + TStringBuf(others).Split(':', path, partitionCountStr); + size_t partitionCount = !partitionCountStr.empty() ? FromString(partitionCountStr) : 1; + if (topicName.empty() || path.empty()) { + ythrow yexception() << "Incorrect PQ file mapping, expected form topic@path[:partitions_count]" << Endl; } - if (!PqFilesMapping.emplace(topicName, filePath).second) { + if (!PqFilesMapping.emplace(topicName, NYql::TDummyTopic("pq", TString(topicName), TString(path), partitionCount)).second) { ythrow yexception() << "Got duplicated topic name: " << topicName; } }); @@ -222,8 +224,8 @@ class TMain : public TMainBase { if (!PqFilesMapping.empty()) { auto fileGateway = MakeIntrusive(); - for (const auto& [topic, file] : PqFilesMapping) { - fileGateway->AddDummyTopic(NYql::TDummyTopic("pq", TString(topic), TString(file))); + for (const auto& [_, topic] : PqFilesMapping) { + fileGateway->AddDummyTopic(topic); } RunnerOptions.FqSettings.PqGateway = std::move(fileGateway); } @@ -246,7 +248,7 @@ class TMain : public TMainBase { private: TExecutionOptions ExecutionOptions; TRunnerOptions RunnerOptions; - std::unordered_map PqFilesMapping; + std::unordered_map PqFilesMapping; }; } // anonymous namespace From c635c804d27c8bf663f6766a0ecec7b939fac689 Mon Sep 17 00:00:00 2001 From: yumkam Date: Wed, 25 Dec 2024 23:08:35 +0300 Subject: [PATCH 10/18] dummy pq emulation: cosmetics (#12992) --- .../dummy/yql_pq_file_topic_client.cpp | 70 +++++++++++-------- 1 file changed, 41 insertions(+), 29 deletions(-) diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp index 3348e3a2b503..d9176551934a 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp @@ -13,14 +13,17 @@ namespace NYql { class TFileTopicReadSession : public NYdb::NTopic::IReadSession { -constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); +constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); public: - TFileTopicReadSession(TFile file, NYdb::NTopic::TPartitionSession::TPtr session, const TString& producerId = ""): - File_(std::move(file)), Session_(std::move(session)), ProducerId_(producerId), - FilePoller_([this] () { - PollFileForChanges(); - }), Counters_() + TFileTopicReadSession(TFile file, NYdb::NTopic::TPartitionSession::TPtr session, const TString& producerId = "") + : File_(std::move(file)) + , Session_(std::move(session)) + , ProducerId_(producerId) + , FilePoller_([this] () { + PollFileForChanges(); + }) + , Counters_() { Pool_.Start(1); } @@ -37,7 +40,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); Y_UNUSED(maxByteSize); TVector res; - for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { + for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { res.push_back(std::move(*event)); } return res; @@ -60,13 +63,15 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); bool Close(TDuration timeout = TDuration::Max()) override { Y_UNUSED(timeout); + // TODO send TSessionClosedEvent + // XXX (... but if we stop queues, nobody will receive it, needs rethinking) EventsQ_.Stop(); Pool_.Stop(); if (FilePoller_.joinable()) { FilePoller_.join(); } - return true; + return true; // TODO incorrect if EventQ_ was non-empty } NYdb::NTopic::TReaderCounters::TPtr GetCounters() const override { @@ -86,11 +91,11 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); } private: - using TMessageInformation = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation; - using TMessage = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage; + using TMessageInformation = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation; + using TMessage = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage; - TMessageInformation MakeNextMessageInformation(size_t offset, size_t uncompressedSize, const TString& messageGroupId = "") { - auto now = TInstant::Now(); + TMessageInformation MakeNextMessageInformation(size_t offset, size_t uncompressedSize, const TString& messageGroupId = "") { + auto now = TInstant::Now(); TMessageInformation msgInfo( offset, ProducerId_, @@ -104,7 +109,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); ); return msgInfo; } - + TMessage MakeNextMessage(const TString& msgBuff) { TMessage msg(msgBuff, nullptr, MakeNextMessageInformation(MsgOffset_, msgBuff.size()), Session_); return msg; @@ -146,12 +151,14 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); ui64 SeqNo_ = 0; }; -class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb::NTopic::TContinuationTokenIssuer { +class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb::NTopic::TContinuationTokenIssuer { public: - TFileTopicWriteSession(TFile file): - File_(std::move(file)), FileWriter_([this] () { - PushToFile(); - }), Counters_() + explicit TFileTopicWriteSession(TFile file) + : File_(std::move(file)) + , FileWriter_([this] () { + PushToFile(); + }) + , Counters_() { Pool_.Start(1); EventsQ_.Push(NYdb::NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); @@ -170,7 +177,7 @@ class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb TVector GetEvents(bool block, TMaybe maxEventsCount) override { TVector res; - for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { + for (auto event = EventsQ_.Pop(block); !event.Empty() && res.size() < maxEventsCount.GetOrElse(std::numeric_limits::max()); event = EventsQ_.Pop(/*block=*/ false)) { res.push_back(std::move(*event)); } return res; @@ -180,10 +187,10 @@ class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb return NThreading::MakeFuture(SeqNo_); } - void Write(NYdb::NTopic::TContinuationToken&&, NYdb::NTopic::TWriteMessage&& message, + void Write(NYdb::NTopic::TContinuationToken&&, NYdb::NTopic::TWriteMessage&& message, NYdb::NTable::TTransaction* tx) override { Y_UNUSED(tx); - + auto size = message.Data.size(); EventsMsgQ_.Push(TOwningWriteMessage(std::move(message)), size); } @@ -238,6 +245,8 @@ class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb bool Close(TDuration timeout = TDuration::Max()) override { Y_UNUSED(timeout); + // TODO send TSessionClosedEvent + // XXX (... but if we stop queues, nobody will receive it, needs rethinking) EventsQ_.Stop(); EventsMsgQ_.Stop(); Pool_.Stop(); @@ -245,7 +254,7 @@ class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb if (FileWriter_.joinable()) { FileWriter_.join(); } - return true; + return true; // TODO incorrect if Event*Q_ was non-empty } NYdb::NTopic::TWriterCounters::TPtr GetCounters() override { @@ -289,16 +298,19 @@ class TFileTopicWriteSession : public NYdb::NTopic::IWriteSession, private NYdb } } } - + TFile File_; - + // We acquire ownership of messages immediately // TODO: remove extra message copying to and from queue struct TOwningWriteMessage { TString content; NYdb::NTopic::TWriteMessage msg; - - TOwningWriteMessage(NYdb::NTopic::TWriteMessage&& msg): content(msg.Data), msg(std::move(msg)) { + + explicit TOwningWriteMessage(NYdb::NTopic::TWriteMessage&& msg) + : content(msg.Data) + , msg(std::move(msg)) + { msg.Data = content; } }; @@ -368,7 +380,7 @@ NYdb::TAsyncStatus TFileTopicClient::DropTopic(const TString& path, const NYdb:: return NThreading::MakeFuture(NYdb::TStatus(NYdb::EStatus::SUCCESS, {})); } -NYdb::NTopic::TAsyncDescribeTopicResult TFileTopicClient::DescribeTopic(const TString& path, +NYdb::NTopic::TAsyncDescribeTopicResult TFileTopicClient::DescribeTopic(const TString& path, const NYdb::NTopic::TDescribeTopicSettings& settings) { Y_UNUSED(path); Y_UNUSED(settings); @@ -377,7 +389,7 @@ NYdb::NTopic::TAsyncDescribeTopicResult TFileTopicClient::DescribeTopic(const TS return NThreading::MakeFuture(NYdb::NTopic::TDescribeTopicResult(std::move(success), {})); } -NYdb::NTopic::TAsyncDescribeConsumerResult TFileTopicClient::DescribeConsumer(const TString& path, const TString& consumer, +NYdb::NTopic::TAsyncDescribeConsumerResult TFileTopicClient::DescribeConsumer(const TString& path, const TString& consumer, const NYdb::NTopic::TDescribeConsumerSettings& settings) { Y_UNUSED(path); Y_UNUSED(consumer); @@ -387,7 +399,7 @@ NYdb::NTopic::TAsyncDescribeConsumerResult TFileTopicClient::DescribeConsumer(co return NThreading::MakeFuture(NYdb::NTopic::TDescribeConsumerResult(std::move(success), {})); } -NYdb::NTopic::TAsyncDescribePartitionResult TFileTopicClient::DescribePartition(const TString& path, i64 partitionId, +NYdb::NTopic::TAsyncDescribePartitionResult TFileTopicClient::DescribePartition(const TString& path, i64 partitionId, const NYdb::NTopic::TDescribePartitionSettings& settings) { Y_UNUSED(path); Y_UNUSED(partitionId); From dd9b8b089a6954820ce10d2439657d51801c884a Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Thu, 20 Feb 2025 17:44:51 +0500 Subject: [PATCH 11/18] YQ-3561 add features for FQ-run (#14826) --- .../pq/gateway/dummy/yql_pq_dummy_gateway.h | 1 + .../dummy/yql_pq_file_topic_client.cpp | 15 ++- .../gateway/dummy/yql_pq_file_topic_client.h | 6 +- ydb/tests/tools/fqrun/.gitignore | 2 + ydb/tests/tools/fqrun/README.md | 11 +- .../tools/fqrun/configuration/as_config.conf | 46 +++++++ ydb/tests/tools/fqrun/flame_graph.sh | 25 ++++ ydb/tests/tools/fqrun/fqrun.cpp | 65 +++++++++- ydb/tests/tools/fqrun/src/common.h | 4 + ydb/tests/tools/fqrun/src/fq_setup.cpp | 12 ++ ydb/tests/tools/fqrun/src/ya.make | 2 +- ydb/tests/tools/fqrun/ya.make | 8 +- ydb/tests/tools/kqprun/README.md | 4 +- ydb/tests/tools/kqprun/kqprun.cpp | 84 +------------ ydb/tests/tools/kqprun/runlib/application.cpp | 119 +++++++++++++----- ydb/tests/tools/kqprun/runlib/application.h | 20 +++ ydb/tests/tools/kqprun/runlib/utils.cpp | 30 +++++ ydb/tests/tools/kqprun/runlib/utils.h | 24 +++- ydb/tests/tools/kqprun/runlib/ya.make | 4 +- ydb/tests/tools/kqprun/ya.make | 1 - 20 files changed, 351 insertions(+), 132 deletions(-) create mode 100644 ydb/tests/tools/fqrun/configuration/as_config.conf create mode 100755 ydb/tests/tools/fqrun/flame_graph.sh diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h index a39fc4f3882e..fef31d7dc09b 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h @@ -26,6 +26,7 @@ struct TDummyTopic { TString TopicName; TMaybe Path; size_t PartitionsCount; + bool CancelOnFileFinish = false; }; // Dummy Pq gateway for tests. diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp index d9176551934a..9a587aca63c1 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp @@ -16,7 +16,7 @@ class TFileTopicReadSession : public NYdb::NTopic::IReadSession { constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); public: - TFileTopicReadSession(TFile file, NYdb::NTopic::TPartitionSession::TPtr session, const TString& producerId = "") + TFileTopicReadSession(TFile file, NYdb::NTopic::TPartitionSession::TPtr session, const TString& producerId = "", bool cancelOnFileFinish = false) : File_(std::move(file)) , Session_(std::move(session)) , ProducerId_(producerId) @@ -24,6 +24,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); PollFileForChanges(); }) , Counters_() + , CancelOnFileFinish_(cancelOnFileFinish) { Pool_.Start(1); } @@ -123,7 +124,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); size_t size = 0; ui64 maxBatchRowSize = 100; - while (size_t read = fi.ReadLine(rawMsg)) { + while (fi.ReadLine(rawMsg)) { msgs.emplace_back(MakeNextMessage(rawMsg)); MsgOffset_++; if (!maxBatchRowSize--) { @@ -133,6 +134,8 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); } if (!msgs.empty()) { EventsQ_.Push(NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent(msgs, {}, Session_), size); + } else if (CancelOnFileFinish_) { + EventsQ_.Push(NYdb::NTopic::TSessionClosedEvent(NYdb::EStatus::CANCELLED, {NYdb::NIssue::TIssue("PQ file topic was finished")}), size); } Sleep(FILE_POLL_PERIOD); @@ -145,6 +148,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); TString ProducerId_; std::thread FilePoller_; NYdb::NTopic::TReaderCounters::TPtr Counters_; + bool CancelOnFileFinish_ = false; TThreadPool Pool_; size_t MsgOffset_ = 0; @@ -336,6 +340,10 @@ struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession { } }; +TFileTopicClient::TFileTopicClient(THashMap topics) + : Topics_(std::move(topics)) +{} + std::shared_ptr TFileTopicClient::CreateReadSession(const NYdb::NTopic::TReadSessionSettings& settings) { Y_ENSURE(!settings.Topics_.empty()); const auto& topic = settings.Topics_.front(); @@ -358,7 +366,8 @@ std::shared_ptr TFileTopicClient::CreateReadSession( ui64 sessionId = 0; return std::make_shared( TFile(*filePath, EOpenMode::TEnum::RdOnly), - MakeIntrusive(sessionId, topicPath, partitionId) + MakeIntrusive(sessionId, TString{topicPath}, partitionId), + "", topicsIt->second.CancelOnFileFinish ); } diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h index f80426726159..7b4ccae0d85c 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.h @@ -4,7 +4,7 @@ namespace NYql { struct TFileTopicClient : public ITopicClient { - TFileTopicClient(THashMap topics): Topics_(topics) {} + explicit TFileTopicClient(THashMap topics); NYdb::TAsyncStatus CreateTopic(const TString& path, const NYdb::NTopic::TCreateTopicSettings& settings = {}) override; @@ -32,5 +32,7 @@ struct TFileTopicClient : public ITopicClient { private: THashMap Topics_; + bool CancelOnFileFinish_ = false; }; -} \ No newline at end of file + +} diff --git a/ydb/tests/tools/fqrun/.gitignore b/ydb/tests/tools/fqrun/.gitignore index 99dcb6bb049f..ee3ed5154ff0 100644 --- a/ydb/tests/tools/fqrun/.gitignore +++ b/ydb/tests/tools/fqrun/.gitignore @@ -5,3 +5,5 @@ sync_dir *.conf *.parquet *.json +*.svg +*.txt diff --git a/ydb/tests/tools/fqrun/README.md b/ydb/tests/tools/fqrun/README.md index cd133168dcd4..dcaee95ccc7c 100644 --- a/ydb/tests/tools/fqrun/README.md +++ b/ydb/tests/tools/fqrun/README.md @@ -2,6 +2,15 @@ Tool can be used to execute streaming queries by using FQ proxy infrastructure. +For profiling memory allocations build fqrun with ya make flags `-D PROFILE_MEMORY_ALLOCATIONS -D CXXFLAGS=-DPROFILE_MEMORY_ALLOCATIONS`. + +## Scripts + +* `flame_graph.sh` - script for collecting flame graphs in svg format, usage: + ```(bash) + ./flame_graph.sh [graph collection time in seconds] + ``` + ## Examples ### Queries @@ -25,7 +34,7 @@ Tool can be used to execute streaming queries by using FQ proxy infrastructure. ./fqrun -M 32000 ``` - Monitoring endpoint: https://localhost:32000 + Monitoring endpoint: http://localhost:32000 * gRPC endpoint: ```(bash) diff --git a/ydb/tests/tools/fqrun/configuration/as_config.conf b/ydb/tests/tools/fqrun/configuration/as_config.conf new file mode 100644 index 000000000000..c1f7d09e7f6a --- /dev/null +++ b/ydb/tests/tools/fqrun/configuration/as_config.conf @@ -0,0 +1,46 @@ +Executor { + Type: BASIC + Threads: 1 + SpinThreshold: 10 + Name: "System" +} +Executor { + Type: BASIC + Threads: 6 + SpinThreshold: 1 + Name: "User" +} +Executor { + Type: BASIC + Threads: 1 + SpinThreshold: 1 + Name: "Batch" +} +Executor { + Type: IO + Threads: 1 + Name: "IO" +} +Executor { + Type: BASIC + Threads: 2 + SpinThreshold: 10 + Name: "IC" + TimePerMailboxMicroSecs: 100 +} + +Scheduler { + Resolution: 64 + SpinThreshold: 0 + ProgressThreshold: 10000 +} + +SysExecutor: 0 +UserExecutor: 1 +IoExecutor: 3 +BatchExecutor: 2 + +ServiceExecutor { + ServiceName: "Interconnect" + ExecutorId: 4 +} diff --git a/ydb/tests/tools/fqrun/flame_graph.sh b/ydb/tests/tools/fqrun/flame_graph.sh new file mode 100755 index 000000000000..b6f8fe6da4f0 --- /dev/null +++ b/ydb/tests/tools/fqrun/flame_graph.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash + +set -eux + +function cleanup { + sudo rm ./profdata + rm ./profdata.txt +} +trap cleanup EXIT + +if [ $# -gt 1 ]; then + echo "Too many arguments" + exit -1 +fi + +fqrun_pid=$(pgrep -u $USER fqrun) + +sudo perf record -F 50 --call-graph dwarf -g --proc-map-timeout=10000 --pid $fqrun_pid -v -o profdata -- sleep ${1:-'30'} +sudo perf script -i profdata > profdata.txt + +SCRIPT_DIR=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +flame_graph_tool="$SCRIPT_DIR/../../../../contrib/tools/flame-graph/" + +${flame_graph_tool}/stackcollapse-perf.pl profdata.txt | ${flame_graph_tool}/flamegraph.pl > profdata.svg diff --git a/ydb/tests/tools/fqrun/fqrun.cpp b/ydb/tests/tools/fqrun/fqrun.cpp index c5c0133fc4de..4f1b096bde10 100644 --- a/ydb/tests/tools/fqrun/fqrun.cpp +++ b/ydb/tests/tools/fqrun/fqrun.cpp @@ -3,11 +3,16 @@ #include +#include #include #include #include #include +#ifdef PROFILE_MEMORY_ALLOCATIONS +#include +#endif + using namespace NKikimrRun; namespace NFqRun { @@ -156,6 +161,21 @@ class TMain : public TMainBase { } }); + options.AddLongOption("as-cfg", "File with actor system config (TActorSystemConfig), use '-' for default") + .RequiredArgument("file") + .DefaultValue("./configuration/as_config.conf") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + const TString file(option->CurValOrDef()); + if (file == "-") { + return; + } + + RunnerOptions.FqSettings.ActorSystemConfig = NKikimrConfig::TActorSystemConfig(); + if (!google::protobuf::TextFormat::ParseFromString(LoadFile(file), &(*RunnerOptions.FqSettings.ActorSystemConfig))) { + ythrow yexception() << "Bad format of actor system configuration"; + } + }); + options.AddLongOption("emulate-s3", "Enable readings by s3 provider from files, `bucket` value in connection - path to folder with files") .NoArgument() .SetFlag(&RunnerOptions.FqSettings.EmulateS3); @@ -165,17 +185,27 @@ class TMain : public TMainBase { .Handler1([this](const NLastGetopt::TOptsParser* option) { TStringBuf topicName, others; TStringBuf(option->CurVal()).Split('@', topicName, others); + TStringBuf path, partitionCountStr; TStringBuf(others).Split(':', path, partitionCountStr); size_t partitionCount = !partitionCountStr.empty() ? FromString(partitionCountStr) : 1; + if (!partitionCount) { + ythrow yexception() << "Topic partition count should be at least one"; + } if (topicName.empty() || path.empty()) { - ythrow yexception() << "Incorrect PQ file mapping, expected form topic@path[:partitions_count]" << Endl; + ythrow yexception() << "Incorrect PQ file mapping, expected form topic@path[:partitions_count]"; } if (!PqFilesMapping.emplace(topicName, NYql::TDummyTopic("pq", TString(topicName), TString(path), partitionCount)).second) { ythrow yexception() << "Got duplicated topic name: " << topicName; } }); + options.AddLongOption("cancel-on-file-finish", "Cancel emulate YDS topics when topic file finished") + .RequiredArgument("topic") + .Handler1([this](const NLastGetopt::TOptsParser* option) { + TopicsSettings[option->CurVal()].CancelOnFileFinish = true; + }); + // Outputs options.AddLongOption("result-file", "File with query results (use '-' to write in stdout)") @@ -210,6 +240,7 @@ class TMain : public TMainBase { ExecutionOptions.Validate(RunnerOptions); RunnerOptions.FqSettings.YqlToken = GetEnv(YQL_TOKEN_VARIABLE); + RunnerOptions.FqSettings.FunctionRegistry = CreateFunctionRegistry().Get(); auto& gatewayConfig = *RunnerOptions.FqSettings.FqConfig.mutable_gateways(); FillTokens(gatewayConfig.mutable_pq()); @@ -224,14 +255,39 @@ class TMain : public TMainBase { if (!PqFilesMapping.empty()) { auto fileGateway = MakeIntrusive(); - for (const auto& [_, topic] : PqFilesMapping) { + for (auto [_, topic] : PqFilesMapping) { + if (const auto it = TopicsSettings.find(topic.TopicName); it != TopicsSettings.end()) { + topic.CancelOnFileFinish = it->second.CancelOnFileFinish; + TopicsSettings.erase(it); + } fileGateway->AddDummyTopic(topic); } RunnerOptions.FqSettings.PqGateway = std::move(fileGateway); } + if (!TopicsSettings.empty()) { + ythrow yexception() << "Found topic settings for not existing topic: '" << TopicsSettings.begin()->first << "'"; + } + +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.FqSettings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Starting profile memory allocations" << CoutColors.Default() << Endl; + } + NAllocProfiler::StartAllocationSampling(true); +#else + if (ProfileAllocationsOutput) { + ythrow yexception() << "Profile memory allocations disabled, please rebuild fqrun with flag `-D PROFILE_MEMORY_ALLOCATIONS`"; + } +#endif RunScript(ExecutionOptions, RunnerOptions); +#ifdef PROFILE_MEMORY_ALLOCATIONS + if (RunnerOptions.FqSettings.VerboseLevel >= EVerbose::Info) { + Cout << CoutColors.Cyan() << "Finishing profile memory allocations" << CoutColors.Default() << Endl; + } + FinishProfileMemoryAllocations(); +#endif + return 0; } @@ -249,6 +305,11 @@ class TMain : public TMainBase { TExecutionOptions ExecutionOptions; TRunnerOptions RunnerOptions; std::unordered_map PqFilesMapping; + + struct TTopicSettings { + bool CancelOnFileFinish = false; + }; + std::unordered_map TopicsSettings; }; } // anonymous namespace diff --git a/ydb/tests/tools/fqrun/src/common.h b/ydb/tests/tools/fqrun/src/common.h index 53dfbe654e93..457897781891 100644 --- a/ydb/tests/tools/fqrun/src/common.h +++ b/ydb/tests/tools/fqrun/src/common.h @@ -7,6 +7,8 @@ #include #include +#include + namespace NFqRun { constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN"; @@ -27,8 +29,10 @@ struct TFqSetupSettings : public NKikimrRun::TServerSettings { TString YqlToken; NYql::IPqGateway::TPtr PqGateway; + TIntrusivePtr FunctionRegistry; NFq::NConfig::TConfig FqConfig; NKikimrConfig::TLogConfig LogConfig; + std::optional ActorSystemConfig; }; struct TRunnerOptions { diff --git a/ydb/tests/tools/fqrun/src/fq_setup.cpp b/ydb/tests/tools/fqrun/src/fq_setup.cpp index 394bfc5c3e16..a4fbfd104777 100644 --- a/ydb/tests/tools/fqrun/src/fq_setup.cpp +++ b/ydb/tests/tools/fqrun/src/fq_setup.cpp @@ -44,6 +44,14 @@ class TFqSetup::TImpl { serverSettings.SetLoggerInitializer(loggerInitializer); } + void SetFunctionRegistry(NKikimr::Tests::TServerSettings& serverSettings) const { + if (Settings.FunctionRegistry) { + serverSettings.SetFrFactory([this](const NKikimr::NScheme::TTypeRegistry&) { + return Settings.FunctionRegistry.Get(); + }); + } + } + NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) { NKikimr::Tests::TServerSettings serverSettings(PortManager.GetPort()); @@ -52,9 +60,13 @@ class TFqSetup::TImpl { NKikimrConfig::TAppConfig config; *config.MutableLogConfig() = Settings.LogConfig; + if (Settings.ActorSystemConfig) { + *config.MutableActorSystemConfig() = *Settings.ActorSystemConfig; + } serverSettings.SetAppConfig(config); SetLoggerSettings(serverSettings); + SetFunctionRegistry(serverSettings); if (Settings.MonitoringEnabled) { serverSettings.InitKikimrRunConfig(); diff --git a/ydb/tests/tools/fqrun/src/ya.make b/ydb/tests/tools/fqrun/src/ya.make index 6f44cf16c38d..a5a54edf3a87 100644 --- a/ydb/tests/tools/fqrun/src/ya.make +++ b/ydb/tests/tools/fqrun/src/ya.make @@ -9,7 +9,6 @@ SRCS( PEERDIR( library/cpp/colorizer library/cpp/testing/unittest - util ydb/core/fq/libs/config/protos ydb/core/fq/libs/control_plane_proxy/events ydb/core/fq/libs/init @@ -20,6 +19,7 @@ PEERDIR( ydb/library/security ydb/library/yql/providers/pq/provider ydb/tests/tools/kqprun/runlib + yql/essentials/minikql ) YQL_LAST_ABI_VERSION() diff --git a/ydb/tests/tools/fqrun/ya.make b/ydb/tests/tools/fqrun/ya.make index 6553c9748b08..a80046acebed 100644 --- a/ydb/tests/tools/fqrun/ya.make +++ b/ydb/tests/tools/fqrun/ya.make @@ -1,5 +1,10 @@ PROGRAM(fqrun) +IF (PROFILE_MEMORY_ALLOCATIONS) + MESSAGE("Enabled profile memory allocations") + ALLOCATOR(LF_DBG) +ENDIF() + SRCS( fqrun.cpp ) @@ -7,7 +12,8 @@ SRCS( PEERDIR( library/cpp/colorizer library/cpp/getopt - util + library/cpp/lfalloc/alloc_profiler + ydb/core/blob_depot ydb/library/yql/providers/pq/gateway/dummy ydb/tests/tools/fqrun/src ydb/tests/tools/kqprun/runlib diff --git a/ydb/tests/tools/kqprun/README.md b/ydb/tests/tools/kqprun/README.md index e5d82a018c56..10e283294135 100644 --- a/ydb/tests/tools/kqprun/README.md +++ b/ydb/tests/tools/kqprun/README.md @@ -2,7 +2,7 @@ Tool can be used to execute queries by using kikimr provider. -For profiling memory allocations build kqprun with ya make flag `-D PROFILE_MEMORY_ALLOCATIONS`. +For profiling memory allocations build kqprun with ya make flag `-D PROFILE_MEMORY_ALLOCATIONS -D CXXFLAGS=-DPROFILE_MEMORY_ALLOCATIONS`. ## Examples @@ -42,7 +42,7 @@ For profiling memory allocations build kqprun with ya make flag `-D PROFILE_MEMO ./kqprun -M 32000 ``` - Monitoring endpoint: https://localhost:32000 + Monitoring endpoint: http://localhost:32000 * gRPC endpoint: ```(bash) diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index e0bbe707bbea..a9ba125897e1 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -14,9 +14,6 @@ #include #include -#include -#include - #include #include #include @@ -399,62 +396,18 @@ void RunScript(const TExecutionOptions& executionOptions, const TRunnerOptions& } -TIntrusivePtr CreateFunctionRegistry(const TString& udfsDirectory, TVector udfsPaths, bool excludeLinkedUdfs) { - if (!udfsDirectory.empty() || !udfsPaths.empty()) { - NColorizer::TColors colors = NColorizer::AutoColors(Cout); - Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Fetching udfs..." << colors.Default() << Endl; - } - - NKikimr::NMiniKQL::FindUdfsInDir(udfsDirectory, &udfsPaths); - auto functionRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&PrintBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, udfsPaths)->Clone(); - - if (excludeLinkedUdfs) { - for (const auto& wrapper : NYql::NUdf::GetStaticUdfModuleWrapperList()) { - auto [name, ptr] = wrapper(); - if (!functionRegistry->IsLoadedUdfModule(name)) { - functionRegistry->AddModule(TString(NKikimr::NMiniKQL::StaticModulePrefix) + name, name, std::move(ptr)); - } - } - } else { - NKikimr::NMiniKQL::FillStaticModules(*functionRegistry); - } - - return functionRegistry; -} - - class TMain : public TMainBase { using EVerbose = TYdbSetupSettings::EVerbose; inline static const TString YqlToken = GetEnv(YQL_TOKEN_VARIABLE); - inline static IOutputStream* ProfileAllocationsOutput = nullptr; - inline static NColorizer::TColors CoutColors = NColorizer::AutoColors(Cout); TExecutionOptions ExecutionOptions; TRunnerOptions RunnerOptions; std::unordered_map Templates; THashMap TablesMapping; - TVector UdfsPaths; - TString UdfsDirectory; - bool ExcludeLinkedUdfs = false; bool EmulateYt = false; -#ifdef PROFILE_MEMORY_ALLOCATIONS -public: - static void FinishProfileMemoryAllocations() { - if (ProfileAllocationsOutput) { - NAllocProfiler::StopAllocationSampling(*ProfileAllocationsOutput); - } else { - TString output; - TStringOutput stream(output); - NAllocProfiler::StopAllocationSampling(stream); - - Cout << CoutColors.Red() << "Warning: profile memory allocations output is not specified, please use flag `--profile-output` for writing profile info (dump size " << NKikimr::NBlobDepot::FormatByteSize(output.size()) << ")" << CoutColors.Default() << Endl; - } - } -#endif - protected: void RegisterOptions(NLastGetopt::TOpts& options) override { options.SetTitle("KqpRun -- tool to execute queries by using kikimr provider (instead of dq provider in DQrun tool)"); @@ -535,18 +488,6 @@ class TMain : public TMainBase { } }); - options.AddLongOption('u', "udf", "Load shared library with UDF by given path") - .RequiredArgument("file") - .EmplaceTo(&UdfsPaths); - - options.AddLongOption("udfs-dir", "Load all shared libraries with UDFs found in given directory") - .RequiredArgument("directory") - .StoreResult(&UdfsDirectory); - - options.AddLongOption("exclude-linked-udfs", "Exclude linked udfs when same udf passed from -u or --udfs-dir") - .NoArgument() - .SetFlag(&ExcludeLinkedUdfs); - // Outputs TChoices traceOpt({ @@ -641,10 +582,6 @@ class TMain : public TMainBase { RunnerOptions.ScriptQueryTimelineFiles.emplace_back(file); }); - options.AddLongOption("profile-output", "File with profile memory allocations output (use '-' to write in stdout)") - .RequiredArgument("file") - .StoreMappedResultT(&ProfileAllocationsOutput, &GetDefaultOutput); - // Pipeline settings TChoices executionCase({ @@ -842,7 +779,7 @@ class TMain : public TMainBase { } RunnerOptions.YdbSettings.YqlToken = YqlToken; - RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(UdfsDirectory, UdfsPaths, ExcludeLinkedUdfs).Get(); + RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry().Get(); auto& appConfig = RunnerOptions.YdbSettings.AppConfig; if (ExecutionOptions.ResultsRowsLimit) { @@ -862,7 +799,7 @@ class TMain : public TMainBase { } #ifdef PROFILE_MEMORY_ALLOCATIONS - if (RunnerOptions.YdbSettings.VerboseLevel >= 1) { + if (RunnerOptions.YdbSettings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors.Cyan() << "Starting profile memory allocations" << CoutColors.Default() << Endl; } NAllocProfiler::StartAllocationSampling(true); @@ -875,7 +812,7 @@ class TMain : public TMainBase { RunScript(ExecutionOptions, RunnerOptions); #ifdef PROFILE_MEMORY_ALLOCATIONS - if (RunnerOptions.YdbSettings.VerboseLevel >= 1) { + if (RunnerOptions.YdbSettings.VerboseLevel >= EVerbose::Info) { Cout << CoutColors.Cyan() << "Finishing profile memory allocations" << CoutColors.Default() << Endl; } FinishProfileMemoryAllocations(); @@ -900,17 +837,6 @@ class TMain : public TMainBase { } }; -#ifdef PROFILE_MEMORY_ALLOCATIONS -void InterruptHandler(int) { - NColorizer::TColors colors = NColorizer::AutoColors(Cerr); - - Cout << colors.Red() << "Execution interrupted, finishing profile memory allocations..." << colors.Default() << Endl; - TMain::FinishProfileMemoryAllocations(); - - abort(); -} -#endif - } // anonymous namespace } // namespace NKqpRun @@ -918,10 +844,6 @@ void InterruptHandler(int) { int main(int argc, const char* argv[]) { SetupSignalActions(); -#ifdef PROFILE_MEMORY_ALLOCATIONS - signal(SIGINT, &NKqpRun::InterruptHandler); -#endif - try { NKqpRun::TMain().Run(argc, argv); } catch (...) { diff --git a/ydb/tests/tools/kqprun/runlib/application.cpp b/ydb/tests/tools/kqprun/runlib/application.cpp index 9a35d584f19e..6b27760b5b51 100644 --- a/ydb/tests/tools/kqprun/runlib/application.cpp +++ b/ydb/tests/tools/kqprun/runlib/application.cpp @@ -1,15 +1,50 @@ #include "application.h" #include "utils.h" +#include #include #include #include +#include + +#include +#include + +#ifdef PROFILE_MEMORY_ALLOCATIONS +#include +#endif namespace NKikimrRun { +#ifdef PROFILE_MEMORY_ALLOCATIONS +void TMainBase::FinishProfileMemoryAllocations() { + if (ProfileAllocationsOutput) { + NAllocProfiler::StopAllocationSampling(*ProfileAllocationsOutput); + } else { + TString output; + TStringOutput stream(output); + NAllocProfiler::StopAllocationSampling(stream); + + Cout << CoutColors.Red() << "Warning: profile memory allocations output is not specified, please use flag `--profile-output` for writing profile info (dump size " << NKikimr::NBlobDepot::FormatByteSize(output.size()) << ")" << CoutColors.Default() << Endl; + } +} +#endif + void TMainBase::RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettings& settings) { + options.AddLongOption('u', "udf", "Load shared library with UDF by given path") + .RequiredArgument("file") + .EmplaceTo(&UdfsPaths); + + options.AddLongOption("udfs-dir", "Load all shared libraries with UDFs found in given directory") + .RequiredArgument("directory") + .StoreResult(&UdfsDirectory); + + options.AddLongOption("exclude-linked-udfs", "Exclude linked udfs when same udf passed from -u or --udfs-dir") + .NoArgument() + .SetFlag(&ExcludeLinkedUdfs); + options.AddLongOption("log-file", "File with execution logs (writes in stderr if empty)") .RequiredArgument("file") .StoreResult(&settings.LogOutputFile) @@ -19,41 +54,11 @@ void TMainBase::RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettin } }); - TChoices logPriority({ - {"emerg", NActors::NLog::EPriority::PRI_EMERG}, - {"alert", NActors::NLog::EPriority::PRI_ALERT}, - {"crit", NActors::NLog::EPriority::PRI_CRIT}, - {"error", NActors::NLog::EPriority::PRI_ERROR}, - {"warn", NActors::NLog::EPriority::PRI_WARN}, - {"notice", NActors::NLog::EPriority::PRI_NOTICE}, - {"info", NActors::NLog::EPriority::PRI_INFO}, - {"debug", NActors::NLog::EPriority::PRI_DEBUG}, - {"trace", NActors::NLog::EPriority::PRI_TRACE}, - }); - options.AddLongOption("log-default", "Default log priority") - .RequiredArgument("priority") - .Choices(logPriority.GetChoices()) - .StoreMappedResultT(&DefaultLogPriority, logPriority); + RegisterLogOptions(options); - options.AddLongOption("log", "Component log priority in format = (e. g. KQP_YQL=trace)") - .RequiredArgument("component priority") - .Handler1([this, logPriority](const NLastGetopt::TOptsParser* option) { - TStringBuf component; - TStringBuf priority; - TStringBuf(option->CurVal()).Split('=', component, priority); - if (component.empty() || priority.empty()) { - ythrow yexception() << "Incorrect log setting, expected form component=priority, e. g. KQP_YQL=trace"; - } - - if (!logPriority.Contains(TString(priority))) { - ythrow yexception() << "Incorrect log priority: " << priority; - } - - const auto service = GetLogService(TString(component)); - if (!LogPriorities.emplace(service, logPriority(TString(priority))).second) { - ythrow yexception() << "Got duplicated log service name: " << component; - } - }); + options.AddLongOption("profile-output", "File with profile memory allocations output (use '-' to write in stdout)") + .RequiredArgument("file") + .StoreMappedResultT(&ProfileAllocationsOutput, &GetDefaultOutput); options.AddLongOption('M', "monitoring", "Embedded UI port (use 0 to start on random free port), if used will be run as daemon") .RequiredArgument("uint") @@ -92,6 +97,28 @@ void TMainBase::RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettin }); } +void TMainBase::RegisterLogOptions(NLastGetopt::TOpts& options) { + options.AddLongOption("log-default", "Default log priority") + .RequiredArgument("priority") + .StoreMappedResultT(&DefaultLogPriority, GetLogPrioritiesMap("log-default")); + + options.AddLongOption("log", "Component log priority in format = (e. g. KQP_YQL=trace)") + .RequiredArgument("component priority") + .Handler1([this, logPriority = GetLogPrioritiesMap("log")](const NLastGetopt::TOptsParser* option) { + TStringBuf component; + TStringBuf priority; + TStringBuf(option->CurVal()).Split('=', component, priority); + if (component.empty() || priority.empty()) { + ythrow yexception() << "Incorrect log setting, expected form component=priority, e. g. KQP_YQL=trace"; + } + + const auto service = GetLogService(TString(component)); + if (!LogPriorities.emplace(service, logPriority(TString(priority))).second) { + ythrow yexception() << "Got duplicated log service name: " << component; + } + }); +} + void TMainBase::FillLogConfig(NKikimrConfig::TLogConfig& config) const { if (DefaultLogPriority) { config.SetDefaultLevel(*DefaultLogPriority); @@ -110,4 +137,28 @@ IOutputStream* TMainBase::GetDefaultOutput(const TString& file) { return nullptr; } +TIntrusivePtr TMainBase::CreateFunctionRegistry() const { + if (!UdfsDirectory.empty() || !UdfsPaths.empty()) { + NColorizer::TColors colors = NColorizer::AutoColors(Cout); + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Fetching udfs..." << colors.Default() << Endl; + } + + auto paths = UdfsPaths; + NKikimr::NMiniKQL::FindUdfsInDir(UdfsDirectory, &paths); + auto functionRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&PrintBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, paths)->Clone(); + + if (ExcludeLinkedUdfs) { + for (const auto& wrapper : NYql::NUdf::GetStaticUdfModuleWrapperList()) { + auto [name, ptr] = wrapper(); + if (!functionRegistry->IsLoadedUdfModule(name)) { + functionRegistry->AddModule(TString(NKikimr::NMiniKQL::StaticModulePrefix) + name, name, std::move(ptr)); + } + } + } else { + NKikimr::NMiniKQL::FillStaticModules(*functionRegistry); + } + + return functionRegistry; +} + } // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/application.h b/ydb/tests/tools/kqprun/runlib/application.h index 88eb2e289011..fd67eb2ac0d9 100644 --- a/ydb/tests/tools/kqprun/runlib/application.h +++ b/ydb/tests/tools/kqprun/runlib/application.h @@ -2,6 +2,7 @@ #include "settings.h" +#include #include #include @@ -10,21 +11,40 @@ #include #include +#include + namespace NKikimrRun { class TMainBase : public TMainClassArgs { +#ifdef PROFILE_MEMORY_ALLOCATIONS +public: + static void FinishProfileMemoryAllocations(); +#endif + protected: void RegisterKikimrOptions(NLastGetopt::TOpts& options, TServerSettings& settings); + virtual void RegisterLogOptions(NLastGetopt::TOpts& options); + void FillLogConfig(NKikimrConfig::TLogConfig& config) const; static IOutputStream* GetDefaultOutput(const TString& file); + TIntrusivePtr CreateFunctionRegistry() const; + +protected: + inline static NColorizer::TColors CoutColors = NColorizer::AutoColors(Cout); + inline static IOutputStream* ProfileAllocationsOutput = nullptr; + private: inline static std::vector> FileHolders; std::optional DefaultLogPriority; std::unordered_map LogPriorities; + + TString UdfsDirectory; + TVector UdfsPaths; + bool ExcludeLinkedUdfs; }; } // namespace NKikimrRun diff --git a/ydb/tests/tools/kqprun/runlib/utils.cpp b/ydb/tests/tools/kqprun/runlib/utils.cpp index 208cf4daf045..f8c37d7b67fe 100644 --- a/ydb/tests/tools/kqprun/runlib/utils.cpp +++ b/ydb/tests/tools/kqprun/runlib/utils.cpp @@ -1,4 +1,5 @@ #include "utils.h" +#include "application.h" #include #include @@ -49,6 +50,17 @@ void FloatingPointExceptionHandler(int) { abort(); } +#ifdef PROFILE_MEMORY_ALLOCATIONS +void InterruptHandler(int) { + NColorizer::TColors colors = NColorizer::AutoColors(Cerr); + + Cout << colors.Red() << "Execution interrupted, finishing profile memory allocations..." << colors.Default() << Endl; + TMainBase::FinishProfileMemoryAllocations(); + + abort(); +} +#endif + } // nonymous namespace @@ -233,10 +245,28 @@ void InitLogSettings(const NKikimrConfig::TLogConfig& logConfig, NActors::TTestA } } +TChoices GetLogPrioritiesMap(const TString& optionName) { + return TChoices({ + {"emerg", NActors::NLog::EPriority::PRI_EMERG}, + {"alert", NActors::NLog::EPriority::PRI_ALERT}, + {"crit", NActors::NLog::EPriority::PRI_CRIT}, + {"error", NActors::NLog::EPriority::PRI_ERROR}, + {"warn", NActors::NLog::EPriority::PRI_WARN}, + {"notice", NActors::NLog::EPriority::PRI_NOTICE}, + {"info", NActors::NLog::EPriority::PRI_INFO}, + {"debug", NActors::NLog::EPriority::PRI_DEBUG}, + {"trace", NActors::NLog::EPriority::PRI_TRACE}, + }, optionName, false); +} + void SetupSignalActions() { std::set_terminate(&TerminateHandler); signal(SIGSEGV, &SegmentationFaultHandler); signal(SIGFPE, &FloatingPointExceptionHandler); + +#ifdef PROFILE_MEMORY_ALLOCATIONS + signal(SIGINT, &InterruptHandler); +#endif } void PrintResultSet(EResultOutputFormat format, IOutputStream& output, const Ydb::ResultSet& resultSet) { diff --git a/ydb/tests/tools/kqprun/runlib/utils.h b/ydb/tests/tools/kqprun/runlib/utils.h index 9e4d69c14501..4bf873e48cae 100644 --- a/ydb/tests/tools/kqprun/runlib/utils.h +++ b/ydb/tests/tools/kqprun/runlib/utils.h @@ -34,12 +34,26 @@ struct TRequestResult { template class TChoices { public: - explicit TChoices(std::map choicesMap) + explicit TChoices(std::map choicesMap, const TString& optionName = "", bool checkRegister = true) : ChoicesMap(std::move(choicesMap)) + , OptionName(optionName) + , CheckRegister(checkRegister) {} - TResult operator()(const TString& choice) const { - return ChoicesMap.at(choice); + TResult operator()(TString choice) const { + if (!CheckRegister) { + std::for_each(choice.begin(), choice.vend(), [](char& c) { c = std::tolower(c); }); + } + + const auto it = ChoicesMap.find(choice); + if (it == ChoicesMap.end()) { + auto error = yexception() << "Value '" << choice << "' is not allowed " << (OptionName ? TStringBuilder() << "for option " << OptionName : TStringBuilder()) << ", available variants:"; + for (const auto& [value, _] : ChoicesMap) { + error << " " << value; + } + throw error; + } + return it->second; } TVector GetChoices() const { @@ -57,6 +71,8 @@ class TChoices { private: const std::map ChoicesMap; + const TString OptionName; + const bool CheckRegister; }; class TStatsPrinter { @@ -88,6 +104,8 @@ void ModifyLogPriorities(std::unordered_map GetLogPrioritiesMap(const TString& optionName); + void SetupSignalActions(); void PrintResultSet(EResultOutputFormat format, IOutputStream& output, const Ydb::ResultSet& resultSet); diff --git a/ydb/tests/tools/kqprun/runlib/ya.make b/ydb/tests/tools/kqprun/runlib/ya.make index 352a3a423105..39bd1328bfe7 100644 --- a/ydb/tests/tools/kqprun/runlib/ya.make +++ b/ydb/tests/tools/kqprun/runlib/ya.make @@ -9,7 +9,6 @@ PEERDIR( library/cpp/colorizer library/cpp/getopt library/cpp/json - util ydb/core/base ydb/core/blob_depot ydb/core/fq/libs/compute/common @@ -20,7 +19,10 @@ PEERDIR( ydb/public/api/protos ydb/public/lib/json_value ydb/public/lib/ydb_cli/common + yql/essentials/minikql + yql/essentials/minikql/invoke_builtins yql/essentials/public/issue + yql/essentials/public/udf ) YQL_LAST_ABI_VERSION() diff --git a/ydb/tests/tools/kqprun/ya.make b/ydb/tests/tools/kqprun/ya.make index 81689d3bd5a3..b39159f0ea40 100644 --- a/ydb/tests/tools/kqprun/ya.make +++ b/ydb/tests/tools/kqprun/ya.make @@ -3,7 +3,6 @@ PROGRAM(kqprun) IF (PROFILE_MEMORY_ALLOCATIONS) MESSAGE("Enabled profile memory allocations") ALLOCATOR(LF_DBG) - CFLAGS(-D PROFILE_MEMORY_ALLOCATIONS) ENDIF() SRCS( From 05259837307ce572d84007af530140108cda4a8e Mon Sep 17 00:00:00 2001 From: Dmitry Kardymon Date: Wed, 26 Feb 2025 09:47:03 +0300 Subject: [PATCH 12/18] YQ-4108 Add topic session thread num settings (#14532) --- ydb/core/fq/libs/actors/pending_fetcher.cpp | 12 +++--- ydb/core/fq/libs/actors/proxy.h | 2 +- ydb/core/fq/libs/actors/run_actor.cpp | 12 ++---- .../libs/compute/common/run_actor_params.cpp | 4 +- .../fq/libs/compute/common/run_actor_params.h | 4 +- ydb/core/fq/libs/config/protos/common.proto | 2 + ydb/core/fq/libs/init/init.cpp | 40 +++++++++++++++---- ydb/core/fq/libs/init/init.h | 2 +- .../fq/libs/read_rule/read_rule_creator.cpp | 12 +++++- .../fq/libs/read_rule/read_rule_creator.h | 2 + .../fq/libs/read_rule/read_rule_deleter.cpp | 12 +++++- .../fq/libs/read_rule/read_rule_deleter.h | 2 + .../fq/libs/row_dispatcher/topic_session.cpp | 5 +-- .../pq/async_io/dq_pq_read_actor.cpp | 2 +- .../pq/async_io/dq_pq_write_actor.cpp | 2 +- .../pq/gateway/dummy/yql_pq_dummy_gateway.cpp | 24 +++++++++++ .../pq/gateway/dummy/yql_pq_dummy_gateway.h | 7 +++- .../pq/gateway/native/yql_pq_gateway.cpp | 32 +++++++++++++-- .../pq/gateway/native/yql_pq_gateway.h | 7 +++- .../providers/pq/provider/yql_pq_gateway.h | 21 ++++++++++ ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp | 6 +++ ydb/tests/tools/fq_runner/kikimr_runner.py | 7 ++++ ydb/tests/tools/fqrun/fqrun.cpp | 2 +- ydb/tests/tools/fqrun/src/common.h | 2 +- ydb/tests/tools/fqrun/src/fq_setup.cpp | 2 +- 25 files changed, 180 insertions(+), 45 deletions(-) diff --git a/ydb/core/fq/libs/actors/pending_fetcher.cpp b/ydb/core/fq/libs/actors/pending_fetcher.cpp index c9d437a361f7..99cef5053e0c 100644 --- a/ydb/core/fq/libs/actors/pending_fetcher.cpp +++ b/ydb/core/fq/libs/actors/pending_fetcher.cpp @@ -158,7 +158,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { const TString& tenantName, NActors::TMon* monitoring, std::shared_ptr s3ActorsFactory, - NYql::IPqGateway::TPtr defaultPqGateway + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ) : YqSharedResources(yqSharedResources) , CredentialsProviderFactory(credentialsProviderFactory) @@ -181,7 +181,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { , Monitoring(monitoring) , ComputeConfig(config.GetCompute()) , S3ActorsFactory(std::move(s3ActorsFactory)) - , DefaultPqGateway(std::move(defaultPqGateway)) + , PqGatewayFactory(std::move(pqGatewayFactory)) { Y_ENSURE(GetYqlDefaultModuleResolverWithContext(ModuleResolver)); } @@ -475,7 +475,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { std::map(task.parameters().begin(), task.parameters().end()), S3ActorsFactory, ComputeConfig.GetWorkloadManagerConfig(task.scope()), - DefaultPqGateway + PqGatewayFactory ); auto runActorId = @@ -551,7 +551,7 @@ class TPendingFetcher : public NActors::TActorBootstrapped { NActors::TMon* Monitoring; TComputeConfig ComputeConfig; std::shared_ptr S3ActorsFactory; - NYql::IPqGateway::TPtr DefaultPqGateway; + NYql::IPqGatewayFactory::TPtr PqGatewayFactory; }; @@ -572,7 +572,7 @@ NActors::IActor* CreatePendingFetcher( const TString& tenantName, NActors::TMon* monitoring, std::shared_ptr s3ActorsFactory, - NYql::IPqGateway::TPtr defaultPqGateway) + NYql::IPqGatewayFactory::TPtr pqGatewayFactory) { return new TPendingFetcher( yqSharedResources, @@ -591,7 +591,7 @@ NActors::IActor* CreatePendingFetcher( tenantName, monitoring, std::move(s3ActorsFactory), - defaultPqGateway); + std::move(pqGatewayFactory)); } TActorId MakePendingFetcherId(ui32 nodeId) { diff --git a/ydb/core/fq/libs/actors/proxy.h b/ydb/core/fq/libs/actors/proxy.h index 9f457859189b..7ebe1f378bf6 100644 --- a/ydb/core/fq/libs/actors/proxy.h +++ b/ydb/core/fq/libs/actors/proxy.h @@ -55,7 +55,7 @@ NActors::IActor* CreatePendingFetcher( const TString& tenantName, NActors::TMon* monitoring, std::shared_ptr s3ActorsFactory, - NYql::IPqGateway::TPtr defaultPqGateway + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ); NActors::IActor* CreateRunActor( diff --git a/ydb/core/fq/libs/actors/run_actor.cpp b/ydb/core/fq/libs/actors/run_actor.cpp index f18b666a4423..c1bba77dee44 100644 --- a/ydb/core/fq/libs/actors/run_actor.cpp +++ b/ydb/core/fq/libs/actors/run_actor.cpp @@ -571,6 +571,7 @@ class TRunActor : public NActors::TActorBootstrapped { SelfId(), Params.QueryId, Params.YqSharedResources->UserSpaceYdbDriver, + Params.PqGatewayFactory->CreatePqGateway(), Params.Resources.topic_consumers(), PrepareReadRuleCredentials() ) @@ -1435,6 +1436,7 @@ class TRunActor : public NActors::TActorBootstrapped { SelfId(), Params.QueryId, Params.YqSharedResources->UserSpaceYdbDriver, + Params.PqGatewayFactory->CreatePqGateway(), Params.Resources.topic_consumers(), PrepareReadRuleCredentials() ) @@ -1970,14 +1972,8 @@ class TRunActor : public NActors::TActorBootstrapped { } { - NYql::TPqGatewayServices pqServices( - Params.YqSharedResources->UserSpaceYdbDriver, - Params.PqCmConnections, - Params.CredentialsFactory, - std::make_shared(gatewaysConfig.GetPq()), - Params.FunctionRegistry - ); - const auto pqGateway = Params.DefaultPqGateway ? Params.DefaultPqGateway : NYql::CreatePqNativeGateway(pqServices); + auto pqGateway = Params.PqGatewayFactory->CreatePqGateway(); + pqGateway->UpdateClusterConfigs(std::make_shared(gatewaysConfig.GetPq())); dataProvidersInit.push_back(GetPqDataProviderInitializer(pqGateway, false, dbResolver)); } diff --git a/ydb/core/fq/libs/compute/common/run_actor_params.cpp b/ydb/core/fq/libs/compute/common/run_actor_params.cpp index 82071360cab9..2af6a14c8293 100644 --- a/ydb/core/fq/libs/compute/common/run_actor_params.cpp +++ b/ydb/core/fq/libs/compute/common/run_actor_params.cpp @@ -61,7 +61,7 @@ TRunActorParams::TRunActorParams( std::map&& queryParameters, std::shared_ptr s3ActorsFactory, const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager, - NYql::IPqGateway::TPtr defaultPqGateway + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ) : YqSharedResources(yqSharedResources) , CredentialsProviderFactory(credentialsProviderFactory) @@ -118,7 +118,7 @@ TRunActorParams::TRunActorParams( , QueryParameters(std::move(queryParameters)) , S3ActorsFactory(std::move(s3ActorsFactory)) , WorkloadManager(workloadManager) - , DefaultPqGateway(defaultPqGateway) + , PqGatewayFactory(std::move(pqGatewayFactory)) { } diff --git a/ydb/core/fq/libs/compute/common/run_actor_params.h b/ydb/core/fq/libs/compute/common/run_actor_params.h index 1ff1db905cd5..f49419a6334d 100644 --- a/ydb/core/fq/libs/compute/common/run_actor_params.h +++ b/ydb/core/fq/libs/compute/common/run_actor_params.h @@ -81,7 +81,7 @@ struct TRunActorParams { // TODO2 : Change name std::map&& queryParameters, std::shared_ptr s3ActorsFactory, const ::NFq::NConfig::TWorkloadManagerConfig& workloadManager, - NYql::IPqGateway::TPtr defaultPqGateway + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ); TRunActorParams(const TRunActorParams& params) = default; @@ -147,7 +147,7 @@ struct TRunActorParams { // TODO2 : Change name std::map QueryParameters; std::shared_ptr S3ActorsFactory; ::NFq::NConfig::TWorkloadManagerConfig WorkloadManager; - NYql::IPqGateway::TPtr DefaultPqGateway; + NYql::IPqGatewayFactory::TPtr PqGatewayFactory; }; } /* NFq */ diff --git a/ydb/core/fq/libs/config/protos/common.proto b/ydb/core/fq/libs/config/protos/common.proto index 4d89d11a35aa..b4b7a974b7c6 100644 --- a/ydb/core/fq/libs/config/protos/common.proto +++ b/ydb/core/fq/libs/config/protos/common.proto @@ -32,4 +32,6 @@ message TCommonConfig { bool ShowQueryTimeline = 16; uint64 MaxQueryTimelineSize = 17; // default: 200KB string PqReconnectPeriod = 18; // default: disabled + uint32 TopicClientHandlersExecutorThreadsNum = 19; // default: 0 that means use default from TopicClientSettings (1) + uint32 TopicClientCompressionExecutorThreadsNum = 20; // default: 0 that means use default from TopicClientSettings (2) } diff --git a/ydb/core/fq/libs/init/init.cpp b/ydb/core/fq/libs/init/init.cpp index 3bdc089b3677..85bd1296a8d7 100644 --- a/ydb/core/fq/libs/init/init.cpp +++ b/ydb/core/fq/libs/init/init.cpp @@ -56,6 +56,17 @@ namespace NFq { using namespace NKikimr; +NYdb::NTopic::TTopicClientSettings GetCommonTopicClientSettings(const NFq::NConfig::TCommonConfig& config) { + NYdb::NTopic::TTopicClientSettings settings; + if (config.GetTopicClientHandlersExecutorThreadsNum()) { + settings.DefaultHandlersExecutor(NYdb::NTopic::CreateThreadPoolExecutor(config.GetTopicClientHandlersExecutorThreadsNum())); + } + if (config.GetTopicClientCompressionExecutorThreadsNum()) { + settings.DefaultCompressionExecutor(NYdb::NTopic::CreateThreadPoolExecutor(config.GetTopicClientCompressionExecutorThreadsNum())); + } + return settings; +} + void Init( const NFq::NConfig::TConfig& protoConfig, ui32 nodeId, @@ -67,7 +78,7 @@ void Init( const std::function& folderServiceFactory, ui32 icPort, const std::vector& additionalCompNodeFactories, - NYql::IPqGateway::TPtr defaultPqGateway + NYql::IPqGatewayFactory::TPtr pqGatewayFactory ) { Y_ABORT_UNLESS(iyqSharedResources, "No YQ shared resources created"); @@ -190,14 +201,18 @@ void Init( credentialsFactory = NYql::CreateSecuredServiceAccountCredentialsOverTokenAccessorFactory(tokenAccessorConfig.GetEndpoint(), tokenAccessorConfig.GetUseSsl(), caContent, tokenAccessorConfig.GetConnectionPoolSize()); } + auto commonTopicClientSettings = GetCommonTopicClientSettings(protoConfig.GetCommon()); + if (protoConfig.GetRowDispatcher().GetEnabled()) { NYql::TPqGatewayServices pqServices( yqSharedResources->UserSpaceYdbDriver, nullptr, nullptr, std::make_shared(), - nullptr); - + nullptr, + nullptr, + commonTopicClientSettings + ); auto rowDispatcher = NFq::NewRowDispatcherService( protoConfig.GetRowDispatcher(), NKikimr::CreateYdbCredentialsProviderFactory, @@ -205,7 +220,7 @@ void Init( credentialsFactory, tenant, yqCounters->GetSubgroup("subsystem", "row_dispatcher"), - defaultPqGateway ? defaultPqGateway : CreatePqNativeGateway(pqServices), + pqGatewayFactory ? pqGatewayFactory->CreatePqGateway() : CreatePqNativeGateway(pqServices), appData->Mon, appData->Counters); actorRegistrator(NFq::RowDispatcherServiceActorId(), rowDispatcher.release()); @@ -224,9 +239,11 @@ void Init( pqCmConnections, credentialsFactory, std::make_shared(protoConfig.GetGateways().GetPq()), - appData->FunctionRegistry + appData->FunctionRegistry, + nullptr, + commonTopicClientSettings ); - auto pqGateway = defaultPqGateway ? defaultPqGateway : NYql::CreatePqNativeGateway(std::move(pqServices)); + auto pqGateway = pqGatewayFactory ? pqGatewayFactory->CreatePqGateway() : NYql::CreatePqNativeGateway(std::move(pqServices)); RegisterDqPqReadActorFactory(*asyncIoFactory, yqSharedResources->UserSpaceYdbDriver, credentialsFactory, pqGateway, yqCounters->GetSubgroup("subsystem", "DqSourceTracker"), protoConfig.GetCommon().GetPqReconnectPeriod()); @@ -330,6 +347,15 @@ void Init( } if (protoConfig.GetPendingFetcher().GetEnabled()) { + NYql::TPqGatewayServices pqServices( + yqSharedResources->UserSpaceYdbDriver, + pqCmConnections, + credentialsFactory, + std::make_shared(protoConfig.GetGateways().GetPq()), + appData->FunctionRegistry, + nullptr, + commonTopicClientSettings + ); auto fetcher = CreatePendingFetcher( yqSharedResources, NKikimr::CreateYdbCredentialsProviderFactory, @@ -347,7 +373,7 @@ void Init( tenant, appData->Mon, s3ActorsFactory, - defaultPqGateway + pqGatewayFactory ? pqGatewayFactory : NYql::CreatePqNativeGatewayFactory(pqServices) ); actorRegistrator(MakePendingFetcherId(nodeId), fetcher); diff --git a/ydb/core/fq/libs/init/init.h b/ydb/core/fq/libs/init/init.h index 517557b04d8d..ad1bca9a431f 100644 --- a/ydb/core/fq/libs/init/init.h +++ b/ydb/core/fq/libs/init/init.h @@ -38,7 +38,7 @@ void Init( const std::function& folderServiceFactory, ui32 icPort, const std::vector& additionalCompNodeFactories, - NYql::IPqGateway::TPtr defaultPqGateway = nullptr + NYql::IPqGatewayFactory::TPtr pqGatewayFactory = nullptr ); } // NFq diff --git a/ydb/core/fq/libs/read_rule/read_rule_creator.cpp b/ydb/core/fq/libs/read_rule/read_rule_creator.cpp index e926018a2be5..fb1f111ad9e0 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_creator.cpp +++ b/ydb/core/fq/libs/read_rule/read_rule_creator.cpp @@ -71,6 +71,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const Fq::Private::TopicConsumer& topicConsumer, std::shared_ptr credentialsProvider, ui64 index @@ -79,6 +80,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped , QueryId(std::move(queryId)) , TopicConsumer(topicConsumer) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , TopicClient(YdbDriver, GetTopicClientSettings(std::move(credentialsProvider))) , Index(index) { @@ -183,7 +185,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped private: NYdb::NTopic::TTopicClientSettings GetTopicClientSettings(std::shared_ptr credentialsProvider) { - return NYdb::NTopic::TTopicClientSettings() + return PqGateway->GetTopicClientSettings() .Database(TopicConsumer.database()) .DiscoveryEndpoint(TopicConsumer.cluster_endpoint()) .CredentialsProviderFactory(std::move(credentialsProvider)) @@ -196,6 +198,7 @@ class TSingleReadRuleCreator : public TActorBootstrapped const TString QueryId; const Fq::Private::TopicConsumer TopicConsumer; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; NYdb::NTopic::TTopicClient TopicClient; ui64 Index = 0; NYdb::NTopic::IRetryPolicy::IRetryState::TPtr RetryState; @@ -210,12 +213,14 @@ class TReadRuleCreator : public TActorBootstrapped { NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials ) : Owner(owner) , QueryId(std::move(queryId)) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , TopicConsumers(VectorFromProto(topicConsumers)) , Credentials(std::move(credentials)) { @@ -232,7 +237,7 @@ class TReadRuleCreator : public TActorBootstrapped { Results.reserve(TopicConsumers.size()); for (size_t i = 0; i < TopicConsumers.size(); ++i) { LOG_D("Create read rule creation actor for `" << TopicConsumers[i].topic_path() << "` [" << i << "]"); - Children.push_back(Register(new TSingleReadRuleCreator(SelfId(), QueryId, YdbDriver, TopicConsumers[i], Credentials[i], i))); + Children.push_back(Register(new TSingleReadRuleCreator(SelfId(), QueryId, YdbDriver, PqGateway, TopicConsumers[i], Credentials[i], i))); } } @@ -281,6 +286,7 @@ class TReadRuleCreator : public TActorBootstrapped { const NActors::TActorId Owner; const TString QueryId; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; const TVector TopicConsumers; const TVector> Credentials; size_t ResultsGot = 0; @@ -295,6 +301,7 @@ NActors::IActor* MakeReadRuleCreatorActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials ) @@ -303,6 +310,7 @@ NActors::IActor* MakeReadRuleCreatorActor( owner, std::move(queryId), std::move(ydbDriver), + pqGateway, topicConsumers, std::move(credentials) ); diff --git a/ydb/core/fq/libs/read_rule/read_rule_creator.h b/ydb/core/fq/libs/read_rule/read_rule_creator.h index 51873fb98272..87b4b4ede5d2 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_creator.h +++ b/ydb/core/fq/libs/read_rule/read_rule_creator.h @@ -4,6 +4,7 @@ #include #include +#include namespace NFq { @@ -11,6 +12,7 @@ NActors::IActor* MakeReadRuleCreatorActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials // For each topic ); diff --git a/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp b/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp index 7a3837fbbe02..970d87f838a0 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp +++ b/ydb/core/fq/libs/read_rule/read_rule_deleter.cpp @@ -67,6 +67,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, Fq::Private::TopicConsumer topic, std::shared_ptr credentialsProvider, ui64 index, @@ -76,6 +77,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped , QueryId(std::move(queryId)) , Topic(std::move(topic)) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , TopicClient(YdbDriver, GetTopicClientSettings(std::move(credentialsProvider))) , Index(index) , MaxRetries(maxRetries) @@ -158,7 +160,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped private: NYdb::NTopic::TTopicClientSettings GetTopicClientSettings(std::shared_ptr credentialsProvider) { - return NYdb::NTopic::TTopicClientSettings() + return PqGateway->GetTopicClientSettings() .Database(Topic.database()) .DiscoveryEndpoint(Topic.cluster_endpoint()) .CredentialsProviderFactory(std::move(credentialsProvider)) @@ -171,6 +173,7 @@ class TSingleReadRuleDeleter : public TActorBootstrapped const TString QueryId; const Fq::Private::TopicConsumer Topic; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; NYdb::NTopic::TTopicClient TopicClient; ui64 Index = 0; const size_t MaxRetries; @@ -184,6 +187,7 @@ class TReadRuleDeleter : public TActorBootstrapped { NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials, size_t maxRetries @@ -191,6 +195,7 @@ class TReadRuleDeleter : public TActorBootstrapped { : Owner(owner) , QueryId(std::move(queryId)) , YdbDriver(std::move(ydbDriver)) + , PqGateway(pqGateway) , Topics(VectorFromProto(topicConsumers)) , Credentials(std::move(credentials)) , MaxRetries(maxRetries) @@ -206,7 +211,7 @@ class TReadRuleDeleter : public TActorBootstrapped { Results.reserve(Topics.size()); for (size_t i = 0; i < Topics.size(); ++i) { LOG_D("Create read rule deleter actor for `" << Topics[i].topic_path() << "` [" << i << "]"); - Children.push_back(Register(new TSingleReadRuleDeleter(SelfId(), QueryId, YdbDriver, Topics[i], Credentials[i], i, MaxRetries))); + Children.push_back(Register(new TSingleReadRuleDeleter(SelfId(), QueryId, YdbDriver, PqGateway, Topics[i], Credentials[i], i, MaxRetries))); } } @@ -257,6 +262,7 @@ class TReadRuleDeleter : public TActorBootstrapped { const NActors::TActorId Owner; const TString QueryId; NYdb::TDriver YdbDriver; + NYql::IPqGateway::TPtr PqGateway; const TVector Topics; const TVector> Credentials; const size_t MaxRetries; @@ -272,6 +278,7 @@ NActors::IActor* MakeReadRuleDeleterActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials, // For each topic size_t maxRetries @@ -281,6 +288,7 @@ NActors::IActor* MakeReadRuleDeleterActor( owner, std::move(queryId), std::move(ydbDriver), + pqGateway, topicConsumers, std::move(credentials), maxRetries diff --git a/ydb/core/fq/libs/read_rule/read_rule_deleter.h b/ydb/core/fq/libs/read_rule/read_rule_deleter.h index bcf332b0f81e..a23e23260248 100644 --- a/ydb/core/fq/libs/read_rule/read_rule_deleter.h +++ b/ydb/core/fq/libs/read_rule/read_rule_deleter.h @@ -4,6 +4,7 @@ #include #include +#include namespace NFq { @@ -11,6 +12,7 @@ NActors::IActor* MakeReadRuleDeleterActor( NActors::TActorId owner, TString queryId, NYdb::TDriver ydbDriver, + const NYql::IPqGateway::TPtr& pqGateway, const ::google::protobuf::RepeatedPtrField& topicConsumers, TVector> credentials, // For each topic size_t maxRetries = 15 diff --git a/ydb/core/fq/libs/row_dispatcher/topic_session.cpp b/ydb/core/fq/libs/row_dispatcher/topic_session.cpp index d3fe3fc4729b..976937172a0d 100644 --- a/ydb/core/fq/libs/row_dispatcher/topic_session.cpp +++ b/ydb/core/fq/libs/row_dispatcher/topic_session.cpp @@ -409,12 +409,11 @@ void TTopicSession::SubscribeOnNextEvent() { } NYdb::NTopic::TTopicClientSettings TTopicSession::GetTopicClientSettings(const NYql::NPq::NProto::TDqPqTopicSource& sourceParams) const { - NYdb::NTopic::TTopicClientSettings opts; - opts.Database(Database) + return PqGateway->GetTopicClientSettings() + .Database(Database) .DiscoveryEndpoint(Endpoint) .SslCredentials(NYdb::TSslCredentials(sourceParams.GetUseSsl())) .CredentialsProviderFactory(CredentialsProviderFactory); - return opts; } NYql::ITopicClient& TTopicSession::GetTopicClient(const NYql::NPq::NProto::TDqPqTopicSource& sourceParams) { diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp index f7ccfdf8c664..f9d5e968b290 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp @@ -159,7 +159,7 @@ class TDqPqReadActor : public NActors::TActor, public NYql::NDq: } NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const { - NYdb::NTopic::TTopicClientSettings opts; + NYdb::NTopic::TTopicClientSettings opts = PqGateway->GetTopicClientSettings(); opts.Database(SourceParams.GetDatabase()) .DiscoveryEndpoint(SourceParams.GetEndpoint()) .SslCredentials(NYdb::TSslCredentials(SourceParams.GetUseSsl())) diff --git a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp index 21ef29fcc48a..8d0085b9587e 100644 --- a/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp +++ b/ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.cpp @@ -312,7 +312,7 @@ class TDqPqWriteActor : public NActors::TActor, public IDqCompu } NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() { - return NYdb::NTopic::TTopicClientSettings() + return PqGateway->GetTopicClientSettings() .Database(SinkParams.GetDatabase()) .DiscoveryEndpoint(SinkParams.GetEndpoint()) .SslCredentials(NYdb::TSslCredentials(SinkParams.GetUseSsl())) diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp index 3253d0f47d19..f4b56a3ea280 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp @@ -80,4 +80,28 @@ void TDummyPqGateway::UpdateClusterConfigs( Y_UNUSED(secure); } +void TDummyPqGateway::UpdateClusterConfigs(const TPqGatewayConfigPtr& config) { + Y_UNUSED(config); +} + +NYdb::NTopic::TTopicClientSettings TDummyPqGateway::GetTopicClientSettings() const { + return NYdb::NTopic::TTopicClientSettings(); +} + +class TPqFileGatewayFactory : public IPqGatewayFactory { +public: + TPqFileGatewayFactory(const TDummyPqGateway::TPtr pqFileGateway) + : PqFileGateway(pqFileGateway) {} + + IPqGateway::TPtr CreatePqGateway() override { + return PqFileGateway; + } +private: + const TDummyPqGateway::TPtr PqFileGateway; +}; + +IPqGatewayFactory::TPtr CreatePqFileGatewayFactory(const TDummyPqGateway::TPtr pqFileGateway) { + return MakeIntrusive(pqFileGateway); +} + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h index fef31d7dc09b..084dd07f572c 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h @@ -58,8 +58,11 @@ class TDummyPqGateway : public IPqGateway { const TString& endpoint, const TString& database, bool secure) override; - + + void UpdateClusterConfigs(const TPqGatewayConfigPtr& config) override; + ITopicClient::TPtr GetTopicClient(const NYdb::TDriver& driver, const NYdb::NTopic::TTopicClientSettings& settings) override; + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const override; using TClusterNPath = std::pair; private: @@ -71,4 +74,6 @@ class TDummyPqGateway : public IPqGateway { IPqGateway::TPtr CreatePqFileGateway(); +IPqGatewayFactory::TPtr CreatePqFileGatewayFactory(const TDummyPqGateway::TPtr pqFileGateway); + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp index 8c94f126ecda..063557322848 100644 --- a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp +++ b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.cpp @@ -40,10 +40,12 @@ class TPqNativeGateway : public IPqGateway { const TString& database, bool secure) override; + void UpdateClusterConfigs(const TPqGatewayConfigPtr& config) override; + ITopicClient::TPtr GetTopicClient(const NYdb::TDriver& driver, const NYdb::NTopic::TTopicClientSettings& settings) override; + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const override; private: - void InitClusterConfigs(); TPqSession::TPtr GetExistingSession(const TString& sessionId) const; private: @@ -56,6 +58,7 @@ class TPqNativeGateway : public IPqGateway { NYdb::TDriver YdbDriver; TPqClusterConfigsMapPtr ClusterConfigs; THashMap Sessions; + TMaybe CommonTopicClientSettings; }; TPqNativeGateway::TPqNativeGateway(const TPqGatewayServices& services) @@ -65,14 +68,15 @@ TPqNativeGateway::TPqNativeGateway(const TPqGatewayServices& services) , CredentialsFactory(services.CredentialsFactory) , CmConnections(services.CmConnections) , YdbDriver(services.YdbDriver) + , CommonTopicClientSettings(services.CommonTopicClientSettings) { Y_UNUSED(FunctionRegistry); - InitClusterConfigs(); + UpdateClusterConfigs(Config); } -void TPqNativeGateway::InitClusterConfigs() { +void TPqNativeGateway::UpdateClusterConfigs(const TPqGatewayConfigPtr& config) { ClusterConfigs = std::make_shared(); - for (const auto& cfg : Config->GetClusterMapping()) { + for (const auto& cfg : config->GetClusterMapping()) { auto& config = (*ClusterConfigs)[cfg.GetName()]; config = cfg; } @@ -144,8 +148,28 @@ ITopicClient::TPtr TPqNativeGateway::GetTopicClient(const NYdb::TDriver& driver, return MakeIntrusive(driver, settings); } +NYdb::NTopic::TTopicClientSettings TPqNativeGateway::GetTopicClientSettings() const { + return CommonTopicClientSettings ? *CommonTopicClientSettings : NYdb::NTopic::TTopicClientSettings(); +} + TPqNativeGateway::~TPqNativeGateway() { Sessions.clear(); } +class TPqNativeGatewayFactory : public IPqGatewayFactory { +public: + TPqNativeGatewayFactory(const NYql::TPqGatewayServices& services) + : Services(services) {} + + IPqGateway::TPtr CreatePqGateway() override { + return CreatePqNativeGateway(Services); + } + const NYql::TPqGatewayServices Services; +}; + +IPqGatewayFactory::TPtr CreatePqNativeGatewayFactory(const NYql::TPqGatewayServices& services) { + return MakeIntrusive(services); +} + + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h index 1294abfe6a34..32d28a2549a6 100644 --- a/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h +++ b/ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h @@ -26,6 +26,7 @@ struct TPqGatewayServices { ISecuredServiceAccountCredentialsFactory::TPtr CredentialsFactory; ::NPq::NConfigurationManager::IConnections::TPtr CmConnections; NYdb::TDriver YdbDriver; + TMaybe CommonTopicClientSettings; TPqGatewayServices( NYdb::TDriver driver, @@ -33,17 +34,21 @@ struct TPqGatewayServices { ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory, TPqGatewayConfigPtr config, const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry, - IMetricsRegistryPtr metrics = nullptr) + IMetricsRegistryPtr metrics = nullptr, + TMaybe commonTopicClientSettings = Nothing()) : FunctionRegistry(functionRegistry) , Config(std::move(config)) , Metrics(std::move(metrics)) , CredentialsFactory(std::move(credentialsFactory)) , CmConnections(std::move(cmConnections)) , YdbDriver(std::move(driver)) + , CommonTopicClientSettings(commonTopicClientSettings) { } }; IPqGateway::TPtr CreatePqNativeGateway(const TPqGatewayServices& services); +IPqGatewayFactory::TPtr CreatePqNativeGatewayFactory(const NYql::TPqGatewayServices& services); + } // namespace NYql diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h b/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h index e9dd09de6009..6cfd71d929ab 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h +++ b/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h @@ -6,13 +6,23 @@ #include #include +#include #include #include #include +namespace NKikimr { +namespace NMiniKQL { +class IFunctionRegistry; +} // namespace NMiniKQL +} + namespace NYql { +class TPqGatewayConfig; +using TPqGatewayConfigPtr = std::shared_ptr; + struct IPqGateway : public TThrRefBase { using TPtr = TIntrusivePtr; @@ -36,6 +46,17 @@ struct IPqGateway : public TThrRefBase { const TString& endpoint, const TString& database, bool secure) = 0; + + virtual void UpdateClusterConfigs(const TPqGatewayConfigPtr& config) = 0; + + virtual NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const = 0; +}; + +struct IPqGatewayFactory : public TThrRefBase { + using TPtr = TIntrusivePtr; + + virtual ~IPqGatewayFactory() = default; + virtual IPqGateway::TPtr CreatePqGateway() = 0; }; } // namespace NYql diff --git a/ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp b/ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp index ff4dba4de4f3..5f41cf2b9ac3 100644 --- a/ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp +++ b/ydb/tests/fq/pq_async_io/mock_pq_gateway.cpp @@ -148,6 +148,8 @@ class TMockPqGateway : public IMockPqGateway { const TString& /*database*/, bool /*secure*/) override {} + void UpdateClusterConfigs(const TPqGatewayConfigPtr& /*config*/) override {}; + NYql::ITopicClient::TPtr GetTopicClient(const NYdb::TDriver& /*driver*/, const NYdb::NTopic::TTopicClientSettings& /*settings*/) override { return MakeIntrusive(this); } @@ -163,6 +165,10 @@ class TMockPqGateway : public IMockPqGateway { GetEventQueue(topic)->Push(std::move(e), size); } + NYdb::NTopic::TTopicClientSettings GetTopicClientSettings() const override { + return NYdb::NTopic::TTopicClientSettings(); + } + private: std::unordered_map> Queues; NActors::TTestActorRuntime& Runtime; diff --git a/ydb/tests/tools/fq_runner/kikimr_runner.py b/ydb/tests/tools/fq_runner/kikimr_runner.py index 6af8dac66a35..9ed1d2fcfd1d 100644 --- a/ydb/tests/tools/fq_runner/kikimr_runner.py +++ b/ydb/tests/tools/fq_runner/kikimr_runner.py @@ -478,6 +478,13 @@ def fill_config(self, control_plane): fq_config['test_connection'] = {'enabled': True} fq_config['common']['keep_internal_errors'] = True + fq_config['common']['ydb_driver_config'] = {} + fq_config['common']['ydb_driver_config']['network_threads_num'] = 1 + fq_config['common']['ydb_driver_config']['client_threads_num'] = 1 + + fq_config['common']['topic_client_handlers_executor_threads_num'] = 1 + fq_config['common']['topic_client_compression_executor_threads_num'] = 1 + if self.mvp_mock_port is not None: fq_config['common']['ydb_mvp_cloud_endpoint'] = "localhost:" + str(self.mvp_mock_port) diff --git a/ydb/tests/tools/fqrun/fqrun.cpp b/ydb/tests/tools/fqrun/fqrun.cpp index 4f1b096bde10..99f597d95c58 100644 --- a/ydb/tests/tools/fqrun/fqrun.cpp +++ b/ydb/tests/tools/fqrun/fqrun.cpp @@ -262,7 +262,7 @@ class TMain : public TMainBase { } fileGateway->AddDummyTopic(topic); } - RunnerOptions.FqSettings.PqGateway = std::move(fileGateway); + RunnerOptions.FqSettings.PqGatewayFactory = CreatePqFileGatewayFactory(fileGateway); } if (!TopicsSettings.empty()) { ythrow yexception() << "Found topic settings for not existing topic: '" << TopicsSettings.begin()->first << "'"; diff --git a/ydb/tests/tools/fqrun/src/common.h b/ydb/tests/tools/fqrun/src/common.h index 457897781891..5c139d83f5cc 100644 --- a/ydb/tests/tools/fqrun/src/common.h +++ b/ydb/tests/tools/fqrun/src/common.h @@ -28,7 +28,7 @@ struct TFqSetupSettings : public NKikimrRun::TServerSettings { EVerbose VerboseLevel = EVerbose::Info; TString YqlToken; - NYql::IPqGateway::TPtr PqGateway; + NYql::IPqGatewayFactory::TPtr PqGatewayFactory; TIntrusivePtr FunctionRegistry; NFq::NConfig::TConfig FqConfig; NKikimrConfig::TLogConfig LogConfig; diff --git a/ydb/tests/tools/fqrun/src/fq_setup.cpp b/ydb/tests/tools/fqrun/src/fq_setup.cpp index a4fbfd104777..b3f38ee69942 100644 --- a/ydb/tests/tools/fqrun/src/fq_setup.cpp +++ b/ydb/tests/tools/fqrun/src/fq_setup.cpp @@ -150,7 +150,7 @@ class TFqSetup::TImpl { NFq::Init( fqConfig, GetRuntime()->GetNodeId(), actorRegistrator, &GetRuntime()->GetAppData(), - Settings.DomainName, nullptr, YqSharedResources, folderServiceFactory, 0, {}, Settings.PqGateway + Settings.DomainName, nullptr, YqSharedResources, folderServiceFactory, 0, {}, Settings.PqGatewayFactory ); YqSharedResources->Init(GetRuntime()->GetActorSystem(0)); } From 807748fa6772318135ddb967a5e23e2cbd8d95b5 Mon Sep 17 00:00:00 2001 From: Pisarenko Grigoriy Date: Thu, 27 Feb 2025 15:24:18 +0500 Subject: [PATCH 13/18] YQ-3561 supported in memory cp storage (#15031) --- .../control_plane_storage.h | 8 +- .../control_plane_storage/events/events.h | 12 + .../in_memory_control_plane_storage.cpp | 1036 +++++++++++++---- .../internal/nodes_health_check.cpp | 11 +- .../internal/task_get.cpp | 163 +-- .../internal/task_ping.cpp | 693 ++++++----- .../internal/task_result_write.cpp | 16 +- .../control_plane_storage/internal/utils.cpp | 15 +- .../control_plane_storage/internal/utils.h | 14 +- .../ydb_control_plane_storage.cpp | 16 +- .../ydb_control_plane_storage_bindings.cpp | 68 +- .../ydb_control_plane_storage_connections.cpp | 67 +- .../ydb_control_plane_storage_impl.h | 351 ++++-- .../ydb_control_plane_storage_queries.cpp | 373 +++--- ydb/core/fq/libs/init/init.cpp | 11 +- .../in_memory_control_plane_storage_ut.cpp | 123 +- ydb/tests/fq/control_plane_storage/ya.make | 3 + .../ydb_test_bootstrap.h | 8 + 18 files changed, 1982 insertions(+), 1006 deletions(-) diff --git a/ydb/core/fq/libs/control_plane_storage/control_plane_storage.h b/ydb/core/fq/libs/control_plane_storage/control_plane_storage.h index aa0c2f478eda..99fad28211ae 100644 --- a/ydb/core/fq/libs/control_plane_storage/control_plane_storage.h +++ b/ydb/core/fq/libs/control_plane_storage/control_plane_storage.h @@ -40,7 +40,13 @@ namespace NFq { NActors::TActorId ControlPlaneStorageServiceActorId(ui32 nodeId = 0); -NActors::IActor* CreateInMemoryControlPlaneStorageServiceActor(const NConfig::TControlPlaneStorageConfig& config); +NActors::IActor* CreateInMemoryControlPlaneStorageServiceActor( + const NConfig::TControlPlaneStorageConfig& config, + const NYql::TS3GatewayConfig& s3Config, + const NConfig::TCommonConfig& common, + const NConfig::TComputeConfig& computeConfig, + const ::NMonitoring::TDynamicCounterPtr& counters, + const TString& tenantName); NActors::IActor* CreateYdbControlPlaneStorageServiceActor( const NConfig::TControlPlaneStorageConfig& config, diff --git a/ydb/core/fq/libs/control_plane_storage/events/events.h b/ydb/core/fq/libs/control_plane_storage/events/events.h index e6e730d1e9d7..314cd3f588d2 100644 --- a/ydb/core/fq/libs/control_plane_storage/events/events.h +++ b/ydb/core/fq/libs/control_plane_storage/events/events.h @@ -392,6 +392,7 @@ struct TEvControlPlaneStorage { // internal messages struct TEvWriteResultDataRequest : NActors::TEventLocal { + using TProto = Fq::Private::WriteTaskResultRequest; TEvWriteResultDataRequest() = default; @@ -411,6 +412,8 @@ struct TEvControlPlaneStorage { struct TEvWriteResultDataResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::WriteTaskResultResult; + explicit TEvWriteResultDataResponse( const Fq::Private::WriteTaskResultResult& record) : Record(record) @@ -434,6 +437,7 @@ struct TEvControlPlaneStorage { }; struct TEvGetTaskRequest : NActors::TEventLocal { + using TProto = Fq::Private::GetTaskRequest; TEvGetTaskRequest() = default; @@ -454,6 +458,8 @@ struct TEvControlPlaneStorage { struct TEvGetTaskResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::GetTaskResult; + explicit TEvGetTaskResponse( const Fq::Private::GetTaskResult& record) : Record(record) @@ -499,6 +505,7 @@ struct TEvControlPlaneStorage { }; struct TEvPingTaskRequest : NActors::TEventLocal { + using TProto = Fq::Private::PingTaskRequest; TEvPingTaskRequest() = default; @@ -519,6 +526,8 @@ struct TEvControlPlaneStorage { struct TEvPingTaskResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::PingTaskResult; + explicit TEvPingTaskResponse( const Fq::Private::PingTaskResult& record) : Record(record) @@ -542,6 +551,7 @@ struct TEvControlPlaneStorage { }; struct TEvNodesHealthCheckRequest : NActors::TEventLocal { + using TProto = Fq::Private::NodesHealthCheckRequest; TEvNodesHealthCheckRequest() = default; @@ -561,6 +571,8 @@ struct TEvControlPlaneStorage { struct TEvNodesHealthCheckResponse : NActors::TEventLocal { static constexpr bool Auditable = false; + using TProto = Fq::Private::NodesHealthCheckResult; + explicit TEvNodesHealthCheckResponse( const Fq::Private::NodesHealthCheckResult& record) : Record(record) diff --git a/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp b/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp index f9089b0cac4b..7f715d2ec9f0 100644 --- a/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp +++ b/ydb/core/fq/libs/control_plane_storage/in_memory_control_plane_storage.cpp @@ -1,59 +1,173 @@ #include "control_plane_storage.h" -#include "util.h" +#include "ydb_control_plane_storage_impl.h" -#include -#include +#include -#include - -#include -#include +#include namespace NFq { -class TInMemoryControlPlaneStorageActor : public NActors::TActor { - struct TKey { +class TInMemoryControlPlaneStorageActor : public NActors::TActor, + public TControlPlaneStorageBase { + struct TScopeKey { TString Scope; TString Id; - bool operator<(const TKey& other) const - { - return tie(Scope, Id) < tie(other.Scope, other.Id); - } + std::strong_ordering operator<=>(const TScopeKey& other) const = default; }; - struct TConfig { - NConfig::TControlPlaneStorageConfig Proto; - TDuration IdempotencyKeyTtl; - TDuration AutomaticQueriesTtl; - TDuration ResultSetsTtl; - TDuration AnalyticsRetryCounterUpdateTime; - TDuration StreamingRetryCounterUpdateTime; - TDuration TaskLeaseTtl; - - TConfig(const NConfig::TControlPlaneStorageConfig& config) - : Proto(FillDefaultParameters(config)) - , IdempotencyKeyTtl(GetDuration(Proto.GetIdempotencyKeysTtl(), TDuration::Minutes(10))) - , AutomaticQueriesTtl(GetDuration(Proto.GetAutomaticQueriesTtl(), TDuration::Days(1))) - , ResultSetsTtl(GetDuration(Proto.GetResultSetsTtl(), TDuration::Days(1))) - , TaskLeaseTtl(GetDuration(Proto.GetTaskLeaseTtl(), TDuration::Seconds(30))) - { - } + struct TQueries { + using TKey = TScopeKey; + + struct TValue { + FederatedQuery::Query Query; + FederatedQuery::Internal::QueryInternal QueryInternal; + TString LastJobId; + TString User; + TString ResultId; + TInstant ResultExpireAt; + ui64 Generation = 0; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; }; - TConfig Config; - TMap Queries; - TMap Connections; - TMap IdempotencyKeys; // idempotency_key -> created_at + struct TPendingQueries { + struct TKey { + TString Tenant; + TString Scope; + TString QueryId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + TRetryLimiter RetryLimiter; + TString Owner; + TInstant AssignedUntil; + TInstant LastSeenAt; + }; + + TMap Values; + }; + + struct TJobs { + struct TKey { + TString Scope; + TString QueryId; + TString JobId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + FederatedQuery::Job Job; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; - static constexpr int64_t InitialRevision = 1; + struct TConnections { + using TKey = TScopeKey; + using TEntity = FederatedQuery::Connection; + + struct TValue { + TEntity Connection; + TString User; + + const TEntity& GetEntity() const { + return Connection; + } + }; + + TMap Values; + }; + + struct TBindings { + using TKey = TScopeKey; + using TEntity = FederatedQuery::Binding; + + struct TValue { + TEntity Binding; + TString User; + + const TEntity& GetEntity() const { + return Binding; + } + }; + + TMap Values; + }; + + struct TIdempotencyKeys { + using TKey = TScopeKey; + + struct TValue { + TString Response; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; + + struct TResultSets { + struct TKey { + TString ResultId; + i32 ResultSetId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + TVector Rows; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; + + struct TNodes { + static constexpr TDuration TTL = TDuration::Seconds(15); + + struct TKey { + TString Tenant; + ui32 NodeId; + + std::strong_ordering operator<=>(const TKey& other) const = default; + }; + + struct TValue { + Fq::Private::NodeInfo Node; + TInstant ExpireAt = TInstant::Zero(); + }; + + TMap Values; + }; + + using TBase = TControlPlaneStorageBase; + + TQueries Queries; + TPendingQueries PendingQueries; + TJobs Jobs; + TConnections Connections; + TBindings Bindings; + TResultSets ResultSets; + TIdempotencyKeys IdempotencyKeys; + TNodes Nodes; public: - TInMemoryControlPlaneStorageActor(const NConfig::TControlPlaneStorageConfig& config) + TInMemoryControlPlaneStorageActor( + const NConfig::TControlPlaneStorageConfig& config, + const NYql::TS3GatewayConfig& s3Config, + const NConfig::TCommonConfig& common, + const NConfig::TComputeConfig& computeConfig, + const ::NMonitoring::TDynamicCounterPtr& counters, + const TString& tenantName) : TActor(&TThis::StateFunc) - , Config(config) - { - } + , TBase(config, s3Config, common, computeConfig, counters, tenantName) + {} static constexpr char ActorName[] = "YQ_CONTROL_PLANE_STORAGE"; @@ -83,75 +197,258 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActorGet()->Request; - CPS_LOG_D("CreateQueryRequest: " << request.DebugString()); - CleanupIndempotencyKeys(); - auto now = TInstant::Now(); - const TString idempotencyKey = request.idempotency_key(); - if (idempotencyKey && IdempotencyKeys.contains(idempotencyKey)) { - CPS_LOG_D("CreateQueryRequest, idempotency key already exist: " << request.DebugString()); - NYql::TIssue issue = MakeErrorIssue(TIssuesIds::BAD_REQUEST, "idempotency key already exist"); - Send(ev->Sender, new TEvControlPlaneStorage::TEvCreateQueryResponse(NYql::TIssues{issue}), 0, ev->Cookie); - return; - } - - NYql::TIssues issues = ValidateCreateQueryRequest(ev); - if (issues) { - CPS_LOG_D("CreateQueryRequest, validation failed: " << request.DebugString() << " error: " << issues.ToString()); - Send(ev->Sender, new TEvControlPlaneStorage::TEvCreateQueryResponse(issues), 0, ev->Cookie); - return; - } - - const TString user = ev->Get()->User; - const TString scope = ev->Get()->Scope; - const TString queryId = CreateGuidAsString(); - FederatedQuery::Query query; - FederatedQuery::QueryContent& content = *query.mutable_content(); - content = request.content(); - FederatedQuery::QueryMeta& meta = *query.mutable_meta(); - FederatedQuery::CommonMeta& common = *meta.mutable_common(); - common.set_id(queryId); - common.set_created_by(user); - auto timestamp = NProtoInterop::CastToProto(now); - *common.mutable_created_at() = timestamp; - common.set_revision(InitialRevision); - - Queries[{scope, queryId}] = query; - - if (!idempotencyKey) { - IdempotencyKeys[idempotencyKey] = now; - } - - CPS_LOG_D("CreateQueryRequest, success: " << request.DebugString() << " query_id: " << queryId); - FederatedQuery::CreateQueryResult result; - result.set_query_id(queryId); - Send(ev->Sender, new TEvControlPlaneStorage::TEvCreateQueryResponse(result, TAuditDetails{}), 0, ev->Cookie); - } - - void Handle(TEvControlPlaneStorage::TEvListQueriesRequest::TPtr& ev) - { + template + class TCommonRequestContext { + public: + using TResultType = TPrepareResponseResultType; + using TResponse = TResultType::Type; + using TAuditDetails = TResultType::TResponseAuditDetails; + + TCommonRequestContext(TInMemoryControlPlaneStorageActor& self, TEvRequest::TPtr& ev, const TString& cloudId, + const TString& scope, const TString& logPrefix, const TString& requestStr, const TString& responseStr) + : StartTime(TInstant::Now()) + , Event(*ev->Get()) + , Request(Event.Request) + , LogPrefix(logPrefix) + , RequestCounters(self.Counters.GetCounters(cloudId, scope, TYPE_SCOPE, TYPE_COMMON)) + , Self(self) + , EventPtr(ev) + , RequestStr(requestStr) + , ResponseStr(responseStr) + { + Self.Cleanup(); + CPS_LOG_I(RequestStr); + CPS_LOG_T(RequestStr << ":" << LogPrefix); + + RequestCounters.IncInFly(); + RequestCounters.Common->RequestBytes->Add(Event.GetByteSize()); + } + + TCommonRequestContext(TInMemoryControlPlaneStorageActor& self, TEvRequest::TPtr& ev, const TString& requestStr, const TString& responseStr) + : TCommonRequestContext(self, ev, "", "", GetLogPrefix(ev), requestStr, responseStr) + {} + + virtual bool Validate() { + if (const auto& issues = Self.ValidateRequest(EventPtr)) { + Fail("query validation", issues); + return false; + } + return true; + } + + bool IsFailed() const { + return Failed; + } + + void Fail(const TString& logInfo, const NYql::TIssues& issues) const { + Y_ABORT_UNLESS(!Failed, "Can not fail twice"); + CPS_LOG_W(RequestStr << ":" << LogPrefix << logInfo << " FAILED: " << issues.ToOneLineString()); + Self.SendResponseIssues(EventPtr->Sender, issues, EventPtr->Cookie, TInstant::Now() - StartTime, RequestCounters); + } + + virtual ~TCommonRequestContext() { + if (Failed) { + return; + } + + Self.SendResponse( + TStringBuilder() << RequestStr << " - " << ResponseStr, + NActors::TActivationContext::ActorSystem(), + NThreading::MakeFuture(NYdb::TStatus(NYdb::EStatus::SUCCESS, {})), + Self.SelfId(), + EventPtr, + StartTime, + RequestCounters, + [response = Response] { return response; }, + Self.Config->Proto.GetEnableDebugMode() ? std::make_shared() : TDebugInfoPtr{}); + } + + static TString GetLogPrefix(TEvRequest::TPtr& ev) { + return TStringBuilder() << "{" << ev->Get()->Request.DebugString() << "} "; + } + + public: + const TInstant StartTime; + const TEvRequest& Event; + const TEvRequest::TProto Request; + const TString LogPrefix; + TRequestCounters RequestCounters; + TResponse Response; + + protected: + TInMemoryControlPlaneStorageActor& Self; + TEvRequest::TPtr& EventPtr; + const TString RequestStr; + const TString ResponseStr; + bool Failed = false; + }; + + template + class TRequestContext : public TCommonRequestContext { + using TBase = TCommonRequestContext; + + Y_HAS_MEMBER(idempotency_key); + static constexpr bool HasIdempotencyKey = THasidempotency_key::value; + + public: + TRequestContext(TInMemoryControlPlaneStorageActor& self, TEvRequest::TPtr& ev, const TString& requestStr, const TString& responseStr) + : TBase( + self, ev, ev->Get()->CloudId, ev->Get()->Scope, + TStringBuilder() << TBase::GetLogPrefix(ev) << MakeUserInfo(ev->Get()->User, ev->Get()->Token), + requestStr, responseStr + ) + , CloudId(TBase::Event.CloudId) + , Scope(TBase::Event.Scope) + , User(TBase::Event.User) + , Token(TBase::Event.Token) + , Permissions(TBase::Event.Permissions) + , Quotas(TBase::Event.Quotas) + , TenantInfo(TBase::Event.TenantInfo) + , ComputeDatabase(TBase::Event.ComputeDatabase) + { + if constexpr (!std::is_same_v) { + TBase::Response.second.CloudId = CloudId; + } + } + + bool Validate() override { + if (!TBase::Validate()) { + return false; + } + if constexpr (HasIdempotencyKey) { + if (const TString& idempotencyKey = TBase::Request.idempotency_key()) { + if (const auto& value = TBase::Self.GetEntity(TBase::Self.IdempotencyKeys, {Scope, idempotencyKey})) { + if (!TBase::Response.first.ParseFromString(value->Response)) { + TBase::RequestCounters.Common->ParseProtobufError->Inc(); + TBase::Fail("idempotency key parse", {NYql::TIssue("INTERNAL ERROR. Error parsing proto message for idempotency key request. Please contact internal support")}); + } else { + TBase::Response.second.IdempotencyResult = true; + } + return false; + } + } + } + return true; + } + + ~TRequestContext() override { + if (TBase::Failed) { + return; + } + + if constexpr (HasIdempotencyKey) { + if (const TString& idempotencyKey = TBase::Request.idempotency_key()) { + this->Self.AddEntity(this->Self.IdempotencyKeys, {this->Scope, idempotencyKey}, { + .Response = TBase::Response.first.SerializeAsString(), + .ExpireAt = TBase::StartTime + this->Self.Config->IdempotencyKeyTtl + }); + } + } + } + + public: + const TString CloudId; + const TString Scope; + const TString User; + const TString Token; + const TPermissions Permissions; + const TMaybe Quotas; + const TTenantInfo::TPtr TenantInfo; + const FederatedQuery::Internal::ComputeDatabaseInternal ComputeDatabase; + }; + +#define HANDLE_CPS_REQUEST_IMPL(TEvRequest, TEvResponse, TContext, RTS_COUNTERS_ENUM, RTC_COUNTERS_ENUM) \ + using TContext##TEvRequest = TContext< \ + TEvControlPlaneStorage::TEvRequest, TEvControlPlaneStorage::TEvResponse, \ + RTS_COUNTERS_ENUM, RTC_COUNTERS_ENUM>; \ + void Handle(TEvControlPlaneStorage::TEvRequest::TPtr& ev) { \ + TContext##TEvRequest ctx(*this, ev, #TEvRequest, #TEvResponse); \ + if (!ctx.Validate()) { \ + return; \ + } \ + try { \ + Process##TRequest(ctx); \ + } catch (...) { \ + const auto& backtrace = TBackTrace::FromCurrentException().PrintToString(); \ + const auto logError = TStringBuilder() << "pocess "#TEvRequest" call, back trace:\n" << backtrace; \ + ctx.Fail(logError, {NYql::TIssue(CurrentExceptionMessage())}); \ + } \ + } \ + void Process##TRequest(TContext##TEvRequest& ctx) + +#define HANDLE_CPS_REQUEST(TEvRequest, TEvResponse, COUNTERS_ENUM) HANDLE_CPS_REQUEST_IMPL(TEvRequest, TEvResponse, TRequestContext, RTS_##COUNTERS_ENUM, RTC_##COUNTERS_ENUM) + + HANDLE_CPS_REQUEST(TEvCreateQueryRequest, TEvCreateQueryResponse, CREATE_QUERY) { + const auto& [query, job] = GetCreateQueryProtos(ctx.Request, ctx.User, ctx.StartTime); + if (query.ByteSizeLong() > Config->Proto.GetMaxRequestSize()) { + return ctx.Fail("query size validation", {NYql::TIssue(TStringBuilder() << "incoming request exceeded the size limit: " << query.ByteSizeLong() << " of " << Config->Proto.GetMaxRequestSize() << ". Please shorten your request")}); + } + ctx.Response.second.After = query; + + const TString& queryId = query.meta().common().id(); + ctx.Response.first.set_query_id(queryId); + + auto queryInternal = GetQueryInternalProto(ctx.Request, ctx.CloudId, ctx.Token, ctx.Quotas); + const auto queryType = ctx.Request.content().type(); + if (ctx.Request.execute_mode() != FederatedQuery::SAVE) { + *queryInternal.mutable_compute_connection() = ctx.ComputeDatabase.connection(); + FillConnectionsAndBindings( + queryInternal, + queryType, + GetEntities(Connections, ctx.Scope, ctx.User), + GetEntitiesWithVisibilityPriority(Connections, ctx.Scope, ctx.User), + GetEntitiesWithVisibilityPriority(Bindings, ctx.Scope, ctx.User) + ); + } + if (queryInternal.ByteSizeLong() > Config->Proto.GetMaxRequestSize()) { + return ctx.Fail("query internal size validation", {NYql::TIssue(TStringBuilder() << "the size of all connections and bindings in the project exceeded the limit: " << queryInternal.ByteSizeLong() << " of " << Config->Proto.GetMaxRequestSize() << ". Please reduce the number of connections and bindings")}); + } + + const auto& jobId = job.meta().id(); + if (ctx.Request.execute_mode() != FederatedQuery::SAVE) { + AddEntity(Jobs, {ctx.Scope, queryId, jobId}, {job}); + + TRetryLimiter retryLimiter; + retryLimiter.Assign(0, ctx.StartTime, 0.0); + + AddEntity(PendingQueries, { + .Tenant = ctx.TenantInfo->Assign(ctx.CloudId, ctx.Scope, queryType, TenantName), + .Scope = ctx.Scope, + .QueryId = queryId + }, {.RetryLimiter = retryLimiter}); + } + + AddEntity(Queries, {ctx.Scope, queryId}, { + .Query = query, + .QueryInternal = queryInternal, + .LastJobId = jobId, + .User = ctx.User + }); + } + + void Handle(TEvControlPlaneStorage::TEvListQueriesRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListQueriesRequest::TPtr, FederatedQuery::ListQueriesResult, TEvControlPlaneStorage::TEvListQueriesResponse>(ev, "ListQueriesRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr, - FederatedQuery::DescribeQueryResult, - TEvControlPlaneStorage::TEvDescribeQueryResponse>(ev, "DescribeQueryRequest"); + HANDLE_CPS_REQUEST(TEvDescribeQueryRequest, TEvDescribeQueryResponse, DESCRIBE_QUERY) { + const auto& query = GetEntity(Queries, {ctx.Scope, ctx.Request.query_id()}); + if (!query) { + return ctx.Fail("find query", {NYql::TIssue("Query does not exist")}); + } + + *ctx.Response.mutable_query() = query->Query; + FillDescribeQueryResult(ctx.Response, query->QueryInternal, ctx.User, ctx.Permissions); } - void Handle(TEvControlPlaneStorage::TEvModifyQueryRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvModifyQueryRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvModifyQueryRequest::TPtr, FederatedQuery::ModifyQueryResult, @@ -159,8 +456,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ModifyQueryRequest"); } - void Handle(TEvControlPlaneStorage::TEvDeleteQueryRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDeleteQueryRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvDeleteQueryRequest::TPtr, FederatedQuery::DeleteQueryResult, @@ -168,8 +465,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "DeleteQueryRequest"); } - void Handle(TEvControlPlaneStorage::TEvControlQueryRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvControlQueryRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvControlQueryRequest::TPtr, FederatedQuery::ControlQueryResult, @@ -177,49 +474,98 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ControlQueryRequest"); } - void Handle(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr, - FederatedQuery::GetResultDataResult, - TEvControlPlaneStorage::TEvGetResultDataResponse>(ev, "GetResultDataRequest"); + HANDLE_CPS_REQUEST(TEvGetResultDataRequest, TEvGetResultDataResponse, GET_RESULT_DATA) { + const auto& query = GetEntity(Queries, {ctx.Scope, ctx.Request.query_id()}); + if (!query) { + return ctx.Fail("find query", {NYql::TIssue("Query does not exist")}); + } + + if (!HasViewAccess(GetResultDataReadPerimssions(ctx.Event), query->Query.content().acl().visibility(), query->User, ctx.User)) { + return ctx.Fail("check ACL", {NYql::TIssue("Permission denied")}); + } + + const auto resultSetIndex = ctx.Request.result_set_index(); + const auto& resultSetMeta = query->Query.result_set_meta(); + if (resultSetIndex >= resultSetMeta.size()) { + return ctx.Fail("check result set index", {NYql::TIssue(TStringBuilder() << "Result set index out of bound: " << resultSetIndex << " >= " << resultSetMeta.size())}); + } + + const auto expireAt = query->ResultExpireAt; + if (query->Query.meta().status() != FederatedQuery::QueryMeta::COMPLETED || !expireAt) { + return ctx.Fail("check status", {NYql::TIssue("Result doesn't exist")}); + } + + if (expireAt < TInstant::Now()) { + return ctx.Fail("check expiration", {NYql::TIssue("Result removed by TTL")}); + } + + const auto& result = GetEntity(ResultSets, {query->ResultId, ctx.Request.result_set_index()}); + if (!result) { + return ctx.Fail("get result", {NYql::TIssue("INTERNAL ERROR. Failed to find result set")}); + } + + auto& resultSet = *ctx.Response.mutable_result_set(); + *resultSet.mutable_columns() = resultSetMeta[resultSetIndex].column(); + + const i64 offset = ctx.Request.offset(); + const i64 numberRows = result->Rows.size(); + for (i64 rowId = offset; rowId < offset + ctx.Request.limit() && rowId < numberRows; ++rowId) { + *resultSet.add_rows() = result->Rows[rowId]; + } } - void Handle(TEvControlPlaneStorage::TEvListJobsRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvListJobsRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListJobsRequest::TPtr, FederatedQuery::ListJobsResult, TEvControlPlaneStorage::TEvListJobsResponse>(ev, "ListJobsRequest"); } - void Handle(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) - { - SendEmptyAuditResponse< - TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr, - FederatedQuery::CreateConnectionResult, - TEvControlPlaneStorage::TEvCreateConnectionResponse, - TAuditDetails>(ev, "CreateConnectionRequest"); + HANDLE_CPS_REQUEST(TEvCreateConnectionRequest, TEvCreateConnectionResponse, CREATE_CONNECTION) { + const auto& content = ctx.Request.content(); + const auto visibility = content.acl().visibility(); + const auto& name = content.name(); + if (!CheckConnectionOrBindingName(Connections, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Connection with the same name already exists. Please choose another name")}); + } + if (!CheckConnectionOrBindingName(Bindings, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Binding with the same name already exists. Please choose another name")}); + } + if (GetNumberEntitiesByScope(Connections, ctx.Scope) >= Config->Proto.GetMaxCountConnections()) { + return ctx.Fail("check number", {NYql::TIssue(TStringBuilder() << "Too many connections in folder: " << Config->Proto.GetMaxCountConnections() << ". Please remove unused connections")}); + } + + const auto& [connection, _] = GetCreateConnectionProtos(ctx.Request, ctx.CloudId, ctx.User, ctx.StartTime); + ctx.Response.second.After = connection; + + const TString& connectionId = connection.meta().id(); + ctx.Response.first.set_connection_id(connectionId); + + AddEntity(Connections, {ctx.Scope, connectionId}, { + .Connection = connection, + .User = ctx.User + }); } - void Handle(TEvControlPlaneStorage::TEvListConnectionsRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvListConnectionsRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListConnectionsRequest::TPtr, FederatedQuery::ListConnectionsResult, TEvControlPlaneStorage::TEvListConnectionsResponse>(ev, "ListConnectionsRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeConnectionRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDescribeConnectionRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvDescribeConnectionRequest::TPtr, FederatedQuery::DescribeConnectionResult, TEvControlPlaneStorage::TEvDescribeConnectionResponse>(ev, "DescribeConnectionRequest"); } - void Handle(TEvControlPlaneStorage::TEvModifyConnectionRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvModifyConnectionRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvModifyConnectionRequest::TPtr, FederatedQuery::ModifyConnectionResult, @@ -227,8 +573,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ModifyConnectionRequest"); } - void Handle(TEvControlPlaneStorage::TEvDeleteConnectionRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDeleteConnectionRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvDeleteConnectionRequest::TPtr, FederatedQuery::DeleteConnectionResult, @@ -236,33 +582,64 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "DeleteConnectionRequest"); } - void Handle(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) - { - SendEmptyAuditResponse< - TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr, - FederatedQuery::CreateBindingResult, - TEvControlPlaneStorage::TEvCreateBindingResponse, - TAuditDetails>(ev, "CreateBindingRequest"); + HANDLE_CPS_REQUEST(TEvCreateBindingRequest, TEvCreateBindingResponse, CREATE_BINDING) { + const auto& content = ctx.Request.content(); + const auto visibility = content.acl().visibility(); + const auto& name = content.name(); + if (!CheckConnectionOrBindingName(Connections, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Connection with the same name already exists. Please choose another name")}); + } + if (!CheckConnectionOrBindingName(Bindings, ctx.Scope, ctx.User, visibility, name)) { + return ctx.Fail("check name", {NYql::TIssue("Binding with the same name already exists. Please choose another name")}); + } + if (GetNumberEntitiesByScope(Bindings, ctx.Scope) >= Config->Proto.GetMaxCountBindings()) { + return ctx.Fail("check number", {NYql::TIssue(TStringBuilder() << "Too many bindings in folder: " << Config->Proto.GetMaxCountBindings() << ". Please remove unused bindings")}); + } + + const auto& connection = GetEntity(Connections, {ctx.Scope, content.connection_id()}); + if (!connection) { + return ctx.Fail("check connection", {NYql::TIssue("Connection for binding not found")}); + } + + const auto connectionVisibility = connection->Connection.content().acl().visibility(); + if (content.acl().visibility() == FederatedQuery::Acl::SCOPE && connectionVisibility == FederatedQuery::Acl::PRIVATE) { + return ctx.Fail("check connection ACL", {NYql::TIssue("Binding with SCOPE visibility cannot refer to connection with PRIVATE visibility")}); + } + + if (!HasManageAccess(GetCreateBindingPerimssions(ctx.Event), connectionVisibility, connection->User, ctx.User)) { + return ctx.Fail("check connection ACL", {NYql::TIssue("Permission denied for binding connection")}); + } + + const auto& [binding, _] = GetCreateBindingProtos(ctx.Request, ctx.CloudId, ctx.User, ctx.StartTime); + ctx.Response.second.After = binding; + + const TString& bindingId = binding.meta().id(); + ctx.Response.first.set_binding_id(bindingId); + + AddEntity(Bindings, {ctx.Scope, bindingId}, { + .Binding = binding, + .User = ctx.User + }); } - void Handle(TEvControlPlaneStorage::TEvListBindingsRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvListBindingsRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvListBindingsRequest::TPtr, FederatedQuery::ListBindingsResult, TEvControlPlaneStorage::TEvListBindingsResponse>(ev, "ListBindingsRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeBindingRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDescribeBindingRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvDescribeBindingRequest::TPtr, FederatedQuery::DescribeBindingResult, TEvControlPlaneStorage::TEvDescribeBindingResponse>(ev, "DescribeBindingRequest"); } - void Handle(TEvControlPlaneStorage::TEvModifyBindingRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvModifyBindingRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvModifyBindingRequest::TPtr, FederatedQuery::ModifyBindingResult, @@ -270,8 +647,8 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "ModifyBindingRequest"); } - void Handle(TEvControlPlaneStorage::TEvDeleteBindingRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDeleteBindingRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyAuditResponse< TEvControlPlaneStorage::TEvDeleteBindingRequest::TPtr, FederatedQuery::DeleteBindingResult, @@ -279,53 +656,109 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor>(ev, "DeleteBindingRequest"); } - void Handle(TEvControlPlaneStorage::TEvDescribeJobRequest::TPtr& ev) - { + void Handle(TEvControlPlaneStorage::TEvDescribeJobRequest::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); SendEmptyResponse< TEvControlPlaneStorage::TEvDescribeJobRequest::TPtr, FederatedQuery::DescribeJobResult, TEvControlPlaneStorage::TEvDescribeJobResponse>(ev, "DescribeJobRequest"); } - void Handle(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr, - NYql::TIssues, - TEvControlPlaneStorage::TEvWriteResultDataResponse>(ev, "WriteResultDataRequest"); + HANDLE_CPS_REQUEST_IMPL(TEvWriteResultDataRequest, TEvWriteResultDataResponse, TCommonRequestContext, RTS_MAX, RTC_WRITE_RESULT_DATA) { + ctx.Response.set_request_id(ctx.Request.request_id()); + + const auto offset = ctx.Request.offset(); + const auto& newRows = ctx.Request.result_set().rows(); + + auto& [resultRows, expireAt] = ResultSets.Values[{ctx.Request.result_id().value(), static_cast(ctx.Request.result_set_id())}]; + expireAt = NProtoInterop::CastFromProto(ctx.Request.deadline()); + + resultRows.resize(std::max(offset + newRows.size(), resultRows.size())); + for (size_t i = offset; const auto& row : newRows) { + resultRows[i++] = row; + } } - void Handle(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) - { - CPS_LOG_I("GetTaskRequest"); - Fq::Private::GetTaskResult result; - auto event = std::make_unique(result); - NActors::TActivationContext::ActorSystem()->Send(new IEventHandle(ev->Sender, SelfId(), event.release(), 0, ev->Cookie)); + HANDLE_CPS_REQUEST_IMPL(TEvGetTaskRequest, TEvGetTaskResponse, TCommonRequestContext, RTS_MAX, RTC_GET_TASK) { + const auto& tasksInternal = GetActiveTasks(ctx); + + TVector tasks; + tasks.reserve(tasksInternal.size()); + for (const auto& taskInternal : tasksInternal) { + if (const auto& task = AssignTask(ctx, taskInternal)) { + tasks.emplace_back(*task); + } + if (ctx.IsFailed()) { + return; + } + } + + FillGetTaskResult(ctx.Response, tasks); } - void Handle(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvPingTaskRequest::TPtr, - Fq::Private::PingTaskResult, - TEvControlPlaneStorage::TEvPingTaskResponse>(ev, "PingTaskRequest"); + HANDLE_CPS_REQUEST_IMPL(TEvPingTaskRequest, TEvPingTaskResponse, TCommonRequestContext, RTS_MAX, RTC_PING_TASK) { + const auto& scope = ctx.Request.scope(); + const auto& queryId = ctx.Request.query_id().value(); + + auto query = GetEntity(Queries, {scope, queryId}); + if (!query) { + return ctx.Fail("get query", {NYql::TIssue("INTERNAL ERROR. Query for ping task not found")}); + } + + auto job = GetEntity(Jobs, {scope, queryId, query->LastJobId}); + if (!job) { + return ctx.Fail("get job", {NYql::TIssue("INTERNAL ERROR. Job for ping task not found")}); + } + + auto pendingQuery = GetEntity(PendingQueries, {ctx.Request.tenant(), scope, queryId}); + if (!pendingQuery) { + return ctx.Fail("get pending query", {NYql::TIssue("INTERNAL ERROR. Pending query for ping task not found")}); + } + + auto resuest = ctx.Request; + auto finalStatus = std::make_shared(); + + TDuration backoff = Config->TaskLeaseTtl; + TInstant expireAt = ctx.StartTime + Config->AutomaticQueriesTtl; + UpdateTaskInfo(TActivationContext::ActorSystem(), resuest, finalStatus, query->Query, query->QueryInternal, job->Job, pendingQuery->Owner, pendingQuery->RetryLimiter, backoff, expireAt); + PingTask(ctx, *query, *job, *pendingQuery, backoff, expireAt); + + if (IsTerminalStatus(ctx.Request.status())) { + FillQueryStatistics(finalStatus, query->Query, query->QueryInternal, pendingQuery->RetryLimiter); + } + Send(SelfId(), new TEvControlPlaneStorage::TEvFinalStatusReport( + queryId, finalStatus->JobId, finalStatus->CloudId, scope, std::move(finalStatus->FinalStatistics), + finalStatus->Status, finalStatus->StatusCode, finalStatus->QueryType, finalStatus->Issues, finalStatus->TransientIssues)); } - void Handle(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) - { - SendEmptyResponse< - TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr, - Fq::Private::NodesHealthCheckResult, - TEvControlPlaneStorage::TEvNodesHealthCheckResponse>(ev, "NodesHealthCheckRequest"); + HANDLE_CPS_REQUEST_IMPL(TEvNodesHealthCheckRequest, TEvNodesHealthCheckResponse, TCommonRequestContext, RTS_MAX, RTC_NODES_HEALTH_CHECK) { + const auto& tenant = ctx.Request.tenant(); + const auto& node = ctx.Request.node(); + + AddEntity(Nodes, {tenant, node.node_id()}, { + .Node = node, + .ExpireAt = ctx.StartTime + TNodes::TTL + }); + + for (const auto& [key, value] : Nodes.Values) { + if (key.Tenant == tenant) { + *ctx.Response.add_nodes() = value.Node; + } + } } +#undef HANDLE_CPS_REQUEST +#undef HANDLE_CPS_REQUEST_IMPL + void Handle(NActors::NMon::TEvHttpInfo::TPtr& ev) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("Unimplemented " << __LINE__); TStringStream str; Send(ev->Sender, new NActors::NMon::TEvHttpInfoRes(str.Str())); } template void SendEmptyResponse(TRequest& ev, std::string logText) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("SendEmptyResponse"); CPS_LOG_I(logText); TResult result = {}; @@ -335,6 +768,7 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActor void SendEmptyAuditResponse(TRequest& ev, std::string logText) { + LOG_YQ_CONTROL_PLANE_STORAGE_CRIT("SendEmptyAuditResponse"); CPS_LOG_I(logText); TResult result = {}; @@ -343,50 +777,250 @@ class TInMemoryControlPlaneStorageActor : public NActors::TActorSend(new IEventHandle(ev->Sender, SelfId(), event.release(), 0, ev->Cookie)); } - NYql::TIssues ValidateCreateQueryRequest(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) - { - NYql::TIssues issues; - const FederatedQuery::CreateQueryRequest& request = ev->Get()->Request; - const TString user = ev->Get()->User; - const TString scope = ev->Get()->Scope; - const FederatedQuery::QueryContent& query = request.content(); +private: + template + static std::optional GetEntity(const TTable& table, const TTable::TKey& key) { + const auto it = table.Values.find(key); + if (it == table.Values.end()) { + return std::nullopt; + } + return it->second; + } + + template + static bool AddEntity(TTable& table, const TTable::TKey& key, const TTable::TValue& value) { + return table.Values.emplace(key, value).second; + } + + template + auto GetScopeRange(const TMap& table, const TString& scope) const { + const auto startIt = table.lower_bound({scope, ""}); + + std::ranges::subrange range(startIt, table.end()); + return range | std::views::take_while([scope](auto element) { + return element.first.Scope == scope; + }); + } + + template + auto GetVisibleRange(const TMap& table, const TString& scope, const TString& user, std::optional visibility = std::nullopt) const { + auto range = GetScopeRange(table, scope); + return range | std::views::filter([user, visibility, ignorePrivate = Config->Proto.GetIgnorePrivateSources()](const auto& element) { + const auto entityVisibility = element.second.GetEntity().content().acl().visibility(); + if (ignorePrivate && entityVisibility == FederatedQuery::Acl::PRIVATE) { + return false; + } + if (visibility && entityVisibility != *visibility) { + return false; + } + return entityVisibility == FederatedQuery::Acl::SCOPE || element.second.User == user; + }); + } + + template + TVector GetEntities(const TTable& table, const TString& scope, const TString& user) const { + auto range = GetVisibleRange(table.Values, scope, user) | std::views::transform([](const auto& element) { + return element.second.GetEntity(); + }) | std::views::common; + return {range.begin(), range.end()}; + } + + template + THashMap GetEntitiesWithVisibilityPriority(const TTable& table, const TString& scope, const TString& user) const { + THashMap entities; + for (const auto& [_, value] : GetVisibleRange(table.Values, scope, user)) { + const auto& entity = value.GetEntity(); + const auto visibility = entity.content().acl().visibility(); + const TString& name = entity.content().name(); + if (auto it = entities.find(name); it != entities.end()) { + if (visibility == FederatedQuery::Acl::PRIVATE) { + it->second = entity; + } + } else { + entities[name] = entity; + } + } + return entities; + } + + template + bool CheckConnectionOrBindingName(const TTable& table, const TString& scope, const TString& user, FederatedQuery::Acl::Visibility visibility, const TString& name) const { + auto range = GetVisibleRange(table.Values, scope, user, visibility) | std::views::transform([](const auto& element) { + return element.second.GetEntity().content().name(); + }); + return std::ranges::find(range, name) == range.end(); + } + + template + ui64 GetNumberEntitiesByScope(const TTable& table, const TString& scope) const { + auto range = GetScopeRange(table.Values, scope) | std::views::common; + return std::distance(range.begin(), range.end()); + } + +private: + void Cleanup() { + CleanupTable(Queries); + CleanupTable(Jobs); + CleanupTable(ResultSets); + CleanupTable(IdempotencyKeys); + CleanupTable(Nodes); + } + + template + static void CleanupTable(TTable& table) { + const auto now = TInstant::Now(); + std::erase_if(table.Values, [now](const auto& item) { + const auto expireAt = item.second.ExpireAt; + return expireAt && expireAt < now; + }); + } + +private: + // Get / Ping task utils + + struct TTaskInternal { + TTask Task; + TString Owner; + TRetryLimiter RetryLimiter; + TString TenantName; + bool ShouldAbortTask; + }; + + TVector GetActiveTasks(const TCommonRequestContextTEvGetTaskRequest& ctx) const { + const ui64 tasksBatchSize = Config->Proto.GetTasksBatchSize(); + const TString& tenantName = ctx.Request.tenant(); + + TVector tasks; + tasks.reserve(std::min(tasksBatchSize, PendingQueries.Values.size())); + for (const auto& [key, query] : PendingQueries.Values) { + if (key.Tenant != tenantName || query.AssignedUntil >= ctx.StartTime) { + continue; + } + + TTaskInternal& taskInternal = tasks.emplace_back(); + taskInternal.Owner = ctx.Request.owner_id(); + taskInternal.TenantName = tenantName; + taskInternal.RetryLimiter = query.RetryLimiter; + + auto& task = taskInternal.Task; + task.Scope = key.Scope; + task.QueryId = key.QueryId; + + if (query.Owner) { + CPS_LOG_T("Task (Query): " << task.QueryId << " Lease TIMEOUT, RetryCounterUpdatedAt " << taskInternal.RetryLimiter.RetryCounterUpdatedAt << " LastSeenAt: " << query.LastSeenAt); + taskInternal.ShouldAbortTask = !taskInternal.RetryLimiter.UpdateOnRetry(query.LastSeenAt, Config->TaskLeaseRetryPolicy, ctx.StartTime); + } + task.RetryCount = taskInternal.RetryLimiter.RetryCount; + + CPS_LOG_T("Task (Query): " << task.QueryId << " RetryRate: " << taskInternal.RetryLimiter.RetryRate << " RetryCounter: " << taskInternal.RetryLimiter.RetryCount << " At: " << taskInternal.RetryLimiter.RetryCounterUpdatedAt << (taskInternal.ShouldAbortTask ? " ABORTED" : "")); + + if (tasks.size() >= tasksBatchSize) { + break; + } + } + + std::shuffle(tasks.begin(), tasks.end(), std::default_random_engine(TInstant::Now().MicroSeconds())); + const ui64 numTasksProportion = Config->Proto.GetNumTasksProportion(); + tasks.resize((tasks.size() + numTasksProportion - 1) / numTasksProportion); + + return tasks; + } + + std::optional AssignTask(const TCommonRequestContextTEvGetTaskRequest& ctx, TTaskInternal taskInternal) { + auto& task = taskInternal.Task; - TString error; - if (!request.validate(error)) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, error)); + const auto& query = GetEntity(Queries, {task.Scope, task.QueryId}); + if (!query) { + ctx.Fail("task build", {NYql::TIssue(TStringBuilder() << "INTERNAL ERROR. Not found query for task in scope " << task.Scope << " with query id " << task.QueryId)}); + return std::nullopt; } - if (query.type() == FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "type field is not specified")); + task.Generation = query->Generation + 1; + task.Query = query->Query; + task.Internal = query->QueryInternal; + task.Deadline = TInstant::Now() + (task.Query.content().automatic() ? std::min(Config->AutomaticQueriesTtl, Config->ResultSetsTtl) : Config->ResultSetsTtl); + *task.Internal.mutable_result_ttl() = NProtoInterop::CastToProto(Config->ResultSetsTtl); + + if (Config->Proto.GetDisableCurrentIam()) { + task.Internal.clear_token(); } - if (query.acl().visibility() == FederatedQuery::Acl::VISIBILITY_UNSPECIFIED) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "acl.visibility field is not specified")); + if (taskInternal.ShouldAbortTask) { + AddTransientIssues(task.Query.mutable_transient_issue(), {NYql::TIssue("Query was aborted by system due to high failure rate")}); + task.Query.mutable_meta()->set_status(FederatedQuery::QueryMeta::ABORTING_BY_SYSTEM); } - if (request.ByteSize() > static_cast(Config.Proto.GetMaxRequestSize())) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "Request size exceeded " + ToString(request.ByteSize()) + " out of " + ToString(Config.Proto.GetMaxRequestSize()) + " bytes")); + if (const auto tenantInfo = ctx.Event.TenantInfo) { + const TString& newTenant = tenantInfo->Assign(task.Internal.cloud_id(), task.Scope, task.Query.content().type(), taskInternal.TenantName); + if (newTenant != taskInternal.TenantName) { + UpdateTaskState(ctx, taskInternal, newTenant); + return std::nullopt; + } + if (tenantInfo->TenantState.Value(taskInternal.TenantName, TenantState::Active) != TenantState::Active) { + return std::nullopt; + } } - const uint64_t countQueries = count_if(Queries.begin(), Queries.end(), [scope](const auto& item) { - const auto& [key, value] = item; - return key.Scope == scope; - }); + UpdateTaskState(ctx, taskInternal); + + return task; + } - if (countQueries > Config.Proto.GetMaxCountQueries()) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "The count of the queries exceeds the limit of " + ToString(countQueries) + " out of " + ToString(Config.Proto.GetMaxCountQueries()))); + void UpdateTaskState(const TCommonRequestContextTEvGetTaskRequest& ctx, const TTaskInternal& taskInternal, const TString& newTenant = "") { + const auto& task = taskInternal.Task; + + const auto queryIt = Queries.Values.find({task.Scope, task.QueryId}); + if (queryIt != Queries.Values.end()) { + queryIt->second.Query = task.Query; + queryIt->second.QueryInternal = task.Internal; + queryIt->second.Generation = task.Generation; } - return issues; + const auto pendingIt = PendingQueries.Values.find({taskInternal.TenantName, task.Scope, task.QueryId}); + if (pendingIt != PendingQueries.Values.end()) { + pendingIt->second.Owner = taskInternal.Owner; + pendingIt->second.RetryLimiter = taskInternal.RetryLimiter; + pendingIt->second.AssignedUntil = ctx.StartTime + Config->TaskLeaseTtl; + pendingIt->second.LastSeenAt = ctx.StartTime; + if (newTenant) { + const auto value = pendingIt->second; + PendingQueries.Values.erase(pendingIt); + AddEntity(PendingQueries, {newTenant, task.Scope, task.QueryId}, value); + } + } } - void CleanupIndempotencyKeys() - { - auto now = TInstant::Now(); - erase_if(IdempotencyKeys, [this, now](const auto& item) { - auto const& [idempotencyKey, timestamp] = item; - return timestamp + Config.IdempotencyKeyTtl < now; - }); + void PingTask(const TCommonRequestContextTEvPingTaskRequest& ctx, const TQueries::TValue& query, const TJobs::TValue& job, const TPendingQueries::TValue& pendingQuery, TDuration backoff, TInstant expireAt) { + const auto& scope = ctx.Request.scope(); + const auto& queryId = ctx.Request.query_id().value(); + const auto status = query.Query.meta().status(); + + const auto pendingIt = PendingQueries.Values.find({ctx.Request.tenant(), scope, queryId}); + if (pendingIt != PendingQueries.Values.end()) { + if (IsTerminalStatus(status)) { + PendingQueries.Values.erase(pendingIt); + } else { + pendingIt->second = pendingQuery; + pendingIt->second.AssignedUntil = ctx.StartTime + backoff; + } + } + + const bool automaticFinished = IsTerminalStatus(status) && query.Query.content().automatic(); + const auto jobIt = Jobs.Values.find({scope, queryId, job.Job.meta().id()}); + if (jobIt != Jobs.Values.end()) { + jobIt->second = job; + jobIt->second.ExpireAt = automaticFinished ? expireAt : TInstant::Zero(); + } + + const auto queryIt = Queries.Values.find({scope, queryId}); + if (queryIt != Queries.Values.end()) { + queryIt->second = query; + if (ctx.Request.has_result_id()) { + queryIt->second.ResultId = ctx.Request.result_id().value(); + } + queryIt->second.ResultExpireAt = status == FederatedQuery::QueryMeta::COMPLETED ? NProtoInterop::CastFromProto(ctx.Request.deadline()) : TInstant::Zero(); + queryIt->second.ExpireAt = automaticFinished ? expireAt : TInstant::Zero(); + } } }; @@ -395,8 +1029,14 @@ NActors::TActorId ControlPlaneStorageServiceActorId(ui32 nodeId) { return NActors::TActorId(nodeId, name); } -NActors::IActor* CreateInMemoryControlPlaneStorageServiceActor(const NConfig::TControlPlaneStorageConfig& config) { - return new TInMemoryControlPlaneStorageActor(config); +NActors::IActor* CreateInMemoryControlPlaneStorageServiceActor( + const NConfig::TControlPlaneStorageConfig& config, + const NYql::TS3GatewayConfig& s3Config, + const NConfig::TCommonConfig& common, + const NConfig::TComputeConfig& computeConfig, + const ::NMonitoring::TDynamicCounterPtr& counters, + const TString& tenantName) { + return new TInMemoryControlPlaneStorageActor(config, s3Config, common, computeConfig, counters, tenantName); } } // NFq diff --git a/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp b/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp index 652098e00de3..ac815ff3b37d 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/nodes_health_check.cpp @@ -1,9 +1,17 @@ #include "utils.h" +#include #include namespace NFq { +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) const { + const auto& request = ev->Get()->Request; + const auto& node = request.node(); + + return ValidateNodesHealthCheck(request.tenant(), node.instance_id(), node.hostname()); +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -29,8 +37,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvNodesHealth CPS_LOG_T("NodesHealthCheckRequest: {" << request.DebugString() << "}"); - NYql::TIssues issues = ValidateNodesHealthCheck(tenant, instanceId, hostName); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("NodesHealthCheckRequest: {" << request.DebugString() << "} validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); diff --git a/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp b/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp index 41bad0bb970d..3031cac8df01 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/task_get.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -247,6 +248,85 @@ TDuration ExtractLimit(const TTask& task) { return executionLimit; } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) const +{ + const auto& request = ev->Get()->Request; + NYql::TIssues issues = ValidateGetTask(request.owner_id(), request.host()); + + if (!ev->Get()->TenantInfo) { + issues.AddIssue(MakeErrorIssue(TIssuesIds::NOT_READY, "Control Plane is not ready yet. Please retry later.")); + } + + return issues; +} + +void TControlPlaneStorageBase::FillGetTaskResult(Fq::Private::GetTaskResult& result, const TVector& tasks) const +{ + for (const auto& task : tasks) { + const auto& queryType = task.Query.content().type(); + if (queryType != FederatedQuery::QueryContent::ANALYTICS && queryType != FederatedQuery::QueryContent::STREAMING) { //TODO: fix + ythrow yexception() + << "query type " + << FederatedQuery::QueryContent::QueryType_Name(queryType) + << " unsupported"; + } + + auto* newTask = result.add_tasks(); + newTask->set_query_type(queryType); + newTask->set_query_syntax(task.Query.content().syntax()); + newTask->set_execute_mode(task.Query.meta().execute_mode()); + newTask->set_state_load_mode(task.Internal.state_load_mode()); + auto* queryId = newTask->mutable_query_id(); + queryId->set_value(task.Query.meta().common().id()); + newTask->set_streaming(queryType == FederatedQuery::QueryContent::STREAMING); + newTask->set_text(task.Query.content().text()); + *newTask->mutable_connection() = task.Internal.connection(); + *newTask->mutable_binding() = task.Internal.binding(); + newTask->set_user_token(task.Internal.token()); + newTask->set_user_id(task.Query.meta().common().created_by()); + newTask->set_generation(task.Generation); + newTask->set_status(task.Query.meta().status()); + *newTask->mutable_created_topic_consumers() = task.Internal.created_topic_consumers(); + newTask->mutable_sensor_labels()->insert({"cloud_id", task.Internal.cloud_id()}); + newTask->mutable_sensor_labels()->insert({"scope", task.Scope}); + newTask->set_automatic(task.Query.content().automatic()); + newTask->set_query_name(task.Query.content().name()); + *newTask->mutable_deadline() = NProtoInterop::CastToProto(task.Deadline); + newTask->mutable_disposition()->CopyFrom(task.Internal.disposition()); + newTask->set_result_limit(task.Internal.result_limit()); + *newTask->mutable_execution_limit() = NProtoInterop::CastToProto(ExtractLimit(task)); + *newTask->mutable_request_started_at() = task.Query.meta().started_at(); + *newTask->mutable_request_submitted_at() = task.Query.meta().submitted_at(); + + newTask->set_restart_count(task.RetryCount); + auto* jobId = newTask->mutable_job_id(); + jobId->set_value(task.Query.meta().last_job_id()); + + for (const auto& connection: task.Internal.connection()) { + const auto serviceAccountId = ExtractServiceAccountId(connection); + if (!serviceAccountId) { + continue; + } + auto* account = newTask->add_service_accounts(); + account->set_value(serviceAccountId); + } + + *newTask->mutable_dq_graph() = task.Internal.dq_graph(); + newTask->set_dq_graph_index(task.Internal.dq_graph_index()); + *newTask->mutable_dq_graph_compressed() = task.Internal.dq_graph_compressed(); + + *newTask->mutable_result_set_meta() = task.Query.result_set_meta(); + newTask->set_scope(task.Scope); + *newTask->mutable_resources() = task.Internal.resources(); + + newTask->set_execution_id(task.Internal.execution_id()); + newTask->set_operation_id(task.Internal.operation_id()); + *newTask->mutable_compute_connection() = task.Internal.compute_connection(); + *newTask->mutable_result_ttl() = task.Internal.result_ttl(); + *newTask->mutable_parameters() = task.Query.content().parameters(); + } +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -262,13 +342,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequ CPS_LOG_T("GetTaskRequest: {" << request.DebugString() << "}"); - NYql::TIssues issues = ValidateGetTask(owner, hostName); - - if (!ev->Get()->TenantInfo) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::NOT_READY, "Control Plane is not ready yet. Please retry later.")); - } - - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("GetTaskRequest: {" << request.DebugString() << "} FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -364,15 +438,17 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequ responseTasks=responseTasks] (const auto& readFuture) mutable { try { - if (!readFuture.GetValue().IsSuccess()) + if (!readFuture.GetValue().IsSuccess()) { return readFuture; + } } catch (...) { return readFuture; } auto pickTaskParams = prepareParams(*resultSets); - if (pickTaskParams.empty()) + if (pickTaskParams.empty()) { return readFuture; + } auto debugInfos = std::make_shared>(pickTaskParams.size()); if (Config->Proto.GetEnableDebugMode()) { @@ -410,74 +486,9 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetTaskRequ }); }); - auto prepare = [response] { + auto prepare = [this, response] { Fq::Private::GetTaskResult result; - const auto& tasks = std::get<0>(*response); - - for (const auto& task : tasks) { - const auto& queryType = task.Query.content().type(); - if (queryType != FederatedQuery::QueryContent::ANALYTICS && queryType != FederatedQuery::QueryContent::STREAMING) { //TODO: fix - ythrow yexception() - << "query type " - << FederatedQuery::QueryContent::QueryType_Name(queryType) - << " unsupported"; - } - - auto* newTask = result.add_tasks(); - newTask->set_query_type(queryType); - newTask->set_query_syntax(task.Query.content().syntax()); - newTask->set_execute_mode(task.Query.meta().execute_mode()); - newTask->set_state_load_mode(task.Internal.state_load_mode()); - auto* queryId = newTask->mutable_query_id(); - queryId->set_value(task.Query.meta().common().id()); - newTask->set_streaming(queryType == FederatedQuery::QueryContent::STREAMING); - newTask->set_text(task.Query.content().text()); - *newTask->mutable_connection() = task.Internal.connection(); - *newTask->mutable_binding() = task.Internal.binding(); - newTask->set_user_token(task.Internal.token()); - newTask->set_user_id(task.Query.meta().common().created_by()); - newTask->set_generation(task.Generation); - newTask->set_status(task.Query.meta().status()); - *newTask->mutable_created_topic_consumers() = task.Internal.created_topic_consumers(); - newTask->mutable_sensor_labels()->insert({"cloud_id", task.Internal.cloud_id()}); - newTask->mutable_sensor_labels()->insert({"scope", task.Scope}); - newTask->set_automatic(task.Query.content().automatic()); - newTask->set_query_name(task.Query.content().name()); - *newTask->mutable_deadline() = NProtoInterop::CastToProto(task.Deadline); - newTask->mutable_disposition()->CopyFrom(task.Internal.disposition()); - newTask->set_result_limit(task.Internal.result_limit()); - *newTask->mutable_execution_limit() = NProtoInterop::CastToProto(ExtractLimit(task)); - *newTask->mutable_request_started_at() = task.Query.meta().started_at(); - *newTask->mutable_request_submitted_at() = task.Query.meta().submitted_at(); - - newTask->set_restart_count(task.RetryCount); - auto* jobId = newTask->mutable_job_id(); - jobId->set_value(task.Query.meta().last_job_id()); - - for (const auto& connection: task.Internal.connection()) { - const auto serviceAccountId = ExtractServiceAccountId(connection); - if (!serviceAccountId) { - continue; - } - auto* account = newTask->add_service_accounts(); - account->set_value(serviceAccountId); - } - - *newTask->mutable_dq_graph() = task.Internal.dq_graph(); - newTask->set_dq_graph_index(task.Internal.dq_graph_index()); - *newTask->mutable_dq_graph_compressed() = task.Internal.dq_graph_compressed(); - - *newTask->mutable_result_set_meta() = task.Query.result_set_meta(); - newTask->set_scope(task.Scope); - *newTask->mutable_resources() = task.Internal.resources(); - - newTask->set_execution_id(task.Internal.execution_id()); - newTask->set_operation_id(task.Internal.operation_id()); - *newTask->mutable_compute_connection() = task.Internal.compute_connection(); - *newTask->mutable_result_ttl() = task.Internal.result_ttl(); - *newTask->mutable_parameters() = task.Query.content().parameters(); - } - + FillGetTaskResult(result, std::get<0>(*response)); return result; }; auto success = SendResponse diff --git a/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp b/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp index 976cceeca87b..b98a0e6b549d 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/task_ping.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -45,36 +46,15 @@ THashMap DeserializeFlatStats(const google::protobuf::RepeatedPtrF } -struct TPingTaskParams { - TString Query; - TParams Params; - const std::function(const TVector&)> Prepare; - std::shared_ptr> MeteringRecords; -}; - -struct TFinalStatus { - FederatedQuery::QueryMeta::ComputeStatus Status = FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED; - NYql::NDqProto::StatusIds::StatusCode StatusCode = NYql::NDqProto::StatusIds::UNSPECIFIED; - FederatedQuery::QueryContent::QueryType QueryType = FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED; - NYql::TIssues Issues; - NYql::TIssues TransientIssues; - StatsValuesList FinalStatistics; - TString CloudId; - TString JobId; -}; - -TPingTaskParams ConstructHardPingTask( +TYdbControlPlaneStorageActor::TPingTaskParams TYdbControlPlaneStorageActor::ConstructHardPingTask( const Fq::Private::PingTaskRequest& request, std::shared_ptr response, - const TString& tablePathPrefix, const TDuration& automaticQueriesTtl, const TDuration& taskLeaseTtl, - const THashMap& retryPolicies, ::NMonitoring::TDynamicCounterPtr rootCounters, - uint64_t maxRequestSize, bool dumpRawStatistics, const std::shared_ptr& finalStatus, - const TRequestCommonCountersPtr& commonCounters) { + const std::shared_ptr& finalStatus, const TRequestCommonCountersPtr& commonCounters) const { auto scope = request.scope(); auto query_id = request.query_id().value(); - auto counters = rootCounters->GetSubgroup("scope", scope)->GetSubgroup("query_id", query_id); + auto counters = Counters.Counters->GetSubgroup("scope", scope)->GetSubgroup("query_id", query_id); - TSqlQueryBuilder readQueryBuilder(tablePathPrefix, "HardPingTask(read)"); + TSqlQueryBuilder readQueryBuilder(YdbConnection->TablePathPrefix, "HardPingTask(read)"); readQueryBuilder.AddString("tenant", request.tenant()); readQueryBuilder.AddString("scope", scope); readQueryBuilder.AddString("query_id", query_id); @@ -146,297 +126,12 @@ TPingTaskParams ConstructHardPingTask( ); } - TMaybe queryStatus; - if (request.status() != FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED) { - queryStatus = request.status(); - } - TMaybe issues; - if (request.issues().size() > 0) { - NYql::TIssues requestIssues; - NYql::IssuesFromMessage(request.issues(), requestIssues); - issues = requestIssues; - } - TMaybe transientIssues; - if (request.transient_issues().size() > 0) { - NYql::TIssues requestTransientIssues; - NYql::IssuesFromMessage(request.transient_issues(), requestTransientIssues); - transientIssues = requestTransientIssues; - } // running query us locked for lease period - TDuration backoff = taskLeaseTtl; - - if (request.resign_query()) { - if (request.status_code() == NYql::NDqProto::StatusIds::UNSPECIFIED && internal.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { - request.set_status_code(internal.pending_status_code()); - internal.clear_pending_status_code(); - internal.clear_execution_id(); - internal.clear_operation_id(); - } - - TRetryPolicyItem policy(0, 0, TDuration::Seconds(1), TDuration::Zero()); - auto it = retryPolicies.find(request.status_code()); - auto policyFound = it != retryPolicies.end(); - if (policyFound) { - policy = it->second; - } - - auto now = TInstant::Now(); - auto executionDeadline = TInstant::Max(); - - auto submittedAt = NProtoInterop::CastFromProto(query.meta().submitted_at()); - auto executionTtl = NProtoInterop::CastFromProto(internal.execution_ttl()); - if (submittedAt && executionTtl) { - executionDeadline = submittedAt + executionTtl; - } - - if (retryLimiter.UpdateOnRetry(now, policy) && now < executionDeadline) { - queryStatus.Clear(); - // failing query is throttled for backoff period - backoff = policy.BackoffPeriod * (retryLimiter.RetryRate + 1); - owner = ""; - if (!transientIssues) { - transientIssues.ConstructInPlace(); - } - TStringBuilder builder; - builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) - << " and will be restarted (RetryCount: " << retryLimiter.RetryCount << ")" - << " at " << now; - transientIssues->AddIssue(NYql::TIssue(builder)); - } else { - // failure query should be processed instantly - queryStatus = FederatedQuery::QueryMeta::FAILING; - backoff = TDuration::Zero(); - TStringBuilder builder; - builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()); - if (policy.RetryCount) { - builder << " (" << retryLimiter.LastError << ")"; - } - builder << " at " << now; - - // in case of problems with finalization, do not change the issues - if (query.meta().status() == FederatedQuery::QueryMeta::FAILING || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_SYSTEM || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_USER) { - if (issues) { - transientIssues->AddIssues(*issues); - } - transientIssues->AddIssue(NYql::TIssue(builder)); - } else { - if (!issues) { - issues.ConstructInPlace(); - } - auto issue = NYql::TIssue(builder); - if (query.issue().size() > 0 && request.issues().empty()) { - NYql::TIssues queryIssues; - NYql::IssuesFromMessage(query.issue(), queryIssues); - for (auto& subIssue : queryIssues) { - issue.AddSubIssue(MakeIntrusive(subIssue)); - } - } - if (transientIssues) { - for (auto& subIssue : *transientIssues) { - issue.AddSubIssue(MakeIntrusive(subIssue)); - } - transientIssues.Clear(); - } - issues->AddIssue(issue); - } - } - CPS_LOG_AS_D(*actorSystem, "PingTaskRequest (resign): " << (!policyFound ? " DEFAULT POLICY" : "") << (owner ? " FAILURE " : " ") << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) << " " << retryLimiter.RetryCount << " " << retryLimiter.RetryCounterUpdatedAt << " " << backoff); - } - - if (queryStatus) { - query.mutable_meta()->set_status(*queryStatus); - job.mutable_query_meta()->set_status(*queryStatus); - } - - if (request.status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { - internal.set_status_code(request.status_code()); - } - - if (request.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { - internal.set_pending_status_code(request.pending_status_code()); - } - - if (issues) { - NYql::IssuesToMessage(*issues, query.mutable_issue()); - NYql::IssuesToMessage(*issues, job.mutable_issue()); - } - - if (transientIssues) { - AddTransientIssues(query.mutable_transient_issue(), std::move(*transientIssues)); - } - - if (request.internal_issues().size()) { - *internal.mutable_internal_issue() = request.internal_issues(); - } - - if (request.statistics()) { - TString statistics = request.statistics(); - if (request.flat_stats_size() == 0) { - internal.clear_statistics(); - // TODO: remove once V1 and V2 stats go the same way - PackStatisticsToProtobuf(*internal.mutable_statistics(), statistics, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); - } - - // global dumpRawStatistics will be removed with YQv1 - if (!dumpRawStatistics && !request.dump_raw_statistics()) { - try { - statistics = GetPrettyStatistics(statistics); - } catch (const std::exception&) { - // LOG_AS? - CPS_LOG_E("Error on statistics prettification: " << CurrentExceptionMessage()); - } - } - *query.mutable_statistics()->mutable_json() = statistics; - *job.mutable_statistics()->mutable_json() = statistics; - } - - if (request.current_load()) { - internal.set_current_load(request.current_load()); - } - - if (request.timeline()) { - internal.set_timeline(request.timeline()); - } - - if (request.flat_stats_size() != 0) { - internal.clear_statistics(); - auto stats = DeserializeFlatStats(request.flat_stats()); - PackStatisticsToProtobuf(*internal.mutable_statistics(), stats, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); - } - - if (!request.result_set_meta().empty()) { - // we will overwrite result_set_meta's COMPLETELY - *query.mutable_result_set_meta() = request.result_set_meta(); - *job.mutable_result_set_meta() = request.result_set_meta(); - } - - if (request.ast()) { - query.mutable_ast()->set_data(request.ast()); - job.mutable_ast()->set_data(request.ast()); - } - - if (request.plan()) { - query.mutable_plan()->set_json(request.plan()); - job.mutable_plan()->set_json(request.plan()); - } - - if (request.ast_compressed().data()) { - internal.mutable_ast_compressed()->set_method(request.ast_compressed().method()); - internal.mutable_ast_compressed()->set_data(request.ast_compressed().data()); - // todo: keep AST compressed in JobInternal - // job.mutable_ast()->set_data(request.ast()); - } - - if (request.plan_compressed().data()) { - internal.mutable_plan_compressed()->set_method(request.plan_compressed().method()); - internal.mutable_plan_compressed()->set_data(request.plan_compressed().data()); - // todo: keep plan compressed in JobInternal - // job.mutable_plan()->set_json(request.plan()); - } - - if (request.has_started_at()) { - *query.mutable_meta()->mutable_started_at() = request.started_at(); - *job.mutable_query_meta()->mutable_started_at() = request.started_at(); - } - - if (request.has_finished_at()) { - *query.mutable_meta()->mutable_finished_at() = request.finished_at(); - *job.mutable_query_meta()->mutable_finished_at() = request.finished_at(); - if (!query.meta().has_started_at()) { - *query.mutable_meta()->mutable_started_at() = request.finished_at(); - *job.mutable_query_meta()->mutable_started_at() = request.finished_at(); - } - } - - TInstant expireAt = TInstant::Now() + automaticQueriesTtl; - if (IsTerminalStatus(query.meta().status()) && query.content().automatic()) { - *query.mutable_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); - *job.mutable_query_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); - *job.mutable_expire_at() = NProtoInterop::CastToProto(expireAt); - } - - if (query.meta().status() == FederatedQuery::QueryMeta::COMPLETED) { - *query.mutable_meta()->mutable_result_expire_at() = request.deadline(); - } - - if (request.state_load_mode()) { - internal.set_state_load_mode(request.state_load_mode()); - if (request.state_load_mode() == FederatedQuery::FROM_LAST_CHECKPOINT) { // Saved checkpoint - query.mutable_meta()->set_has_saved_checkpoints(true); - } - } - - if (request.has_disposition()) { - *internal.mutable_disposition() = request.disposition(); - } - - if (request.status() && IsTerminalStatus(request.status())) { - internal.clear_created_topic_consumers(); - // internal.clear_dq_graph(); keep for debug - internal.clear_dq_graph_index(); - // internal.clear_execution_id(); keep for debug - // internal.clear_operation_id(); keep for debug - } - - if (!request.created_topic_consumers().empty()) { - std::set mergedConsumers; - for (auto&& c : *internal.mutable_created_topic_consumers()) { - mergedConsumers.emplace(std::move(c)); - } - for (const auto& c : request.created_topic_consumers()) { - mergedConsumers.emplace(c); - } - internal.clear_created_topic_consumers(); - for (auto&& c : mergedConsumers) { - *internal.add_created_topic_consumers() = std::move(c); - } - } - - if (!request.execution_id().empty()) { - internal.set_execution_id(request.execution_id()); - } - - if (!request.operation_id().empty()) { - internal.set_operation_id(request.operation_id()); - } + TDuration backoff = Config->TaskLeaseTtl; + TInstant expireAt = TInstant::Now() + Config->AutomaticQueriesTtl; + UpdateTaskInfo(actorSystem, request, finalStatus, query, internal, job, owner, retryLimiter, backoff, expireAt); - if (!request.dq_graph().empty()) { - *internal.mutable_dq_graph() = request.dq_graph(); - } - - if (!request.dq_graph_compressed().empty()) { - *internal.mutable_dq_graph_compressed() = request.dq_graph_compressed(); - } - - if (request.dq_graph_index()) { - internal.set_dq_graph_index(request.dq_graph_index()); - } - - if (request.has_resources()) { - *internal.mutable_resources() = request.resources(); - } - - if (job.ByteSizeLong() > maxRequestSize) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Job proto exceeded the size limit: " << job.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(job).ToString(); - } - - if (query.ByteSizeLong() > maxRequestSize) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Query proto exceeded the size limit: " << query.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(query).ToString(); - } - - if (internal.ByteSizeLong() > maxRequestSize) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "QueryInternal proto exceeded the size limit: " << internal.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(internal).ToString(); - } - - finalStatus->Status = query.meta().status(); - finalStatus->QueryType = query.content().type(); - finalStatus->StatusCode = internal.status_code(); - finalStatus->CloudId = internal.cloud_id(); - finalStatus->JobId = jobId; - NYql::IssuesFromMessage(query.issue(), finalStatus->Issues); - NYql::IssuesFromMessage(query.transient_issue(), finalStatus->TransientIssues); - - TSqlQueryBuilder writeQueryBuilder(tablePathPrefix, "HardPingTask(write)"); + TSqlQueryBuilder writeQueryBuilder(YdbConnection->TablePathPrefix, "HardPingTask(write)"); writeQueryBuilder.AddString("tenant", request.tenant()); writeQueryBuilder.AddString("scope", request.scope()); writeQueryBuilder.AddString("job_id", jobId); @@ -531,18 +226,7 @@ TPingTaskParams ConstructHardPingTask( // YQv2 may not provide statistics with terminal status, use saved one statistics = query.statistics().json(); } - finalStatus->FinalStatistics = ExtractStatisticsFromProtobuf(internal.statistics()); - finalStatus->FinalStatistics.push_back(std::make_pair("IsAutomatic", query.content().automatic())); - if (query.content().name().Contains("DataLens YQ query")) { - finalStatus->FinalStatistics.push_back(std::make_pair("IsDataLens", 1)); - } else if (query.content().name().Contains("Audit-trails")) { - finalStatus->FinalStatistics.push_back(std::make_pair("IsAuditTrails", 1)); - } else if (query.content().name().Contains("Query from YDB SDK")) { - finalStatus->FinalStatistics.push_back(std::make_pair("IsSDK", 1)); - } - finalStatus->FinalStatistics.push_back(std::make_pair("RetryCount", retryLimiter.RetryCount)); - finalStatus->FinalStatistics.push_back(std::make_pair("RetryRate", retryLimiter.RetryRate * 100)); - finalStatus->FinalStatistics.push_back(std::make_pair("Load", internal.current_load())); + FillQueryStatistics(finalStatus, query, internal, retryLimiter); auto records = GetMeteringRecords(statistics, isBillable, jobId, request.scope(), HostName()); meteringRecords->swap(records); @@ -558,10 +242,10 @@ TPingTaskParams ConstructHardPingTask( return {readQuery.Sql, readQuery.Params, prepareParams, meteringRecords}; } -TPingTaskParams ConstructSoftPingTask( +TYdbControlPlaneStorageActor::TPingTaskParams TYdbControlPlaneStorageActor::ConstructSoftPingTask( const Fq::Private::PingTaskRequest& request, std::shared_ptr response, - const TString& tablePathPrefix, const TDuration& taskLeaseTtl, const TRequestCommonCountersPtr& commonCounters) { - TSqlQueryBuilder readQueryBuilder(tablePathPrefix, "SoftPingTask(read)"); + const TRequestCommonCountersPtr& commonCounters) const { + TSqlQueryBuilder readQueryBuilder(YdbConnection->TablePathPrefix, "SoftPingTask(read)"); readQueryBuilder.AddString("tenant", request.tenant()); readQueryBuilder.AddString("scope", request.scope()); readQueryBuilder.AddString("query_id", request.query_id().value()); @@ -603,11 +287,11 @@ TPingTaskParams ConstructSoftPingTask( } } - TInstant ttl = TInstant::Now() + taskLeaseTtl; + TInstant ttl = TInstant::Now() + Config->TaskLeaseTtl; response->set_action(internal.action()); *response->mutable_expired_at() = google::protobuf::util::TimeUtil::MillisecondsToTimestamp(ttl.MilliSeconds()); - TSqlQueryBuilder writeQueryBuilder(tablePathPrefix, "SoftPingTask(write)"); + TSqlQueryBuilder writeQueryBuilder(YdbConnection->TablePathPrefix, "SoftPingTask(write)"); writeQueryBuilder.AddTimestamp("now", TInstant::Now()); writeQueryBuilder.AddTimestamp("ttl", ttl); writeQueryBuilder.AddString("tenant", request.tenant()); @@ -626,6 +310,330 @@ TPingTaskParams ConstructSoftPingTask( return {readQuery.Sql, readQuery.Params, prepareParams, std::shared_ptr>{}}; } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) const { + const Fq::Private::PingTaskRequest& request = ev->Get()->Request; + + NYql::TIssues issues = ValidatePingTask(request.scope(), request.query_id().value(), request.owner_id(), NProtoInterop::CastFromProto(request.deadline()), Config->ResultSetsTtl); + + const auto tenantInfo = ev->Get()->TenantInfo; + if (tenantInfo && tenantInfo->TenantState.Value(request.tenant(), TenantState::Active) == TenantState::Idle) { + issues.AddIssue("Tenant is idle, no processing is allowed"); + } + + return issues; +} + +void TControlPlaneStorageBase::UpdateTaskInfo( + NActors::TActorSystem* actorSystem, Fq::Private::PingTaskRequest& request, const std::shared_ptr& finalStatus, FederatedQuery::Query& query, + FederatedQuery::Internal::QueryInternal& internal, FederatedQuery::Job& job, TString& owner, + TRetryLimiter& retryLimiter, TDuration& backoff, TInstant& expireAt) const +{ + TMaybe queryStatus; + if (request.status() != FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED) { + queryStatus = request.status(); + } + TMaybe issues; + if (request.issues().size() > 0) { + NYql::TIssues requestIssues; + NYql::IssuesFromMessage(request.issues(), requestIssues); + issues = requestIssues; + } + TMaybe transientIssues; + if (request.transient_issues().size() > 0) { + NYql::TIssues requestTransientIssues; + NYql::IssuesFromMessage(request.transient_issues(), requestTransientIssues); + transientIssues = requestTransientIssues; + } + + if (request.resign_query()) { + if (request.status_code() == NYql::NDqProto::StatusIds::UNSPECIFIED && internal.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { + request.set_status_code(internal.pending_status_code()); + internal.clear_pending_status_code(); + internal.clear_execution_id(); + internal.clear_operation_id(); + } + + TRetryPolicyItem policy(0, 0, TDuration::Seconds(1), TDuration::Zero()); + auto it = Config->RetryPolicies.find(request.status_code()); + auto policyFound = it != Config->RetryPolicies.end(); + if (policyFound) { + policy = it->second; + } + + auto now = TInstant::Now(); + auto executionDeadline = TInstant::Max(); + + auto submittedAt = NProtoInterop::CastFromProto(query.meta().submitted_at()); + auto executionTtl = NProtoInterop::CastFromProto(internal.execution_ttl()); + if (submittedAt && executionTtl) { + executionDeadline = submittedAt + executionTtl; + } + + if (retryLimiter.UpdateOnRetry(now, policy) && now < executionDeadline) { + queryStatus.Clear(); + // failing query is throttled for backoff period + backoff = policy.BackoffPeriod * (retryLimiter.RetryRate + 1); + owner = ""; + if (!transientIssues) { + transientIssues.ConstructInPlace(); + } + TStringBuilder builder; + builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) + << " and will be restarted (RetryCount: " << retryLimiter.RetryCount << ")" + << " at " << now; + transientIssues->AddIssue(NYql::TIssue(builder)); + } else { + // failure query should be processed instantly + queryStatus = FederatedQuery::QueryMeta::FAILING; + backoff = TDuration::Zero(); + TStringBuilder builder; + builder << "Query failed with code " << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()); + if (policy.RetryCount) { + builder << " (" << retryLimiter.LastError << ")"; + } + builder << " at " << now; + + // in case of problems with finalization, do not change the issues + if (query.meta().status() == FederatedQuery::QueryMeta::FAILING || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_SYSTEM || query.meta().status() == FederatedQuery::QueryMeta::ABORTING_BY_USER) { + if (issues) { + transientIssues->AddIssues(*issues); + } + transientIssues->AddIssue(NYql::TIssue(builder)); + } else { + if (!issues) { + issues.ConstructInPlace(); + } + auto issue = NYql::TIssue(builder); + if (query.issue().size() > 0 && request.issues().empty()) { + NYql::TIssues queryIssues; + NYql::IssuesFromMessage(query.issue(), queryIssues); + for (auto& subIssue : queryIssues) { + issue.AddSubIssue(MakeIntrusive(subIssue)); + } + } + if (transientIssues) { + for (auto& subIssue : *transientIssues) { + issue.AddSubIssue(MakeIntrusive(subIssue)); + } + transientIssues.Clear(); + } + issues->AddIssue(issue); + } + } + CPS_LOG_AS_D(*actorSystem, "PingTaskRequest (resign): " << (!policyFound ? " DEFAULT POLICY" : "") << (owner ? " FAILURE " : " ") << NYql::NDqProto::StatusIds_StatusCode_Name(request.status_code()) << " " << retryLimiter.RetryCount << " " << retryLimiter.RetryCounterUpdatedAt << " " << backoff); + } + + if (queryStatus) { + query.mutable_meta()->set_status(*queryStatus); + job.mutable_query_meta()->set_status(*queryStatus); + } + + if (request.status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { + internal.set_status_code(request.status_code()); + } + + if (request.pending_status_code() != NYql::NDqProto::StatusIds::UNSPECIFIED) { + internal.set_pending_status_code(request.pending_status_code()); + } + + if (issues) { + NYql::IssuesToMessage(*issues, query.mutable_issue()); + NYql::IssuesToMessage(*issues, job.mutable_issue()); + } + + if (transientIssues) { + AddTransientIssues(query.mutable_transient_issue(), std::move(*transientIssues)); + } + + if (request.internal_issues().size()) { + *internal.mutable_internal_issue() = request.internal_issues(); + } + + if (request.statistics()) { + TString statistics = request.statistics(); + if (request.flat_stats_size() == 0) { + internal.clear_statistics(); + // TODO: remove once V1 and V2 stats go the same way + PackStatisticsToProtobuf(*internal.mutable_statistics(), statistics, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); + } + + // global dumpRawStatistics will be removed with YQv1 + if (!Config->Proto.GetDumpRawStatistics() && !request.dump_raw_statistics()) { + try { + statistics = GetPrettyStatistics(statistics); + } catch (const std::exception&) { + CPS_LOG_AS_E(*actorSystem, "Error on statistics prettification: " << CurrentExceptionMessage()); + } + } + *query.mutable_statistics()->mutable_json() = statistics; + *job.mutable_statistics()->mutable_json() = statistics; + } + + if (request.current_load()) { + internal.set_current_load(request.current_load()); + } + + if (request.timeline()) { + internal.set_timeline(request.timeline()); + } + + if (request.flat_stats_size() != 0) { + internal.clear_statistics(); + auto stats = DeserializeFlatStats(request.flat_stats()); + PackStatisticsToProtobuf(*internal.mutable_statistics(), stats, TInstant::Now() - NProtoInterop::CastFromProto(job.meta().created_at())); + } + + if (!request.result_set_meta().empty()) { + // we will overwrite result_set_meta's COMPLETELY + *query.mutable_result_set_meta() = request.result_set_meta(); + *job.mutable_result_set_meta() = request.result_set_meta(); + } + + if (request.ast()) { + query.mutable_ast()->set_data(request.ast()); + job.mutable_ast()->set_data(request.ast()); + } + + if (request.plan()) { + query.mutable_plan()->set_json(request.plan()); + job.mutable_plan()->set_json(request.plan()); + } + + if (request.ast_compressed().data()) { + internal.mutable_ast_compressed()->set_method(request.ast_compressed().method()); + internal.mutable_ast_compressed()->set_data(request.ast_compressed().data()); + // todo: keep AST compressed in JobInternal + // job.mutable_ast()->set_data(request.ast()); + } + + if (request.plan_compressed().data()) { + internal.mutable_plan_compressed()->set_method(request.plan_compressed().method()); + internal.mutable_plan_compressed()->set_data(request.plan_compressed().data()); + // todo: keep plan compressed in JobInternal + // job.mutable_plan()->set_json(request.plan()); + } + + if (request.has_started_at()) { + *query.mutable_meta()->mutable_started_at() = request.started_at(); + *job.mutable_query_meta()->mutable_started_at() = request.started_at(); + } + + if (request.has_finished_at()) { + *query.mutable_meta()->mutable_finished_at() = request.finished_at(); + *job.mutable_query_meta()->mutable_finished_at() = request.finished_at(); + if (!query.meta().has_started_at()) { + *query.mutable_meta()->mutable_started_at() = request.finished_at(); + *job.mutable_query_meta()->mutable_started_at() = request.finished_at(); + } + } + + if (IsTerminalStatus(query.meta().status()) && query.content().automatic()) { + *query.mutable_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); + *job.mutable_query_meta()->mutable_expire_at() = NProtoInterop::CastToProto(expireAt); + *job.mutable_expire_at() = NProtoInterop::CastToProto(expireAt); + } + + if (query.meta().status() == FederatedQuery::QueryMeta::COMPLETED) { + *query.mutable_meta()->mutable_result_expire_at() = request.deadline(); + } + + if (request.state_load_mode()) { + internal.set_state_load_mode(request.state_load_mode()); + if (request.state_load_mode() == FederatedQuery::FROM_LAST_CHECKPOINT) { // Saved checkpoint + query.mutable_meta()->set_has_saved_checkpoints(true); + } + } + + if (request.has_disposition()) { + *internal.mutable_disposition() = request.disposition(); + } + + if (request.status() && IsTerminalStatus(request.status())) { + internal.clear_created_topic_consumers(); + // internal.clear_dq_graph(); keep for debug + internal.clear_dq_graph_index(); + // internal.clear_execution_id(); keep for debug + // internal.clear_operation_id(); keep for debug + } + + if (!request.created_topic_consumers().empty()) { + std::set mergedConsumers; + for (auto&& c : *internal.mutable_created_topic_consumers()) { + mergedConsumers.emplace(std::move(c)); + } + for (const auto& c : request.created_topic_consumers()) { + mergedConsumers.emplace(c); + } + internal.clear_created_topic_consumers(); + for (auto&& c : mergedConsumers) { + *internal.add_created_topic_consumers() = std::move(c); + } + } + + if (!request.execution_id().empty()) { + internal.set_execution_id(request.execution_id()); + } + + if (!request.operation_id().empty()) { + internal.set_operation_id(request.operation_id()); + } + + if (!request.dq_graph().empty()) { + *internal.mutable_dq_graph() = request.dq_graph(); + } + + if (!request.dq_graph_compressed().empty()) { + *internal.mutable_dq_graph_compressed() = request.dq_graph_compressed(); + } + + if (request.dq_graph_index()) { + internal.set_dq_graph_index(request.dq_graph_index()); + } + + if (request.has_resources()) { + *internal.mutable_resources() = request.resources(); + } + + const auto maxRequestSize = Config->Proto.GetMaxRequestSize(); + if (job.ByteSizeLong() > maxRequestSize) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Job proto exceeded the size limit: " << job.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(job).ToString(); + } + + if (query.ByteSizeLong() > maxRequestSize) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Query proto exceeded the size limit: " << query.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(query).ToString(); + } + + if (internal.ByteSizeLong() > maxRequestSize) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "QueryInternal proto exceeded the size limit: " << internal.ByteSizeLong() << " of " << maxRequestSize << " " << TSizeFormatPrinter(internal).ToString(); + } + + finalStatus->Status = query.meta().status(); + finalStatus->QueryType = query.content().type(); + finalStatus->StatusCode = internal.status_code(); + finalStatus->CloudId = internal.cloud_id(); + finalStatus->JobId = job.meta().id(); + NYql::IssuesFromMessage(query.issue(), finalStatus->Issues); + NYql::IssuesFromMessage(query.transient_issue(), finalStatus->TransientIssues); +} + +void TControlPlaneStorageBase::FillQueryStatistics( + const std::shared_ptr& finalStatus, const FederatedQuery::Query& query, + const FederatedQuery::Internal::QueryInternal& internal, const TRetryLimiter& retryLimiter) const +{ + finalStatus->FinalStatistics = ExtractStatisticsFromProtobuf(internal.statistics()); + finalStatus->FinalStatistics.push_back(std::make_pair("IsAutomatic", query.content().automatic())); + if (query.content().name().Contains("DataLens YQ query")) { + finalStatus->FinalStatistics.push_back(std::make_pair("IsDataLens", 1)); + } else if (query.content().name().Contains("Audit-trails")) { + finalStatus->FinalStatistics.push_back(std::make_pair("IsAuditTrails", 1)); + } else if (query.content().name().Contains("Query from YDB SDK")) { + finalStatus->FinalStatistics.push_back(std::make_pair("IsSDK", 1)); + } + finalStatus->FinalStatistics.push_back(std::make_pair("RetryCount", retryLimiter.RetryCount)); + finalStatus->FinalStatistics.push_back(std::make_pair("RetryRate", retryLimiter.RetryRate * 100)); + finalStatus->FinalStatistics.push_back(std::make_pair("Load", internal.current_load())); +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -635,20 +643,10 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskReq requestCounters.IncInFly(); requestCounters.Common->RequestBytes->Add(ev->Get()->GetByteSize()); const TString queryId = request.query_id().value(); - const TString owner = request.owner_id(); - const TInstant deadline = NProtoInterop::CastFromProto(request.deadline()); - const TString tenant = request.tenant(); CPS_LOG_T("PingTaskRequest: {" << request.DebugString() << "}"); - NYql::TIssues issues = ValidatePingTask(scope, queryId, owner, deadline, Config->ResultSetsTtl); - - auto tenantInfo = ev->Get()->TenantInfo; - if (tenantInfo && tenantInfo->TenantState.Value(tenant, TenantState::Active) == TenantState::Idle) { - issues.AddIssue("Tenant is idle, no processing is allowed"); - } - - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("PingTaskRequest: {" << request.DebugString() << "} validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -660,10 +658,8 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskReq std::shared_ptr finalStatus = std::make_shared(); auto pingTaskParams = DoesPingTaskUpdateQueriesTable(request) ? - ConstructHardPingTask(request, response, YdbConnection->TablePathPrefix, Config->AutomaticQueriesTtl, - Config->TaskLeaseTtl, Config->RetryPolicies, Counters.Counters, Config->Proto.GetMaxRequestSize(), - Config->Proto.GetDumpRawStatistics(), finalStatus, requestCounters.Common) : - ConstructSoftPingTask(request, response, YdbConnection->TablePathPrefix, Config->TaskLeaseTtl, requestCounters.Common); + ConstructHardPingTask(request, response, finalStatus, requestCounters.Common) : + ConstructSoftPingTask(request, response, requestCounters.Common); auto debugInfo = Config->Proto.GetEnableDebugMode() ? std::make_shared() : TDebugInfoPtr{}; auto result = ReadModifyWrite(pingTaskParams.Query, pingTaskParams.Params, pingTaskParams.Prepare, requestCounters, debugInfo); auto prepare = [response] { return *response; }; @@ -696,7 +692,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvPingTaskReq }); } -void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvFinalStatusReport::TPtr& ev) { +void TControlPlaneStorageBase::Handle(TEvControlPlaneStorage::TEvFinalStatusReport::TPtr& ev) { const auto& event = *ev->Get(); if (!IsTerminalStatus(event.Status)) { return; @@ -719,11 +715,10 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvFinalStatus Counters.GetFinalStatusCounters(event.CloudId, event.Scope)->IncByStatus(event.Status); - Statistics statistics{event.Statistics}; + TStatistics statistics{event.Statistics}; LOG_YQ_AUDIT_SERVICE_INFO("FinalStatus: cloud id: [" << event.CloudId << "], scope: [" << event.Scope << "], query id: [" << event.QueryId << "], job id: [" << event.JobId << "], query type: [" << FederatedQuery::QueryContent::QueryType_Name(event.QueryType) << "], " << statistics << ", " << "status: " << FederatedQuery::QueryMeta::ComputeStatus_Name(event.Status)); } - } // NFq diff --git a/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp b/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp index c2a39dbfe05b..469f8eab3230 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/task_result_write.cpp @@ -1,7 +1,19 @@ #include "utils.h" +#include + namespace NFq { +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) const { + const auto& request = ev->Get()->Request; + return ValidateWriteResultData( + request.result_id().value(), + request.result_set(), + NProtoInterop::CastFromProto(request.deadline()), + Config->ResultSetsTtl + ); +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -17,11 +29,9 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvWriteResult const Ydb::ResultSet& resultSet = request.result_set(); const int byteSize = resultSet.ByteSize(); - CPS_LOG_T("WriteResultDataRequest: " << resultId << " " << resultSetId << " " << startRowId << " " << resultSet.ByteSize() << " " << deadline); - NYql::TIssues issues = ValidateWriteResultData(resultId, resultSet, deadline, Config->ResultSetsTtl); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_D("WriteResultDataRequest, validation failed: " << resultId << " " << resultSetId << " " << startRowId << " " << resultSet.DebugString() << " " << deadline << " error: " << issues.ToString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); diff --git a/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp b/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp index 3c9281e86eda..ab7c8a3617e0 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp +++ b/ydb/core/fq/libs/control_plane_storage/internal/utils.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -143,15 +144,15 @@ std::vector GetMeteringRecords(const TString& statistics, bool billable } auto now = Now(); - result.emplace_back(TBillRecord() + result.emplace_back(NKikimr::TBillRecord() .Id(jobId + "_i") .Schema("yq.ingress.v1") .FolderId(TScope(scope).ParseFolder()) .SourceWt(now) .SourceId(sourceId) - .Usage(TBillRecord::TUsage() - .Type(TBillRecord::TUsage::EType::Delta) - .Unit(TBillRecord::TUsage::EUnit::MByte) + .Usage(NKikimr::TBillRecord::TUsage() + .Type(NKikimr::TBillRecord::TUsage::EType::Delta) + .Unit(NKikimr::TBillRecord::TUsage::EUnit::MByte) .Quantity(ingressMBytes) .Start(now) .Finish(now) @@ -311,8 +312,8 @@ void PackStatisticsToProtobuf(google::protobuf::RepeatedPtrField& statsProto) { - StatsValuesList statPairs; +TStatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPtrField& statsProto) { + TStatsValuesList statPairs; statPairs.reserve(statsProto.size()); for (const auto& stat : statsProto) { statPairs.emplace_back(stat.name(), stat.value()); @@ -320,7 +321,7 @@ StatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPt return statPairs; } -TStringBuilder& operator<<(TStringBuilder& builder, const Statistics& statistics) { +TStringBuilder& operator<<(TStringBuilder& builder, const TStatistics& statistics) { bool first = true; builder << '{'; for (const auto& [field, value] : statistics.Stats) { diff --git a/ydb/core/fq/libs/control_plane_storage/internal/utils.h b/ydb/core/fq/libs/control_plane_storage/internal/utils.h index 8a35170c094c..d52cf84ec42f 100644 --- a/ydb/core/fq/libs/control_plane_storage/internal/utils.h +++ b/ydb/core/fq/libs/control_plane_storage/internal/utils.h @@ -3,11 +3,13 @@ #include #include +#include #include #include -#include +#include +#include #include namespace NFq { @@ -41,17 +43,17 @@ void PackStatisticsToProtobuf(google::protobuf::RepeatedPtrField>; +using TStatsValuesList = std::vector>; -StatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPtrField& statsProto); +TStatsValuesList ExtractStatisticsFromProtobuf(const google::protobuf::RepeatedPtrField& statsProto); -struct Statistics { +struct TStatistics { operator bool() const noexcept { return !Stats.empty(); } - const StatsValuesList& Stats; + const TStatsValuesList& Stats; }; -TStringBuilder& operator<<(TStringBuilder& builder, const Statistics& statistics); +TStringBuilder& operator<<(TStringBuilder& builder, const TStatistics& statistics); void AddTransientIssues(::google::protobuf::RepeatedPtrField< ::Ydb::Issue::IssueMessage>* protoIssues, NYql::TIssues&& issues); diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp index 362cb5d7ec6d..a5057a16804c 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage.cpp @@ -32,6 +32,20 @@ TYdbSdkRetryPolicy::TPtr MakeCreateSchemaRetryPolicy() { } // namespace +TControlPlaneStorageBase::TControlPlaneStorageBase( + const NConfig::TControlPlaneStorageConfig& config, + const NYql::TS3GatewayConfig& s3Config, + const NConfig::TCommonConfig& common, + const NConfig::TComputeConfig& computeConfig, + const ::NMonitoring::TDynamicCounterPtr& counters, + const TString& tenantName) + : TBase(config, s3Config, common, computeConfig) + , Counters(counters, *Config) + , FailedStatusCodeCounters(MakeIntrusive("FinalFailedStatusCode", counters->GetSubgroup("component", "QueryDiagnostic"))) + , TenantName(tenantName) +{ +} + void TYdbControlPlaneStorageActor::Bootstrap() { CPS_LOG_I("Starting ydb control plane storage service. Actor id: " << SelfId()); NLwTraceMonPage::ProbeRegistry().AddProbesList(LWTRACE_GET_PROBES(YQ_CONTROL_PLANE_STORAGE_PROVIDER)); @@ -323,7 +337,7 @@ void TYdbControlPlaneStorageActor::AfterTablesCreated() { // Schedule(TDuration::Zero(), new NActors::TEvents::TEvWakeup()); } -bool TControlPlaneStorageUtils::IsSuperUser(const TString& user) +bool TControlPlaneStorageUtils::IsSuperUser(const TString& user) const { return AnyOf(Config->Proto.GetSuperUsers(), [&user](const auto& superUser) { return superUser == user; diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp index 75f5557ea197..2bc787fc4522 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_bindings.cpp @@ -8,6 +8,46 @@ namespace NFq { +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) const +{ + NYql::TIssues issues = ValidateBinding(ev); + + const auto& event = *ev->Get(); + const auto& permissions = GetCreateBindingPerimssions(event); + if (event.Request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { + issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a binding with these parameters. Please receive a permission yq.resources.managePublic")); + } + + return issues; +} + +TPermissions TControlPlaneStorageBase::GetCreateBindingPerimssions(const TEvControlPlaneStorage::TEvCreateBindingRequest& event) const +{ + TPermissions permissions = Config->Proto.GetEnablePermissions() + ? event.Permissions + : TPermissions{TPermissions::MANAGE_PUBLIC}; + if (IsSuperUser(event.User)) { + permissions.SetAll(); + } + return permissions; +} + +std::pair TControlPlaneStorageBase::GetCreateBindingProtos( + const FederatedQuery::CreateBindingRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const +{ + const TString& bindingId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::BINDING); + + FederatedQuery::Binding binding; + FederatedQuery::BindingContent& content = *binding.mutable_content(); + content = request.content(); + *binding.mutable_meta() = CreateCommonMeta(bindingId, user, startTime, InitialRevision); + + FederatedQuery::Internal::BindingInternal bindingInternal; + bindingInternal.set_cloud_id(cloudId); + + return {binding, bindingInternal}; +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -19,28 +59,21 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi requestCounters.Common->RequestBytes->Add(event.GetByteSize()); const TString user = event.User; const TString token = event.Token; - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::MANAGE_PUBLIC}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } const FederatedQuery::CreateBindingRequest& request = event.Request; - const TString bindingId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::BINDING); int byteSize = request.ByteSize(); const TString connectionId = request.content().connection_id(); const TString idempotencyKey = request.idempotency_key(); + const auto [binding, bindingInternal] = GetCreateBindingProtos(request, cloudId, user, startTime); + const auto& content = binding.content(); + const TString& bindingId = binding.meta().id(); + CPS_LOG_T(MakeLogPrefix(scope, user, bindingId) << "CreateBindingRequest: " << NKikimr::MaskTicket(token) << " " << request.DebugString()); - NYql::TIssues issues = ValidateBinding(ev); - if (request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a binding with these parameters. Please receive a permission yq.resources.managePublic")); - } - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_D(MakeLogPrefix(scope, user, bindingId) << "CreateBindingRequest, validation failed: " << NKikimr::MaskTicket(token) << " " @@ -52,14 +85,6 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi return; } - FederatedQuery::Binding binding; - FederatedQuery::BindingContent& content = *binding.mutable_content(); - content = request.content(); - *binding.mutable_meta() = CreateCommonMeta(bindingId, user, startTime, InitialRevision); - - FederatedQuery::Internal::BindingInternal bindingInternal; - bindingInternal.set_cloud_id(cloudId); - std::shared_ptr>> response = std::make_shared>>(); response->first.set_binding_id(bindingId); response->second.After.ConstructInPlace().CopyFrom(binding); @@ -111,7 +136,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi scope, connectionId, "Connection " + connectionId + " does not exist or permission denied. Please check the id connection or your access rights", - permissions, + GetCreateBindingPerimssions(event), user, content.acl().visibility(), YdbConnection->TablePathPrefix); @@ -122,7 +147,6 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateBindi user, YdbConnection->TablePathPrefix); - TVector validators; if (idempotencyKey) { validators.push_back(CreateIdempotencyKeyValidator(scope, idempotencyKey, response, YdbConnection->TablePathPrefix, requestCounters.Common->ParseProtobufError)); diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp index 426ce7594725..76ebf8ddd451 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_connections.cpp @@ -30,6 +30,46 @@ void PrepareSensitiveFields(::FederatedQuery::Connection& connection, bool extra } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) const +{ + NYql::TIssues issues = ValidateConnection(ev); + + const auto& event = *ev->Get(); + const auto& permissions = GetCreateConnectionPerimssions(event); + if (event.Request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { + issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a connection with these parameters. Please receive a permission yq.resources.managePublic")); + } + + return issues; +} + +TPermissions TControlPlaneStorageBase::GetCreateConnectionPerimssions(const TEvControlPlaneStorage::TEvCreateConnectionRequest& event) const +{ + TPermissions permissions = Config->Proto.GetEnablePermissions() + ? event.Permissions + : TPermissions{TPermissions::MANAGE_PUBLIC}; + if (IsSuperUser(event.User)) { + permissions.SetAll(); + } + return permissions; +} + +std::pair TControlPlaneStorageBase::GetCreateConnectionProtos( + const FederatedQuery::CreateConnectionRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const +{ + const TString& connectionId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::CONNECTION); + + FederatedQuery::Connection connection; + FederatedQuery::ConnectionContent& content = *connection.mutable_content(); + content = request.content(); + *connection.mutable_meta() = CreateCommonMeta(connectionId, user, startTime, InitialRevision); + + FederatedQuery::Internal::ConnectionInternal connectionInternal; + connectionInternal.set_cloud_id(cloudId); + + return {connection, connectionInternal}; +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -43,25 +83,18 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConne const TString user = event.User; const TString token = event.Token; const int byteSize = request.ByteSize(); - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::MANAGE_PUBLIC}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } const TString idempotencyKey = request.idempotency_key(); - const TString connectionId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::CONNECTION); + + const auto [connection, connectionInternal] = GetCreateConnectionProtos(request, cloudId, user, startTime); + const auto& content = connection.content(); + const TString& connectionId = connection.meta().id(); CPS_LOG_T(MakeLogPrefix(scope, user, connectionId) << "CreateConnectionRequest: " << NKikimr::MaskTicket(token) << " " << request.DebugString()); - NYql::TIssues issues = ValidateConnection(ev); - if (request.content().acl().visibility() == FederatedQuery::Acl::SCOPE && !permissions.Check(TPermissions::MANAGE_PUBLIC)) { - issues.AddIssue(MakeErrorIssue(TIssuesIds::ACCESS_DENIED, "Permission denied to create a connection with these parameters. Please receive a permission yq.resources.managePublic")); - } - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_D(MakeLogPrefix(scope, user, connectionId) << "CreateConnectionRequest, validation failed: " << NKikimr::MaskTicket(token) << " " @@ -73,14 +106,6 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConne return; } - FederatedQuery::Connection connection; - FederatedQuery::ConnectionContent& content = *connection.mutable_content(); - content = request.content(); - *connection.mutable_meta() = CreateCommonMeta(connectionId, user, startTime, InitialRevision); - - FederatedQuery::Internal::ConnectionInternal connectionInternal; - connectionInternal.set_cloud_id(cloudId); - std::shared_ptr>> response = std::make_shared>>(); response->first.set_connection_id(connectionId); response->second.After.ConstructInPlace().CopyFrom(connection); @@ -141,7 +166,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateConne auto overridBindingValidator = CreateConnectionOverrideBindingValidator( scope, content.name(), - permissions, + GetCreateConnectionPerimssions(event), user, YdbConnection->TablePathPrefix); validators.push_back(overridBindingValidator); diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h index dc212439650d..3c947e638930 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h @@ -8,6 +8,7 @@ #include "request_validators.h" #include "util.h" #include "validators.h" +#include #include #include @@ -312,10 +313,10 @@ class TControlPlaneStorageUtils { /* * Utility */ - bool IsSuperUser(const TString& user); + bool IsSuperUser(const TString& user) const; template - NYql::TIssues ValidateConnection(T& ev, bool passwordRequired = true) + NYql::TIssues ValidateConnection(T& ev, bool passwordRequired = true) const { return ::NFq::ValidateConnection(ev, Config->Proto.GetMaxRequestSize(), Config->AvailableConnections, Config->Proto.GetDisableCurrentIam(), @@ -323,19 +324,19 @@ class TControlPlaneStorageUtils { } template - NYql::TIssues ValidateBinding(T& ev) + NYql::TIssues ValidateBinding(T& ev) const { return ::NFq::ValidateBinding(ev, Config->Proto.GetMaxRequestSize(), Config->AvailableBindings, Config->GeneratorPathsLimit); } template - NYql::TIssues ValidateQuery(const T& ev) + NYql::TIssues ValidateQuery(const T& ev) const { return ::NFq::ValidateQuery(ev, Config->Proto.GetMaxRequestSize()); } template - NYql::TIssues ValidateEvent(const P& ev) + NYql::TIssues ValidateEvent(const P& ev) const { return ::NFq::ValidateEvent

(ev, Config->Proto.GetMaxRequestSize()); } @@ -370,10 +371,8 @@ class TControlPlaneStorageUtils { static constexpr int64_t InitialRevision = 1; }; -class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped, - public TDbRequester, - public TControlPlaneStorageUtils -{ +class TControlPlaneStorageBase : public TControlPlaneStorageUtils { +protected: enum ERequestTypeScope { RTS_CREATE_QUERY, RTS_LIST_QUERIES, @@ -558,6 +557,10 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped cacheVal; ScopeCounters.Get(key, &cacheVal); @@ -580,13 +583,217 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped GetCreateQueryProtos( + const FederatedQuery::CreateQueryRequest& request, const TString& user, TInstant startTime) const; + + FederatedQuery::Internal::QueryInternal GetQueryInternalProto( + const FederatedQuery::CreateQueryRequest& request, const TString& cloudId, const TString& token, + const TMaybe& quotas) const; + + void FillConnectionsAndBindings( + FederatedQuery::Internal::QueryInternal& queryInternal, FederatedQuery::QueryContent::QueryType queryType, + const TVector& allConnections, + const THashMap& visibleConnections, + const THashMap& visibleBindings) const; + + // Describe query request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) const; + + void FillDescribeQueryResult( + FederatedQuery::DescribeQueryResult& result, FederatedQuery::Internal::QueryInternal internal, + const TString& user, TPermissions permissions) const; + + // Get result data request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) const; + + TPermissions GetResultDataReadPerimssions(const TEvControlPlaneStorage::TEvGetResultDataRequest& event) const; + + // Create connection request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvCreateConnectionRequest::TPtr& ev) const; + + TPermissions GetCreateConnectionPerimssions(const TEvControlPlaneStorage::TEvCreateConnectionRequest& event) const; + + std::pair GetCreateConnectionProtos( + const FederatedQuery::CreateConnectionRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const; + + // Create binding request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvCreateBindingRequest::TPtr& ev) const; + + TPermissions GetCreateBindingPerimssions(const TEvControlPlaneStorage::TEvCreateBindingRequest& event) const; + + std::pair GetCreateBindingProtos( + const FederatedQuery::CreateBindingRequest& request, const TString& cloudId, const TString& user, TInstant startTime) const; + + // Write result data request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvWriteResultDataRequest::TPtr& ev) const; + + // Get task request + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvGetTaskRequest::TPtr& ev) const; + + void FillGetTaskResult(Fq::Private::GetTaskResult& result, const TVector& tasks) const; + + // Ping task request + + struct TFinalStatus { + FederatedQuery::QueryMeta::ComputeStatus Status = FederatedQuery::QueryMeta::COMPUTE_STATUS_UNSPECIFIED; + NYql::NDqProto::StatusIds::StatusCode StatusCode = NYql::NDqProto::StatusIds::UNSPECIFIED; + FederatedQuery::QueryContent::QueryType QueryType = FederatedQuery::QueryContent::QUERY_TYPE_UNSPECIFIED; + NYql::TIssues Issues; + NYql::TIssues TransientIssues; + TStatsValuesList FinalStatistics; + TString CloudId; + TString JobId; + }; + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvPingTaskRequest::TPtr& ev) const; + + void UpdateTaskInfo( + NActors::TActorSystem* actorSystem, Fq::Private::PingTaskRequest& request, const std::shared_ptr& finalStatus, FederatedQuery::Query& query, + FederatedQuery::Internal::QueryInternal& internal, FederatedQuery::Job& job, TString& owner, + TRetryLimiter& retryLimiter, TDuration& backoff, TInstant& expireAt) const; + + void FillQueryStatistics( + const std::shared_ptr& finalStatus, const FederatedQuery::Query& query, + const FederatedQuery::Internal::QueryInternal& internal, const TRetryLimiter& retryLimiter) const; + + void Handle(TEvControlPlaneStorage::TEvFinalStatusReport::TPtr& ev); + + // Node health check + + NYql::TIssues ValidateRequest(TEvControlPlaneStorage::TEvNodesHealthCheckRequest::TPtr& ev) const; + +protected: + // Should not be used from callbacks + template + void SendResponseIssues(TActorId sender, const NYql::TIssues& issues, ui64 cookie, const TDuration& delta, TRequestCounters requestCounters) { + std::unique_ptr event(new T{issues}); + requestCounters.Common->ResponseBytes->Add(event->GetByteSize()); + TActivationContext::Send(sender, std::move(event), 0, cookie); + requestCounters.DecInFly(); + requestCounters.IncError(); + requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); + } + + template + TFuture SendResponse(const TString& name, + NActors::TActorSystem* actorSystem, + const TAsyncStatus& status, + TActorId self, + const RequestEventPtr& ev, + const TInstant& startTime, + const TRequestCounters& requestCounters, + const std::function::Type()>& prepare, + TDebugInfoPtr debugInfo) + { + return status.Apply([=, requestCounters=requestCounters](const auto& future) mutable { + NYql::TIssues internalIssues; + NYql::TIssues issues; + Result result; + typename TPrepareResponseResultType::TResponseAuditDetails auditDetails; // void* for nonauditable events + + try { + TStatus status = future.GetValue(); + if (status.IsSuccess()) { + if constexpr (ResponseEvent::Auditable) { + auto p = prepare(); + result = std::move(p.first); + auditDetails = std::move(p.second); + } else { + result = prepare(); + } + } else { + issues.AddIssues(status.GetIssues()); + internalIssues.AddIssues(status.GetIssues()); + } + } catch (const NYql::TCodeLineException& exception) { + NYql::TIssue issue = MakeErrorIssue(exception.Code, exception.GetRawMessage()); + issues.AddIssue(issue); + NYql::TIssue internalIssue = MakeErrorIssue(exception.Code, CurrentExceptionMessage()); + internalIssues.AddIssue(internalIssue); + } catch (const std::exception& exception) { + NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, exception.what()); + issues.AddIssue(issue); + NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); + internalIssues.AddIssue(internalIssue); + } catch (...) { + NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); + issues.AddIssue(issue); + NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); + internalIssues.AddIssue(internalIssue); + } + + const auto& request = ev->Get()->Request; + size_t responseByteSize = 0; + if (issues) { + CPS_LOG_AS_W(*actorSystem, name << ": {" << TrimForLogs(request.DebugString()) << "} ERROR: " << internalIssues.ToOneLineString()); + auto event = std::make_unique(issues); + event->DebugInfo = debugInfo; + responseByteSize = event->GetByteSize(); + actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); + requestCounters.IncError(); + for (const auto& issue : issues) { + NYql::WalkThroughIssues(issue, true, [&requestCounters](const NYql::TIssue& err, ui16 level) { + Y_UNUSED(level); + requestCounters.Common->Issues->GetCounter(ToString(err.GetCode()), true)->Inc(); + }); + } + } else { + CPS_LOG_AS_T(*actorSystem, name << ": {" << TrimForLogs(result.DebugString()) << "} SUCCESS"); + std::unique_ptr event; + if constexpr (ResponseEvent::Auditable) { + event = std::make_unique(result, auditDetails); + } else { + event = std::make_unique(result); + } + event->DebugInfo = debugInfo; + responseByteSize = event->GetByteSize(); + actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); + requestCounters.IncOk(); + } + requestCounters.DecInFly(); + requestCounters.Common->ResponseBytes->Add(responseByteSize); + TDuration delta = TInstant::Now() - startTime; + requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); + return MakeFuture(!issues); + }); + } + + static ui64 GetExecutionLimitMills(FederatedQuery::QueryContent::QueryType queryType, const TMaybe& quotas); +}; + +class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped, + public TDbRequester, + public TControlPlaneStorageBase +{ + using TBase = TControlPlaneStorageBase; ::NFq::TYqSharedResources::TPtr YqSharedResources; NKikimr::TYdbCredentialsProviderFactory CredProviderFactory; - TString TenantName; // Query Quota THashMap QueryQuotas; @@ -606,12 +813,9 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped("FinalFailedStatusCode", counters->GetSubgroup("component", "QueryDiagnostic"))) + : TBase(config, s3Config, common, computeConfig, counters, tenantName) , YqSharedResources(yqSharedResources) , CredProviderFactory(credProviderFactory) - , TenantName(tenantName) { } @@ -654,7 +858,7 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped void HandleRateLimiterImpl(TEventPtr& ev); @@ -736,90 +938,6 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped - TFuture SendResponse(const TString& name, - NActors::TActorSystem* actorSystem, - const TAsyncStatus& status, - TActorId self, - const RequestEventPtr& ev, - const TInstant& startTime, - const TRequestCounters& requestCounters, - const std::function::Type()>& prepare, - TDebugInfoPtr debugInfo) - { - return status.Apply([=, requestCounters=requestCounters](const auto& future) mutable { - NYql::TIssues internalIssues; - NYql::TIssues issues; - Result result; - typename TPrepareResponseResultType::TResponseAuditDetails auditDetails; // void* for nonauditable events - - try { - TStatus status = future.GetValue(); - if (status.IsSuccess()) { - if constexpr (ResponseEvent::Auditable) { - auto p = prepare(); - result = std::move(p.first); - auditDetails = std::move(p.second); - } else { - result = prepare(); - } - } else { - issues.AddIssues(status.GetIssues()); - internalIssues.AddIssues(status.GetIssues()); - } - } catch (const NYql::TCodeLineException& exception) { - NYql::TIssue issue = MakeErrorIssue(exception.Code, exception.GetRawMessage()); - issues.AddIssue(issue); - NYql::TIssue internalIssue = MakeErrorIssue(exception.Code, CurrentExceptionMessage()); - internalIssues.AddIssue(internalIssue); - } catch (const std::exception& exception) { - NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, exception.what()); - issues.AddIssue(issue); - NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); - internalIssues.AddIssue(internalIssue); - } catch (...) { - NYql::TIssue issue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); - issues.AddIssue(issue); - NYql::TIssue internalIssue = MakeErrorIssue(TIssuesIds::INTERNAL_ERROR, CurrentExceptionMessage()); - internalIssues.AddIssue(internalIssue); - } - - const auto& request = ev->Get()->Request; - size_t responseByteSize = 0; - if (issues) { - CPS_LOG_AS_W(*actorSystem, name << ": {" << TrimForLogs(request.DebugString()) << "} ERROR: " << internalIssues.ToOneLineString()); - auto event = std::make_unique(issues); - event->DebugInfo = debugInfo; - responseByteSize = event->GetByteSize(); - actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); - requestCounters.IncError(); - for (const auto& issue : issues) { - NYql::WalkThroughIssues(issue, true, [&requestCounters](const NYql::TIssue& err, ui16 level) { - Y_UNUSED(level); - requestCounters.Common->Issues->GetCounter(ToString(err.GetCode()), true)->Inc(); - }); - } - } else { - CPS_LOG_AS_T(*actorSystem, name << ": {" << TrimForLogs(result.DebugString()) << "} SUCCESS"); - std::unique_ptr event; - if constexpr (ResponseEvent::Auditable) { - event = std::make_unique(result, auditDetails); - } else { - event = std::make_unique(result); - } - event->DebugInfo = debugInfo; - responseByteSize = event->GetByteSize(); - actorSystem->Send(new IEventHandle(ev->Sender, self, event.release(), 0, ev->Cookie)); - requestCounters.IncOk(); - } - requestCounters.DecInFly(); - requestCounters.Common->ResponseBytes->Add(responseByteSize); - TDuration delta = TInstant::Now() - startTime; - requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); - return MakeFuture(!issues); - }); - } - template TFuture SendResponseTuple(const TString& name, NActors::TActorSystem* actorSystem, @@ -890,20 +1008,6 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped - void SendResponseIssues(const TActorId sender, - const NYql::TIssues& issues, - ui64 cookie, - const TDuration& delta, - TRequestCounters requestCounters) { - std::unique_ptr event(new T{issues}); - requestCounters.Common->ResponseBytes->Add(event->GetByteSize()); - Send(sender, event.release(), 0, cookie); - requestCounters.DecInFly(); - requestCounters.IncError(); - requestCounters.Common->LatencyMs->Collect(delta.MilliSeconds()); - } - struct TPickTaskParams { TString ReadQuery; TParams ReadParams; @@ -920,9 +1024,20 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped& validators = {}, TTxSettings transactionMode = TTxSettings::SerializableRW()); - ui64 GetExecutionLimitMills( - FederatedQuery::QueryContent_QueryType queryType, - const TMaybe& quotas); + struct TPingTaskParams { + TString Query; + TParams Params; + const std::function(const std::vector&)> Prepare; + std::shared_ptr> MeteringRecords; + }; + + TPingTaskParams ConstructHardPingTask( + const Fq::Private::PingTaskRequest& request, std::shared_ptr response, + const std::shared_ptr& finalStatus, const TRequestCommonCountersPtr& commonCounters) const; + + TPingTaskParams ConstructSoftPingTask( + const Fq::Private::PingTaskRequest& request, std::shared_ptr response, + const TRequestCommonCountersPtr& commonCounters) const; }; } diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp index fa2e67f2e27c..94b2353ba082 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp @@ -54,36 +54,17 @@ FederatedQuery::IamAuth::IdentityCase GetIamAuth(const FederatedQuery::Connectio namespace NFq { -void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) const { - TInstant startTime = TInstant::Now(); const TEvControlPlaneStorage::TEvCreateQueryRequest& event = *ev->Get(); - const TString cloudId = event.CloudId; const FederatedQuery::CreateQueryRequest& request = event.Request; - const FederatedQuery::Internal::ComputeDatabaseInternal& computeDatabase = event.ComputeDatabase; - ui64 resultLimit = 0; - if (event.Quotas) { - if (auto it = event.Quotas->find(QUOTA_QUERY_RESULT_LIMIT); it != event.Quotas->end()) { - resultLimit = it->second.Limit.Value; - } - } - auto queryType = request.content().type(); - ui64 executionLimitMills = GetExecutionLimitMills(queryType, event.Quotas); - const TString scope = event.Scope; - TRequestCounters requestCounters = Counters.GetCounters(cloudId, scope, RTS_CREATE_QUERY, RTC_CREATE_QUERY); - requestCounters.IncInFly(); - requestCounters.Common->RequestBytes->Add(event.GetByteSize()); - const TString user = event.User; - const TString token = event.Token; + TPermissions permissions = Config->Proto.GetEnablePermissions() ? event.Permissions : TPermissions{TPermissions::QUERY_INVOKE | TPermissions::MANAGE_PUBLIC}; - if (IsSuperUser(user)) { + if (IsSuperUser(event.User)) { permissions.SetAll(); } - const size_t byteSize = request.ByteSizeLong(); - const TString queryId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::QUERY); - CPS_LOG_T("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); NYql::TIssues issues = ValidateQuery(ev); if (request.execute_mode() != FederatedQuery::SAVE && !permissions.Check(TPermissions::QUERY_INVOKE)) { @@ -97,10 +78,9 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "Streaming disposition \"from_last_checkpoint\" is not allowed in CreateQuery request")); } - auto tenant = ev->Get()->TenantInfo->Assign(cloudId, scope, request.content().type(), TenantName); - if (event.Quotas) { TQuotaMap::const_iterator it = event.Quotas->end(); + const auto queryType = request.content().type(); if (queryType == FederatedQuery::QueryContent::ANALYTICS) { it = event.Quotas->find(QUOTA_ANALYTICS_COUNT_LIMIT); } else if (queryType == FederatedQuery::QueryContent::STREAMING) { @@ -120,24 +100,22 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery issues.AddIssue(MakeErrorIssue(TIssuesIds::BAD_REQUEST, "VCPU rate limit can't be less than zero")); } - if (issues) { - CPS_LOG_W("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); - const TDuration delta = TInstant::Now() - startTime; - SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); - LWPROBE(CreateQueryRequest, scope, user, delta, byteSize, false); - return; - } - - const TString idempotencyKey = request.idempotency_key(); - const TString jobId = request.execute_mode() == FederatedQuery::SAVE ? "" : GetEntityIdAsString(Config->IdsPrefix, EEntityType::JOB); + return issues; +} +std::pair TControlPlaneStorageBase::GetCreateQueryProtos( + const FederatedQuery::CreateQueryRequest& request, const TString& user, TInstant startTime) const +{ FederatedQuery::Query query; FederatedQuery::QueryContent& content = *query.mutable_content() = request.content(); FederatedQuery::QueryMeta& meta = *query.mutable_meta(); + + const TString queryId = GetEntityIdAsString(Config->IdsPrefix, EEntityType::QUERY); FederatedQuery::CommonMeta& common = *meta.mutable_common() = CreateCommonMeta(queryId, user, startTime, InitialRevision); meta.set_execute_mode(request.execute_mode()); meta.set_status(request.execute_mode() == FederatedQuery::SAVE ? FederatedQuery::QueryMeta::COMPLETED : FederatedQuery::QueryMeta::STARTING); + const TString jobId = request.execute_mode() == FederatedQuery::SAVE ? "" : GetEntityIdAsString(Config->IdsPrefix, EEntityType::JOB); FederatedQuery::Job job; if (request.execute_mode() != FederatedQuery::SAVE) { meta.set_last_job_query_revision(InitialRevision); @@ -154,6 +132,112 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery *job.mutable_parameters() = content.parameters(); } + return {query, job}; +} + +FederatedQuery::Internal::QueryInternal TControlPlaneStorageBase::GetQueryInternalProto( + const FederatedQuery::CreateQueryRequest& request, const TString& cloudId, const TString& token, + const TMaybe& quotas) const +{ + FederatedQuery::Internal::QueryInternal queryInternal; + if (!Config->Proto.GetDisableCurrentIam()) { + queryInternal.set_token(token); + } + + queryInternal.set_cloud_id(cloudId); + queryInternal.set_state_load_mode(FederatedQuery::StateLoadMode::EMPTY); + queryInternal.mutable_disposition()->CopyFrom(request.disposition()); + ui64 resultLimit = 0; + if (quotas) { + if (const auto it = quotas->find(QUOTA_QUERY_RESULT_LIMIT); it != quotas->end()) { + resultLimit = it->second.Limit.Value; + } + } + queryInternal.set_result_limit(resultLimit); + *queryInternal.mutable_execution_ttl() = NProtoInterop::CastToProto(TDuration::MilliSeconds(GetExecutionLimitMills(request.content().type(), quotas))); + + return queryInternal; +} + +void TControlPlaneStorageBase::FillConnectionsAndBindings( + FederatedQuery::Internal::QueryInternal& queryInternal, FederatedQuery::QueryContent::QueryType queryType, + const TVector& allConnections, + const THashMap& visibleConnections, + const THashMap& visibleBindings) const +{ + TSet disabledConnections; + for (const auto& connection : allConnections) { + const auto connectionCase = connection.content().setting().connection_case(); + if (!Config->AvailableConnections.contains(connectionCase)) { + disabledConnections.insert(connection.meta().id()); + continue; + } + if ((queryType == FederatedQuery::QueryContent::STREAMING) && !Config->AvailableStreamingConnections.contains(connectionCase)) { + disabledConnections.insert(connection.meta().id()); + continue; + } + if (GetIamAuth(connection) == FederatedQuery::IamAuth::kCurrentIam && Config->Proto.GetDisableCurrentIam()) { + disabledConnections.insert(connection.meta().id()); + continue; + } + } + + TSet connectionIds; + for (const auto& [_, connection] : visibleConnections) { + if (disabledConnections.contains(connection.meta().id())) { + continue; + } + *queryInternal.add_connection() = connection; + connectionIds.insert(connection.meta().id()); + } + + for (const auto& [_, binding] : visibleBindings) { + if (!Config->AvailableBindings.contains(binding.content().setting().binding_case())) { + continue; + } + if (disabledConnections.contains(binding.content().connection_id())) { + continue; + } + + *queryInternal.add_binding() = binding; + if (!connectionIds.contains(binding.content().connection_id())) { + ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Unable to resolve connection for binding " << binding.meta().id() << ", name " << binding.content().name() << ", connection id " << binding.content().connection_id(); + } + } +} + +void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQueryRequest::TPtr& ev) +{ + TInstant startTime = TInstant::Now(); + const TEvControlPlaneStorage::TEvCreateQueryRequest& event = *ev->Get(); + const TString cloudId = event.CloudId; + const FederatedQuery::CreateQueryRequest& request = event.Request; + const FederatedQuery::Internal::ComputeDatabaseInternal& computeDatabase = event.ComputeDatabase; + const auto queryType = request.content().type(); + const TString scope = event.Scope; + TRequestCounters requestCounters = Counters.GetCounters(cloudId, scope, RTS_CREATE_QUERY, RTC_CREATE_QUERY); + requestCounters.IncInFly(); + requestCounters.Common->RequestBytes->Add(event.GetByteSize()); + const TString user = event.User; + const TString token = event.Token; + const size_t byteSize = request.ByteSizeLong(); + CPS_LOG_T("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); + + auto tenant = ev->Get()->TenantInfo->Assign(cloudId, scope, queryType, TenantName); + + if (const auto& issues = ValidateRequest(ev)) { + CPS_LOG_W("CreateQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); + const TDuration delta = TInstant::Now() - startTime; + SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); + LWPROBE(CreateQueryRequest, scope, user, delta, byteSize, false); + return; + } + + const TString idempotencyKey = request.idempotency_key(); + const auto [query, job] = GetCreateQueryProtos(request, user, startTime); + const TString queryId = query.meta().common().id(); + const TString jobId = job.meta().id(); + std::shared_ptr>> response = std::make_shared>>(); response->first.set_query_id(queryId); response->second.CloudId = cloudId; @@ -179,7 +263,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery ); } - auto prepareParams = [=, as=TActivationContext::ActorSystem(), commonCounters=requestCounters.Common](const TVector& resultSets) mutable { + auto prepareParams = [=, as=TActivationContext::ActorSystem(), commonCounters=requestCounters.Common, quotas=event.Quotas](const TVector& resultSets) mutable { const size_t countSets = (idempotencyKey ? 1 : 0) + (request.execute_mode() != FederatedQuery::SAVE ? 2 : 0); if (resultSets.size() != countSets) { ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Result set size is not equal to " << countSets << " but equal " << resultSets.size() << ". Please contact internal support"; @@ -197,63 +281,18 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery } } - FederatedQuery::Internal::QueryInternal queryInternal; - if (!Config->Proto.GetDisableCurrentIam()) { - queryInternal.set_token(token); - } - - queryInternal.set_cloud_id(cloudId); - queryInternal.set_state_load_mode(FederatedQuery::StateLoadMode::EMPTY); - queryInternal.mutable_disposition()->CopyFrom(request.disposition()); - queryInternal.set_result_limit(resultLimit); - *queryInternal.mutable_execution_ttl() = NProtoInterop::CastToProto(TDuration::MilliSeconds(executionLimitMills)); + auto queryInternal = GetQueryInternalProto(request, cloudId, token, quotas); if (request.execute_mode() != FederatedQuery::SAVE) { // TODO: move to run actor priority selection *queryInternal.mutable_compute_connection() = computeDatabase.connection(); - TSet disabledConnections; - for (const auto& connection: GetEntities(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters)) { - auto connectionCase = connection.content().setting().connection_case(); - if (!Config->AvailableConnections.contains(connectionCase)) { - disabledConnections.insert(connection.meta().id()); - continue; - } - if ((queryType == FederatedQuery::QueryContent::STREAMING) && !Config->AvailableStreamingConnections.contains(connectionCase)) { - disabledConnections.insert(connection.meta().id()); - continue; - } - - if (GetIamAuth(connection) == FederatedQuery::IamAuth::kCurrentIam && Config->Proto.GetDisableCurrentIam()) { - disabledConnections.insert(connection.meta().id()); - continue; - } - } - - TSet connectionIds; - auto connections = GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters); - for (const auto& [_, connection]: connections) { - if (disabledConnections.contains(connection.meta().id())) { - continue; - } - *queryInternal.add_connection() = connection; - connectionIds.insert(connection.meta().id()); - } - - auto bindings = GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 1], BINDING_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters); - for (const auto& [_, binding]: bindings) { - if (!Config->AvailableBindings.contains(binding.content().setting().binding_case())) { - continue; - } - - if (disabledConnections.contains(binding.content().connection_id())) { - continue; - } - - *queryInternal.add_binding() = binding; - if (!connectionIds.contains(binding.content().connection_id())) { - ythrow NYql::TCodeLineException(TIssuesIds::BAD_REQUEST) << "Unable to resolve connection for binding " << binding.meta().id() << ", name " << binding.content().name() << ", connection id " << binding.content().connection_id(); - } - } + FillConnectionsAndBindings( + queryInternal, + queryType, + GetEntities(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters), + GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters), + GetEntitiesWithVisibilityPriority(resultSets[resultSets.size() - 1], BINDING_COLUMN_NAME, Config->Proto.GetIgnorePrivateSources(), commonCounters) + ); } if (query.ByteSizeLong() > Config->Proto.GetMaxRequestSize()) { @@ -510,6 +549,73 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvListQueries }); } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) const +{ + return ValidateEvent(ev); +} + +void TControlPlaneStorageBase::FillDescribeQueryResult( + FederatedQuery::DescribeQueryResult& result, FederatedQuery::Internal::QueryInternal internal, + const TString& user, TPermissions permissions) const +{ + const auto lastJobId = result.query().meta().last_job_id(); + result.mutable_query()->mutable_meta()->set_last_job_id(lastJobId + "-" + result.query().meta().common().id()); + + permissions = Config->Proto.GetEnablePermissions() + ? permissions + : TPermissions{TPermissions::VIEW_PUBLIC | TPermissions::VIEW_AST | TPermissions::VIEW_QUERY_TEXT}; + if (IsSuperUser(user)) { + permissions.SetAll(); + } + + const auto queryVisibility = result.query().content().acl().visibility(); + const auto queryUser = result.query().meta().common().created_by(); + const bool hasViewAccess = HasViewAccess(permissions, queryVisibility, queryUser, user); + if (!hasViewAccess) { + ythrow NYql::TCodeLineException(TIssuesIds::ACCESS_DENIED) << "Query does not exist or permission denied. Please check the id of the query or your access rights"; + } + + // decompress plan + if (internal.plan_compressed().data()) { // todo: remove this if after migration + TCompressor compressor(internal.plan_compressed().method()); + result.mutable_query()->mutable_plan()->set_json(compressor.Decompress(internal.plan_compressed().data())); + if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { + if (result.query().plan().json().size() > 1000) { + // modifing plan this way should definitely reduce query msg size + result.mutable_query()->mutable_plan()->set_json(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping plan of size " << result.query().plan().json().size() << " bytes"); + } + } + } + + auto timeline = internal.timeline(); + if (timeline) { + result.mutable_query()->mutable_timeline()->set_svg(timeline); + } + + if (!permissions.Check(TPermissions::VIEW_AST)) { + result.mutable_query()->clear_ast(); + } else { + // decompress AST + if (internal.ast_compressed().data()) { // todo: remove this if after migration + TCompressor compressor(internal.ast_compressed().method()); + result.mutable_query()->mutable_ast()->set_data(compressor.Decompress(internal.ast_compressed().data())); + } + if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { + if (result.query().ast().data().size() > 1000) { + // modifing AST this way should definitely reduce query msg size + result.mutable_query()->mutable_ast()->set_data(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping AST of size " << result.query().ast().data().size() << " bytes"); + } + } + } + if (!permissions.Check(TPermissions::VIEW_QUERY_TEXT)) { + result.mutable_query()->mutable_content()->clear_text(); + } + + if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { + ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Resulting query of size " << result.query().ByteSizeLong() << " bytes is too big"; + } +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQueryRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -521,20 +627,13 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQue requestCounters.Common->RequestBytes->Add(event.GetByteSize()); const TString user = event.User; const TString token = event.Token; - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::VIEW_PUBLIC | TPermissions::VIEW_AST | TPermissions::VIEW_QUERY_TEXT}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } const FederatedQuery::DescribeQueryRequest& request = event.Request; const TString queryId = request.query_id(); const int byteSize = request.ByteSize(); CPS_LOG_T("DescribeQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); - NYql::TIssues issues = ValidateEvent(ev); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("DescribeQueryRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -553,7 +652,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQue const auto query = queryBuilder.Build(); auto debugInfo = Config->Proto.GetEnableDebugMode() ? std::make_shared() : TDebugInfoPtr{}; auto [result, resultSets] = Read(query.Sql, query.Params, requestCounters, debugInfo); - auto prepare = [resultSets=resultSets, user, permissions, commonCounters=requestCounters.Common] { + auto prepare = [this, resultSets=resultSets, user, permissions=event.Permissions, commonCounters=requestCounters.Common] { if (resultSets->size() != 1) { ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Result set size is not equal to 1 but equal " << resultSets->size() << ". Please contact internal support"; } @@ -569,61 +668,13 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeQue ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Error parsing proto message for query. Please contact internal support"; } - const auto lastJobId = result.query().meta().last_job_id(); - result.mutable_query()->mutable_meta()->set_last_job_id(lastJobId + "-" + result.query().meta().common().id()); - - const auto queryVisibility = result.query().content().acl().visibility(); - const auto queryUser = result.query().meta().common().created_by(); - const bool hasViewAccess = HasViewAccess(permissions, queryVisibility, queryUser, user); - if (!hasViewAccess) { - ythrow NYql::TCodeLineException(TIssuesIds::ACCESS_DENIED) << "Query does not exist or permission denied. Please check the id of the query or your access rights"; - } - FederatedQuery::Internal::QueryInternal internal; if (!internal.ParseFromString(*parser.ColumnParser(INTERNAL_COLUMN_NAME).GetOptionalString())) { commonCounters->ParseProtobufError->Inc(); ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Error parsing proto message for query internal. Please contact internal support"; } - // decompress plan - if (internal.plan_compressed().data()) { // todo: remove this if after migration - TCompressor compressor(internal.plan_compressed().method()); - result.mutable_query()->mutable_plan()->set_json(compressor.Decompress(internal.plan_compressed().data())); - if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { - if (result.query().plan().json().size() > 1000) { - // modifing plan this way should definitely reduce query msg size - result.mutable_query()->mutable_plan()->set_json(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping plan of size " << result.query().plan().json().size() << " bytes"); - } - } - } - - auto timeline = internal.timeline(); - if (timeline) { - result.mutable_query()->mutable_timeline()->set_svg(timeline); - } - - if (!permissions.Check(TPermissions::VIEW_AST)) { - result.mutable_query()->clear_ast(); - } else { - // decompress AST - if (internal.ast_compressed().data()) { // todo: remove this if after migration - TCompressor compressor(internal.ast_compressed().method()); - result.mutable_query()->mutable_ast()->set_data(compressor.Decompress(internal.ast_compressed().data())); - } - if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { - if (result.query().ast().data().size() > 1000) { - // modifing AST this way should definitely reduce query msg size - result.mutable_query()->mutable_ast()->set_data(TStringBuilder() << "Message is too big: " << result.query().ByteSizeLong() << " bytes, dropping AST of size " << result.query().ast().data().size() << " bytes"); - } - } - } - if (!permissions.Check(TPermissions::VIEW_QUERY_TEXT)) { - result.mutable_query()->mutable_content()->clear_text(); - } - - if (result.query().ByteSizeLong() > GRPC_MESSAGE_SIZE_LIMIT) { - ythrow NYql::TCodeLineException(TIssuesIds::INTERNAL_ERROR) << "Resulting query of size " << result.query().ByteSizeLong() << " bytes is too big"; - } + FillDescribeQueryResult(result, internal, user, permissions); return result; }; @@ -1484,6 +1535,21 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvControlQuer }); } +NYql::TIssues TControlPlaneStorageBase::ValidateRequest(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) const +{ + return ValidateEvent(ev); +} + +TPermissions TControlPlaneStorageBase::GetResultDataReadPerimssions(const TEvControlPlaneStorage::TEvGetResultDataRequest& event) const { + TPermissions permissions = Config->Proto.GetEnablePermissions() + ? event.Permissions + : TPermissions{TPermissions::VIEW_PUBLIC}; + if (IsSuperUser(event.User)) { + permissions.SetAll(); + } + return permissions; +} + void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetResultDataRequest::TPtr& ev) { TInstant startTime = TInstant::Now(); @@ -1501,17 +1567,11 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvGetResultDa const TString queryId = request.query_id(); const int byteSize = event.Request.ByteSize(); const TString token = event.Token; - TPermissions permissions = Config->Proto.GetEnablePermissions() - ? event.Permissions - : TPermissions{TPermissions::VIEW_PUBLIC}; - if (IsSuperUser(user)) { - permissions.SetAll(); - } + const TPermissions permissions = GetResultDataReadPerimssions(event); const int64_t limit = request.limit(); CPS_LOG_T("GetResultDataRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token)); - NYql::TIssues issues = ValidateEvent(ev); - if (issues) { + if (const auto& issues = ValidateRequest(ev)) { CPS_LOG_W("GetResultDataRequest: {" << request.DebugString() << "} " << MakeUserInfo(user, token) << "validation FAILED: " << issues.ToOneLineString()); const TDuration delta = TInstant::Now() - startTime; SendResponseIssues(ev->Sender, issues, ev->Cookie, delta, requestCounters); @@ -1851,10 +1911,7 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvDescribeJob }); } -ui64 TYdbControlPlaneStorageActor::GetExecutionLimitMills( - FederatedQuery::QueryContent_QueryType queryType, - const TMaybe& quotas) { - +ui64 TControlPlaneStorageBase::GetExecutionLimitMills(FederatedQuery::QueryContent::QueryType queryType, const TMaybe& quotas) { if (!quotas) { return 0; } diff --git a/ydb/core/fq/libs/init/init.cpp b/ydb/core/fq/libs/init/init.cpp index 85bd1296a8d7..cd9afe4dcbc0 100644 --- a/ydb/core/fq/libs/init/init.cpp +++ b/ydb/core/fq/libs/init/init.cpp @@ -88,14 +88,21 @@ void Init( const auto clientCounters = yqCounters->GetSubgroup("subsystem", "ClientMetrics"); if (protoConfig.GetControlPlaneStorage().GetEnabled()) { + const auto counters = yqCounters->GetSubgroup("subsystem", "ControlPlaneStorage"); auto controlPlaneStorage = protoConfig.GetControlPlaneStorage().GetUseInMemory() - ? NFq::CreateInMemoryControlPlaneStorageServiceActor(protoConfig.GetControlPlaneStorage()) + ? NFq::CreateInMemoryControlPlaneStorageServiceActor( + protoConfig.GetControlPlaneStorage(), + protoConfig.GetGateways().GetS3(), + protoConfig.GetCommon(), + protoConfig.GetCompute(), + counters, + tenant) : NFq::CreateYdbControlPlaneStorageServiceActor( protoConfig.GetControlPlaneStorage(), protoConfig.GetGateways().GetS3(), protoConfig.GetCommon(), protoConfig.GetCompute(), - yqCounters->GetSubgroup("subsystem", "ControlPlaneStorage"), + counters, yqSharedResources, NKikimr::CreateYdbCredentialsProviderFactory, tenant); diff --git a/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp b/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp index c7266237bd7f..de0bdb869de4 100644 --- a/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp +++ b/ydb/tests/fq/control_plane_storage/in_memory_control_plane_storage_ut.cpp @@ -1,3 +1,5 @@ +#include + #include #include @@ -12,74 +14,111 @@ #include #include +#include + namespace NFq { using namespace NActors; using namespace NKikimr; +using namespace NFqRun; +using namespace NKikimrRun; namespace { -//////////////////////////////////////////////////////////////////////////////// +class TTestFuxture : public NUnitTest::TBaseFixture { +protected: + using TBase = NUnitTest::TBaseFixture; -using TRuntimePtr = std::unique_ptr; + static constexpr TDuration WAIT_TIMEOUT = TDuration::Seconds(15); -struct TTestBootstrap { - NConfig::TControlPlaneStorageConfig Config; - TRuntimePtr Runtime; +public: + void SetUp(NUnitTest::TTestContext& ctx) override { + TBase::SetUp(ctx); - const TDuration RequestTimeout = TDuration::Seconds(10); + SetupSignalActions(); + EnableYDBBacktraceFormat(); - TTestBootstrap() - { - Runtime = PrepareTestActorRuntime(); - } + TFqSetupSettings settings; + settings.VerboseLevel = TFqSetupSettings::EVerbose::Max; - std::pair CreateQuery() - { - TActorId sender = Runtime->AllocateEdgeActor(); - FederatedQuery::CreateQueryRequest proto; + settings.LogConfig.SetDefaultLevel(NLog::EPriority::PRI_WARN); + ModifyLogPriorities({{NKikimrServices::EServiceKikimr::YQ_CONTROL_PLANE_STORAGE, NLog::EPriority::PRI_DEBUG}}, settings.LogConfig); - FederatedQuery::QueryContent& content = *proto.mutable_content(); - content.set_name("my_query_1"); - content.set_text("SELECT 1;"); - content.set_type(FederatedQuery::QueryContent::ANALYTICS); - content.mutable_acl()->set_visibility(FederatedQuery::Acl::SCOPE); + auto& cpStorageConfig = *settings.FqConfig.MutableControlPlaneStorage(); + cpStorageConfig.SetEnabled(true); + cpStorageConfig.SetUseInMemory(true); - auto request = std::make_unique("yandexcloud://my_cloud_1/my_folder_1", proto, "user@staff", "", "mock_cloud", TPermissions{}, TQuotaMap{}, std::make_shared(), FederatedQuery::Internal::ComputeDatabaseInternal{}); - Runtime->Send(new IEventHandle(ControlPlaneStorageServiceActorId(), sender, request.release())); + auto& privateApiConfig = *settings.FqConfig.MutablePrivateApi(); + privateApiConfig.SetEnabled(true); + privateApiConfig.SetLoopback(true); + + settings.FqConfig.SetEnableDynamicNameservice(true); + settings.FqConfig.MutableControlPlaneProxy()->SetEnabled(true); + settings.FqConfig.MutableDbPool()->SetEnabled(true); + settings.FqConfig.MutableNodesManager()->SetEnabled(true); + settings.FqConfig.MutablePendingFetcher()->SetEnabled(true); + settings.FqConfig.MutablePrivateProxy()->SetEnabled(true); + settings.FqConfig.MutableGateways()->SetEnabled(true); + settings.FqConfig.MutableResourceManager()->SetEnabled(true); + settings.FqConfig.MutableCommon()->SetIdsPrefix("ut"); + + FqSetup = std::make_unique(settings); + } - TAutoPtr handle; - TEvControlPlaneStorage::TEvCreateQueryResponse* event = Runtime->GrabEdgeEvent(handle); - return {event->Result, event->Issues}; +protected: + static void CheckSuccess(const TRequestResult& result) { + UNIT_ASSERT_VALUES_EQUAL_C(result.Status, Ydb::StatusIds::SUCCESS, result.Issues.ToOneLineString()); } -private: - TRuntimePtr PrepareTestActorRuntime() { - TRuntimePtr runtime(new TTestBasicRuntime(1)); - runtime->SetLogPriority(NKikimrServices::YQ_CONTROL_PLANE_STORAGE, NLog::PRI_DEBUG); + TExecutionMeta WaitQueryExecution(const TString& queryId, TDuration timeout = WAIT_TIMEOUT) const { + using EStatus = FederatedQuery::QueryMeta; - auto controlPlaneProxy = CreateInMemoryControlPlaneStorageServiceActor(Config); + const TInstant start = TInstant::Now(); + while (TInstant::Now() - start <= timeout) { + TExecutionMeta meta; + CheckSuccess(FqSetup->DescribeQuery(queryId, meta)); - runtime->AddLocalService( - ControlPlaneStorageServiceActorId(), - TActorSetupCmd(controlPlaneProxy, TMailboxType::Simple, 0)); + if (!IsIn({EStatus::FAILED, EStatus::COMPLETED, EStatus::ABORTED_BY_USER, EStatus::ABORTED_BY_SYSTEM}, meta.Status)) { + Cerr << "Wait query execution " << TInstant::Now() - start << ": " << EStatus::ComputeStatus_Name(meta.Status) << "\n"; + Sleep(TDuration::Seconds(1)); + continue; + } - SetupTabletServices(*runtime); + UNIT_ASSERT_C(meta.Status == EStatus::COMPLETED, "issues: " << meta.Issues.ToOneLineString() << ", transient issues: " << meta.TransientIssues.ToOneLineString()); + return meta; + } - return runtime; + UNIT_ASSERT_C(false, "Waiting query execution timeout. Spent time " << TInstant::Now() - start << " exceeds limit " << timeout); + return {}; } + +protected: + std::unique_ptr FqSetup; }; -} +} // anonymous namespace + +Y_UNIT_TEST_SUITE(InMemoryControlPlaneStorage) { + Y_UNIT_TEST_F(ExecuteSimpleQuery, TTestFuxture) { + TString queryId; + CheckSuccess(FqSetup->StreamRequest({ + .Query = "SELECT 42 AS result_value" + }, queryId)); + UNIT_ASSERT(queryId); -Y_UNIT_TEST_SUITE(CreateQueryRequest) { - Y_UNIT_TEST(ShouldCreateSimpleQuery) - { - TTestBootstrap bootstrap; - const auto [result, issues] = bootstrap.CreateQuery(); - UNIT_ASSERT(!issues); + const auto meta = WaitQueryExecution(queryId); + UNIT_ASSERT_VALUES_EQUAL(meta.ResultSetSizes.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(meta.ResultSetSizes[0], 1); + + Ydb::ResultSet resultSet; + CheckSuccess(FqSetup->FetchQueryResults(queryId, 0, resultSet)); + + NYdb::TResultSetParser parser(resultSet); + UNIT_ASSERT_VALUES_EQUAL(parser.ColumnsCount(), 1); + UNIT_ASSERT_VALUES_EQUAL(parser.RowsCount(), 1); + UNIT_ASSERT(parser.TryNextRow()); + UNIT_ASSERT_VALUES_EQUAL(parser.ColumnParser("result_value").GetInt32(), 42); } } - } // NFq diff --git a/ydb/tests/fq/control_plane_storage/ya.make b/ydb/tests/fq/control_plane_storage/ya.make index 889ba53f9d36..b8ee3698c58d 100644 --- a/ydb/tests/fq/control_plane_storage/ya.make +++ b/ydb/tests/fq/control_plane_storage/ya.make @@ -14,6 +14,7 @@ ENDIF() PEERDIR( library/cpp/retry library/cpp/testing/unittest + ydb/core/base ydb/core/external_sources ydb/core/fq/libs/actors/logging ydb/core/fq/libs/init @@ -22,6 +23,8 @@ PEERDIR( ydb/core/fq/libs/rate_limiter/events ydb/core/testlib/default ydb/library/security + ydb/tests/tools/fqrun/src + ydb/tests/tools/kqprun/runlib ) INCLUDE(${ARCADIA_ROOT}/ydb/public/tools/ydb_recipe/recipe.inc) diff --git a/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h b/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h index 7620c35a1578..a80b7cef0084 100644 --- a/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h +++ b/ydb/tests/fq/control_plane_storage/ydb_test_bootstrap.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -31,6 +32,9 @@ #include #include #include + +#include + #include #include @@ -41,6 +45,7 @@ namespace NFq { using namespace NActors; using namespace NKikimr; +using namespace NKikimrRun; ////////////////////////////////////////////////////// @@ -96,6 +101,9 @@ struct TTestBootstrap { TTestBootstrap(std::string tablePrefix, const NConfig::TControlPlaneStorageConfig& config = {}) : Config(config) { + SetupSignalActions(); + EnableYDBBacktraceFormat(); + Cerr << "Netstat: " << Exec("netstat --all --program") << Endl; Cerr << "Process stat: " << Exec("ps aux") << Endl; Cerr << "YDB receipt endpoint: " << GetEnv("YDB_ENDPOINT") << ", database: " << GetEnv("YDB_DATABASE") << Endl; From 6200b659fb44c6d6b32cf02652c3c8d0f21e6343 Mon Sep 17 00:00:00 2001 From: Gigoriy Pisarenko Date: Thu, 27 Feb 2025 12:50:27 +0000 Subject: [PATCH 14/18] Fixed include file --- .../libs/control_plane_storage/ydb_control_plane_storage_impl.h | 2 +- ydb/library/yql/providers/pq/provider/yql_pq_gateway.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h index 3c947e638930..434c9405446b 100644 --- a/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h +++ b/ydb/core/fq/libs/control_plane_storage/ydb_control_plane_storage_impl.h @@ -1027,7 +1027,7 @@ class TYdbControlPlaneStorageActor : public NActors::TActorBootstrapped(const std::vector&)> Prepare; + const std::function(const TVector&)> Prepare; std::shared_ptr> MeteringRecords; }; diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h b/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h index 6cfd71d929ab..cbe70bb3b3c4 100644 --- a/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h +++ b/ydb/library/yql/providers/pq/provider/yql_pq_gateway.h @@ -6,7 +6,7 @@ #include #include -#include +#include #include #include From 2a8487506868fd5b75221604ab7962245bfb2f96 Mon Sep 17 00:00:00 2001 From: Gigoriy Pisarenko Date: Thu, 27 Feb 2025 15:08:20 +0000 Subject: [PATCH 15/18] Fixed build 2 --- ydb/tests/tools/kqprun/runlib/utils.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb/tests/tools/kqprun/runlib/utils.cpp b/ydb/tests/tools/kqprun/runlib/utils.cpp index f8c37d7b67fe..abd1865dde83 100644 --- a/ydb/tests/tools/kqprun/runlib/utils.cpp +++ b/ydb/tests/tools/kqprun/runlib/utils.cpp @@ -7,7 +7,7 @@ #include #include -#include +#include #include #include #include From 5415636e1a80672777c87251c0ce399ccdaecb8f Mon Sep 17 00:00:00 2001 From: Gigoriy Pisarenko Date: Thu, 27 Feb 2025 16:07:41 +0000 Subject: [PATCH 16/18] Fixed build 3 --- .../yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp index 9a587aca63c1..e85243df7e58 100644 --- a/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp +++ b/ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp @@ -135,7 +135,7 @@ constexpr static auto FILE_POLL_PERIOD = TDuration::MilliSeconds(5); if (!msgs.empty()) { EventsQ_.Push(NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent(msgs, {}, Session_), size); } else if (CancelOnFileFinish_) { - EventsQ_.Push(NYdb::NTopic::TSessionClosedEvent(NYdb::EStatus::CANCELLED, {NYdb::NIssue::TIssue("PQ file topic was finished")}), size); + EventsQ_.Push(NYdb::NTopic::TSessionClosedEvent(NYdb::EStatus::CANCELLED, {NYql::TIssue("PQ file topic was finished")}), size); } Sleep(FILE_POLL_PERIOD); From 55709883211854d25da851af8bcf3d1291160da1 Mon Sep 17 00:00:00 2001 From: Gigoriy Pisarenko Date: Fri, 28 Feb 2025 13:21:35 +0000 Subject: [PATCH 17/18] Fixed build --- ydb/core/testlib/test_client.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp index fc49d26427ea..e67daa9777c6 100644 --- a/ydb/core/testlib/test_client.cpp +++ b/ydb/core/testlib/test_client.cpp @@ -58,6 +58,7 @@ #include #include #include +#include #include #include #include From b86bbc51c95cd86774c7d2d4f510bf7648cc7807 Mon Sep 17 00:00:00 2001 From: Gigoriy Pisarenko Date: Tue, 4 Mar 2025 07:22:41 +0000 Subject: [PATCH 18/18] Fixed build 2 --- ydb/tests/tools/fqrun/src/fq_setup.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ydb/tests/tools/fqrun/src/fq_setup.cpp b/ydb/tests/tools/fqrun/src/fq_setup.cpp index b3f38ee69942..e870d3cff211 100644 --- a/ydb/tests/tools/fqrun/src/fq_setup.cpp +++ b/ydb/tests/tools/fqrun/src/fq_setup.cpp @@ -11,6 +11,8 @@ #include #include +#include + using namespace NKikimrRun; namespace NFqRun {