YQ-3912 Add read_group ( ~ connection_id) to pq cluster config (#12064)

kardymonds · web-flow · commit e7cb027263be · 2024-12-03T13:25:06.000+03:00
diff --git a/ydb/core/fq/libs/actors/clusters_from_connections.cpp b/ydb/core/fq/libs/actors/clusters_from_connections.cpp
@@ -41,7 +41,8 @@ void FillClusterAuth(TClusterConfig& clusterCfg,
 void FillPqClusterConfig(NYql::TPqClusterConfig& clusterConfig,
         const TString& name, bool useBearerForYdb,
         const TString& authToken, const THashMap<TString, TString>& accountIdSignatures,
-        const FederatedQuery::DataStreams& ds) {
+        const FederatedQuery::DataStreams& ds,
+        const TString& readGroup) {
     clusterConfig.SetName(name);
     if (ds.endpoint()) {
         clusterConfig.SetEndpoint(ds.endpoint());
@@ -52,6 +53,7 @@ void FillPqClusterConfig(NYql::TPqClusterConfig& clusterConfig,
     clusterConfig.SetAddBearerToToken(useBearerForYdb);
     clusterConfig.SetClusterType(TPqClusterConfig::CT_DATA_STREAMS);
     clusterConfig.SetSharedReading(ds.shared_reading());
+    clusterConfig.SetReadGroup(readGroup);
     FillClusterAuth(clusterConfig, ds.auth(), authToken, accountIdSignatures);
 }
 
@@ -177,13 +179,14 @@ void FillGenericClusterConfig<FederatedQuery::PostgreSQLCluster>(
 
 NYql::TPqClusterConfig CreatePqClusterConfig(const TString& name,
         bool useBearerForYdb, const TString& authToken,
-        const TString& accountSignature, const FederatedQuery::DataStreams& ds) {
+        const TString& accountSignature, const FederatedQuery::DataStreams& ds,
+        const TString& readGroup) {
     NYql::TPqClusterConfig cluster;
     THashMap<TString, TString> accountIdSignatures;
     if (ds.auth().has_service_account()) {
         accountIdSignatures[ds.auth().service_account().id()] = accountSignature;
     }
-    FillPqClusterConfig(cluster, name, useBearerForYdb, authToken, accountIdSignatures, ds);
+    FillPqClusterConfig(cluster, name, useBearerForYdb, authToken, accountIdSignatures, ds, readGroup);
     return cluster;
 }
 
@@ -255,7 +258,7 @@ void AddClustersFromConnections(
         case FederatedQuery::ConnectionSetting::kDataStreams: {
             const auto& ds = conn.content().setting().data_streams();
             auto* clusterCfg = gatewaysConfig.MutablePq()->AddClusterMapping();
-            FillPqClusterConfig(*clusterCfg, connectionName, common.GetUseBearerForYdb(), authToken, accountIdSignatures, ds);
+            FillPqClusterConfig(*clusterCfg, connectionName, common.GetUseBearerForYdb(), authToken, accountIdSignatures, ds, conn.meta().id());
             clusters.emplace(connectionName, PqProviderName);
             break;
         }
diff --git a/ydb/core/fq/libs/actors/clusters_from_connections.h b/ydb/core/fq/libs/actors/clusters_from_connections.h
@@ -6,7 +6,7 @@
 
 namespace NFq {
 
-NYql::TPqClusterConfig CreatePqClusterConfig(const TString& name, bool useBearerForYdb, const TString& authToken, const TString& accountSignature, const FederatedQuery::DataStreams& ds);
+NYql::TPqClusterConfig CreatePqClusterConfig(const TString& name, bool useBearerForYdb, const TString& authToken, const TString& accountSignature, const FederatedQuery::DataStreams& ds, const TString& readGroup);
 
 NYql::TS3ClusterConfig CreateS3ClusterConfig(const TString& name, const TString& authToken, const TString& objectStorageEndpoint, const TString& accountSignature, const FederatedQuery::ObjectStorageConnection& s3);
 
diff --git a/ydb/core/fq/libs/test_connection/test_data_streams.cpp b/ydb/core/fq/libs/test_connection/test_data_streams.cpp
@@ -71,6 +71,7 @@ using namespace NActors;
 
 class TTestDataStreamsConnectionActor : public NActors::TActorBootstrapped<TTestDataStreamsConnectionActor> {
     inline static const TString SessionName = "test_connection_data_streams";
+    inline static const TString ReadGroup = "read_group";
 
     NFq::NConfig::TCommonConfig CommonConfig;
     TActorId Sender;
@@ -117,7 +118,7 @@ class TTestDataStreamsConnectionActor : public NActors::TActorBootstrapped<TTest
         , CmConnections(cmConnections)
         , FunctionRegistry(functionRegistry)
         , DbResolver(dbResolver)
-        , ClusterConfig(CreateClusterConfig(SessionName, CommonConfig, Token, signer, ds))
+        , ClusterConfig(CreateClusterConfig(SessionName, CommonConfig, Token, signer, ds, ReadGroup))
         , StructuredToken(NYql::ComposeStructuredTokenJsonForServiceAccount(ClusterConfig.GetServiceAccountId(), ClusterConfig.GetServiceAccountIdSignature(), ClusterConfig.GetToken()))
     {
         Counters->InFly->Inc();
@@ -248,10 +249,10 @@ class TTestDataStreamsConnectionActor : public NActors::TActorBootstrapped<TTest
         DestroyActor();
     }
 
-    static NYql::TPqClusterConfig CreateClusterConfig(const TString& sessionName, const NFq::NConfig::TCommonConfig& commonConfig, const TString& token, const NFq::TSigner::TPtr& signer, const FederatedQuery::DataStreams& ds) {
+    static NYql::TPqClusterConfig CreateClusterConfig(const TString& sessionName, const NFq::NConfig::TCommonConfig& commonConfig, const TString& token, const NFq::TSigner::TPtr& signer, const FederatedQuery::DataStreams& ds, const TString& readGroup) {
         const auto& auth = ds.auth();
         const TString signedAccountId = signer && auth.has_service_account() ? signer->SignAccountId(auth.service_account().id()) : TString{};
-        return NFq::CreatePqClusterConfig(sessionName, commonConfig.GetUseBearerForYdb(), token, signedAccountId, ds);
+        return NFq::CreatePqClusterConfig(sessionName, commonConfig.GetUseBearerForYdb(), token, signedAccountId, ds, readGroup);
     }
 
     NYql::TPqGatewayServices CreateGatewayServices() {
diff --git a/ydb/library/yql/providers/pq/common/yql_names.h b/ydb/library/yql/providers/pq/common/yql_names.h
@@ -16,5 +16,6 @@ constexpr TStringBuf WatermarksGranularityUsSetting = "WatermarksGranularityUs";
 constexpr TStringBuf WatermarksLateArrivalDelayUsSetting = "WatermarksLateArrivalDelayUs";
 constexpr TStringBuf WatermarksIdlePartitionsSetting = "WatermarksIdlePartitions";
 constexpr TStringBuf ReconnectPeriod = "ReconnectPeriod";
+constexpr TStringBuf ReadGroup = "ReadGroup";
 
 } // namespace NYql
diff --git a/ydb/library/yql/providers/pq/proto/dq_io.proto b/ydb/library/yql/providers/pq/proto/dq_io.proto
@@ -39,6 +39,7 @@ message TDqPqTopicSource {
     bool SharedReading = 15;
     string ReconnectPeriod = 16; // disabled by default, example of a parameter: 5m
     bool EnabledLLVM = 17;
+    string ReadGroup = 18;
 }
 
 message TDqPqTopicSink {
diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp
@@ -229,6 +229,8 @@ class TPqDqIntegration: public TDqIntegrationBase {
                         sharedReading = FromString<bool>(Value(setting));
                     } else if (name == ReconnectPeriod) {
                         srcDesc.SetReconnectPeriod(TString(Value(setting)));
+                    } else if (name == ReadGroup) {
+                        srcDesc.SetReadGroup(TString(Value(setting)));
                     } else if (name == Format) {
                         format = TString(Value(setting));
                     } else if (name == UseSslSetting) {
@@ -348,7 +350,7 @@ class TPqDqIntegration: public TDqIntegrationBase {
         Add(props, SharedReading, ToString(clusterConfiguration->SharedReading), pos, ctx);
         Add(props, ReconnectPeriod, ToString(clusterConfiguration->ReconnectPeriod), pos, ctx);
         Add(props, Format, format, pos, ctx);
-
+        Add(props, ReadGroup, clusterConfiguration->ReadGroup, pos, ctx);
 
         if (clusterConfiguration->UseSsl) {
             Add(props, UseSslSetting, "1", pos, ctx);
diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_settings.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_settings.cpp
@@ -43,6 +43,7 @@ void TPqConfiguration::Init(
         clusterSettings.UseSsl = cluster.GetUseSsl();
         clusterSettings.AddBearerToToken = cluster.GetAddBearerToToken();
         clusterSettings.SharedReading = cluster.GetSharedReading();
+        clusterSettings.ReadGroup = cluster.GetReadGroup();
 
         const TString authToken = typeCtx->Credentials->FindCredentialContent("cluster:default_" + clusterSettings.ClusterName, "default_pq", cluster.GetToken());
         clusterSettings.AuthToken = authToken;
diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_settings.h b/ydb/library/yql/providers/pq/provider/yql_pq_settings.h
@@ -31,6 +31,7 @@ struct TPqClusterConfigurationSettings {
     bool AddBearerToToken = false;
     bool SharedReading = false;
     TString ReconnectPeriod;
+    TString ReadGroup;
 };
 
 struct TPqConfiguration : public TPqSettings, public NCommon::TSettingDispatcher {
diff --git a/ydb/library/yql/tools/dqrun/examples/gateways.conf b/ydb/library/yql/tools/dqrun/examples/gateways.conf
@@ -143,7 +143,8 @@ Pq {
         Database: "local"
         ClusterType: CT_DATA_STREAMS
         UseSsl: True
-        SharedReading:False
+        SharedReading:True
+        ReadGroup: "read_group"
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@ message TDqPqTopicSource {`
`39`	`39`	`bool SharedReading = 15;`
`40`	`40`	`string ReconnectPeriod = 16; // disabled by default, example of a parameter: 5m`
`41`	`41`	`bool EnabledLLVM = 17;`
	`42`	`+ string ReadGroup = 18;`
`42`	`43`	`}`
`43`	`44`
`44`	`45`	`message TDqPqTopicSink {`
Original file line number	Diff line number	Diff line change
`@@ -143,7 +143,8 @@ Pq {`
`143`	`143`	`Database: "local"`
`144`	`144`	`ClusterType: CT_DATA_STREAMS`
`145`	`145`	`UseSsl: True`
`146`		`- SharedReading:False`
	`146`	`+ SharedReading:True`
	`147`	`+ ReadGroup: "read_group"`
`147`	`148`	`}`
`148`	`149`	`}`
`149`	`150`