improve replay execution (#6211)

3 files changed, 51 insertions, 10 deletions

diff --git a/ydb/tools/query_replay_yt/main.cpp b/ydb/tools/query_replay_yt/main.cpp
index c9237e84a24..972e56f7e81 100644
--- a/ydb/tools/query_replay_yt/main.cpp
+++ b/ydb/tools/query_replay_yt/main.cpp
@@ -4,6 +4,7 @@
 #include <ydb/core/client/minikql_compile/mkql_compile_service.h>
 #include <ydb/core/kqp/rm_service/kqp_rm_service.h>
 #include <ydb/library/yql/minikql/invoke_builtins/mkql_builtins.h>
+#include <ydb/library/yql/providers/common/http_gateway/yql_http_gateway.h>
 
 #include <library/cpp/monlib/dynamic_counters/counters.h>
 
@@ -45,6 +46,7 @@ class TQueryReplayMapper
     TIntrusivePtr<NKikimr::NMiniKQL::IMutableFunctionRegistry> FunctionRegistry;
     TIntrusivePtr<NKikimr::NKqp::TModuleResolverState> ModuleResolverState;
 
+    NYql::IHTTPGateway::TPtr HttpGateway;
     TVector<TString> UdfFiles;
     ui32 ActorSystemThreadsCount = 5;
 
@@ -72,6 +74,8 @@ class TQueryReplayMapper
                 return "write_columns_mismatch";
             case TQueryReplayEvents::UncategorizedPlanMismatch:
                 return "uncategorized_plan_mismatch";
+            case TQueryReplayEvents::MissingTableMetadata:
+                return "missing_table_metadata";
             default:
                 return "unspecified";
         }
@@ -107,6 +111,7 @@ public:
         ActorSystem->Start();
         ActorSystem->Register(NKikimr::NKqp::CreateKqpResourceManagerActor({}, nullptr));
         ModuleResolverState = MakeIntrusive<NKikimr::NKqp::TModuleResolverState>();
+        HttpGateway = NYql::IHTTPGateway::Make();
         Y_ABORT_UNLESS(GetYqlDefaultModuleResolver(ModuleResolverState->ExprCtx, ModuleResolverState->ModuleResolver));
     }
 
@@ -120,7 +125,16 @@ public:
                 json.InsertValue(key, NJson::TJsonValue(child.AsString()));
             }
 
-            auto compileActorId = ActorSystem->Register(CreateQueryCompiler(ModuleResolverState, FunctionRegistry.Get()));
+            TString queryType = row["query_type"].AsString();
+            if (queryType == "QUERY_TYPE_AST_SCAN") {
+                continue;
+            }
+
+            if (queryType == "QUERY_TYPE_SQL_GENERIC_SCRIPT") {
+                continue;
+            }
+
+            auto compileActorId = ActorSystem->Register(CreateQueryCompiler(ModuleResolverState, FunctionRegistry.Get(), HttpGateway));
 
             auto future = ActorSystem->Ask<TQueryReplayEvents::TEvCompileResponse>(
                 compileActorId,
@@ -132,7 +146,7 @@ public:
 
             TString failReason = GetFailReason(status);
 
-            if (failReason == "unspecified") {
+            if (failReason == "unspecified" || status == TQueryReplayEvents::MissingTableMetadata) {
                 continue;
             }
 
diff --git a/ydb/tools/query_replay_yt/query_compiler.cpp b/ydb/tools/query_replay_yt/query_compiler.cpp
index 64cd4aecfd0..d3e410973d5 100644
--- a/ydb/tools/query_replay_yt/query_compiler.cpp
+++ b/ydb/tools/query_replay_yt/query_compiler.cpp
@@ -137,6 +137,7 @@ struct TMetadataInfoHolder {
 class TStaticTableMetadataLoader: public NYql::IKikimrGateway::IKqpTableMetadataLoader, public NYql::IDbSchemeResolver {
     TActorSystem* ActorSystem;
     std::shared_ptr<TMetadataInfoHolder> TableMetadata;
+    bool IsMissingTableMetadata = false;
 
 public:
     TStaticTableMetadataLoader(TActorSystem* actorSystem, std::shared_ptr<TMetadataInfoHolder>& tableMetadata)
@@ -152,6 +153,10 @@ public:
         Y_UNUSED(database);
         Y_UNUSED(userToken);
         auto ptr = TableMetadata->find(table);
+        if (ptr == TableMetadata->end()) {
+            IsMissingTableMetadata = true;
+        }
+
         Y_ENSURE(ptr != TableMetadata->end());
 
         NYql::IKikimrGateway::TTableMetadataResult result;
@@ -199,6 +204,10 @@ public:
         return results;
     }
 
+    bool HasMissingTableMetadata() const {
+        return IsMissingTableMetadata;
+    }
+
     virtual NThreading::TFuture<TTableResults> ResolveTables(const TVector<TTable>& tables) override {
         return NThreading::MakeFuture(Resolve(tables));
     }
@@ -215,13 +224,17 @@ public:
 
 class TReplayCompileActor: public TActorBootstrapped<TReplayCompileActor> {
 public:
-    TReplayCompileActor(TIntrusivePtr<TModuleResolverState> moduleResolverState, const NMiniKQL::IFunctionRegistry* functionRegistry)
+    TReplayCompileActor(TIntrusivePtr<TModuleResolverState> moduleResolverState, const NMiniKQL::IFunctionRegistry* functionRegistry,
+        NYql::IHTTPGateway::TPtr httpGateway)
         : ModuleResolverState(moduleResolverState)
         , KqpSettings()
         , Config(MakeIntrusive<TKikimrConfiguration>())
         , FunctionRegistry(functionRegistry)
+        , HttpGateway(std::move(httpGateway))
     {
         Config->EnableKqpScanQueryStreamLookup = true;
+        Config->PredicateExtract20 = true;
+        Config->EnablePreparedDdl = true;
     }
 
     void Bootstrap() {
@@ -495,7 +508,13 @@ private:
         Y_UNUSED(queryPlan);
         if (status != Ydb::StatusIds::SUCCESS) {
             ev->Success = false;
-            ev->Status = status == Ydb::StatusIds::TIMEOUT ? TQueryReplayEvents::CompileTimeout : TQueryReplayEvents::CompileError;
+            if (MetadataLoader->HasMissingTableMetadata()) {
+                ev->Status = TQueryReplayEvents::MissingTableMetadata;
+            } else if (status == Ydb::StatusIds::TIMEOUT) {
+                ev->Status = TQueryReplayEvents::CompileTimeout;
+            } else {
+                ev->Status = TQueryReplayEvents::CompileError;
+            }
             ev->Message = issues.ToString();
             Cerr << "Failed to compile query: " << ev->Message << Endl;
             WriteJsonData("-repro.txt", ReplayDetails);
@@ -594,16 +613,17 @@ private:
 	    Config->_KqpYqlSyntaxVersion = syntax;
         Config->FreezeDefaults();
 
-        std::shared_ptr<NYql::IKikimrGateway::IKqpTableMetadataLoader> loader = make_shared<TStaticTableMetadataLoader>(TlsActivationContext->ActorSystem(), TableMetadata);
+        MetadataLoader = make_shared<TStaticTableMetadataLoader>(TlsActivationContext->ActorSystem(), TableMetadata);
+        std::shared_ptr<NYql::IKikimrGateway::IKqpTableMetadataLoader> loader = MetadataLoader;
 
         auto c = MakeIntrusive<NMonitoring::TDynamicCounters>();
         auto counters = MakeIntrusive<TKqpRequestCounters>();
         counters->Counters = new TKqpCounters(c);
         counters->TxProxyMon = new NTxProxy::TTxProxyMon(c);
 
-        Gateway = CreateKikimrIcGateway(Query->Cluster, NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY, Query->Database, std::move(loader),
+        Gateway = CreateKikimrIcGateway(Query->Cluster, queryType, Query->Database, std::move(loader),
             TlsActivationContext->ExecutorThread.ActorSystem, SelfId().NodeId(), counters);
-        auto federatedQuerySetup = std::make_optional<TKqpFederatedQuerySetup>({NYql::IHTTPGateway::Make(), nullptr, nullptr, nullptr, {}, {}, {}, nullptr, nullptr});
+        auto federatedQuerySetup = std::make_optional<TKqpFederatedQuerySetup>({HttpGateway, nullptr, nullptr, nullptr, {}, {}, {}, nullptr, nullptr});
         KqpHost = CreateKqpHost(Gateway, Query->Cluster, Query->Database, Config, ModuleResolverState->ModuleResolver,
             federatedQuerySetup, nullptr, GUCSettings, Nothing(), FunctionRegistry, false);
 
@@ -652,10 +672,12 @@ private:
     std::shared_ptr<TMetadataInfoHolder> TableMetadata;
     TActorId MiniKQLCompileService;
     NJson::TJsonValue ReplayDetails;
+    std::shared_ptr<TStaticTableMetadataLoader> MetadataLoader;
+    NYql::IHTTPGateway::TPtr HttpGateway;
 };
 
 IActor* CreateQueryCompiler(TIntrusivePtr<TModuleResolverState> moduleResolverState,
-    const NMiniKQL::IFunctionRegistry* functionRegistry)
+    const NMiniKQL::IFunctionRegistry* functionRegistry, NYql::IHTTPGateway::TPtr httpGateway)
 {
-    return new TReplayCompileActor(moduleResolverState, functionRegistry);
+    return new TReplayCompileActor(moduleResolverState, functionRegistry, httpGateway);
 }
diff --git a/ydb/tools/query_replay_yt/query_replay.h b/ydb/tools/query_replay_yt/query_replay.h
index 10cfff2837d..9a13e92b78a 100644
--- a/ydb/tools/query_replay_yt/query_replay.h
+++ b/ydb/tools/query_replay_yt/query_replay.h
@@ -23,6 +23,10 @@ struct TQueryReplayConfig {
     void ParseConfig(int argc, const char** argv);
 };
 
+namespace NYql {
+    class IHTTPGateway;
+}
+
 using namespace NActors;
 
 THolder<TActorSystemSetup> BuildActorSystemSetup(ui32 threads, ui32 pools = 1);
@@ -46,6 +50,7 @@ struct TQueryReplayEvents {
         ExtraWriting,
         WriteColumnsMismatch,
         UncategorizedPlanMismatch,
+        MissingTableMetadata,
         Unspecified,
     };
 
@@ -73,4 +78,4 @@ struct TQueryReplayEvents {
 };
 
 NActors::IActor* CreateQueryCompiler(TIntrusivePtr<NKikimr::NKqp::TModuleResolverState> moduleResolverState,
-    const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry);
+    const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry, std::shared_ptr<NYql::IHTTPGateway> httpGateway);