Skip to content

Commit 1239efb

Browse files
shnikdshnikd
authored
Merge 2ce9291 into eafe9a1 
2 parents eafe9a1 + 2ce9291 commit 1239efb

File tree

19 files changed

+241
-19
lines changed

19 files changed

+241
-19
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
@@ -276,7 +276,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
276276
cachePolicy,
277277
nullptr, // operationParams
278278
settings,
279-
req->pool_id());
279+
req->pool_id(),
280+
req->Getcollect_full_diagnostics());
280281

281282
if (!ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release(), 0, 0, Span_.GetTraceId())) {
282283
NYql::TIssues issues;
@@ -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: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,10 @@ 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(),
150+
"",
151+
req->Getcollect_full_diagnostics());
149152

150153
ReportCostInfo_ = req->operation_params().report_cost_info() == Ydb::FeatureFlag::ENABLED;
151154

@@ -203,6 +206,7 @@ class TExecuteDataQueryRPC : public TRpcKqpRequestActor<TExecuteDataQueryRPC, TE
203206
queryMeta.mutable_parameters_types()->insert({queryParameter.GetName(), parameterType});
204207
}
205208
}
209+
queryResult->set_query_full_diagnostics(kqpResponse.GetQueryDiagnostics());
206210
} catch (const std::exception& ex) {
207211
NYql::TIssues issues;
208212
issues.AddIssue(NYql::ExceptionToIssue(ex));

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,8 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
6868
const ::Ydb::Table::QueryCachePolicy* queryCachePolicy,
6969
const ::Ydb::Operations::OperationParams* operationParams,
7070
const TQueryRequestSettings& querySettings = TQueryRequestSettings(),
71-
const TString& poolId = "");
71+
const TString& poolId = "",
72+
std::optional<bool> collectFullDiagnostics = std::nullopt);
7273

7374
TEvQueryRequest() {
7475
Record.MutableRequest()->SetUsePublicResponseDataFormat(true);
@@ -282,7 +283,7 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
282283
}
283284

284285
bool GetCollectDiagnostics() const {
285-
return Record.GetRequest().GetCollectDiagnostics();
286+
return CollectFullDiagnostics.has_value() ? CollectFullDiagnostics.value() : Record.GetRequest().GetCollectDiagnostics();
286287
}
287288

288289
ui32 CalculateSerializedSize() const override {
@@ -395,6 +396,7 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
395396
TIntrusivePtr<TUserRequestContext> UserRequestContext;
396397
TDuration ProgressStatsPeriod;
397398
std::optional<NResourcePool::TPoolSettings> PoolConfig;
399+
std::optional<bool> CollectFullDiagnostics = std::nullopt;
398400
};
399401

400402
struct TEvDataQueryStreamPart: public TEventPB<TEvDataQueryStreamPart,

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ TEvKqp::TEvQueryRequest::TEvQueryRequest(
1919
const ::Ydb::Table::QueryCachePolicy* queryCachePolicy,
2020
const ::Ydb::Operations::OperationParams* operationParams,
2121
const TQueryRequestSettings& querySettings,
22-
const TString& poolId)
22+
const TString& poolId,
23+
std::optional<bool> collectFullDiagnostics)
2324
: RequestCtx(ctx)
2425
, RequestActorId(requestActorId)
2526
, Database(CanonizePath(ctx->GetDatabaseName().GetOrElse("")))
@@ -35,6 +36,7 @@ TEvKqp::TEvQueryRequest::TEvQueryRequest(
3536
, QueryCachePolicy(queryCachePolicy)
3637
, HasOperationParams(operationParams)
3738
, QuerySettings(querySettings)
39+
, CollectFullDiagnostics(collectFullDiagnostics)
3840
{
3941
if (HasOperationParams) {
4042
OperationTimeout = GetDuration(operationParams->operation_timeout());
@@ -107,6 +109,10 @@ void TEvKqp::TEvQueryRequest::PrepareRemote() const {
107109
Record.MutableRequest()->SetIsInternalCall(RequestCtx->IsInternalCall());
108110
Record.MutableRequest()->SetOutputChunkMaxSize(QuerySettings.OutputChunkMaxSize);
109111

112+
if (CollectFullDiagnostics.has_value()) {
113+
Record.MutableRequest()->SetCollectDiagnostics(CollectFullDiagnostics.value());
114+
}
115+
110116
RequestCtx.reset();
111117
}
112118
}

‎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)