Skip to content

Commit f3ee3c3

Browse files
shnikdshnikd
authored
Merge 6d7e7a5 into 95ce0df 
2 parents 95ce0df + 6d7e7a5 commit f3ee3c3

File tree

19 files changed

+237
-17
lines changed

19 files changed

+237
-17
lines changed

‎ydb/core/grpc_services/query/rpc_execute_query.cpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
260260
.SetUseCancelAfter(false)
261261
.SetSyntax(syntax)
262262
.SetSupportStreamTrailingResult(true)
263-
.SetOutputChunkMaxSize(req->response_part_limit_bytes());
263+
.SetOutputChunkMaxSize(req->response_part_limit_bytes())
264+
.SetCollectFullDiagnostics(req->Getcollect_full_diagnostics());
264265

265266
auto ev = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>(
266267
QueryAction,
@@ -394,6 +395,7 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
394395
hasTrailingMessage = true;
395396
response.mutable_tx_meta()->set_id(kqpResponse.GetTxMeta().id());
396397
}
398+
response.set_query_full_diagnostics(kqpResponse.GetQueryDiagnostics());
397399
}
398400

399401
if (hasTrailingMessage) {

‎ydb/core/grpc_services/rpc_execute_data_query.cpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,8 @@ class TExecuteDataQueryRPC : public TRpcKqpRequestActor<TExecuteDataQueryRPC, TE
145145
&req->parameters(),
146146
req->collect_stats(),
147147
req->has_query_cache_policy() ? &req->query_cache_policy() : nullptr,
148-
req->has_operation_params() ? &req->operation_params() : nullptr);
148+
req->has_operation_params() ? &req->operation_params() : nullptr,
149+
NKqp::NPrivateEvents::TQueryRequestSettings().SetCollectFullDiagnostics(req->Getcollect_full_diagnostics()));
149150

150151
ReportCostInfo_ = req->operation_params().report_cost_info() == Ydb::FeatureFlag::ENABLED;
151152

@@ -203,6 +204,7 @@ class TExecuteDataQueryRPC : public TRpcKqpRequestActor<TExecuteDataQueryRPC, TE
203204
queryMeta.mutable_parameters_types()->insert({queryParameter.GetName(), parameterType});
204205
}
205206
}
207+
queryResult->set_query_full_diagnostics(kqpResponse.GetQueryDiagnostics());
206208
} catch (const std::exception& ex) {
207209
NYql::TIssues issues;
208210
issues.AddIssue(NYql::ExceptionToIssue(ex));

‎ydb/core/kqp/common/events/query.h‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,17 @@ struct TQueryRequestSettings {
4545
return *this;
4646
}
4747

48+
TQueryRequestSettings& SetCollectFullDiagnostics(bool flag) {
49+
CollectFullDiagnostics = flag;
50+
return *this;
51+
}
52+
4853
ui64 OutputChunkMaxSize = 0;
4954
bool KeepSession = false;
5055
bool UseCancelAfter = true;
5156
::Ydb::Query::Syntax Syntax = Ydb::Query::Syntax::SYNTAX_UNSPECIFIED;
5257
bool SupportsStreamTrailingResult = false;
58+
bool CollectFullDiagnostics = false;
5359
};
5460

5561
struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents::EvQueryRequest> {
@@ -282,7 +288,7 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
282288
}
283289

284290
bool GetCollectDiagnostics() const {
285-
return Record.GetRequest().GetCollectDiagnostics();
291+
return QuerySettings.CollectFullDiagnostics;
286292
}
287293

288294
ui32 CalculateSerializedSize() const override {

‎ydb/core/kqp/common/kqp_event_impl.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,8 @@ void TEvKqp::TEvQueryRequest::PrepareRemote() const {
107107
Record.MutableRequest()->SetIsInternalCall(RequestCtx->IsInternalCall());
108108
Record.MutableRequest()->SetOutputChunkMaxSize(QuerySettings.OutputChunkMaxSize);
109109

110+
Record.MutableRequest()->SetCollectDiagnostics(QuerySettings.CollectFullDiagnostics);
111+
110112
RequestCtx.reset();
111113
}
112114
}

‎ydb/core/kqp/ut/query/kqp_query_ut.cpp‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,72 @@ Y_UNIT_TEST_SUITE(KqpQuery) {
179179
UNIT_ASSERT_VALUES_EQUAL(counters.RecompileRequestGet()->Val(), 1);
180180
}
181181

182+
Y_UNIT_TEST(ExecuteDataQueryCollectFullDiagnostics) {
183+
auto setting = NKikimrKqp::TKqpSetting();
184+
auto serverSettings = TKikimrSettings()
185+
.SetKqpSettings({setting});
186+
187+
TKikimrRunner kikimr(serverSettings);
188+
auto db = kikimr.GetTableClient();
189+
auto session = db.CreateSession().GetValueSync().GetSession();
190+
191+
{
192+
UNIT_ASSERT(session.ExecuteSchemeQuery(R"(
193+
CREATE TABLE `/Root/TestTable` (
194+
Key Uint64,
195+
Value String,
196+
PRIMARY KEY (Key)
197+
);
198+
)").GetValueSync().IsSuccess());
199+
}
200+
201+
{
202+
const TString query(Q1_(R"(
203+
SELECT * FROM `/Root/TestTable`;
204+
)"));
205+
206+
{
207+
auto settings = TExecDataQuerySettings();
208+
settings.CollectFullDiagnostics(true);
209+
210+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
211+
212+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString().c_str());
213+
214+
UNIT_ASSERT_C(!result.GetDiagnostics().empty(), "Query result diagnostics is empty");
215+
216+
TStringStream in;
217+
in << result.GetDiagnostics();
218+
NJson::TJsonValue value;
219+
ReadJsonTree(&in, &value);
220+
221+
UNIT_ASSERT_C(value.IsMap(), "Incorrect Diagnostics");
222+
UNIT_ASSERT_C(value.Has("query_id"), "Incorrect Diagnostics");
223+
UNIT_ASSERT_C(value.Has("version"), "Incorrect Diagnostics");
224+
UNIT_ASSERT_C(value.Has("query_text"), "Incorrect Diagnostics");
225+
UNIT_ASSERT_C(value.Has("query_parameter_types"), "Incorrect Diagnostics");
226+
UNIT_ASSERT_C(value.Has("table_metadata"), "Incorrect Diagnostics");
227+
UNIT_ASSERT_C(value["table_metadata"].IsArray(), "Incorrect Diagnostics: table_metadata type should be an array");
228+
UNIT_ASSERT_C(value.Has("created_at"), "Incorrect Diagnostics");
229+
UNIT_ASSERT_C(value.Has("query_syntax"), "Incorrect Diagnostics");
230+
UNIT_ASSERT_C(value.Has("query_database"), "Incorrect Diagnostics");
231+
UNIT_ASSERT_C(value.Has("query_cluster"), "Incorrect Diagnostics");
232+
UNIT_ASSERT_C(value.Has("query_plan"), "Incorrect Diagnostics");
233+
UNIT_ASSERT_C(value.Has("query_type"), "Incorrect Diagnostics");
234+
}
235+
236+
{
237+
auto settings = TExecDataQuerySettings();
238+
239+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
240+
241+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString().c_str());
242+
243+
UNIT_ASSERT_C(result.GetDiagnostics().empty(), "Query result diagnostics should be empty, but it's not");
244+
}
245+
}
246+
}
247+
182248
Y_UNIT_TEST(QueryCachePermissionsLoss) {
183249
TKikimrRunner kikimr;
184250
auto db = kikimr.GetTableClient();

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,55 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
272272
}
273273
}
274274

275+
Y_UNIT_TEST(ExecuteCollectFullDiagnostics) {
276+
auto kikimr = DefaultKikimrRunner();
277+
auto db = kikimr.GetQueryClient();
278+
279+
{
280+
TExecuteQuerySettings settings;
281+
settings.CollectFullDiagnostics(true);
282+
283+
auto result = db.ExecuteQuery(R"(
284+
SELECT Key, Value2 FROM TwoShard WHERE Value2 > 0;
285+
)", TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
286+
287+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
288+
UNIT_ASSERT_C(!result.GetDiagnostics().empty(), "Query result diagnostics is empty");
289+
290+
TStringStream in;
291+
in << result.GetDiagnostics();
292+
NJson::TJsonValue value;
293+
ReadJsonTree(&in, &value);
294+
295+
UNIT_ASSERT_C(value.IsMap(), "Incorrect Diagnostics");
296+
UNIT_ASSERT_C(value.Has("query_id"), "Incorrect Diagnostics");
297+
UNIT_ASSERT_C(value.Has("version"), "Incorrect Diagnostics");
298+
UNIT_ASSERT_C(value.Has("query_text"), "Incorrect Diagnostics");
299+
UNIT_ASSERT_C(value.Has("query_parameter_types"), "Incorrect Diagnostics");
300+
UNIT_ASSERT_C(value.Has("table_metadata"), "Incorrect Diagnostics");
301+
UNIT_ASSERT_C(value["table_metadata"].IsArray(), "Incorrect Diagnostics: table_metadata type should be an array");
302+
UNIT_ASSERT_C(value.Has("created_at"), "Incorrect Diagnostics");
303+
UNIT_ASSERT_C(value.Has("query_syntax"), "Incorrect Diagnostics");
304+
UNIT_ASSERT_C(value.Has("query_database"), "Incorrect Diagnostics");
305+
UNIT_ASSERT_C(value.Has("query_cluster"), "Incorrect Diagnostics");
306+
UNIT_ASSERT_C(value.Has("query_plan"), "Incorrect Diagnostics");
307+
UNIT_ASSERT_C(value.Has("query_type"), "Incorrect Diagnostics");
308+
}
309+
310+
{
311+
TExecuteQuerySettings settings;
312+
settings.CollectFullDiagnostics(true);
313+
314+
auto result = db.ExecuteQuery(R"(
315+
SELECT Key, Value2 FROM TwoShard WHERE Value2 > 0;
316+
)", TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
317+
318+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString().c_str());
319+
320+
UNIT_ASSERT_C(result.GetDiagnostics().empty(), "Query result diagnostics should be empty, but it's not");
321+
}
322+
}
323+
275324
void CheckQueryResult(TExecuteQueryResult result) {
276325
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
277326
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);

‎ydb/public/api/protos/ydb_query.proto‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,8 @@ message ExecuteQueryRequest {
172172
int64 response_part_limit_bytes = 9 [(Ydb.value) = "[0; 33554432]"];
173173

174174
string pool_id = 10; // Workload manager pool id
175+
176+
bool collect_full_diagnostics = 11;
175177
}
176178

177179
message ResultSetMeta {
@@ -191,6 +193,9 @@ message ExecuteQueryResponsePart {
191193
Ydb.TableStats.QueryStats exec_stats = 5;
192194

193195
TransactionMeta tx_meta = 6;
196+
197+
// Full query diagnostics
198+
string query_full_diagnostics = 7;
194199
}
195200

196201
message ExecuteScriptRequest {

‎ydb/public/api/protos/ydb_table.proto‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -942,6 +942,7 @@ message ExecuteDataQueryRequest {
942942
QueryCachePolicy query_cache_policy = 5;
943943
Ydb.Operations.OperationParams operation_params = 6;
944944
QueryStatsCollection.Mode collect_stats = 7;
945+
bool collect_full_diagnostics = 8;
945946
}
946947

947948
message ExecuteDataQueryResponse {
@@ -984,6 +985,8 @@ message ExecuteQueryResult {
984985
QueryMeta query_meta = 3;
985986
// Query execution statistics
986987
Ydb.TableStats.QueryStats query_stats = 4;
988+
// Full query diagnostics
989+
string query_full_diagnostics = 5;
987990
}
988991

989992
// Explain data query

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

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,8 @@ 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);
368370

369371
AddOutputFormats(config, {
370372
EDataFormat::Pretty,
@@ -432,6 +434,9 @@ int TCommandExecuteQuery::ExecuteDataQuery(TConfig& config) {
432434
NTable::TExecDataQuerySettings settings;
433435
settings.KeepInQueryCache(true);
434436
settings.CollectQueryStats(ParseQueryStatsModeOrThrow(CollectStatsMode, defaultStatsMode));
437+
if (CollectFullDiagnostics) {
438+
settings.CollectFullDiagnostics(true);
439+
}
435440

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

521531
int TCommandExecuteQuery::ExecuteSchemeQuery(TConfig& config) {
@@ -548,7 +558,7 @@ namespace {
548558
NQuery::TExecuteQuerySettings>;
549559

550560
template <typename TClient>
551-
auto GetSettings(const TString& collectStatsMode, const bool basicStats, std::optional<TDuration> timeout) {
561+
auto GetSettings(const TString& collectStatsMode, const bool basicStats, std::optional<TDuration> timeout, bool collectFullDiagnostics) {
552562
if constexpr (std::is_same_v<TClient, NTable::TTableClient>) {
553563
const auto defaultStatsMode = basicStats
554564
? NTable::ECollectQueryStatsMode::Basic
@@ -558,6 +568,9 @@ namespace {
558568
if (timeout.has_value()) {
559569
settings.ClientTimeout(*timeout);
560570
}
571+
if (collectFullDiagnostics) {
572+
settings.CollectFullDiagnostics(true);
573+
}
561574
return settings;
562575
} else if constexpr (std::is_same_v<TClient, NQuery::TQueryClient>) {
563576
const auto defaultStatsMode = basicStats
@@ -568,6 +581,9 @@ namespace {
568581
if (timeout.has_value()) {
569582
settings.ClientTimeout(*timeout);
570583
}
584+
if (collectFullDiagnostics) {
585+
settings.CollectFullDiagnostics(true);
586+
}
571587
return settings;
572588
}
573589
Y_UNREACHABLE();
@@ -674,7 +690,7 @@ int TCommandExecuteQuery::ExecuteQueryImpl(TConfig& config) {
674690
if (OperationTimeout) {
675691
optTimeout = TDuration::MilliSeconds(FromString<ui64>(OperationTimeout));
676692
}
677-
const auto settings = GetSettings<TClient>(CollectStatsMode, BasicStats, optTimeout);
693+
const auto settings = GetSettings<TClient>(CollectStatsMode, BasicStats, optTimeout, CollectFullDiagnostics);
678694

679695
TAsyncPartIterator<TClient> asyncResult;
680696
SetInterruptHandlers();
@@ -732,6 +748,7 @@ template <typename TIterator>
732748
bool TCommandExecuteQuery::PrintQueryResponse(TIterator& result) {
733749
TMaybe<TString> stats;
734750
TMaybe<TString> fullStats;
751+
TString diagnostics;
735752
{
736753
TResultSetPrinter printer(OutputFormat, &IsInterrupted);
737754

@@ -753,6 +770,8 @@ bool TCommandExecuteQuery::PrintQueryResponse(TIterator& result) {
753770
fullStats = queryStats.GetPlan();
754771
}
755772
}
773+
774+
diagnostics = streamPart.GetDiagnostics();
756775
}
757776
} // TResultSetPrinter destructor should be called before printing stats
758777

@@ -767,6 +786,12 @@ bool TCommandExecuteQuery::PrintQueryResponse(TIterator& result) {
767786
queryPlanPrinter.Print(*fullStats);
768787
}
769788

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

772797
if (IsInterrupted()) {

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

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

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

0 commit comments

Comments
 (0)