Support of mrjob in dqrun (#4243)

Author: vitstn

ydb/library/yql/tools/dqrun/dqrun.cpp

@@ -84,6 +84,7 @@
 #include <library/cpp/getopt/last_getopt.h>
 #include <library/cpp/logger/priority.h>
 #include <library/cpp/protobuf/util/pb_io.h>
+#include <library/cpp/digest/md5/md5.h>
 #include <ydb/library/actors/http/http_proxy.h>
 
 #include <util/generic/string.h>
@@ -164,6 +165,28 @@ void ReadGatewaysConfig(const TString& configFile, TGatewaysConfig* config, THas
     }
 }
 
+void PatchGatewaysConfig(TGatewaysConfig* config, const TString& mrJobBin, const TString& mrJobUdfsDir,
+    size_t numThreads, bool keepTemp)
+{
+    auto ytConfig = config->MutableYt();
+    ytConfig->SetGatewayThreads(numThreads);
+    if (mrJobBin.empty()) {
+        ytConfig->ClearMrJobBin();
+    } else {
+        ytConfig->SetMrJobBin(mrJobBin);
+        ytConfig->SetMrJobBinMd5(MD5::File(mrJobBin));
+    }
+
+    if (mrJobUdfsDir.empty()) {
+        ytConfig->ClearMrJobUdfsDir();
+    } else {
+        ytConfig->SetMrJobUdfsDir(mrJobUdfsDir);
+    }
+    auto attr = ytConfig->MutableDefaultSettings()->Add();
+    attr->SetName("KeepTempTables");
+    attr->SetValue(keepTemp ? "yes" : "no");
+}
+
 TFileStoragePtr CreateFS(const TString& paramsFile, const TString& defYtServer) {
     TFileStorageConfig params;
     LoadFsConfigFromFile(paramsFile ? paramsFile : "../../../../../yql/cfg/local/fs.conf", params);
@@ -454,6 +477,9 @@ int RunMain(int argc, const char* argv[])
     int verbosity = 3;
     bool showLog = false;
     bool emulateYt = false;
+    TString mrJobBin;
+    TString mrJobUdfsDir;
+    size_t numYtThreads = 1;
     TString token = GetEnv("YQL_TOKEN");
     if (!token) {
         TString home = GetEnv("HOME");
@@ -550,6 +576,16 @@ int RunMain(int argc, const char* argv[])
         .Optional()
         .NoArgument()
         .SetFlag(&udfResolverFilterSyscalls);
+    opts.AddLongOption("mrjob-bin", "Path to mrjob binary")
+        .Optional()
+        .StoreResult(&mrJobBin);
+    opts.AddLongOption("mrjob-udfsdir", "Path to udfs for mr jobs")
+        .Optional()
+        .StoreResult(&mrJobUdfsDir);
+    opts.AddLongOption("yt-threads", "YT gateway threads")
+        .Optional()
+        .RequiredArgument("COUNT")
+        .StoreResult(&numYtThreads);
     opts.AddLongOption('v', "verbosity", "Log verbosity level")
         .Optional()
         .RequiredArgument("LEVEL")
@@ -737,6 +773,7 @@ int RunMain(int argc, const char* argv[])
 
     TGatewaysConfig gatewaysConfig;
     ReadGatewaysConfig(gatewaysCfgFile, &gatewaysConfig, sqlFlags);
+    PatchGatewaysConfig(&gatewaysConfig, mrJobBin, mrJobUdfsDir, numYtThreads, res.Has("keep-temp"));
     if (runOptions.AnalyzeQuery) {
         auto* setting = gatewaysConfig.MutableDq()->AddDefaultSettings();
         setting->SetName("AnalyzeQuery");

ydb/library/yql/tools/dqrun/ya.make

@@ -25,6 +25,7 @@ ENDIF()
         library/cpp/logger
         library/cpp/resource
         library/cpp/yson
+        library/cpp/digest/md5
         yt/cpp/mapreduce/interface
         ydb/library/yql/sql/pg
         ydb/library/yql/core/facade