Add EnableAntlr4Parser to query replay (#9471)

8 files changed, 50 insertions, 22 deletions

diff --git a/ydb/core/kqp/host/kqp_host.cpp b/ydb/core/kqp/host/kqp_host.cpp
index 73171d05e1..c98dc7453a 100644
--- a/ydb/core/kqp/host/kqp_host.cpp
+++ b/ydb/core/kqp/host/kqp_host.cpp
@@ -1230,7 +1230,8 @@ private:
                 .SetIsEnablePgConstsToParams(SessionCtx->Config().EnablePgConstsToParams)
                 .SetQueryParameters(query.ParameterTypes)
                 .SetApplicationName(ApplicationName)
-                .SetIsEnablePgSyntax(SessionCtx->Config().FeatureFlags.GetEnablePgSyntax());
+                .SetIsEnablePgSyntax(SessionCtx->Config().FeatureFlags.GetEnablePgSyntax())
+                .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
             NSQLTranslation::TTranslationSettings effectiveSettings;
             auto astRes = ParseQuery(
                 query.Text,
@@ -1298,7 +1299,8 @@ private:
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, query.Text, SessionCtx->Config().BindingsMode, GUCSettings);
         settingsBuilder
             .SetSqlAutoCommit(false)
-            .SetUsePgParser(settings.UsePgParser);
+            .SetUsePgParser(settings.UsePgParser)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(query, /* isSql */ true, *ExprCtx, sqlVersion, settingsBuilder);
 
         return TSplitResult{
@@ -1392,7 +1394,8 @@ private:
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, query.Text, SessionCtx->Config().BindingsMode, GUCSettings);
         settingsBuilder.SetSqlAutoCommit(false)
-            .SetUsePgParser(settings.UsePgParser);
+            .SetUsePgParser(settings.UsePgParser)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(query, isSql, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1452,7 +1455,8 @@ private:
 
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, query.Text, SessionCtx->Config().BindingsMode, GUCSettings);
-        settingsBuilder.SetSqlAutoCommit(false);
+        settingsBuilder.SetSqlAutoCommit(false)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(query, /* isSql */ true, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1480,7 +1484,8 @@ private:
 
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, queryAst.Text, SessionCtx->Config().BindingsMode, GUCSettings);
-        settingsBuilder.SetSqlAutoCommit(false);
+        settingsBuilder.SetSqlAutoCommit(false)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(queryAst, false, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1526,7 +1531,8 @@ private:
         if (!expr) {
             TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, query.Text, SessionCtx->Config().BindingsMode, GUCSettings);
             settingsBuilder.SetSqlAutoCommit(false)
-                .SetUsePgParser(settings.UsePgParser);
+                .SetUsePgParser(settings.UsePgParser)
+                .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
             auto compileResult = CompileYqlQuery(query, /* isSql */ true, ctx, sqlVersion, settingsBuilder);
             if (compileResult.QueryExprs.empty()) {
                 return nullptr;
@@ -1563,7 +1569,8 @@ private:
 
         TMaybe<TSqlVersion> sqlVersion = 1;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, query.Text, SessionCtx->Config().BindingsMode, GUCSettings);
-        settingsBuilder.SetSqlAutoCommit(false);
+        settingsBuilder.SetSqlAutoCommit(false)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(query, true, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1584,7 +1591,8 @@ private:
 
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, queryAst.Text, SessionCtx->Config().BindingsMode, GUCSettings);
-        settingsBuilder.SetSqlAutoCommit(false);
+        settingsBuilder.SetSqlAutoCommit(false)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(queryAst, false, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1611,7 +1619,8 @@ private:
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, script.Text, SessionCtx->Config().BindingsMode, GUCSettings);
         settingsBuilder.SetSqlAutoCommit(true)
-            .SetUsePgParser(settings.UsePgParser);
+            .SetUsePgParser(settings.UsePgParser)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(script, true, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1640,7 +1649,8 @@ private:
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, script.Text, SessionCtx->Config().BindingsMode, GUCSettings);
         settingsBuilder.SetSqlAutoCommit(true)
-            .SetUsePgParser(settings.UsePgParser);
+            .SetUsePgParser(settings.UsePgParser)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(script, true, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1664,7 +1674,8 @@ private:
 
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, script.Text, SessionCtx->Config().BindingsMode, GUCSettings);
-        settingsBuilder.SetSqlAutoCommit(true);
+        settingsBuilder.SetSqlAutoCommit(true)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(script, true, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
@@ -1692,7 +1703,8 @@ private:
 
         TMaybe<TSqlVersion> sqlVersion;
         TKqpTranslationSettingsBuilder settingsBuilder(SessionCtx->Query().Type, SessionCtx->Config()._KqpYqlSyntaxVersion.Get().GetRef(), Cluster, script.Text, SessionCtx->Config().BindingsMode, GUCSettings);
-        settingsBuilder.SetSqlAutoCommit(true);
+        settingsBuilder.SetSqlAutoCommit(true)
+            .SetIsEnableAntlr4Parser(SessionCtx->Config().EnableAntlr4Parser);
         auto compileResult = CompileYqlQuery(script, true, ctx, sqlVersion, settingsBuilder);
         if (compileResult.QueryExprs.empty()) {
             return nullptr;
diff --git a/ydb/core/kqp/host/kqp_translate.cpp b/ydb/core/kqp/host/kqp_translate.cpp
index 8b7d842f88..cd0022be50 100644
--- a/ydb/core/kqp/host/kqp_translate.cpp
+++ b/ydb/core/kqp/host/kqp_translate.cpp
@@ -64,6 +64,10 @@ NSQLTranslation::TTranslationSettings TKqpTranslationSettingsBuilder::Build(NYql
         settings.AutoParametrizeValuesStmt = IsEnablePgConstsToParams;
     }
 
+    if (!settings.PgParser) {
+        settings.Antlr4Parser = IsEnableAntlr4Parser;
+    }
+
     if (QueryType == NYql::EKikimrQueryType::Scan || QueryType == NYql::EKikimrQueryType::Query) {
         SqlVersion = SqlVersion ? *SqlVersion : 1;
     }
diff --git a/ydb/core/kqp/host/kqp_translate.h b/ydb/core/kqp/host/kqp_translate.h
index d71c01d4ce..f4de2d3f12 100644
--- a/ydb/core/kqp/host/kqp_translate.h
+++ b/ydb/core/kqp/host/kqp_translate.h
@@ -67,6 +67,11 @@ public:
         return *this;
     }
 
+    TKqpTranslationSettingsBuilder& SetIsEnableAntlr4Parser(bool value) {
+        IsEnableAntlr4Parser = value;
+        return *this;
+    }
+
 private:
     const NYql::EKikimrQueryType QueryType;
     const ui16 KqpYqlSyntaxVersion;
@@ -79,6 +84,7 @@ private:
     bool IsEnableExternalDataSources = false;
     bool IsEnablePgConstsToParams = false;
     bool IsEnablePgSyntax = false;
+    bool IsEnableAntlr4Parser = false;
     TMaybe<bool> SqlAutoCommit = {};
     TGUCSettings::TPtr GUCSettings;
     TMaybe<TString> ApplicationName = {};
diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.h b/ydb/core/kqp/provider/yql_kikimr_settings.h
index 8e6e706de4..10eecc0fbc 100644
--- a/ydb/core/kqp/provider/yql_kikimr_settings.h
+++ b/ydb/core/kqp/provider/yql_kikimr_settings.h
@@ -171,6 +171,7 @@ struct TKikimrConfiguration : public TKikimrSettings, public NCommon::TSettingDi
     ui32 DefaultCostBasedOptimizationLevel = 4;
     bool EnableConstantFolding = true;
     ui64 DefaultEnableSpillingNodes = 0;
+    bool EnableAntlr4Parser = false;
 
     void SetDefaultEnabledSpillingNodes(const TString& node);
     ui64 GetEnabledSpillingNodes() const;
diff --git a/ydb/tools/query_replay_yt/main.cpp b/ydb/tools/query_replay_yt/main.cpp
index 851de48612..15ca12154d 100644
--- a/ydb/tools/query_replay_yt/main.cpp
+++ b/ydb/tools/query_replay_yt/main.cpp
@@ -57,6 +57,7 @@ class TQueryReplayMapper
     TVector<TString> UdfFiles;
     ui32 ActorSystemThreadsCount = 5;
     NActors::NLog::EPriority YqlLogPriority = NActors::NLog::EPriority::PRI_ERROR;
+    bool EnableAntlr4Parser;
 
 public:
     static TString GetFailReason(const TQueryReplayEvents::TCheckQueryPlanStatus& status) {
@@ -93,13 +94,14 @@ public:
 public:
     TQueryReplayMapper() = default;
 
-    Y_SAVELOAD_JOB(UdfFiles, ActorSystemThreadsCount, YqlLogPriority);
+    Y_SAVELOAD_JOB(UdfFiles, ActorSystemThreadsCount, EnableAntlr4Parser, YqlLogPriority);
 
-    TQueryReplayMapper(TVector<TString> udfFiles, ui32 actorSystemThreadsCount,
+    TQueryReplayMapper(TVector<TString> udfFiles, ui32 actorSystemThreadsCount, bool enableAntlr4Parser,
         NActors::NLog::EPriority yqlLogPriority = NActors::NLog::EPriority::PRI_ERROR)
         : UdfFiles(udfFiles)
         , ActorSystemThreadsCount(actorSystemThreadsCount)
         , YqlLogPriority(yqlLogPriority)
+        , EnableAntlr4Parser(enableAntlr4Parser)
     {}
 
     void Start(NYT::TTableWriter<NYT::TNode>*) override {
@@ -147,12 +149,12 @@ public:
             return nullptr;
         }
 
-        auto compileActorId = ActorSystem->Register(CreateQueryCompiler(ModuleResolverState, FunctionRegistry.Get(), HttpGateway));
+        auto compileActorId = ActorSystem->Register(CreateQueryCompiler(ModuleResolverState, FunctionRegistry.Get(), HttpGateway, EnableAntlr4Parser));
 
         auto future = ActorSystem->Ask<TQueryReplayEvents::TEvCompileResponse>(
             compileActorId,
             THolder(new TQueryReplayEvents::TEvCompileRequest(std::move(json))),
-            TDuration::Seconds(100));
+            TDuration::Seconds(300));
 
         return future.ExtractValueSync();
     }
@@ -236,7 +238,7 @@ int main(int argc, const char** argv) {
     TQueryReplayConfig config;
     config.ParseConfig(argc, argv);
     if (config.QueryFile) {
-        auto fakeMapper = TQueryReplayMapper(config.UdfFiles, config.ActorSystemThreadsCount, config.YqlLogLevel);
+        auto fakeMapper = TQueryReplayMapper(config.UdfFiles, config.ActorSystemThreadsCount, config.EnableAntlr4Parser, config.YqlLogLevel);
         fakeMapper.Start(nullptr);
         Y_DEFER {
             fakeMapper.Finish(nullptr);
@@ -281,7 +283,7 @@ int main(int argc, const char** argv) {
         spec.CoreTablePath(config.CoreTablePath);
     }
 
-    client->Map(spec, new TQueryReplayMapper(config.UdfFiles, config.ActorSystemThreadsCount, config.YqlLogLevel));
+    client->Map(spec, new TQueryReplayMapper(config.UdfFiles, config.ActorSystemThreadsCount, config.EnableAntlr4Parser, config.YqlLogLevel));
 
     return EXIT_SUCCESS;
 }
diff --git a/ydb/tools/query_replay_yt/query_compiler.cpp b/ydb/tools/query_replay_yt/query_compiler.cpp
index b32255ecac..23290a3e83 100644
--- a/ydb/tools/query_replay_yt/query_compiler.cpp
+++ b/ydb/tools/query_replay_yt/query_compiler.cpp
@@ -225,7 +225,7 @@ public:
 class TReplayCompileActor: public TActorBootstrapped<TReplayCompileActor> {
 public:
     TReplayCompileActor(TIntrusivePtr<TModuleResolverState> moduleResolverState, const NMiniKQL::IFunctionRegistry* functionRegistry,
-        NYql::IHTTPGateway::TPtr httpGateway)
+        NYql::IHTTPGateway::TPtr httpGateway, bool enableAntlr4Parser)
         : ModuleResolverState(moduleResolverState)
         , KqpSettings()
         , Config(MakeIntrusive<TKikimrConfiguration>())
@@ -234,6 +234,7 @@ public:
     {
         Config->EnableKqpScanQueryStreamLookup = true;
         Config->EnablePreparedDdl = true;
+        Config->EnableAntlr4Parser = enableAntlr4Parser;
     }
 
     void Bootstrap() {
@@ -679,7 +680,7 @@ private:
 };
 
 IActor* CreateQueryCompiler(TIntrusivePtr<TModuleResolverState> moduleResolverState,
-    const NMiniKQL::IFunctionRegistry* functionRegistry, NYql::IHTTPGateway::TPtr httpGateway)
+    const NMiniKQL::IFunctionRegistry* functionRegistry, NYql::IHTTPGateway::TPtr httpGateway, bool enableAntlr4Parser)
 {
-    return new TReplayCompileActor(moduleResolverState, functionRegistry, httpGateway);
+    return new TReplayCompileActor(moduleResolverState, functionRegistry, httpGateway, enableAntlr4Parser);
 }
diff --git a/ydb/tools/query_replay_yt/query_replay.cpp b/ydb/tools/query_replay_yt/query_replay.cpp
index 4b3a96a3a4..85f7c041d2 100644
--- a/ydb/tools/query_replay_yt/query_replay.cpp
+++ b/ydb/tools/query_replay_yt/query_replay.cpp
@@ -22,6 +22,7 @@ void TQueryReplayConfig::ParseConfig(int argc, const char** argv) {
     opts.AddLongOption("udf-file", "UDFS to load").AppendTo(&UdfFiles);
     opts.AddLongOption("query", "Single query to replay").StoreResult(&QueryFile);
     opts.AddLongOption("log-level", "Yql log level").StoreResult(&YqlLogLevel);
+    opts.AddLongOption("antlr4-parser", "Use Antlr4 parser").NoArgument().SetFlag(&EnableAntlr4Parser);
 
     NLastGetopt::TOptsParseResult parseResult(&opts, argc, argv);
 }
diff --git a/ydb/tools/query_replay_yt/query_replay.h b/ydb/tools/query_replay_yt/query_replay.h
index f6583169f6..bd968add9a 100644
--- a/ydb/tools/query_replay_yt/query_replay.h
+++ b/ydb/tools/query_replay_yt/query_replay.h
@@ -22,6 +22,7 @@ struct TQueryReplayConfig {
     TVector<TString> UdfFiles;
     TString QueryFile;
     NActors::NLog::EPriority YqlLogLevel = NActors::NLog::EPriority::PRI_ERROR;
+    bool EnableAntlr4Parser = false;
 
     void ParseConfig(int argc, const char** argv);
 };
@@ -81,4 +82,4 @@ struct TQueryReplayEvents {
 };
 
 NActors::IActor* CreateQueryCompiler(TIntrusivePtr<NKikimr::NKqp::TModuleResolverState> moduleResolverState,
-    const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry, std::shared_ptr<NYql::IHTTPGateway> httpGateway);
+    const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry, std::shared_ptr<NYql::IHTTPGateway> httpGateway, bool enableAntlr4Parser);