Added log settings flags

GrigoriyPA · GrigoriyPA · commit aa4e3d5f7105 · 2025-01-06T08:42:30.000Z
diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp
@@ -1,7 +1,5 @@
 #include "src/kqp_runner.h"
 
-#include <cstdio>
-
 #include <contrib/libs/protobuf/src/google/protobuf/text_format.h>
 
 #include <library/cpp/colorizer/colors.h>
@@ -15,11 +13,13 @@
 
 #include <ydb/library/aclib/aclib.h>
 #include <ydb/library/yaml_config/yaml_config.h>
+
 #include <yql/essentials/minikql/invoke_builtins/mkql_builtins.h>
+#include <yql/essentials/public/udf/udf_static_registry.h>
+
 #include <yt/yql/providers/yt/gateway/file/yql_yt_file.h>
 #include <yt/yql/providers/yt/gateway/file/yql_yt_file_comp_nodes.h>
 #include <yt/yql/providers/yt/lib/yt_download/yt_download.h>
-#include <yql/essentials/public/udf/udf_static_registry.h>
 
 
 namespace NKqpRun {
@@ -135,7 +135,7 @@ struct TExecutionOptions {
     void ValidateOptionsSizes(const TRunnerOptions& runnerOptions) const {
         const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName) {
             if (checkSize > numberQueries) {
-                ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of queries is " << numberQueries;
+                ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of script queries is " << numberQueries;
             }
         };
 
@@ -431,6 +431,9 @@ class TMain : public TMainClassArgs {
     bool ExcludeLinkedUdfs = false;
     bool EmulateYt = false;
 
+    std::optional<NActors::NLog::EPriority> DefaultLogPriority;
+    std::unordered_map<NKikimrServices::EServiceKikimr, NActors::NLog::EPriority> LogPriorities;
+
     static TString LoadFile(const TString& file) {
         return TFileInput(file).ReadAll();
     }
@@ -466,6 +469,10 @@ class TMain : public TMainClassArgs {
             return choices;
         }
 
+        bool Contains(const TString& choice) const {
+            return ChoicesMap.contains(choice);
+        }
+
     private:
         const std::map<TString, TResult> ChoicesMap;
     };
@@ -483,11 +490,13 @@ class TMain : public TMainClassArgs {
             .Handler1([this](const NLastGetopt::TOptsParser* option) {
                 ExecutionOptions.SchemeQuery = LoadFile(option->CurVal());
             });
+
         options.AddLongOption('p', "script-query", "Script query to execute (typically DML query)")
             .RequiredArgument("file")
             .Handler1([this](const NLastGetopt::TOptsParser* option) {
                 ExecutionOptions.ScriptQueries.emplace_back(LoadFile(option->CurVal()));
             });
+
         options.AddLongOption("templates", "Enable templates for -s and -p queries, such as ${YQL_TOKEN} and ${QUERY_ID}")
             .NoArgument()
             .SetFlag(&ExecutionOptions.UseTemplates);
@@ -499,10 +508,10 @@ class TMain : public TMainClassArgs {
                 TStringBuf filePath;
                 TStringBuf(option->CurVal()).Split('@', tableName, filePath);
                 if (tableName.empty() || filePath.empty()) {
-                    ythrow yexception() << "Incorrect table mapping, expected form table@file, e.g. yt.Root/plato.Input@input.txt";
+                    ythrow yexception() << "Incorrect table mapping, expected form table@file, e. g. yt.Root/plato.Input@input.txt";
                 }
                 if (TablesMapping.contains(tableName)) {
-                    ythrow yexception() << "Got duplicate table name: " << tableName;
+                    ythrow yexception() << "Got duplicated table name: " << tableName;
                 }
                 TablesMapping[tableName] = filePath;
             });
@@ -523,9 +532,11 @@ class TMain : public TMainClassArgs {
         options.AddLongOption('u', "udf", "Load shared library with UDF by given path")
             .RequiredArgument("file")
             .EmplaceTo(&UdfsPaths);
+
         options.AddLongOption("udfs-dir", "Load all shared libraries with UDFs found in given directory")
             .RequiredArgument("directory")
             .StoreResult(&UdfsDirectory);
+
         options.AddLongOption("exclude-linked-udfs", "Exclude linked udfs when same udf passed from -u or --udfs-dir")
             .NoArgument()
             .SetFlag(&ExcludeLinkedUdfs);
@@ -540,13 +551,51 @@ class TMain : public TMainClassArgs {
                     std::remove(file.c_str());
                 }
             });
+
+        TChoices<NActors::NLog::EPriority> logPriority({
+            {"emerg", NActors::NLog::EPriority::PRI_EMERG},
+            {"alert", NActors::NLog::EPriority::PRI_ALERT},
+            {"crit", NActors::NLog::EPriority::PRI_CRIT},
+            {"error", NActors::NLog::EPriority::PRI_ERROR},
+            {"warn", NActors::NLog::EPriority::PRI_WARN},
+            {"notice", NActors::NLog::EPriority::PRI_NOTICE},
+            {"info", NActors::NLog::EPriority::PRI_INFO},
+            {"debug", NActors::NLog::EPriority::PRI_DEBUG},
+            {"trace", NActors::NLog::EPriority::PRI_TRACE},
+        });
+        options.AddLongOption("log-default", "Default log priority")
+            .RequiredArgument("priority")
+            .Choices(logPriority.GetChoices())
+            .StoreMappedResultT<TString>(&DefaultLogPriority, logPriority);
+
+        options.AddLongOption("log", "Component log priority in format <component>=<priority> (e. g. KQP_YQL=trace)")
+            .RequiredArgument("component priority")
+            .Handler1([this, logPriority](const NLastGetopt::TOptsParser* option) {
+                TStringBuf component;
+                TStringBuf priority;
+                TStringBuf(option->CurVal()).Split('=', component, priority);
+                if (component.empty() || priority.empty()) {
+                    ythrow yexception() << "Incorrect log setting, expected form component=priority, e. g. KQP_YQL=trace";
+                }
+
+                const auto service = GetLogService(TString(component));
+                if (LogPriorities.contains(service)) {
+                    ythrow yexception() << "Got duplicated log service name: " << component;
+                }
+
+                if (!logPriority.Contains(TString(priority))) {
+                    ythrow yexception() << "Incorrect log priority: " << priority;
+                }
+                LogPriorities[service] = logPriority(TString(priority));
+            });
+
         TChoices<TRunnerOptions::ETraceOptType> traceOpt({
             {"all", TRunnerOptions::ETraceOptType::All},
             {"scheme", TRunnerOptions::ETraceOptType::Scheme},
             {"script", TRunnerOptions::ETraceOptType::Script},
             {"disabled", TRunnerOptions::ETraceOptType::Disabled}
         });
-        options.AddLongOption('T', "trace-opt", "print AST in the begin of each transformation (use script@<query id> for tracing one -p query)")
+        options.AddLongOption('T', "trace-opt", "Print AST in the begin of each transformation (use script@<query id> for tracing one -p query)")
             .RequiredArgument("trace-opt-query")
             .DefaultValue("disabled")
             .Choices(traceOpt.GetChoices())
@@ -555,9 +604,11 @@ class TMain : public TMainClassArgs {
                 RunnerOptions.YdbSettings.TraceOptEnabled = traceOptType != NKqpRun::TRunnerOptions::ETraceOptType::Disabled;
                 return traceOptType;
             });
-        options.AddLongOption('I', "trace-opt-index", "index of -p query to use --trace-opt, starts from zero")
+
+        options.AddLongOption('I', "trace-opt-index", "Index of -p query to use --trace-opt, starts from zero")
             .RequiredArgument("uint")
             .StoreResult(&RunnerOptions.TraceOptScriptId);
+
         options.AddLongOption("trace-id", "Trace id for -p queries")
             .RequiredArgument("id")
             .EmplaceTo(&ExecutionOptions.TraceIds);
@@ -566,10 +617,12 @@ class TMain : public TMainClassArgs {
             .RequiredArgument("file")
             .DefaultValue("-")
             .StoreMappedResultT<TString>(&RunnerOptions.ResultOutput, &GetDefaultOutput);
+
         options.AddLongOption('L', "result-rows-limit", "Rows limit for script execution results")
             .RequiredArgument("uint")
             .DefaultValue(0)
             .StoreResult(&ExecutionOptions.ResultsRowsLimit);
+
         TChoices<TRunnerOptions::EResultOutputFormat> resultFormat({
             {"rows", TRunnerOptions::EResultOutputFormat::RowsJson},
             {"full-json", TRunnerOptions::EResultOutputFormat::FullJson},
@@ -596,6 +649,7 @@ class TMain : public TMainClassArgs {
             .Handler1([this](const NLastGetopt::TOptsParser* option) {
                 RunnerOptions.ScriptQueryPlanOutputs.emplace_back(GetDefaultOutput(TString(option->CurValOrDef())));
             });
+
         options.AddLongOption("script-statistics", "File with script inprogress statistics")
             .RequiredArgument("file")
             .Handler1([this](const NLastGetopt::TOptsParser* option) {
@@ -605,6 +659,7 @@ class TMain : public TMainClassArgs {
                 }
                 RunnerOptions.InProgressStatisticsOutputFiles.emplace_back(file);
             });
+
         TChoices<NYdb::NConsoleClient::EDataFormat> planFormat({
             {"pretty", NYdb::NConsoleClient::EDataFormat::Pretty},
             {"table", NYdb::NConsoleClient::EDataFormat::PrettyTable},
@@ -641,19 +696,21 @@ class TMain : public TMainClassArgs {
                 TString choice(option->CurValOrDef());
                 ExecutionOptions.ExecutionCases.emplace_back(executionCase(choice));
             });
+
         options.AddLongOption("inflight-limit", "In flight limit for async queries (use 0 for unlimited)")
             .RequiredArgument("uint")
             .DefaultValue(0)
             .StoreResult(&RunnerOptions.YdbSettings.AsyncQueriesSettings.InFlightLimit);
-        TChoices<TAsyncQueriesSettings::EVerbose> verbose({
-            {"each-query", TAsyncQueriesSettings::EVerbose::EachQuery},
-            {"final", TAsyncQueriesSettings::EVerbose::Final}
-        });
 
         options.AddLongOption("verbose", "Common verbose level (max level 2)")
             .RequiredArgument("uint")
             .DefaultValue(1)
             .StoreResult(&RunnerOptions.YdbSettings.VerboseLevel);
+
+        TChoices<TAsyncQueriesSettings::EVerbose> verbose({
+            {"each-query", TAsyncQueriesSettings::EVerbose::EachQuery},
+            {"final", TAsyncQueriesSettings::EVerbose::Final}
+        });
         options.AddLongOption("async-verbose", "Verbose type for async queries")
             .RequiredArgument("type")
             .DefaultValue("each-query")
@@ -690,10 +747,12 @@ class TMain : public TMainClassArgs {
             .RequiredArgument("uint")
             .DefaultValue(ExecutionOptions.LoopCount)
             .StoreResult(&ExecutionOptions.LoopCount);
+
         options.AddLongOption("loop-delay", "Delay in milliseconds between loop steps")
             .RequiredArgument("uint")
             .DefaultValue(0)
             .StoreMappedResultT<ui64>(&ExecutionOptions.LoopDelay, &TDuration::MilliSeconds<ui64>);
+
         options.AddLongOption("continue-after-fail", "Don't not stop requests execution after fails")
             .NoArgument()
             .SetFlag(&ExecutionOptions.ContinueAfterFail);
@@ -803,12 +862,19 @@ class TMain : public TMainClassArgs {
 
         RunnerOptions.YdbSettings.YqlToken = YqlToken;
         RunnerOptions.YdbSettings.FunctionRegistry = CreateFunctionRegistry(UdfsDirectory, UdfsPaths, ExcludeLinkedUdfs).Get();
+
+        auto& appConfig = RunnerOptions.YdbSettings.AppConfig;
         if (ExecutionOptions.ResultsRowsLimit) {
-            RunnerOptions.YdbSettings.AppConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(ExecutionOptions.ResultsRowsLimit);
+            appConfig.MutableQueryServiceConfig()->SetScriptResultRowsLimit(ExecutionOptions.ResultsRowsLimit);
+        }
+
+        if (DefaultLogPriority) {
+            appConfig.MutableLogConfig()->SetDefaultLevel(*DefaultLogPriority);
         }
+        ModifyLogPriorities(LogPriorities, *appConfig.MutableLogConfig());
 
         if (EmulateYt) {
-            const auto& fileStorageConfig = RunnerOptions.YdbSettings.AppConfig.GetQueryServiceConfig().GetFileStorage();
+            const auto& fileStorageConfig = appConfig.GetQueryServiceConfig().GetFileStorage();
             auto fileStorage = WithAsync(CreateFileStorage(fileStorageConfig, {MakeYtDownloader(fileStorageConfig)}));
             auto ytFileServices = NYql::NFile::TYtFileServices::Make(RunnerOptions.YdbSettings.FunctionRegistry.Get(), TablesMapping, fileStorage);
             RunnerOptions.YdbSettings.YtGateway = NYql::CreateYtFileGateway(ytFileServices);
diff --git a/ydb/tests/tools/kqprun/src/common.cpp b/ydb/tests/tools/kqprun/src/common.cpp
@@ -0,0 +1,29 @@
+#include "common.h"
+
+
+namespace NKqpRun {
+
+NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName) {
+    NKikimrServices::EServiceKikimr service;
+    if (!NKikimrServices::EServiceKikimr_Parse(serviceName, &service)) {
+        ythrow yexception() << "Invalid kikimr service name " << serviceName;
+    }
+    return service;
+}
+
+void ModifyLogPriorities(std::unordered_map<NKikimrServices::EServiceKikimr, NActors::NLog::EPriority> logPriorities, NKikimrConfig::TLogConfig& logConfig) {
+    for (auto& entry : *logConfig.MutableEntry()) {
+        const auto it = logPriorities.find(GetLogService(entry.GetComponent()));
+        if (it != logPriorities.end()) {
+            entry.SetLevel(it->second);
+            logPriorities.erase(it);
+        }
+    }
+    for (const auto& [service, priority] : logPriorities) {
+        auto* entry = logConfig.AddEntry();
+        entry->SetComponent(NKikimrServices::EServiceKikimr_Name(service));
+        entry->SetLevel(priority);
+    }
+}
+
+}  // namespace NKqpRun
diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h
@@ -1,14 +1,18 @@
 #pragma once
 
 #include <ydb/core/protos/config.pb.h>
-#include <ydb/public/api/protos/ydb_cms.pb.h>
 #include <ydb/core/protos/kqp.pb.h>
 
+#include <ydb/library/actors/core/log_iface.h>
+#include <ydb/library/services/services.pb.h>
+
+#include <ydb/public/api/protos/ydb_cms.pb.h>
+#include <ydb/public/lib/ydb_cli/common/formats.h>
+
 #include <yql/essentials/minikql/computation/mkql_computation_node.h>
 #include <yql/essentials/minikql/mkql_function_registry.h>
-#include <yt/yql/providers/yt/provider/yql_yt_gateway.h>
 
-#include <ydb/public/lib/ydb_cli/common/formats.h>
+#include <yt/yql/providers/yt/provider/yql_yt_gateway.h>
 
 
 namespace NKqpRun {
@@ -101,11 +105,14 @@ struct TRequestOptions {
 };
 
 template <typename TValue>
-static TValue GetValue(size_t index, const std::vector<TValue>& values, TValue defaultValue) {
+TValue GetValue(size_t index, const std::vector<TValue>& values, TValue defaultValue) {
     if (values.empty()) {
         return defaultValue;
     }
     return values[std::min(index, values.size() - 1)];
 }
 
+NKikimrServices::EServiceKikimr GetLogService(const TString& serviceName);
+void ModifyLogPriorities(std::unordered_map<NKikimrServices::EServiceKikimr, NActors::NLog::EPriority> logPriorities, NKikimrConfig::TLogConfig& logConfig);
+
 }  // namespace NKqpRun
diff --git a/ydb/tests/tools/kqprun/src/ya.make b/ydb/tests/tools/kqprun/src/ya.make
@@ -2,6 +2,7 @@ LIBRARY()
 
 SRCS(
     actors.cpp
+    common.cpp
     kqp_runner.cpp
     ydb_setup.cpp
 )
diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
@@ -150,13 +150,8 @@ class TYdbSetup::TImpl {
                 }
             }
 
-            for (auto setting : Settings_.AppConfig.GetLogConfig().get_arr_entry()) {
-                NKikimrServices::EServiceKikimr service;
-                if (!NKikimrServices::EServiceKikimr_Parse(setting.GetComponent(), &service)) {
-                    ythrow yexception() << "Invalid kikimr service name " << setting.GetComponent();
-                }
-
-                runtime.SetLogPriority(service, NActors::NLog::EPriority(setting.GetLevel()));
+            for (const auto& setting : Settings_.AppConfig.GetLogConfig().get_arr_entry()) {
+                runtime.SetLogPriority(GetLogService(setting.GetComponent()), NActors::NLog::EPriority(setting.GetLevel()));
             }
 
             runtime.SetLogBackendFactory([this]() { return CreateLogBackend(); });
@@ -316,21 +311,7 @@ class TYdbSetup::TImpl {
             return;
         }
 
-        bool found = false;
-        for (auto& entry : *Settings_.AppConfig.MutableLogConfig()->MutableEntry()) {
-            if (entry.GetComponent() == "KQP_YQL") {
-                entry.SetLevel(NActors::NLog::PRI_TRACE);
-                found = true;
-                break;
-            }
-        }
-
-        if (!found) {
-            auto entry = Settings_.AppConfig.MutableLogConfig()->AddEntry();
-            entry->SetComponent("KQP_YQL");
-            entry->SetLevel(NActors::NLog::PRI_TRACE);
-        }
-
+        ModifyLogPriorities({{NKikimrServices::EServiceKikimr::KQP_YQL, NActors::NLog::PRI_TRACE}}, *Settings_.AppConfig.MutableLogConfig());
         NYql::NLog::InitLogger(NActors::CreateNullBackend());
     }
 

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,7 @@ LIBRARY()`
`2`	`2`
`3`	`3`	`SRCS(`
`4`	`4`	`actors.cpp`
	`5`	`+ common.cpp`
`5`	`6`	`kqp_runner.cpp`
`6`	`7`	`ydb_setup.cpp`
`7`	`8`	`)`