Support --test-param for specifying logging level in KQP tests (#1373)

3 files changed, 72 insertions, 24 deletions

diff --git a/ydb/core/kqp/ut/common/kqp_ut_common.cpp b/ydb/core/kqp/ut/common/kqp_ut_common.cpp
index ceda14be13..28022e3576 100644
--- a/ydb/core/kqp/ut/common/kqp_ut_common.cpp
+++ b/ydb/core/kqp/ut/common/kqp_ut_common.cpp
@@ -14,6 +14,8 @@
 #include <ydb/library/yql/minikql/invoke_builtins/mkql_builtins.h>
 #include <ydb/library/yql/utils/yql_panic.h>
 
+#include <library/cpp/testing/common/env.h>
+
 namespace NKikimr {
 namespace NKqp {
 
@@ -442,31 +444,75 @@ void TKikimrRunner::CreateSampleTables() {
 
 }
 
+static TMaybe<NActors::NLog::EPriority> ParseLogLevel(const TString& level) {
+    static const THashMap<TString, NActors::NLog::EPriority> levels = {
+        { "TRACE", NActors::NLog::PRI_TRACE },
+        { "DEBUG", NActors::NLog::PRI_DEBUG },
+        { "INFO", NActors::NLog::PRI_INFO },
+        { "NOTICE", NActors::NLog::PRI_NOTICE },
+        { "WARN", NActors::NLog::PRI_WARN },
+        { "ERROR", NActors::NLog::PRI_ERROR },
+        { "CRIT", NActors::NLog::PRI_CRIT },
+        { "ALERT", NActors::NLog::PRI_ALERT },
+        { "EMERG", NActors::NLog::PRI_EMERG },
+    };
+
+    TString l = level;
+    l.to_upper();
+    const auto levelIt = levels.find(l);
+    if (levelIt != levels.end()) {
+        return levelIt->second;
+    } else {
+        Cerr << "Failed to parse test log level [" << level << "]" << Endl;
+        return Nothing();
+    }
+}
+
+void TKikimrRunner::SetupLogLevelFromTestParam(NKikimrServices::EServiceKikimr service) {
+    if (const TString paramForService = GetTestParam(TStringBuilder() << "KQP_LOG_" << NKikimrServices::EServiceKikimr_Name(service))) {
+        if (const TMaybe<NActors::NLog::EPriority> level = ParseLogLevel(paramForService)) {
+            Server->GetRuntime()->SetLogPriority(service, *level);
+            return;
+        }
+    }
+    if (const TString commonParam = GetTestParam("KQP_LOG")) {
+        if (const TMaybe<NActors::NLog::EPriority> level = ParseLogLevel(commonParam)) {
+            Server->GetRuntime()->SetLogPriority(service, *level);
+        }
+    }
+}
+
 void TKikimrRunner::Initialize(const TKikimrSettings& settings) {
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::FLAT_TX_SCHEMESHARD, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_YQL, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_YQL, NActors::NLog::PRI_INFO);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::TX_DATASHARD, NActors::NLog::PRI_TRACE);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::TX_COORDINATOR, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_COMPUTE, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_TASKS_RUNNER, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_EXECUTER, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::TX_PROXY_SCHEME_CACHE, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::SCHEME_BOARD_REPLICA, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_WORKER, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_SESSION, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::TABLET_EXECUTOR, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_SLOW_LOG, NActors::NLog::PRI_TRACE);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_PROXY, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_COMPILE_SERVICE, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_COMPILE_ACTOR, NActors::NLog::PRI_TRACE);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_COMPILE_REQUEST, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_GATEWAY, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::RPC_REQUEST, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_RESOURCE_MANAGER, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_NODE, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::KQP_BLOBS_STORAGE, NActors::NLog::PRI_DEBUG);
-    // Server->GetRuntime()->SetLogPriority(NKikimrServices::TX_COLUMNSHARD, NActors::NLog::PRI_INFO);
+    // You can enable logging for these services in test using test option:
+    // `--test-param KQP_LOG=<level>`
+    // or `--test-param KQP_LOG_<service>=<level>`
+    // For example:
+    // --test-param KQP_LOG=TRACE
+    // --test-param KQP_LOG_FLAT_TX_SCHEMESHARD=debug
+    SetupLogLevelFromTestParam(NKikimrServices::FLAT_TX_SCHEMESHARD);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_YQL);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_YQL);
+    SetupLogLevelFromTestParam(NKikimrServices::TX_DATASHARD);
+    SetupLogLevelFromTestParam(NKikimrServices::TX_COORDINATOR);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_COMPUTE);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_TASKS_RUNNER);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_EXECUTER);
+    SetupLogLevelFromTestParam(NKikimrServices::TX_PROXY_SCHEME_CACHE);
+    SetupLogLevelFromTestParam(NKikimrServices::SCHEME_BOARD_REPLICA);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_WORKER);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_SESSION);
+    SetupLogLevelFromTestParam(NKikimrServices::TABLET_EXECUTOR);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_SLOW_LOG);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_PROXY);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_COMPILE_SERVICE);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_COMPILE_ACTOR);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_COMPILE_REQUEST);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_GATEWAY);
+    SetupLogLevelFromTestParam(NKikimrServices::RPC_REQUEST);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_RESOURCE_MANAGER);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_NODE);
+    SetupLogLevelFromTestParam(NKikimrServices::KQP_BLOBS_STORAGE);
+    SetupLogLevelFromTestParam(NKikimrServices::TX_COLUMNSHARD);
 
     RunCall([this, domain = settings.DomainRoot]{
         this->Client->InitRootScheme(domain);
diff --git a/ydb/core/kqp/ut/common/kqp_ut_common.h b/ydb/core/kqp/ut/common/kqp_ut_common.h
index 2f05c0d608..10b01f1581 100644
--- a/ydb/core/kqp/ut/common/kqp_ut_common.h
+++ b/ydb/core/kqp/ut/common/kqp_ut_common.h
@@ -182,6 +182,7 @@ private:
     void Initialize(const TKikimrSettings& settings);
     void WaitForKqpProxyInit();
     void CreateSampleTables();
+    void SetupLogLevelFromTestParam(NKikimrServices::EServiceKikimr service);
 
 private:
     THolder<Tests::TServerSettings> ServerSettings;
diff --git a/ydb/core/kqp/ut/common/ya.make b/ydb/core/kqp/ut/common/ya.make
index d42a4a94e5..97281a7aad 100644
--- a/ydb/core/kqp/ut/common/ya.make
+++ b/ydb/core/kqp/ut/common/ya.make
@@ -11,6 +11,7 @@ SRCS(
 )
 
 PEERDIR(
+    library/cpp/testing/common
     ydb/core/kqp/federated_query
     ydb/core/testlib
     ydb/library/yql/public/udf