ydb_recipe: support pq_client_service_type LOGBROKER-7284

ref:7b86b3eb11098a751b1bae531a40dd00276ac633

4 files changed, 24 insertions, 0 deletions

diff --git a/ydb/core/yq/libs/read_rule/read_rule_creator.cpp b/ydb/core/yq/libs/read_rule/read_rule_creator.cpp
index 4ca1a5ca47..c5d6fa4fa2 100644
--- a/ydb/core/yq/libs/read_rule/read_rule_creator.cpp
+++ b/ydb/core/yq/libs/read_rule/read_rule_creator.cpp
@@ -110,6 +110,7 @@ public:
                 .ReadRule(
                     NYdb::NPersQueue::TReadRuleSettings()
                         .ConsumerName(Topic.GetConsumerName())
+                        .ServiceType("yandex-query")
                         .SupportedCodecs({
                             NYdb::NPersQueue::ECodec::RAW,
                             NYdb::NPersQueue::ECodec::GZIP,
diff --git a/ydb/public/tools/lib/cmds/__init__.py b/ydb/public/tools/lib/cmds/__init__.py
index 977f0a4aae..eceadd27c9 100644
--- a/ydb/public/tools/lib/cmds/__init__.py
+++ b/ydb/public/tools/lib/cmds/__init__.py
@@ -241,6 +241,20 @@ def enable_pq(arguments):
     return getattr(arguments, 'enable_pq', False) or enable_datastreams(arguments)
 
 
+def pq_client_service_types(arguments):
+    items = getattr(arguments, 'pq_client_service_types')
+    if not items:
+        types_str = os.getenv('YDB_PQ_CLIENT_SERVICE_TYPES')
+        if types_str:
+            items = types_str.split(',')
+    service_types = []
+    for item in items:
+        item = item.strip()
+        if item:
+            service_types.append(item)
+    return service_types
+
+
 def enable_pqcd(arguments):
     return enable_pq(arguments) and (getattr(arguments, 'enable_pqcd', False) or os.getenv('YDB_ENABLE_PQCD') == 'true')
 
@@ -281,6 +295,7 @@ def deploy(arguments):
         pdisk_store_path=pdisk_store_path,
         domain_name='local',
         enable_pq=enable_pq(arguments),
+        pq_client_service_types=pq_client_service_types(arguments),
         enable_datastreams=enable_datastreams(arguments),
         enable_pqcd=enable_pqcd(arguments),
         load_udfs=True,
diff --git a/ydb/public/tools/ydb_recipe/__main__.py b/ydb/public/tools/ydb_recipe/__main__.py
index 37562c5a85..e21d27e22f 100644
--- a/ydb/public/tools/ydb_recipe/__main__.py
+++ b/ydb/public/tools/ydb_recipe/__main__.py
@@ -12,6 +12,7 @@ if __name__ == "__main__":
     parser.add_argument("--ydb-working-dir", action="store")
     parser.add_argument("--debug-logging", nargs='*')
     parser.add_argument("--enable-pq", action='store_true', default=False)
+    parser.add_argument("--pq-client-service-type", action='append', default=[])
     parser.add_argument("--enable-datastreams", action='store_true', default=False)
     parser.add_argument("--enable-pqcd", action='store_true', default=False)
     parsed, _ = parser.parse_known_args()
@@ -23,6 +24,7 @@ if __name__ == "__main__":
     if parsed.debug_logging:
         arguments.debug_logging = parsed.debug_logging
     arguments.enable_pq = parsed.enable_pq
+    arguments.pq_client_service_types = parsed.pq_client_service_type
     arguments.enable_datastreams = parsed.enable_datastreams
     arguments.enable_pqcd = parsed.enable_pqcd
     declare_recipe(
diff --git a/ydb/tests/library/harness/kikimr_config.py b/ydb/tests/library/harness/kikimr_config.py
index 57c282b9f4..f81ace71de 100644
--- a/ydb/tests/library/harness/kikimr_config.py
+++ b/ydb/tests/library/harness/kikimr_config.py
@@ -88,6 +88,7 @@ class KikimrConfigGenerator(object):
             udfs_path=None,
             output_path=None,
             enable_pq=False,
+            pq_client_service_types=None,
             slot_count=0,
             pdisk_store_path=None,
             version=None,
@@ -178,6 +179,11 @@ class KikimrConfigGenerator(object):
         self.yaml_config["feature_flags"]["enable_public_api_external_blobs"] = enable_public_api_external_blobs
         self.yaml_config["feature_flags"]["enable_mvcc"] = "VALUE_FALSE" if disable_mvcc else "VALUE_TRUE"
         self.yaml_config['pqconfig']['enabled'] = enable_pq
+        if pq_client_service_types:
+            self.yaml_config['pqconfig']['client_service_type'] = []
+            for service_type in pq_client_service_types:
+                self.yaml_config['pqconfig']['client_service_type'].append({'name': service_type})
+
         self.yaml_config['pqconfig']['topics_are_first_class_citizen'] = enable_pq and enable_datastreams
         self.yaml_config['sqs_config']['enable_sqs'] = enable_sqs
         self.yaml_config['pqcluster_discovery_config']['enabled'] = enable_pqcd