Skip to content

Commit c9a507d

Browse files
shnikdshnikd
committed
Initial commit
1 parent 4306c06 commit c9a507d

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
@@ -262,6 +262,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
262262
.SetSupportStreamTrailingResult(true)
263263
.SetOutputChunkMaxSize(req->response_part_limit_bytes());
264264

265+
assert(req->Getcollect_full_diagnostics());
266+
265267
auto ev = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>(
266268
QueryAction,
267269
queryType,
@@ -276,7 +278,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
276278
cachePolicy,
277279
nullptr, // operationParams
278280
settings,
279-
req->pool_id());
281+
req->pool_id(),
282+
req->Getcollect_full_diagnostics());
280283

281284
if (!ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release(), 0, 0, Span_.GetTraceId())) {
282285
NYql::TIssues issues;
@@ -394,6 +397,8 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
394397
hasTrailingMessage = true;
395398
response.mutable_tx_meta()->set_id(kqpResponse.GetTxMeta().id());
396399
}
400+
assert(!kqpResponse.GetQueryDiagnostics().empty());
401+
response.set_query_full_diagnostics(kqpResponse.GetQueryDiagnostics());
397402
}
398403

399404
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
@@ -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: 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(*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)