Skip to content

Commit ab2e3bc

Browse files
shnikdshnikd
committed
Initial commit
1 parent 28180f6 commit ab2e3bc

File tree

17 files changed

+220
-16
lines changed

17 files changed

+220
-16
lines changed

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
267267
.SetSupportStreamTrailingResult(true)
268268
.SetOutputChunkMaxSize(req->response_part_limit_bytes());
269269

270+
assert(req->Getcollect_full_diagnostics());
271+
270272
auto ev = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>(
271273
QueryAction,
272274
queryType,
@@ -281,7 +283,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
281283
cachePolicy,
282284
nullptr, // operationParams
283285
settings,
284-
req->pool_id());
286+
req->pool_id(),
287+
req->Getcollect_full_diagnostics());
285288

286289
ev->SetProgressStatsPeriod(TDuration::MilliSeconds(req->stats_period_ms()));
287290

@@ -422,6 +425,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
422425
hasTrailingMessage = true;
423426
response.mutable_tx_meta()->set_id(kqpResponse.GetTxMeta().id());
424427
}
428+
assert(!kqpResponse.GetQueryDiagnostics().empty());
429+
response.set_query_full_diagnostics(kqpResponse.GetQueryDiagnostics());
425430
}
426431

427432
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: 3 additions & 1 deletion
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+
bool collectFullDiagnostics = false);
7273

7374
TEvQueryRequest() {
7475
Record.MutableRequest()->SetUsePublicResponseDataFormat(true);
@@ -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+
bool CollectFullDiagnostics = false;
398400
};
399401

400402
struct TEvDataQueryStreamPart: public TEventPB<TEvDataQueryStreamPart,

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

Lines changed: 5 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+
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,8 @@ void TEvKqp::TEvQueryRequest::PrepareRemote() const {
107109
Record.MutableRequest()->SetIsInternalCall(RequestCtx->IsInternalCall());
108110
Record.MutableRequest()->SetOutputChunkMaxSize(QuerySettings.OutputChunkMaxSize);
109111

112+
Record.MutableRequest()->SetCollectDiagnostics(CollectFullDiagnostics);
113+
110114
RequestCtx.reset();
111115
}
112116
}

‎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
@@ -650,6 +650,55 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
650650
}
651651
}
652652

653+
Y_UNIT_TEST(ExecuteCollectFullDiagnostics) {
654+
auto kikimr = DefaultKikimrRunner();
655+
auto db = kikimr.GetQueryClient();
656+
657+
{
658+
TExecuteQuerySettings settings;
659+
settings.CollectFullDiagnostics(true);
660+
661+
auto result = db.ExecuteQuery(R"(
662+
SELECT Key, Value2 FROM TwoShard WHERE Value2 > 0;
663+
)", TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
664+
665+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
666+
UNIT_ASSERT_C(!result.GetDiagnostics().empty(), "Query result diagnostics is empty");
667+
668+
TStringStream in;
669+
in << result.GetDiagnostics();
670+
NJson::TJsonValue value;
671+
ReadJsonTree(&in, &value);
672+
673+
UNIT_ASSERT_C(value.IsMap(), "Incorrect Diagnostics");
674+
UNIT_ASSERT_C(value.Has("query_id"), "Incorrect Diagnostics");
675+
UNIT_ASSERT_C(value.Has("version"), "Incorrect Diagnostics");
676+
UNIT_ASSERT_C(value.Has("query_text"), "Incorrect Diagnostics");
677+
UNIT_ASSERT_C(value.Has("query_parameter_types"), "Incorrect Diagnostics");
678+
UNIT_ASSERT_C(value.Has("table_metadata"), "Incorrect Diagnostics");
679+
UNIT_ASSERT_C(value["table_metadata"].IsArray(), "Incorrect Diagnostics: table_metadata type should be an array");
680+
UNIT_ASSERT_C(value.Has("created_at"), "Incorrect Diagnostics");
681+
UNIT_ASSERT_C(value.Has("query_syntax"), "Incorrect Diagnostics");
682+
UNIT_ASSERT_C(value.Has("query_database"), "Incorrect Diagnostics");
683+
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_type"), "Incorrect Diagnostics");
686+
}
687+
688+
{
689+
TExecuteQuerySettings settings;
690+
settings.CollectFullDiagnostics(true);
691+
692+
auto result = db.ExecuteQuery(R"(
693+
SELECT Key, Value2 FROM TwoShard WHERE Value2 > 0;
694+
)", TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
695+
696+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString().c_str());
697+
698+
UNIT_ASSERT_C(result.GetDiagnostics().empty(), "Query result diagnostics should be empty, but it's not");
699+
}
700+
}
701+
653702
void CheckQueryResult(TExecuteQueryResult result) {
654703
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
655704
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
@@ -181,6 +181,8 @@ message ExecuteQueryRequest {
181181
// When query statistics are enabled (stats_mode != STATS_MODE_NONE), by default statistics will be sent only once after query execution is finished.
182182
// In case when stats_period_ms is specified and is non-zero, query statistics will be additionally sent every stats_period_ms milliseconds beginning from the start of query execution.
183183
int64 stats_period_ms = 11 [(Ydb.value) = ">= 0"];
184+
185+
bool collect_full_diagnostics = 12;
184186
}
185187

186188
message ResultSetMeta {
@@ -200,6 +202,9 @@ message ExecuteQueryResponsePart {
200202
Ydb.TableStats.QueryStats exec_stats = 5;
201203

202204
TransactionMeta tx_meta = 6;
205+
206+
// Full query diagnostics
207+
string query_full_diagnostics = 7;
203208
}
204209

205210
message ExecuteScriptRequest {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -926,6 +926,7 @@ message ExecuteDataQueryRequest {
926926
QueryCachePolicy query_cache_policy = 5;
927927
Ydb.Operations.OperationParams operation_params = 6;
928928
QueryStatsCollection.Mode collect_stats = 7;
929+
bool collect_full_diagnostics = 8;
929930
}
930931

931932
message ExecuteDataQueryResponse {
@@ -968,6 +969,8 @@ message ExecuteQueryResult {
968969
QueryMeta query_meta = 3;
969970
// Query execution statistics
970971
Ydb.TableStats.QueryStats query_stats = 4;
972+
// Full query diagnostics
973+
string query_full_diagnostics = 5;
971974
}
972975

973976
// Explain data query

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

Lines changed: 24 additions & 0 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) {
@@ -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();
@@ -753,6 +769,8 @@ bool TCommandExecuteQuery::PrintQueryResponse(TIterator& result) {
753769
fullStats = queryStats.GetPlan();
754770
}
755771
}
772+
773+
if ()
756774
}
757775
} // TResultSetPrinter destructor should be called before printing stats
758776

@@ -767,6 +785,12 @@ bool TCommandExecuteQuery::PrintQueryResponse(TIterator& result) {
767785
queryPlanPrinter.Print(TString{*fullStats});
768786
}
769787

788+
if (CollectFullDiagnostics)
789+
{
790+
TFileOutput file(TStringBuilder() << "diagnostics_" << TGUID::Create().AsGuidString() << ".txt");
791+
file << result.GetDiagnostics();
792+
}
793+
770794
PrintFlameGraph(fullStats);
771795

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