Skip to content

Commit a2a9a5a

Browse files
shnikdshnikd
committed
Fixes
1 parent 560c541 commit a2a9a5a

File tree

6 files changed

+120
-56
lines changed

6 files changed

+120
-56
lines changed

‎ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp

Lines changed: 84 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -663,10 +663,11 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
663663
)", TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
664664

665665
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
666-
UNIT_ASSERT_C(!result.GetDiagnostics().empty(), "Query result diagnostics is empty");
666+
auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
667+
UNIT_ASSERT_C(!stats.query_diagnostics().empty(), "Query result diagnostics is empty");
667668

668669
TStringStream in;
669-
in << result.GetDiagnostics();
670+
in << stats.query_diagnostics();
670671
NJson::TJsonValue value;
671672
ReadJsonTree(&in, &value);
672673

@@ -681,7 +682,7 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
681682
UNIT_ASSERT_C(value.Has("query_syntax"), "Incorrect Diagnostics");
682683
UNIT_ASSERT_C(value.Has("query_database"), "Incorrect Diagnostics");
683684
UNIT_ASSERT_C(value.Has("query_cluster"), "Incorrect Diagnostics");
684-
UNIT_ASSERT_C(value.Has("query_plan"), "Incorrect Diagnostics");
685+
UNIT_ASSERT_C(!value.Has("query_plan"), "Incorrect Diagnostics");
685686
UNIT_ASSERT_C(value.Has("query_type"), "Incorrect Diagnostics");
686687
}
687688

@@ -695,7 +696,86 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
695696

696697
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString().c_str());
697698

698-
UNIT_ASSERT_C(result.GetDiagnostics().empty(), "Query result diagnostics should be empty, but it's not");
699+
auto& stats = NYdb::TProtoAccessor::GetProto(*result.GetStats());
700+
UNIT_ASSERT_C(stats.query_diagnostics().empty(), "Query result diagnostics should be empty, but it's not");
701+
}
702+
}
703+
704+
Y_UNIT_TEST(StreamExecuteCollectFullDiagnostics) {
705+
auto kikimr = DefaultKikimrRunner();
706+
auto db = kikimr.GetQueryClient();
707+
708+
{
709+
TExecuteQuerySettings settings;
710+
settings.StatsMode(EStatsMode::Full);
711+
712+
auto it = db.StreamExecuteQuery(R"(
713+
SELECT Key, Value2 FROM TwoShard WHERE Value2 > 0;
714+
)", TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
715+
UNIT_ASSERT_VALUES_EQUAL_C(it.GetStatus(), EStatus::SUCCESS, it.GetIssues().ToString());
716+
717+
TString statsString;
718+
for (;;) {
719+
auto streamPart = it.ReadNext().GetValueSync();
720+
if (!streamPart.IsSuccess()) {
721+
UNIT_ASSERT_C(streamPart.EOS(), streamPart.GetIssues().ToString());
722+
break;
723+
}
724+
725+
const auto& execStats = streamPart.GetStats();
726+
if (execStats.Defined()) {
727+
auto& stats = NYdb::TProtoAccessor::GetProto(*execStats);
728+
statsString = stats.query_diagnostics();
729+
}
730+
}
731+
732+
UNIT_ASSERT_C(!statsString.empty(), "Query result diagnostics is empty");
733+
734+
TStringStream in;
735+
in << statsString;
736+
NJson::TJsonValue value;
737+
ReadJsonTree(&in, &value);
738+
739+
UNIT_ASSERT_C(value.IsMap(), "Incorrect Diagnostics");
740+
UNIT_ASSERT_C(value.Has("query_id"), "Incorrect Diagnostics");
741+
UNIT_ASSERT_C(value.Has("version"), "Incorrect Diagnostics");
742+
UNIT_ASSERT_C(value.Has("query_text"), "Incorrect Diagnostics");
743+
UNIT_ASSERT_C(value.Has("query_parameter_types"), "Incorrect Diagnostics");
744+
UNIT_ASSERT_C(value.Has("table_metadata"), "Incorrect Diagnostics");
745+
UNIT_ASSERT_C(value["table_metadata"].IsArray(), "Incorrect Diagnostics: table_metadata type should be an array");
746+
UNIT_ASSERT_C(value.Has("created_at"), "Incorrect Diagnostics");
747+
UNIT_ASSERT_C(value.Has("query_syntax"), "Incorrect Diagnostics");
748+
UNIT_ASSERT_C(value.Has("query_database"), "Incorrect Diagnostics");
749+
UNIT_ASSERT_C(value.Has("query_cluster"), "Incorrect Diagnostics");
750+
UNIT_ASSERT_C(!value.Has("query_plan"), "Incorrect Diagnostics");
751+
UNIT_ASSERT_C(value.Has("query_type"), "Incorrect Diagnostics");
752+
}
753+
754+
{
755+
TExecuteQuerySettings settings;
756+
settings.StatsMode(EStatsMode::Basic);
757+
758+
auto it = db.StreamExecuteQuery(R"(
759+
SELECT Key, Value2 FROM TwoShard WHERE Value2 > 0;
760+
)", TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
761+
UNIT_ASSERT_VALUES_EQUAL_C(it.GetStatus(), EStatus::SUCCESS, it.GetIssues().ToString());
762+
763+
TString statsString;
764+
for (;;) {
765+
auto streamPart = it.ReadNext().GetValueSync();
766+
if (!streamPart.IsSuccess()) {
767+
UNIT_ASSERT_C(streamPart.EOS(), streamPart.GetIssues().ToString());
768+
break;
769+
}
770+
771+
const auto& execStats = streamPart.GetStats();
772+
if (execStats.Defined()) {
773+
auto& stats = NYdb::TProtoAccessor::GetProto(*execStats);
774+
statsString = stats.query_diagnostics();
775+
}
776+
}
777+
778+
UNIT_ASSERT_C(statsString.empty(), "Query result diagnostics should be empty, but it's not");
699779
}
700780
}
701781

‎ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp

Lines changed: 2 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -365,8 +365,6 @@ void TCommandExecuteQuery::Config(TConfig& config) {
365365
config.Opts->AddLongOption('q', "query", "Text of query to execute").RequiredArgument("[String]").StoreResult(&Query);
366366
config.Opts->AddLongOption('f', "file", "Path to file with query text to execute")
367367
.RequiredArgument("PATH").StoreResult(&QueryFile);
368-
config.Opts->AddLongOption("collect-diagnostics", "Collects diagnostics and saves it to file")
369-
.StoreTrue(&CollectFullDiagnostics);
370368

371369
AddOutputFormats(config, {
372370
EDataFormat::Pretty,
@@ -434,9 +432,6 @@ int TCommandExecuteQuery::ExecuteDataQuery(TConfig& config) {
434432
NTable::TExecDataQuerySettings settings;
435433
settings.KeepInQueryCache(true);
436434
settings.CollectQueryStats(ParseQueryStatsModeOrThrow(CollectStatsMode, defaultStatsMode));
437-
if (CollectFullDiagnostics) {
438-
settings.CollectFullDiagnostics(true);
439-
}
440435

441436
NTable::TTxSettings txSettings;
442437
if (TxMode) {
@@ -521,11 +516,6 @@ void TCommandExecuteQuery::PrintDataQueryResponse(NTable::TDataQueryResult& resu
521516
{
522517
Cout << Endl << "Flame graph is available for full or profile stats only" << Endl;
523518
}
524-
if (CollectFullDiagnostics)
525-
{
526-
TFileOutput file(TStringBuilder() << "diagnostics_" << TGUID::Create().AsGuidString() << ".txt");
527-
file << result.GetDiagnostics();
528-
}
529519
}
530520

531521
int TCommandExecuteQuery::ExecuteSchemeQuery(TConfig& config) {
@@ -558,7 +548,7 @@ namespace {
558548
NQuery::TExecuteQuerySettings>;
559549

560550
template <typename TClient>
561-
auto GetSettings(const TString& collectStatsMode, const bool basicStats, std::optional<TDuration> timeout, bool collectFullDiagnostics) {
551+
auto GetSettings(const TString& collectStatsMode, const bool basicStats, std::optional<TDuration> timeout) {
562552
if constexpr (std::is_same_v<TClient, NTable::TTableClient>) {
563553
const auto defaultStatsMode = basicStats
564554
? NTable::ECollectQueryStatsMode::Basic
@@ -568,9 +558,6 @@ namespace {
568558
if (timeout.has_value()) {
569559
settings.ClientTimeout(*timeout);
570560
}
571-
if (collectFullDiagnostics) {
572-
settings.CollectFullDiagnostics(true);
573-
}
574561
return settings;
575562
} else if constexpr (std::is_same_v<TClient, NQuery::TQueryClient>) {
576563
const auto defaultStatsMode = basicStats
@@ -581,9 +568,6 @@ namespace {
581568
if (timeout.has_value()) {
582569
settings.ClientTimeout(*timeout);
583570
}
584-
if (collectFullDiagnostics) {
585-
settings.CollectFullDiagnostics(true);
586-
}
587571
return settings;
588572
}
589573
Y_UNREACHABLE();
@@ -690,7 +674,7 @@ int TCommandExecuteQuery::ExecuteQueryImpl(TConfig& config) {
690674
if (OperationTimeout) {
691675
optTimeout = TDuration::MilliSeconds(FromString<ui64>(OperationTimeout));
692676
}
693-
const auto settings = GetSettings<TClient>(CollectStatsMode, BasicStats, optTimeout, CollectFullDiagnostics);
677+
const auto settings = GetSettings<TClient>(CollectStatsMode, BasicStats, optTimeout);
694678

695679
TAsyncPartIterator<TClient> asyncResult;
696680
SetInterruptHandlers();
@@ -770,8 +754,6 @@ bool TCommandExecuteQuery::PrintQueryResponse(TIterator& result) {
770754
fullStats = queryStats.GetPlan();
771755
}
772756
}
773-
774-
diagnostics = streamPart.GetDiagnostics();
775757
}
776758
} // TResultSetPrinter destructor should be called before printing stats
777759

@@ -786,12 +768,6 @@ bool TCommandExecuteQuery::PrintQueryResponse(TIterator& result) {
786768
queryPlanPrinter.Print(TString{*fullStats});
787769
}
788770

789-
if (CollectFullDiagnostics)
790-
{
791-
TFileOutput file(TStringBuilder() << "diagnostics_" << TGUID::Create().AsGuidString() << ".txt");
792-
file << diagnostics;
793-
}
794-
795771
PrintFlameGraph(fullStats);
796772

797773
if (IsInterrupted()) {
@@ -897,10 +873,6 @@ int TCommandExplain::Run(TConfig& config) {
897873
settings.Explain(true);
898874
}
899875

900-
if (CollectFullDiagnostics) {
901-
settings.CollectFullDiagnostics(true);
902-
}
903-
904876
auto result = client.StreamExecuteScanQuery(Query, settings).GetValueSync();
905877
NStatusHelpers::ThrowOnErrorOrPrintIssues(result);
906878

@@ -917,13 +889,6 @@ int TCommandExplain::Run(TConfig& config) {
917889
planJson = proto.query_plan();
918890
ast = proto.query_ast();
919891
}
920-
if (tablePart.HasDiagnostics()) {
921-
diagnostics = tablePart.ExtractDiagnostics();
922-
}
923-
}
924-
925-
if (CollectFullDiagnostics) {
926-
SaveDiagnosticsToFile(diagnostics);
927892
}
928893

929894
if (IsInterrupted()) {

‎ydb/public/lib/ydb_cli/commands/ydb_service_table.h

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,6 @@ class TCommandExecuteQuery : public TTableCommand, TCommandQueryBase, TCommandWi
123123
TString TxMode;
124124
TString QueryType;
125125
bool BasicStats = false;
126-
bool CollectFullDiagnostics = false;
127126
};
128127

129128
class TCommandExplain : public TTableCommand, public TCommandWithOutput, TCommandQueryBase, TInterruptibleCommand {

‎ydb/public/lib/ydb_cli/commands/ydb_sql.cpp

Lines changed: 32 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
#include "ydb_sql.h"
22

33
#include <library/cpp/json/json_reader.h>
4+
#include <library/cpp/json/json_writer.h>
5+
#include <library/cpp/json/json_prettifier.h>
46
#include <ydb/public/lib/json_value/ydb_json_value.h>
57
#include <ydb-cpp-sdk/library/operation_id/operation_id.h>
68
#include <ydb/public/lib/ydb_cli/common/interactive.h>
@@ -9,7 +11,6 @@
911
#include <ydb/public/lib/ydb_cli/common/query_stats.h>
1012
#include <ydb/public/lib/ydb_cli/common/waiting_bar.h>
1113
#include <ydb-cpp-sdk/client/proto/accessor.h>
12-
#include <util/generic/guid.h>
1314
#include <util/generic/queue.h>
1415
#include <google/protobuf/text_format.h>
1516

@@ -41,11 +42,11 @@ void TCommandSql::Config(TConfig& config) {
4142
.StoreTrue(&ExplainAnalyzeMode);
4243
config.Opts->AddLongOption("stats", "Execution statistics collection mode [none, basic, full, profile]")
4344
.RequiredArgument("[String]").StoreResult(&CollectStatsMode);
45+
config.Opts->AddLongOption("diagnostics-file", "Path to file where the diagnostics will be saved.")
46+
.RequiredArgument("[String]").StoreResult(&DiagnosticsFile);
4447
config.Opts->AddLongOption("syntax", "Query syntax [yql, pg]")
4548
.RequiredArgument("[String]").DefaultValue("yql").StoreResult(&Syntax)
4649
.Hidden();
47-
config.Opts->AddLongOption("collect-diagnostics", "Collects diagnostics and saves it to file")
48-
.StoreTrue(&CollectFullDiagnostics);
4950

5051
AddOutputFormats(config, {
5152
EDataFormat::Pretty,
@@ -149,10 +150,6 @@ int TCommandSql::RunCommand(TConfig& config) {
149150
throw TMisuseException() << "Unknow syntax option \"" << Syntax << "\"";
150151
}
151152

152-
if (CollectFullDiagnostics) {
153-
settings.CollectFullDiagnostics(true);
154-
}
155-
156153
if (!Parameters.empty() || InputParamStream) {
157154
// Execute query with parameters
158155
THolder<TParamsBuilder> paramBuilder;
@@ -190,7 +187,7 @@ int TCommandSql::PrintResponse(NQuery::TExecuteQueryIterator& result) {
190187
std::optional<std::string> stats;
191188
std::optional<std::string> plan;
192189
std::optional<std::string> ast;
193-
TString diagnostics;
190+
TMaybe<TString> meta;
194191
{
195192
TResultSetPrinter printer(OutputFormat, &IsInterrupted);
196193

@@ -212,9 +209,8 @@ int TCommandSql::PrintResponse(NQuery::TExecuteQueryIterator& result) {
212209
if (queryStats.GetPlan()) {
213210
plan = queryStats.GetPlan();
214211
}
212+
diagnostics = queryStats.GetDiagnostics();
215213
}
216-
217-
diagnostics = streamPart.GetDiagnostics();
218214
}
219215
} // TResultSetPrinter destructor should be called before printing stats
220216

@@ -232,6 +228,10 @@ int TCommandSql::PrintResponse(NQuery::TExecuteQueryIterator& result) {
232228
Cout << Endl << "Statistics:" << Endl << *stats;
233229
}
234230

231+
if (diagnostics) {
232+
Cout << Endl << "Diagnostics:" << Endl << NJson::PrettifyJson(*diagnostics, true) << Endl;;
233+
}
234+
235235
if (plan) {
236236
if (!ExplainMode && !ExplainAnalyzeMode
237237
&& (OutputFormat == EDataFormat::Default || OutputFormat == EDataFormat::Pretty)) {
@@ -245,9 +245,28 @@ int TCommandSql::PrintResponse(NQuery::TExecuteQueryIterator& result) {
245245
queryPlanPrinter.Print(TString{*plan});
246246
}
247247

248-
if (CollectFullDiagnostics) {
249-
TFileOutput file(TStringBuilder() << "diagnostics_" << TGUID::Create().AsGuidString() << ".txt");
250-
file << diagnostics;
248+
if (!DiagnosticsFile.empty()) {
249+
TFileOutput file(DiagnosticsFile);
250+
251+
NJson::TJsonValue diagnosticsJson(NJson::JSON_MAP);
252+
253+
if (stats) {
254+
diagnosticsJson.InsertValue("stats", *stats);
255+
}
256+
if (ast) {
257+
diagnosticsJson.InsertValue("ast", *ast);
258+
}
259+
if (plan) {
260+
NJson::TJsonValue planJson;
261+
NJson::ReadJsonTree(*plan, &planJson, true);
262+
diagnosticsJson.InsertValue("plan", planJson);
263+
}
264+
if (diagnostics) {
265+
NJson::TJsonValue metaJson;
266+
NJson::ReadJsonTree(*diagnostics, &metaJson, true);
267+
diagnosticsJson.InsertValue("meta", metaJson);
268+
}
269+
file << NJson::PrettifyJson(NJson::WriteJson(diagnosticsJson, true), false);
251270
}
252271

253272
if (IsInterrupted()) {

‎ydb/public/lib/ydb_cli/commands/ydb_sql.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,13 @@ class TCommandSql : public TYdbCommand, public TCommandWithOutput, public TComma
2828
int PrintResponse(NQuery::TExecuteQueryIterator& result);
2929

3030
TString CollectStatsMode;
31+
TString DiagnosticsFile;
3132
TString Query;
3233
TString QueryFile;
3334
TString Syntax;
3435
bool ExplainMode = false;
3536
bool ExplainAnalyzeMode = false;
3637
bool ExplainAst = false;
37-
bool CollectFullDiagnostics = false;
3838
};
3939

4040
}

‎ydb/public/sdk/cpp/client/ydb_query/stats.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ TString TExecStats::ToString(bool withPlan) const {
2929
if (!withPlan) {
3030
proto.clear_query_plan();
3131
proto.clear_query_ast();
32+
proto.clear_query_diagnostics();
3233
}
3334

3435
TString res;

0 commit comments

Comments
 (0)