Skip to content

Commit 400ce66

Browse files
GrigoriyPAGrigoriyPA
committed
Added script trace opt id
1 parent 9740792 commit 400ce66

File tree

4 files changed

+80
-42
lines changed

4 files changed

+80
-42
lines changed

‎ydb/tests/tools/kqprun/kqprun.cpp

Lines changed: 64 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@
2222
#include <yql/essentials/public/udf/udf_static_registry.h>
2323

2424

25+
namespace NKqpRun {
26+
27+
namespace {
28+
2529
struct TExecutionOptions {
2630
enum class EExecutionCase {
2731
GenericScript,
@@ -77,7 +81,7 @@ struct TExecutionOptions {
7781
return GetValue(index, ScriptQueryActions, NKikimrKqp::EQueryAction::QUERY_ACTION_EXECUTE);
7882
}
7983

80-
NKqpRun::TRequestOptions GetSchemeQueryOptions() const {
84+
TRequestOptions GetSchemeQueryOptions() const {
8185
TString sql = SchemeQuery;
8286
if (UseTemplates) {
8387
ReplaceYqlTokenTemplate(sql);
@@ -94,7 +98,7 @@ struct TExecutionOptions {
9498
};
9599
}
96100

97-
NKqpRun::TRequestOptions GetScriptQueryOptions(size_t index, size_t queryId, TInstant startTime) const {
101+
TRequestOptions GetScriptQueryOptions(size_t index, size_t queryId, TInstant startTime) const {
98102
Y_ABORT_UNLESS(index < ScriptQueries.size());
99103

100104
TString sql = ScriptQueries[index];
@@ -110,11 +114,12 @@ struct TExecutionOptions {
110114
.PoolId = GetValue(index, PoolIds, TString()),
111115
.UserSID = GetValue(index, UserSIDs, TString(BUILTIN_ACL_ROOT)),
112116
.Database = GetValue(index, Databases, TString()),
113-
.Timeout = GetValue(index, Timeouts, TDuration::Zero())
117+
.Timeout = GetValue(index, Timeouts, TDuration::Zero()),
118+
.QueryId = queryId
114119
};
115120
}
116121

117-
void Validate(const NKqpRun::TRunnerOptions& runnerOptions) const {
122+
void Validate(const TRunnerOptions& runnerOptions) const {
118123
if (!SchemeQuery && ScriptQueries.empty() && !runnerOptions.YdbSettings.MonitoringEnabled && !runnerOptions.YdbSettings.GrpcEnabled) {
119124
ythrow yexception() << "Nothing to execute and is not running as daemon";
120125
}
@@ -123,7 +128,7 @@ struct TExecutionOptions {
123128
ValidateSchemeQueryOptions(runnerOptions);
124129
ValidateScriptExecutionOptions(runnerOptions);
125130
ValidateAsyncOptions(runnerOptions.YdbSettings.AsyncQueriesSettings);
126-
ValidateTraceOpt(runnerOptions.TraceOptType);
131+
ValidateTraceOpt(runnerOptions);
127132
}
128133

129134
private:
@@ -143,7 +148,7 @@ struct TExecutionOptions {
143148
checker(Timeouts.size(), "timeouts");
144149
}
145150

146-
void ValidateSchemeQueryOptions(const NKqpRun::TRunnerOptions& runnerOptions) const {
151+
void ValidateSchemeQueryOptions(const TRunnerOptions& runnerOptions) const {
147152
if (SchemeQuery) {
148153
return;
149154
}
@@ -152,7 +157,7 @@ struct TExecutionOptions {
152157
}
153158
}
154159

155-
void ValidateScriptExecutionOptions(const NKqpRun::TRunnerOptions& runnerOptions) const {
160+
void ValidateScriptExecutionOptions(const TRunnerOptions& runnerOptions) const {
156161
if (runnerOptions.YdbSettings.SameSession && HasExecutionCase(EExecutionCase::AsyncQuery)) {
157162
ythrow yexception() << "Same session can not be used with async quries";
158163
}
@@ -194,7 +199,7 @@ struct TExecutionOptions {
194199
}
195200
}
196201

197-
void ValidateAsyncOptions(const NKqpRun::TAsyncQueriesSettings& asyncQueriesSettings) const {
202+
void ValidateAsyncOptions(const TAsyncQueriesSettings& asyncQueriesSettings) const {
198203
if (asyncQueriesSettings.InFlightLimit && !HasExecutionCase(EExecutionCase::AsyncQuery)) {
199204
ythrow yexception() << "In flight limit can not be used without async queries";
200205
}
@@ -205,28 +210,43 @@ struct TExecutionOptions {
205210
}
206211
}
207212

208-
void ValidateTraceOpt(NKqpRun::TRunnerOptions::ETraceOptType traceOptType) const {
209-
switch (traceOptType) {
210-
case NKqpRun::TRunnerOptions::ETraceOptType::Scheme: {
213+
void ValidateTraceOpt(const TRunnerOptions& runnerOptions) const {
214+
NColorizer::TColors colors = NColorizer::AutoColors(Cout);
215+
switch (runnerOptions.TraceOptType) {
216+
case TRunnerOptions::ETraceOptType::Scheme: {
211217
if (!SchemeQuery) {
212218
ythrow yexception() << "Trace opt type scheme cannot be used without scheme query";
213219
}
214220
break;
215221
}
216-
case NKqpRun::TRunnerOptions::ETraceOptType::Script: {
222+
case TRunnerOptions::ETraceOptType::Script: {
217223
if (ScriptQueries.empty()) {
218224
ythrow yexception() << "Trace opt type script cannot be used without script queries";
219225
}
220226
}
221-
case NKqpRun::TRunnerOptions::ETraceOptType::All: {
227+
case TRunnerOptions::ETraceOptType::All: {
222228
if (!SchemeQuery && ScriptQueries.empty()) {
223229
ythrow yexception() << "Trace opt type all cannot be used without any queries";
224230
}
225231
}
226-
case NKqpRun::TRunnerOptions::ETraceOptType::Disabled: {
232+
case TRunnerOptions::ETraceOptType::Disabled: {
227233
break;
228234
}
229235
}
236+
237+
if (const auto traceOptId = runnerOptions.TraceOptScriptId) {
238+
if (runnerOptions.TraceOptType != TRunnerOptions::ETraceOptType::Script) {
239+
ythrow yexception() << "Trace opt id allowed only for trace opt type script (used " << runnerOptions.TraceOptType << ")";
240+
}
241+
242+
const ui64 numberScripts = ScriptQueries.size() * LoopCount;
243+
if (*traceOptId >= numberScripts) {
244+
ythrow yexception() << "Invalid trace opt id " << *traceOptId << ", it should be less than number of scipt queries " << numberScripts;
245+
}
246+
if (numberScripts == 1) {
247+
Cout << colors.Red() << "Warning: trace opt id is not necessary for single script mode, use -T script" << Endl;
248+
}
249+
}
230250
}
231251

232252
private:
@@ -239,8 +259,8 @@ struct TExecutionOptions {
239259
}
240260

241261
static void ReplaceYqlTokenTemplate(TString& sql) {
242-
const TString variableName = TStringBuilder() << "${" << NKqpRun::YQL_TOKEN_VARIABLE << "}";
243-
if (const TString& yqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE)) {
262+
const TString variableName = TStringBuilder() << "${" << YQL_TOKEN_VARIABLE << "}";
263+
if (const TString& yqlToken = GetEnv(YQL_TOKEN_VARIABLE)) {
244264
SubstGlobal(sql, variableName, yqlToken);
245265
} else if (sql.Contains(variableName)) {
246266
ythrow yexception() << "Failed to replace ${YQL_TOKEN} template, please specify YQL_TOKEN environment variable\n";
@@ -249,7 +269,7 @@ struct TExecutionOptions {
249269
};
250270

251271

252-
void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TExecutionOptions& executionOptions, NKqpRun::TKqpRunner& runner) {
272+
void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TExecutionOptions& executionOptions, TKqpRunner& runner) {
253273
NColorizer::TColors colors = NColorizer::AutoColors(Cout);
254274

255275
switch (executionOptions.GetExecutionCase(index)) {
@@ -292,7 +312,7 @@ void RunArgumentQuery(size_t index, size_t queryId, TInstant startTime, const TE
292312
}
293313

294314

295-
void RunArgumentQueries(const TExecutionOptions& executionOptions, NKqpRun::TKqpRunner& runner) {
315+
void RunArgumentQueries(const TExecutionOptions& executionOptions, TKqpRunner& runner) {
296316
NColorizer::TColors colors = NColorizer::AutoColors(Cout);
297317

298318
if (executionOptions.SchemeQuery) {
@@ -354,11 +374,11 @@ void RunAsDaemon() {
354374
}
355375

356376

357-
void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunnerOptions& runnerOptions) {
377+
void RunScript(const TExecutionOptions& executionOptions, const TRunnerOptions& runnerOptions) {
358378
NColorizer::TColors colors = NColorizer::AutoColors(Cout);
359379

360380
Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization of kqp runner..." << colors.Default() << Endl;
361-
NKqpRun::TKqpRunner runner(runnerOptions);
381+
TKqpRunner runner(runnerOptions);
362382

363383
try {
364384
RunArgumentQueries(executionOptions, runner);
@@ -403,11 +423,11 @@ TIntrusivePtr<NKikimr::NMiniKQL::IMutableFunctionRegistry> CreateFunctionRegistr
403423

404424

405425
class TMain : public TMainClassArgs {
406-
inline static const TString YqlToken = GetEnv(NKqpRun::YQL_TOKEN_VARIABLE);
426+
inline static const TString YqlToken = GetEnv(YQL_TOKEN_VARIABLE);
407427
inline static std::vector<std::unique_ptr<TFileOutput>> FileHolders;
408428

409429
TExecutionOptions ExecutionOptions;
410-
NKqpRun::TRunnerOptions RunnerOptions;
430+
TRunnerOptions RunnerOptions;
411431

412432
THashMap<TString, TString> TablesMapping;
413433
TVector<TString> UdfsPaths;
@@ -524,13 +544,13 @@ class TMain : public TMainClassArgs {
524544
std::remove(file.c_str());
525545
}
526546
});
527-
TChoices<NKqpRun::TRunnerOptions::ETraceOptType> traceOpt({
528-
{"all", NKqpRun::TRunnerOptions::ETraceOptType::All},
529-
{"scheme", NKqpRun::TRunnerOptions::ETraceOptType::Scheme},
530-
{"script", NKqpRun::TRunnerOptions::ETraceOptType::Script},
531-
{"disabled", NKqpRun::TRunnerOptions::ETraceOptType::Disabled}
547+
TChoices<TRunnerOptions::ETraceOptType> traceOpt({
548+
{"all", TRunnerOptions::ETraceOptType::All},
549+
{"scheme", TRunnerOptions::ETraceOptType::Scheme},
550+
{"script", TRunnerOptions::ETraceOptType::Script},
551+
{"disabled", TRunnerOptions::ETraceOptType::Disabled}
532552
});
533-
options.AddLongOption('T', "trace-opt", "print AST in the begin of each transformation")
553+
options.AddLongOption('T', "trace-opt", "print AST in the begin of each transformation (use script@<query id> for tracing one -p query)")
534554
.RequiredArgument("trace-opt-query")
535555
.DefaultValue("disabled")
536556
.Choices(traceOpt.GetChoices())
@@ -539,6 +559,9 @@ class TMain : public TMainClassArgs {
539559
RunnerOptions.YdbSettings.TraceOptEnabled = traceOptType != NKqpRun::TRunnerOptions::ETraceOptType::Disabled;
540560
return traceOptType;
541561
});
562+
options.AddLongOption('I', "trace-opt-index", "index of -p query to use --trace-opt, starts from zero")
563+
.RequiredArgument("uint")
564+
.StoreResult(&RunnerOptions.TraceOptScriptId);
542565
options.AddLongOption("trace-id", "Trace id for -p queries")
543566
.RequiredArgument("id")
544567
.EmplaceTo(&ExecutionOptions.TraceIds);
@@ -551,10 +574,10 @@ class TMain : public TMainClassArgs {
551574
.RequiredArgument("uint")
552575
.DefaultValue(0)
553576
.StoreResult(&ExecutionOptions.ResultsRowsLimit);
554-
TChoices<NKqpRun::TRunnerOptions::EResultOutputFormat> resultFormat({
555-
{"rows", NKqpRun::TRunnerOptions::EResultOutputFormat::RowsJson},
556-
{"full-json", NKqpRun::TRunnerOptions::EResultOutputFormat::FullJson},
557-
{"full-proto", NKqpRun::TRunnerOptions::EResultOutputFormat::FullProto}
577+
TChoices<TRunnerOptions::EResultOutputFormat> resultFormat({
578+
{"rows", TRunnerOptions::EResultOutputFormat::RowsJson},
579+
{"full-json", TRunnerOptions::EResultOutputFormat::FullJson},
580+
{"full-proto", TRunnerOptions::EResultOutputFormat::FullProto}
558581
});
559582
options.AddLongOption('R', "result-format", "Script query result format")
560583
.RequiredArgument("result-format")
@@ -620,9 +643,9 @@ class TMain : public TMainClassArgs {
620643
.RequiredArgument("uint")
621644
.DefaultValue(0)
622645
.StoreResult(&RunnerOptions.YdbSettings.AsyncQueriesSettings.InFlightLimit);
623-
TChoices<NKqpRun::TAsyncQueriesSettings::EVerbose> verbose({
624-
{"each-query", NKqpRun::TAsyncQueriesSettings::EVerbose::EachQuery},
625-
{"final", NKqpRun::TAsyncQueriesSettings::EVerbose::Final}
646+
TChoices<TAsyncQueriesSettings::EVerbose> verbose({
647+
{"each-query", TAsyncQueriesSettings::EVerbose::EachQuery},
648+
{"final", TAsyncQueriesSettings::EVerbose::Final}
626649
});
627650
options.AddLongOption("async-verbose", "Verbose type for async queries")
628651
.RequiredArgument("type")
@@ -814,13 +837,16 @@ void SegmentationFaultHandler(int) {
814837
abort();
815838
}
816839

840+
} // anonymous namespace
841+
842+
} // namespace NKqpRun
817843

818844
int main(int argc, const char* argv[]) {
819-
std::set_terminate(KqprunTerminateHandler);
820-
signal(SIGSEGV, &SegmentationFaultHandler);
845+
std::set_terminate(NKqpRun::KqprunTerminateHandler);
846+
signal(SIGSEGV, &NKqpRun::SegmentationFaultHandler);
821847

822848
try {
823-
TMain().Run(argc, argv);
849+
NKqpRun::TMain().Run(argc, argv);
824850
} catch (...) {
825851
NColorizer::TColors colors = NColorizer::AutoColors(Cerr);
826852

‎ydb/tests/tools/kqprun/src/common.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ struct TRunnerOptions {
8080
EResultOutputFormat ResultOutputFormat = EResultOutputFormat::RowsJson;
8181
NYdb::NConsoleClient::EDataFormat PlanOutputFormat = NYdb::NConsoleClient::EDataFormat::Default;
8282
ETraceOptType TraceOptType = ETraceOptType::Disabled;
83+
std::optional<size_t> TraceOptScriptId;
8384

8485
TDuration ScriptCancelAfter;
8586

@@ -95,6 +96,7 @@ struct TRequestOptions {
9596
TString UserSID;
9697
TString Database;
9798
TDuration Timeout;
99+
size_t QueryId = 0;
98100
};
99101

100102
} // namespace NKqpRun

‎ydb/tests/tools/kqprun/src/kqp_runner.cpp

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ class TKqpRunner::TImpl {
120120
}
121121

122122
bool ExecuteScript(const TRequestOptions& script) {
123-
StartScriptTraceOpt();
123+
StartScriptTraceOpt(script.QueryId);
124124

125125
TRequestResult status = YdbSetup_.ScriptRequest(script, ExecutionOperation_);
126126

@@ -136,7 +136,7 @@ class TKqpRunner::TImpl {
136136
}
137137

138138
bool ExecuteQuery(const TRequestOptions& query, EQueryType queryType) {
139-
StartScriptTraceOpt();
139+
StartScriptTraceOpt(query.QueryId);
140140
StartTime_ = TInstant::Now();
141141

142142
TString queryTypeStr;
@@ -267,6 +267,8 @@ class TKqpRunner::TImpl {
267267
Sleep(getOperationPeriod);
268268
}
269269

270+
TYdbSetup::StopTraceOpt();
271+
270272
PrintScriptAst(ExecutionMeta_.Ast);
271273
PrintScriptProgress(ExecutionMeta_.Plan);
272274
PrintScriptPlan(ExecutionMeta_.Plan);
@@ -290,8 +292,14 @@ class TKqpRunner::TImpl {
290292
}
291293
}
292294

293-
void StartScriptTraceOpt() const {
294-
if (Options_.TraceOptType == TRunnerOptions::ETraceOptType::All || Options_.TraceOptType == TRunnerOptions::ETraceOptType::Script) {
295+
void StartScriptTraceOpt(size_t queryId) const {
296+
bool startTraceOpt = Options_.TraceOptType == TRunnerOptions::ETraceOptType::All;
297+
298+
if (Options_.TraceOptType == TRunnerOptions::ETraceOptType::Script) {
299+
startTraceOpt |= !Options_.TraceOptScriptId || *Options_.TraceOptScriptId == queryId;
300+
}
301+
302+
if (startTraceOpt) {
295303
YdbSetup_.StartTraceOpt();
296304
}
297305
}

‎ydb/tests/tools/kqprun/src/ya.make

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ PEERDIR(
1010
ydb/core/testlib
1111
)
1212

13+
GENERATE_ENUM_SERIALIZATION(common.h)
14+
1315
YQL_LAST_ABI_VERSION()
1416

1517
END()

0 commit comments

Comments
 (0)