Yq 3560 Add row dispatcher to dqrun (#9697)

Co-authored-by: Fiodar Miron <fedor-miron@ydb.tech>
Co-authored-by: Fiodar Miron <61616792+fedor-miron@users.noreply.github.com>

Author: kardymonds

ydb/core/fq/libs/control_plane_storage/ya.make

@@ -20,6 +20,7 @@ PEERDIR(
     library/cpp/lwtrace
     library/cpp/protobuf/interop
     ydb/core/base
+    ydb/core/external_sources
     ydb/core/fq/libs/actors/logging
     ydb/core/fq/libs/common
     ydb/core/fq/libs/config
@@ -33,13 +34,13 @@ PEERDIR(
     ydb/core/fq/libs/shared_resources
     ydb/core/fq/libs/ydb
     ydb/core/mon
+    ydb/library/db_pool
     ydb/library/security
+    ydb/library/yql/providers/s3/path_generator
+    ydb/library/yql/public/issue
     ydb/public/api/protos
     ydb/public/sdk/cpp/client/ydb_scheme
     ydb/public/sdk/cpp/client/ydb_table
-    ydb/library/db_pool
-    ydb/library/yql/providers/s3/path_generator
-    ydb/library/yql/public/issue
 )
 
 YQL_LAST_ABI_VERSION()

ydb/core/fq/libs/init/init.cpp

@@ -207,7 +207,7 @@ void Init(
         NYql::NDq::TS3ReadActorFactoryConfig readActorFactoryCfg = NYql::NDq::CreateReadActorFactoryConfig(protoConfig.GetGateways().GetS3());
 
         RegisterDqInputTransformLookupActorFactory(*asyncIoFactory);
-        
+
         NYql::TPqGatewayServices pqServices(
             yqSharedResources->UserSpaceYdbDriver,
             pqCmConnections,

ydb/library/yql/providers/dq/local_gateway/yql_dq_gateway_local.cpp

@@ -31,7 +31,8 @@ class TLocalServiceHolder {
         NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, int threads,
         IMetricsRegistryPtr metricsRegistry,
         const std::function<IActor*(void)>& metricsPusherFactory,
-        bool withSpilling)
+        bool withSpilling,
+        TVector<std::pair<TActorId, TActorSetupCmd>>&& additionalLocalServices)
         : MetricsRegistry(metricsRegistry
             ? metricsRegistry
             : CreateMetricsRegistry(GetSensorsGroupFor(NSensorComponent::kDq))
@@ -89,6 +90,9 @@ class TLocalServiceHolder {
                 NDq::MakeDqLocalFileSpillingServiceID(nodeId),
                 TActorSetupCmd(spillingActor, TMailboxType::Simple, 0));
         }
+        for (auto& [actorId, setupCmd] : additionalLocalServices) {
+            ServiceNode->AddLocalService(actorId, std::move(setupCmd));
+        }
 
         auto statsCollector = CreateStatsCollector(1, *ServiceNode->GetSetup(), MetricsRegistry->GetSensors());
 
@@ -248,7 +252,8 @@ THolder<TLocalServiceHolder> CreateLocalServiceHolder(const NKikimr::NMiniKQL::I
     NBus::TBindResult interconnectPort, NBus::TBindResult grpcPort,
     NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, int threads,
     IMetricsRegistryPtr metricsRegistry,
-    const std::function<IActor*(void)>& metricsPusherFactory, bool withSpilling)
+    const std::function<IActor*(void)>& metricsPusherFactory, bool withSpilling,
+    TVector<std::pair<TActorId, TActorSetupCmd>>&& additionalLocalServices)
 {
     return MakeHolder<TLocalServiceHolder>(functionRegistry,
         compFactory,
@@ -260,15 +265,17 @@ THolder<TLocalServiceHolder> CreateLocalServiceHolder(const NKikimr::NMiniKQL::I
         threads,
         metricsRegistry,
         metricsPusherFactory,
-        withSpilling);
+        withSpilling,
+        std::move(additionalLocalServices));
 }
 
 TIntrusivePtr<IDqGateway> CreateLocalDqGateway(const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry,
     NKikimr::NMiniKQL::TComputationNodeFactory compFactory,
     TTaskTransformFactory taskTransformFactory, const TDqTaskPreprocessorFactoryCollection& dqTaskPreprocessorFactories,
     bool withSpilling, NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, int threads,
     IMetricsRegistryPtr metricsRegistry,
-    const std::function<IActor*(void)>& metricsPusherFactory)
+    const std::function<IActor*(void)>& metricsPusherFactory,
+    TVector<std::pair<TActorId, TActorSetupCmd>>&& additionalLocalServices)
 {
     int startPort = 31337;
     TRangeWalker<int> portWalker(startPort, startPort+100);
@@ -287,7 +294,8 @@ TIntrusivePtr<IDqGateway> CreateLocalDqGateway(const NKikimr::NMiniKQL::IFunctio
             threads,
             metricsRegistry,
             metricsPusherFactory,
-            withSpilling),
+            withSpilling,
+            std::move(additionalLocalServices)),
         CreateDqGateway("[::1]", grpcPort.Addr.GetPort()));
 }
 

ydb/library/yql/providers/dq/local_gateway/yql_dq_gateway_local.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <ydb/library/actors/core/actorsystem.h>
 #include <ydb/library/yql/providers/dq/provider/yql_dq_gateway.h>
 #include <ydb/library/yql/providers/dq/interface/yql_dq_task_preprocessor.h>
 #include <ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io_factory.h>
@@ -17,6 +18,7 @@ TIntrusivePtr<IDqGateway> CreateLocalDqGateway(const NKikimr::NMiniKQL::IFunctio
     bool withSpilling,
     NDq::IDqAsyncIoFactory::TPtr = nullptr, int threads = 16,
     IMetricsRegistryPtr metricsRegistry = {},
-    const std::function<NActors::IActor*(void)>& metricsPusherFactory = {});
+    const std::function<NActors::IActor*(void)>& metricsPusherFactory = {},
+    TVector<std::pair<NActors::TActorId, NActors::TActorSetupCmd>>&& additionalLocalServices = {});
 
 } // namespace NYql

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

@@ -81,8 +81,12 @@
 #include <ydb/library/yql/public/result_format/yql_result_format_data.h>
 
 #include <ydb/core/fq/libs/actors/database_resolver.h>
+#include <ydb/core/fq/libs/config/protos/fq_config.pb.h>
 #include <ydb/core/fq/libs/db_id_async_resolver_impl/db_async_resolver_impl.h>
 #include <ydb/core/fq/libs/db_id_async_resolver_impl/mdb_endpoint_generator.h>
+#include <ydb/core/fq/libs/init/init.h>
+#include <ydb/core/fq/libs/row_dispatcher/row_dispatcher_service.h>
+
 #include <ydb/core/util/pb.h>
 
 #include <yt/cpp/mapreduce/interface/init.h>
@@ -180,6 +184,17 @@ void ReadGatewaysConfig(const TString& configFile, TGatewaysConfig* config, THas
     }
 }
 
+void ReadFqConfig(const TString& fqCfgFile, NFq::NConfig::TConfig* fqConfig) {
+    if (fqCfgFile.empty()) {
+        return;
+    }
+    auto configData = TFileInput(fqCfgFile).ReadAll();
+    using ::google::protobuf::TextFormat;
+    if (!TextFormat::ParseFromString(configData, fqConfig)) {
+        ythrow yexception() << "Bad format of fq configuration";
+    }
+}
+
 void PatchGatewaysConfig(TGatewaysConfig* config, const TString& mrJobBin, const TString& mrJobUdfsDir,
     size_t numThreads, bool keepTemp)
 {
@@ -483,9 +498,35 @@ int RunProgram(TProgramPtr program, const TRunOptions& options, const THashMap<T
     return 0;
 }
 
+void InitFq(const NFq::NConfig::TConfig& fqConfig, TVector<std::pair<TActorId, TActorSetupCmd>>& additionalLocalServices) {
+    if (fqConfig.HasRowDispatcher() && fqConfig.GetRowDispatcher().GetEnabled()) {
+        NFq::IYqSharedResources::TPtr iSharedResources = NFq::CreateYqSharedResources(
+            fqConfig,
+            NKikimr::CreateYdbCredentialsProviderFactory,
+            MakeIntrusive<NMonitoring::TDynamicCounters>());
+        NFq::TYqSharedResources::TPtr yqSharedResources = NFq::TYqSharedResources::Cast(iSharedResources);
+        ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory;
+
+        NFq::NConfig::TCommonConfig commonConfig;
+        auto rowDispatcher = NFq::NewRowDispatcherService(
+            fqConfig.GetRowDispatcher(),
+            commonConfig,
+            NKikimr::CreateYdbCredentialsProviderFactory,
+            yqSharedResources,
+            credentialsFactory,
+            "/tenant",
+            MakeIntrusive<NMonitoring::TDynamicCounters>());
+
+        additionalLocalServices.emplace_back(
+            NFq::RowDispatcherServiceActorId(),
+            TActorSetupCmd(rowDispatcher.release(), TMailboxType::Simple, 0));
+    }
+}
+
 int RunMain(int argc, const char* argv[])
 {
     TString gatewaysCfgFile;
+    TString fqCfgFile;
     TString progFile;
     TVector<TString> tablesMappingList;
     THashMap<TString, TString> tablesMapping;
@@ -581,6 +622,10 @@ int RunMain(int argc, const char* argv[])
         .Optional()
         .RequiredArgument("FILE")
         .StoreResult(&gatewaysCfgFile);
+    opts.AddLongOption("fq-cfg", "federated query configuration file")
+        .Optional()
+        .RequiredArgument("FILE")
+        .StoreResult(&fqCfgFile);
     opts.AddLongOption("fs-cfg", "Path to file storage config")
         .Optional()
         .StoreResult(&fileStorageCfg);
@@ -874,6 +919,9 @@ int RunMain(int argc, const char* argv[])
         setting->SetValue("1");
     }
 
+    NFq::NConfig::TConfig fqConfig;
+    ReadFqConfig(fqCfgFile, &fqConfig);
+
     if (res.Has("enable-spilling")) {
         auto* setting = gatewaysConfig.MutableDq()->AddDefaultSettings();
         setting->SetName("SpillingEngine");
@@ -1040,6 +1088,8 @@ int RunMain(int argc, const char* argv[])
             clusters.emplace(to_lower(cluster.GetName()), TString{NYql::SolomonProviderName});
         }
     }
+    TVector<std::pair<TActorId, TActorSetupCmd>> additionalLocalServices;
+    InitFq(fqConfig, additionalLocalServices);
 
     std::function<NActors::IActor*(void)> metricsPusherFactory = {};
 
@@ -1064,7 +1114,7 @@ int RunMain(int argc, const char* argv[])
             bool enableSpilling = res.Has("enable-spilling");
             dqGateway = CreateLocalDqGateway(funcRegistry.Get(), dqCompFactory, dqTaskTransformFactory, dqTaskPreprocessorFactories, enableSpilling,
                 CreateAsyncIoFactory(driver, httpGateway, ytFileServices, genericClient, credentialsFactory, *funcRegistry, requestTimeout, maxRetries, pqGateway), threads,
-                metricsRegistry, metricsPusherFactory);
+                metricsRegistry, metricsPusherFactory, std::move(additionalLocalServices));
         }
 
         dataProvidersInit.push_back(GetDqDataProviderInitializer(&CreateDqExecTransformer, dqGateway, dqCompFactory, {}, storage));

ydb/library/yql/tools/dqrun/examples/fq.conf

@@ -0,0 +1,20 @@
+RowDispatcher {
+  Enabled: true
+  TimeoutBeforeStartSessionSec: 2 
+  MaxSessionUsedMemory: 0
+  SendStatusPeriodSec: 10
+  WithoutConsumer: true
+  Coordinator {
+      CoordinationNodePath: "not_used"
+      Database {
+        Endpoint: "not_used:2135"
+        Database: "/not/used/database"
+        UseLocalMetadataService: true
+        UseSsl: true
+        ClientTimeoutSec: 70
+        OperationTimeoutSec: 60
+        CancelAfterSec: 60
+      }
+  }
+}
+

ydb/library/yql/tools/dqrun/examples/gateways.conf

@@ -68,6 +68,10 @@ Dq {
         Name: "EnableDqReplicate"
         Value: "true"
     }
+    DefaultSettings {
+        Name: "_TableTimeout"
+        Value: "600000"
+    }
 }
 
 Generic {
@@ -133,3 +137,14 @@ SqlCore {
     TranslationFlags: ["FlexibleTypes", "DisableAnsiOptionalAs", "EmitAggApply"]
 }
 
+Pq {
+    ClusterMapping {
+        Name: "pq"
+        Endpoint: "localhost:2135"
+        Database: "local"
+        ClusterType: CT_DATA_STREAMS
+        UseSsl: True
+        SharedReading:False
+    }
+}
+

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

@@ -95,6 +95,8 @@ ENDIF()
         ydb/library/yql/utils/actor_system
         ydb/core/fq/libs/actors
         ydb/core/fq/libs/db_id_async_resolver_impl
+        ydb/core/fq/libs/init
+        ydb/core/external_sources
 
         ydb/library/yql/udfs/common/clickhouse/client
     )