diff --git a/ydb/tests/tools/kqprun/configuration/app_config.conf b/ydb/tests/tools/kqprun/configuration/app_config.conf index ea9d478a13f4..69fcd118a20c 100644 --- a/ydb/tests/tools/kqprun/configuration/app_config.conf +++ b/ydb/tests/tools/kqprun/configuration/app_config.conf @@ -24,6 +24,7 @@ LogConfig { QueryServiceConfig { MdbTransformHost: false + ProgressStatsPeriodMs: 1000 QueryArtifactsCompressionMethod: "zstd_6" ScriptResultRowsLimit: 0 ScriptResultSizeLimit: 10485760 @@ -108,6 +109,8 @@ ResourceBrokerConfig { TableServiceConfig { BindingsMode: BM_DROP CompileTimeoutMs: 600000 + EnableCreateTableAs: true + EnablePerStatementQueryExecution: true SessionsLimitPerNode: 1000 QueryLimits { diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp index ce435e542abc..f3fb818deba4 100644 --- a/ydb/tests/tools/kqprun/kqprun.cpp +++ b/ydb/tests/tools/kqprun/kqprun.cpp @@ -25,7 +25,7 @@ struct TExecutionOptions { YqlScript }; - TString ScriptQuery; + std::vector ScriptQueries; TString SchemeQuery; bool ForgetExecution = false; @@ -35,7 +35,7 @@ struct TExecutionOptions { TString TraceId = "kqprun"; bool HasResults() const { - return ScriptQuery && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE; + return !ScriptQueries.empty() && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE; } }; @@ -53,11 +53,11 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner } } - if (executionOptions.ScriptQuery) { - Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing script..." << colors.Default() << Endl; + for (size_t id = 0; id < executionOptions.ScriptQueries.size(); ++id) { + Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing script" << (executionOptions.ScriptQueries.size() > 1 ? TStringBuilder() << " " << id : TString()) << "..." << colors.Default() << Endl; switch (executionOptions.ClearExecution) { case TExecutionOptions::EClearExecutionCase::Disabled: - if (!runner.ExecuteScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { + if (!runner.ExecuteScript(executionOptions.ScriptQueries[id], executionOptions.ScriptQueryAction, executionOptions.TraceId)) { ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Script execution failed"; } Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Fetching script results..." << colors.Default() << Endl; @@ -73,13 +73,13 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner break; case TExecutionOptions::EClearExecutionCase::GenericQuery: - if (!runner.ExecuteQuery(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { + if (!runner.ExecuteQuery(executionOptions.ScriptQueries[id], executionOptions.ScriptQueryAction, executionOptions.TraceId)) { ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Query execution failed"; } break; case TExecutionOptions::EClearExecutionCase::YqlScript: - if (!runner.ExecuteYqlScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.TraceId)) { + if (!runner.ExecuteYqlScript(executionOptions.ScriptQueries[id], executionOptions.ScriptQueryAction, executionOptions.TraceId)) { ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Yql script execution failed"; } break; @@ -106,6 +106,16 @@ THolder SetupDefaultFileOutput(const TString& filePath, IOutputStre } +template +EnumType GetCaseVariant(const TString& optionName, const TString& caseName, const std::map& casesMap) { + auto it = casesMap.find(caseName); + if (it == casesMap.end()) { + ythrow yexception() << "Option '" << optionName << "' has no case '" << caseName << "'"; + } + return it->second; +} + + void ReplaceTemplate(const TString& variableName, const TString& variableValue, TString& query) { TString variableTemplate = TStringBuilder() << "${" << variableName << "}"; for (size_t position = query.find(variableTemplate); position != TString::npos; position = query.find(variableTemplate, position)) { @@ -143,7 +153,7 @@ void RunMain(int argc, const char* argv[]) { TExecutionOptions executionOptions; NKqpRun::TRunnerOptions runnerOptions; - TString scriptQueryFile; + std::vector scriptQueryFiles; TString schemeQueryFile; TString resultOutputFile = "-"; TString schemeQueryAstFile; @@ -169,7 +179,7 @@ void RunMain(int argc, const char* argv[]) { options.AddLongOption('p', "script-query", "Script query to execute") .Optional() .RequiredArgument("FILE") - .StoreResult(&scriptQueryFile); + .AppendTo(&scriptQueryFiles); options.AddLongOption('s', "scheme-query", "Scheme query to execute") .Optional() .RequiredArgument("FILE") @@ -273,27 +283,30 @@ void RunMain(int argc, const char* argv[]) { // Execution options - if (!schemeQueryFile && !scriptQueryFile) { + if (!schemeQueryFile && scriptQueryFiles.empty()) { ythrow yexception() << "Nothing to execute"; } + if (schemeQueryFile) { executionOptions.SchemeQuery = TFileInput(schemeQueryFile).ReadAll(); ReplaceTemplate(NKqpRun::YQL_TOKEN_VARIABLE, yqlToken, executionOptions.SchemeQuery); } - if (scriptQueryFile) { - executionOptions.ScriptQuery = TFileInput(scriptQueryFile).ReadAll(); + + executionOptions.ScriptQueries.reserve(scriptQueryFiles.size()); + for (const TString& scriptQueryFile : scriptQueryFiles) { + executionOptions.ScriptQueries.emplace_back(TFileInput(scriptQueryFile).ReadAll()); } - executionOptions.ClearExecution = - (clearExecutionType == TStringBuf("query")) ? TExecutionOptions::EClearExecutionCase::GenericQuery - : (clearExecutionType == TStringBuf("yql-script")) ? TExecutionOptions::EClearExecutionCase::YqlScript - : (clearExecutionType == TStringBuf("disabled")) ? TExecutionOptions::EClearExecutionCase::Disabled - : TExecutionOptions::EClearExecutionCase::Disabled; + executionOptions.ClearExecution = GetCaseVariant("clear-execution", clearExecutionType, { + {"query", TExecutionOptions::EClearExecutionCase::GenericQuery}, + {"yql-script", TExecutionOptions::EClearExecutionCase::YqlScript}, + {"disabled", TExecutionOptions::EClearExecutionCase::Disabled} + }); - executionOptions.ScriptQueryAction = - (scriptQueryAction == TStringBuf("execute")) ? NKikimrKqp::QUERY_ACTION_EXECUTE - : (scriptQueryAction == TStringBuf("explain")) ? NKikimrKqp::QUERY_ACTION_EXPLAIN - : NKikimrKqp::QUERY_ACTION_EXECUTE; + executionOptions.ScriptQueryAction = GetCaseVariant("script-action", scriptQueryAction, { + {"execute", NKikimrKqp::QUERY_ACTION_EXECUTE}, + {"explain", NKikimrKqp::QUERY_ACTION_EXPLAIN} + }); // Runner options @@ -302,24 +315,24 @@ void RunMain(int argc, const char* argv[]) { THolder scriptQueryAstFileHolder = SetupDefaultFileOutput(scriptQueryAstFile, runnerOptions.ScriptQueryAstOutput); THolder scriptQueryPlanFileHolder = SetupDefaultFileOutput(scriptQueryPlanFile, runnerOptions.ScriptQueryPlanOutput); - runnerOptions.TraceOptType = - (traceOptType == TStringBuf("all")) ? NKqpRun::TRunnerOptions::ETraceOptType::All - : (traceOptType == TStringBuf("scheme")) ? NKqpRun::TRunnerOptions::ETraceOptType::Scheme - : (traceOptType == TStringBuf("script")) ? NKqpRun::TRunnerOptions::ETraceOptType::Script - : (traceOptType == TStringBuf("disabled")) ? NKqpRun::TRunnerOptions::ETraceOptType::Disabled - : NKqpRun::TRunnerOptions::ETraceOptType::All; + runnerOptions.TraceOptType = GetCaseVariant("trace-opt", traceOptType, { + {"all", NKqpRun::TRunnerOptions::ETraceOptType::All}, + {"scheme", NKqpRun::TRunnerOptions::ETraceOptType::Scheme}, + {"script", NKqpRun::TRunnerOptions::ETraceOptType::Script}, + {"disabled", NKqpRun::TRunnerOptions::ETraceOptType::Disabled} + }); runnerOptions.YdbSettings.TraceOptEnabled = runnerOptions.TraceOptType != NKqpRun::TRunnerOptions::ETraceOptType::Disabled; - runnerOptions.ResultOutputFormat = - (resultOutputFormat == TStringBuf("rows")) ? NKqpRun::TRunnerOptions::EResultOutputFormat::RowsJson - : (resultOutputFormat == TStringBuf("full")) ? NKqpRun::TRunnerOptions::EResultOutputFormat::FullJson - : NKqpRun::TRunnerOptions::EResultOutputFormat::RowsJson; + runnerOptions.ResultOutputFormat = GetCaseVariant("result-format", resultOutputFormat, { + {"rows", NKqpRun::TRunnerOptions::EResultOutputFormat::RowsJson}, + {"full", NKqpRun::TRunnerOptions::EResultOutputFormat::FullJson} + }); - runnerOptions.PlanOutputFormat = - (planOutputFormat == TStringBuf("pretty")) ? NYdb::NConsoleClient::EOutputFormat::Pretty - : (planOutputFormat == TStringBuf("table")) ? NYdb::NConsoleClient::EOutputFormat::PrettyTable - : (planOutputFormat == TStringBuf("json")) ? NYdb::NConsoleClient::EOutputFormat::JsonUnicode - : NYdb::NConsoleClient::EOutputFormat::Default; + runnerOptions.PlanOutputFormat = GetCaseVariant("plan-format", planOutputFormat, { + {"pretty", NYdb::NConsoleClient::EOutputFormat::Pretty}, + {"table", NYdb::NConsoleClient::EOutputFormat::PrettyTable}, + {"json", NYdb::NConsoleClient::EOutputFormat::JsonUnicode}, + }); // Ydb settings diff --git a/ydb/tests/tools/kqprun/src/actors.cpp b/ydb/tests/tools/kqprun/src/actors.cpp index 48f73e4cabfc..c1597255b8b5 100644 --- a/ydb/tests/tools/kqprun/src/actors.cpp +++ b/ydb/tests/tools/kqprun/src/actors.cpp @@ -11,12 +11,13 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped request, NThreading::TPromise promise, - ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector& resultSets) + ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector& resultSets, TString& queryPlan) : Request_(std::move(request)) , Promise_(promise) , ResultRowsLimit_(std::numeric_limits::max()) , ResultSizeLimit_(std::numeric_limits::max()) , ResultSets_(resultSets) + , QueryPlan_(queryPlan) { if (resultRowsLimit) { ResultRowsLimit_ = resultRowsLimit; @@ -36,6 +37,7 @@ class TRunScriptActorMock : public NActors::TActorBootstrappedGet()->Record.GetQueryPlan(); + } + private: THolder Request_; NThreading::TPromise Promise_; ui64 ResultRowsLimit_; ui64 ResultSizeLimit_; std::vector& ResultSets_; + TString& QueryPlan_; }; } // anonymous namespace NActors::IActor* CreateRunScriptActorMock(THolder request, NThreading::TPromise promise, - ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector& resultSets) { - return new TRunScriptActorMock(std::move(request), promise, resultRowsLimit, resultSizeLimit, resultSets); + ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector& resultSets, TString& queryPlan) { + return new TRunScriptActorMock(std::move(request), promise, resultRowsLimit, resultSizeLimit, resultSets, queryPlan); } } // namespace NKqpRun diff --git a/ydb/tests/tools/kqprun/src/actors.h b/ydb/tests/tools/kqprun/src/actors.h index 9e7a251d14ff..6ccbf32eee34 100644 --- a/ydb/tests/tools/kqprun/src/actors.h +++ b/ydb/tests/tools/kqprun/src/actors.h @@ -6,6 +6,6 @@ namespace NKqpRun { NActors::IActor* CreateRunScriptActorMock(THolder request, NThreading::TPromise promise, - ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector& resultSets); + ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector& resultSets, TString& queryPlan); } // namespace NKqpRun diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp index 21c31f74ade2..9997acec0a05 100644 --- a/ydb/tests/tools/kqprun/src/kqp_runner.cpp +++ b/ydb/tests/tools/kqprun/src/kqp_runner.cpp @@ -2,6 +2,7 @@ #include "ydb_setup.h" #include +#include #include #include @@ -74,6 +75,8 @@ class TKqpRunner::TImpl { PrintScriptAst(meta.Ast); + PrintScriptPlan(meta.Plan); + if (!status.IsSuccess()) { Cerr << CerrColors_.Red() << "Failed to execute query, reason:" << CerrColors_.Default() << Endl << status.ToString() << Endl; return false; @@ -83,8 +86,6 @@ class TKqpRunner::TImpl { Cerr << CerrColors_.Red() << "Request finished with issues:" << CerrColors_.Default() << Endl << status.Issues.ToString() << Endl; } - PrintScriptPlan(meta.Plan); - return true; } @@ -130,6 +131,7 @@ class TKqpRunner::TImpl { private: bool WaitScriptExecutionOperation() { + ExecutionMeta_ = TExecutionMeta(); TRequestResult status; while (true) { status = YdbSetup_.GetScriptExecutionOperationRequest(ExecutionOperation_, ExecutionMeta_); @@ -148,6 +150,8 @@ class TKqpRunner::TImpl { PrintScriptAst(ExecutionMeta_.Ast); + PrintScriptPlan(ExecutionMeta_.Plan); + if (!status.IsSuccess() || ExecutionMeta_.ExecutionStatus != NYdb::NQuery::EExecStatus::Completed) { Cerr << CerrColors_.Red() << "Failed to execute script, invalid final status, reason:" << CerrColors_.Default() << Endl << status.ToString() << Endl; return false; @@ -157,8 +161,6 @@ class TKqpRunner::TImpl { Cerr << CerrColors_.Red() << "Request finished with issues:" << CerrColors_.Default() << Endl << status.Issues.ToString() << Endl; } - PrintScriptPlan(ExecutionMeta_.Plan); - return true; } @@ -189,12 +191,20 @@ class TKqpRunner::TImpl { } void PrintScriptPlan(const TString& plan) const { - if (Options_.ScriptQueryPlanOutput) { - Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; + if (!Options_.ScriptQueryPlanOutput || !plan) { + return; + } - NYdb::NConsoleClient::TQueryPlanPrinter printer(Options_.PlanOutputFormat, true, *Options_.ScriptQueryPlanOutput); - printer.Print(plan); + NJson::TJsonValue planJson; + NJson::ReadJsonTree(plan, &planJson, true); + if (!planJson.GetMapSafe().contains("meta")) { + return; } + + Cout << CoutColors_.Cyan() << "Writing script query plan" << CoutColors_.Default() << Endl; + + NYdb::NConsoleClient::TQueryPlanPrinter printer(Options_.PlanOutputFormat, true, *Options_.ScriptQueryPlanOutput); + printer.Print(plan); } void PrintScriptResult(const Ydb::ResultSet& resultSet) const { diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp index 74e1df18d7c3..e6ccb4ec0735 100644 --- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp +++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp @@ -174,7 +174,7 @@ class TYdbSetup::TImpl { NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr SchemeQueryRequest(const TString& query, const TString& traceId) const { auto event = MakeHolder(); - FillSchemeRequest(query, traceId, event->Record); + FillQueryRequest(query, NKikimrKqp::QUERY_TYPE_SQL_DDL, NKikimrKqp::QUERY_ACTION_EXECUTE, traceId, event->Record); return RunKqpProxyRequest(std::move(event)); } @@ -186,21 +186,25 @@ class TYdbSetup::TImpl { return RunKqpProxyRequest(std::move(event)); } - NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr QueryRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, std::vector& resultSets) const { + NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr QueryRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, std::vector& resultSets, TString& queryPlan) const { auto event = MakeHolder(); - FillScriptRequest(query, action, traceId, event->Record); + FillQueryRequest(query, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, action, traceId, event->Record); + + if (auto progressStatsPeriodMs = Settings_.AppConfig.GetQueryServiceConfig().GetProgressStatsPeriodMs()) { + event->SetProgressStatsPeriod(TDuration::MilliSeconds(Settings_.AppConfig.GetQueryServiceConfig().GetProgressStatsPeriodMs())); + } auto promise = NThreading::NewPromise(); auto rowsLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultRowsLimit(); auto sizeLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultSizeLimit(); - GetRuntime()->Register(CreateRunScriptActorMock(std::move(event), promise, rowsLimit, sizeLimit, resultSets)); + GetRuntime()->Register(CreateRunScriptActorMock(std::move(event), promise, rowsLimit, sizeLimit, resultSets, queryPlan)); return promise.GetFuture().GetValueSync(); } NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr YqlScriptRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId) const { auto event = MakeHolder(); - FillYqlScriptRequest(query, action, traceId, event->Record); + FillQueryRequest(query, NKikimrKqp::QUERY_TYPE_SQL_SCRIPT, action, traceId, event->Record); return RunKqpProxyRequest(std::move(event)); } @@ -271,10 +275,6 @@ class TYdbSetup::TImpl { request->SetDatabase(Settings_.DomainName); } - void FillSchemeRequest(const TString& query, const TString& traceId, NKikimrKqp::TEvQueryRequest& event) const { - FillQueryRequest(query, NKikimrKqp::QUERY_TYPE_SQL_DDL, NKikimrKqp::QUERY_ACTION_EXECUTE, traceId, event); - } - void FillScriptRequest(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId, NKikimrKqp::TEvQueryRequest& event) const { FillQueryRequest(script, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_SCRIPT, action, traceId, event); @@ -285,10 +285,6 @@ class TYdbSetup::TImpl { } } - void FillYqlScriptRequest(const TString& script, NKikimrKqp::EQueryAction action, const TString& traceId, NKikimrKqp::TEvQueryRequest& event) const { - FillQueryRequest(script, NKikimrKqp::QUERY_TYPE_SQL_SCRIPT, action, traceId, event); - } - private: TYdbSetupSettings Settings_; @@ -350,11 +346,13 @@ TRequestResult TYdbSetup::ScriptRequest(const TString& script, NKikimrKqp::EQuer TRequestResult TYdbSetup::QueryRequest(const TString& query, NKikimrKqp::EQueryAction action, const TString& traceId, TQueryMeta& meta, std::vector& resultSets) const { resultSets.clear(); - auto queryOperationResponse = Impl_->QueryRequest(query, action, traceId, resultSets)->Get()->Record.GetRef(); + auto queryOperationResponse = Impl_->QueryRequest(query, action, traceId, resultSets, meta.Plan)->Get()->Record.GetRef(); const auto& responseRecord = queryOperationResponse.GetResponse(); meta.Ast = responseRecord.GetQueryAst(); - meta.Plan = responseRecord.GetQueryPlan(); + if (const auto& plan = responseRecord.GetQueryPlan()) { + meta.Plan = plan; + } return TRequestResult(queryOperationResponse.GetYdbStatus(), responseRecord.GetQueryIssues()); } @@ -390,7 +388,9 @@ TRequestResult TYdbSetup::GetScriptExecutionOperationRequest(const TString& oper meta.ExecutionStatus = static_cast(deserializedMeta.exec_status()); meta.ResultSetsCount = deserializedMeta.result_sets_meta_size(); meta.Ast = deserializedMeta.exec_stats().query_ast(); - meta.Plan = deserializedMeta.exec_stats().query_plan(); + if (deserializedMeta.exec_stats().query_plan() != "{}") { + meta.Plan = deserializedMeta.exec_stats().query_plan(); + } } return TRequestResult(scriptExecutionOperation->Get()->Status, scriptExecutionOperation->Get()->Issues);