Skip to content

Commit 053a1d3

Browse files
GrigoriyPAGrigoriyPA
authored
Merge 892cb1f into a08366f 
2 parents a08366f + 892cb1f commit 053a1d3

File tree

9 files changed

+202
-62
lines changed

9 files changed

+202
-62
lines changed

‎ydb/tests/tools/kqprun/.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
11
sync_dir
2+
example
3+
udfs
24
*.log
35
*.sql
6+
*.bin

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

Lines changed: 50 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include "src/kqp_runner.h"
22

3+
#include <cstdio>
4+
35
#include <contrib/libs/protobuf/src/google/protobuf/text_format.h>
46

57
#include <library/cpp/colorizer/colors.h>
@@ -22,7 +24,7 @@ struct TExecutionOptions {
2224
TString ScriptTraceId = "kqprun";
2325

2426
bool HasResults() const {
25-
return ScriptQuery && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE && !ClearExecution;
27+
return ScriptQuery && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE;
2628
}
2729
};
2830

@@ -46,6 +48,9 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner
4648
if (!runner.ExecuteScript(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.ScriptTraceId)) {
4749
ythrow yexception() << "Script execution failed";
4850
}
51+
if (!runner.FetchScriptResults()) {
52+
ythrow yexception() << "Fetch script results failed";
53+
}
4954
} else {
5055
if (!runner.ExecuteQuery(executionOptions.ScriptQuery, executionOptions.ScriptQueryAction, executionOptions.ScriptTraceId)) {
5156
ythrow yexception() << "Query execution failed";
@@ -54,10 +59,7 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner
5459
}
5560

5661
if (executionOptions.HasResults()) {
57-
Cout << colors.Yellow() << "Writing script results..." << colors.Default() << Endl;
58-
if (!runner.WriteScriptResults()) {
59-
ythrow yexception() << "Writing script results failed";
60-
}
62+
runner.PrintScriptResults();
6163
}
6264
}
6365

@@ -74,6 +76,20 @@ THolder<TFileOutput> SetupDefaultFileOutput(const TString& filePath, IOutputStre
7476
}
7577

7678

79+
TIntrusivePtr<NKikimr::NMiniKQL::IMutableFunctionRegistry> CreateFunctionRegistry(const TString& udfsDirectory, TVector<TString> udfsPaths) {
80+
if (!udfsDirectory.empty() || !udfsPaths.empty()) {
81+
NColorizer::TColors colors = NColorizer::AutoColors(Cout);
82+
Cout << colors.Yellow() << "Fetching udfs..." << colors.Default() << Endl;
83+
}
84+
85+
NKikimr::NMiniKQL::FindUdfsInDir(udfsDirectory, &udfsPaths);
86+
auto functionRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&NYql::NBacktrace::KikimrBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, udfsPaths)->Clone();
87+
NKikimr::NMiniKQL::FillStaticModules(*functionRegistry);
88+
89+
return functionRegistry;
90+
}
91+
92+
7793
void RunMain(int argc, const char* argv[]) {
7894
TExecutionOptions executionOptions;
7995
NKqpRun::TRunnerOptions runnerOptions;
@@ -87,9 +103,11 @@ void RunMain(int argc, const char* argv[]) {
87103
TString logFile = "-";
88104
TString appConfigFile = "./configuration/app_config.conf";
89105

106+
TString traceOptType = "disabled";
90107
TString scriptQueryAction = "execute";
91108
TString planOutputFormat = "pretty";
92109
TString resultOutputFormat = "rows";
110+
i64 resultsRowsLimit = 1000;
93111

94112
TVector<TString> udfsPaths;
95113
TString udfsDirectory;
@@ -103,7 +121,7 @@ void RunMain(int argc, const char* argv[]) {
103121
.Optional()
104122
.RequiredArgument("FILE")
105123
.StoreResult(&schemeQueryFile);
106-
options.AddLongOption("app-config", "File with app config (TAppConfig)")
124+
options.AddLongOption('c', "app-config", "File with app config (TAppConfig)")
107125
.Optional()
108126
.RequiredArgument("FILE")
109127
.DefaultValue(appConfigFile)
@@ -135,33 +153,33 @@ void RunMain(int argc, const char* argv[]) {
135153
.NoArgument()
136154
.DefaultValue(executionOptions.ClearExecution)
137155
.SetFlag(&executionOptions.ClearExecution);
138-
options.AddLongOption("trace-opt", "print AST in the begin of each transformation")
156+
options.AddLongOption('T', "trace-opt", "print AST in the begin of each transformation, one of { scheme | script | all }")
139157
.Optional()
140-
.NoArgument()
141-
.DefaultValue(runnerOptions.YdbSettings.TraceOpt)
142-
.SetFlag(&runnerOptions.YdbSettings.TraceOpt);
143-
options.AddLongOption("script-action", "Script query execute action, one of { execute | explain }")
158+
.RequiredArgument("STR")
159+
.DefaultValue(traceOptType)
160+
.StoreResult(&traceOptType);
161+
options.AddLongOption('A', "script-action", "Script query execute action, one of { execute | explain }")
144162
.Optional()
145163
.RequiredArgument("STR")
146164
.DefaultValue(scriptQueryAction)
147165
.StoreResult(&scriptQueryAction);
148-
options.AddLongOption("plan-format", "Script query plan format, one of { pretty | table | json }")
166+
options.AddLongOption('P', "plan-format", "Script query plan format, one of { pretty | table | json }")
149167
.Optional()
150168
.RequiredArgument("STR")
151169
.DefaultValue(planOutputFormat)
152170
.StoreResult(&planOutputFormat);
153-
options.AddLongOption("result-format", "Script query result format, one of { rows | full }")
171+
options.AddLongOption('R', "result-format", "Script query result format, one of { rows | full }")
154172
.Optional()
155173
.RequiredArgument("STR")
156174
.DefaultValue(resultOutputFormat)
157175
.StoreResult(&resultOutputFormat);
158-
options.AddLongOption("result-rows-limit", "Rows limit for script execution results")
176+
options.AddLongOption('L', "result-rows-limit", "Rows limit for script execution results")
159177
.Optional()
160178
.RequiredArgument("INT")
161-
.DefaultValue(runnerOptions.ResultsRowsLimit)
162-
.StoreResult(&runnerOptions.ResultsRowsLimit);
179+
.DefaultValue(resultsRowsLimit)
180+
.StoreResult(&resultsRowsLimit);
163181

164-
options.AddLongOption("udf", "Load shared library with UDF by given path")
182+
options.AddLongOption('u', "udf", "Load shared library with UDF by given path")
165183
.Optional()
166184
.RequiredArgument("FILE")
167185
.AppendTo(&udfsPaths);
@@ -191,15 +209,19 @@ void RunMain(int argc, const char* argv[]) {
191209

192210
// Runner options
193211

194-
if (runnerOptions.ResultsRowsLimit < 0) {
195-
ythrow yexception() << "Results rows limit less than zero";
196-
}
197-
198212
THolder<TFileOutput> resultFileHolder = SetupDefaultFileOutput(resultOutputFile, runnerOptions.ResultOutput);
199213
THolder<TFileOutput> schemeQueryAstFileHolder = SetupDefaultFileOutput(schemeQueryAstFile, runnerOptions.SchemeQueryAstOutput);
200214
THolder<TFileOutput> scriptQueryAstFileHolder = SetupDefaultFileOutput(scriptQueryAstFile, runnerOptions.ScriptQueryAstOutput);
201215
THolder<TFileOutput> scriptQueryPlanFileHolder = SetupDefaultFileOutput(scriptQueryPlanFile, runnerOptions.ScriptQueryPlanOutput);
202216

217+
runnerOptions.TraceOptType =
218+
(traceOptType == TStringBuf("all")) ? NKqpRun::TRunnerOptions::ETraceOptType::All
219+
: (traceOptType == TStringBuf("scheme")) ? NKqpRun::TRunnerOptions::ETraceOptType::Scheme
220+
: (traceOptType == TStringBuf("script")) ? NKqpRun::TRunnerOptions::ETraceOptType::Script
221+
: (traceOptType == TStringBuf("disabled")) ? NKqpRun::TRunnerOptions::ETraceOptType::Disabled
222+
: NKqpRun::TRunnerOptions::ETraceOptType::All;
223+
runnerOptions.YdbSettings.TraceOptEnabled = runnerOptions.TraceOptType != NKqpRun::TRunnerOptions::ETraceOptType::Disabled;
224+
203225
runnerOptions.ResultOutputFormat =
204226
(resultOutputFormat == TStringBuf("rows")) ? NKqpRun::TRunnerOptions::EResultOutputFormat::RowsJson
205227
: (resultOutputFormat == TStringBuf("full")) ? NKqpRun::TRunnerOptions::EResultOutputFormat::FullJson
@@ -215,20 +237,22 @@ void RunMain(int argc, const char* argv[]) {
215237

216238
if (logFile != "-") {
217239
runnerOptions.YdbSettings.LogOutputFile = logFile;
240+
std::remove(logFile.c_str());
218241
}
219242

220243
runnerOptions.YdbSettings.YqlToken = GetEnv("YQL_TOKEN");
221-
222-
NKikimr::NMiniKQL::FindUdfsInDir(udfsDirectory, &udfsPaths);
223-
auto functionRegistry = NKikimr::NMiniKQL::CreateFunctionRegistry(&NYql::NBacktrace::KikimrBackTrace, NKikimr::NMiniKQL::CreateBuiltinRegistry(), false, udfsPaths)->Clone();
224-
NKikimr::NMiniKQL::FillStaticModules(*functionRegistry);
225-
runnerOptions.YdbSettings.FunctionRegistry = functionRegistry.Get();
244+
runnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(udfsDirectory, udfsPaths).Get();
226245

227246
TString appConfigData = TFileInput(appConfigFile).ReadAll();
228247
if (!google::protobuf::TextFormat::ParseFromString(appConfigData, &runnerOptions.YdbSettings.AppConfig)) {
229248
ythrow yexception() << "Bad format of app configuration";
230249
}
231250

251+
if (resultsRowsLimit < 0) {
252+
ythrow yexception() << "Results rows limit less than zero";
253+
}
254+
runnerOptions.YdbSettings.AppConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(resultsRowsLimit);
255+
232256
RunScript(executionOptions, runnerOptions);
233257
}
234258

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

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,19 @@ namespace {
99

1010
class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMock> {
1111
public:
12-
TRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request, NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise, ui64 resultSizeLimit)
12+
TRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
13+
NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise,
14+
ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets)
1315
: Request_(std::move(request))
1416
, Promise_(promise)
17+
, ResultRowsLimit_(std::numeric_limits<ui64>::max())
1518
, ResultSizeLimit_(std::numeric_limits<i64>::max())
19+
, ResultSets_(resultSets)
1620
{
17-
if (resultSizeLimit && resultSizeLimit < std::numeric_limits<i64>::max()) {
21+
if (resultRowsLimit) {
22+
ResultRowsLimit_ = resultRowsLimit;
23+
}
24+
if (resultSizeLimit) {
1825
ResultSizeLimit_ = resultSizeLimit;
1926
}
2027
}
@@ -36,6 +43,28 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
3643
response->Record.SetSeqNo(ev->Get()->Record.GetSeqNo());
3744
response->Record.SetFreeSpace(ResultSizeLimit_);
3845

46+
auto resultSetIndex = ev->Get()->Record.GetQueryResultIndex();
47+
if (resultSetIndex >= ResultSets_.size()) {
48+
ResultSets_.resize(resultSetIndex + 1);
49+
}
50+
51+
if (!ResultSets_[resultSetIndex].truncated()) {
52+
for (auto& row : *ev->Get()->Record.MutableResultSet()->mutable_rows()) {
53+
if (static_cast<ui64>(ResultSets_[resultSetIndex].rows_size()) >= ResultRowsLimit_) {
54+
ResultSets_[resultSetIndex].set_truncated(true);
55+
break;
56+
}
57+
58+
if (ResultSets_[resultSetIndex].ByteSizeLong() + row.ByteSizeLong() > ResultSizeLimit_) {
59+
ResultSets_[resultSetIndex].set_truncated(true);
60+
break;
61+
}
62+
63+
*ResultSets_[resultSetIndex].add_rows() = std::move(row);
64+
}
65+
*ResultSets_[resultSetIndex].mutable_columns() = ev->Get()->Record.GetResultSet().columns();
66+
}
67+
3968
Send(ev->Sender, response.Release());
4069
}
4170

@@ -47,13 +76,17 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
4776
private:
4877
THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> Request_;
4978
NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> Promise_;
50-
i64 ResultSizeLimit_;
79+
ui64 ResultRowsLimit_;
80+
ui64 ResultSizeLimit_;
81+
std::vector<Ydb::ResultSet>& ResultSets_;
5182
};
5283

5384
} // anonymous namespace
5485

55-
NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request, NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise, ui64 resultSizeLimit) {
56-
return new TRunScriptActorMock(std::move(request), promise, resultSizeLimit);
86+
NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
87+
NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise,
88+
ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets) {
89+
return new TRunScriptActorMock(std::move(request), promise, resultRowsLimit, resultSizeLimit, resultSets);
5790
}
5891

5992
} // namespace NKqpRun

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44

55
namespace NKqpRun {
66

7-
NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request, NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise, ui64 resultSizeLimit);
7+
NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
8+
NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise,
9+
ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets);
810

911
} // namespace NKqpRun

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

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,30 +12,36 @@ namespace NKqpRun {
1212
struct TYdbSetupSettings {
1313
TString DomainName = "Root";
1414

15-
bool TraceOpt = false;
15+
bool TraceOptEnabled = false;
1616
TMaybe<TString> LogOutputFile;
1717

1818
TString YqlToken;
19-
NKikimr::NMiniKQL::IFunctionRegistry* FunctionRegistry = nullptr;
19+
TIntrusivePtr<NKikimr::NMiniKQL::IMutableFunctionRegistry> FunctionRegistry = nullptr;
2020
NKikimrConfig::TAppConfig AppConfig;
2121
};
2222

2323

2424
struct TRunnerOptions {
25+
enum class ETraceOptType {
26+
Disabled,
27+
Scheme,
28+
Script,
29+
All,
30+
};
31+
2532
enum class EResultOutputFormat {
2633
RowsJson, // Rows in json format
2734
FullJson, // Columns, rows and types in json format
2835
};
2936

30-
i64 ResultsRowsLimit = 1000;
31-
3237
IOutputStream* ResultOutput = &Cout;
3338
IOutputStream* SchemeQueryAstOutput = nullptr;
3439
IOutputStream* ScriptQueryAstOutput = nullptr;
3540
IOutputStream* ScriptQueryPlanOutput = nullptr;
3641

3742
EResultOutputFormat ResultOutputFormat = EResultOutputFormat::RowsJson;
3843
NYdb::NConsoleClient::EOutputFormat PlanOutputFormat = NYdb::NConsoleClient::EOutputFormat::Default;
44+
ETraceOptType TraceOptType = ETraceOptType::Disabled;
3945

4046
TYdbSetupSettings YdbSettings;
4147
};

0 commit comments

Comments
 (0)