Use im avtorization in perf tests (#28842)

Author: iddqdex

ydb/tests/olap/lib/ydb_cli.py

@@ -40,11 +40,14 @@ def get_cli_path() -> str:
 
     @staticmethod
     def get_cli_command() -> list[str]:
-        return [
+        result = [
             YdbCliHelper.get_cli_path(),
             '-e', YdbCluster.ydb_endpoint,
             '-d', f'/{YdbCluster.ydb_database}'
         ]
+        if YdbCluster.ydb_iam_file:
+            result += ['--sa-key-file', YdbCluster.ydb_iam_file]
+        return result
 
     class QueryPlan:
         def __init__(self) -> None:

ydb/tests/olap/lib/ydb_cluster.py

@@ -66,7 +66,8 @@ def slot(self) -> str:
     _cluster_info = None
     ydb_endpoint = get_external_param('ydb-endpoint', 'grpc://ydb-olap-testing-vla-0002.search.yandex.net:2135')
     ydb_database = get_external_param('ydb-db', 'olap-testing/kikimr/testing/acceptance-2').lstrip('/')
-    ydb_mon_port = 8765
+    ydb_mon_port = int(get_external_param('ydb-mon-port', 8765))
+    ydb_iam_file = get_external_param('ydb-iam-file', os.getenv('YDB_IAM_FILE'))
     _tables_path = get_external_param('tables-path', 'olap_yatests').rstrip('/')
     _monitoring_urls: list[YdbCluster.MonitoringUrl] = None
     _dyn_nodes_count: Optional[int] = None
@@ -108,6 +109,8 @@ def _get_service_url(cls):
     def get_cluster_nodes(cls, path: Optional[str] = None, db_only: bool = False,
                           role: Optional[YdbCluster.Node.Role] = None
                           ) -> list[YdbCluster.Node]:
+        if cls.ydb_mon_port == 0:
+            return []
 
         # Получаем таймаут из переменной окружения, как в wait_ydb_alive
         timeout = int(os.getenv('WAIT_CLUSTER_ALIVE_TIMEOUT', 2 * 60))  # По умолчанию 20 минут
@@ -203,7 +206,7 @@ def get_cluster_info(cls):
     @staticmethod
     def _create_ydb_driver(endpoint, database, oauth=None, iam_file=None):
         credentials = None
-        LOGGER.info(f"Connecting to {endpoint} to {database} ydb_access_token is set {oauth is not None}")
+        LOGGER.info(f"Connecting to {endpoint} to {database} ydb_access_token is set {oauth is not None}, iam file path is {iam_file}")
 
         if oauth is not None:
             credentials = ydb.AccessTokenCredentials(oauth)
@@ -239,7 +242,7 @@ def reset(cls, ydb_endpoint, ydb_database, ydb_mon_port, dyn_nodes_count):
     def get_ydb_driver(cls):
         if cls._ydb_driver is None:
             cls._ydb_driver = cls._create_ydb_driver(
-                cls.ydb_endpoint, cls.ydb_database, oauth=os.getenv('OLAP_YDB_OAUTH', None)
+                cls.ydb_endpoint, cls.ydb_database, oauth=os.getenv('OLAP_YDB_OAUTH', None), iam_file=cls.ydb_iam_file
             )
         return cls._ydb_driver
 
@@ -294,13 +297,13 @@ def _join_errors(log_level: int, errors: list[str]):
     @allure.step('Execute scan query')
     def execute_single_result_query(cls, query, timeout=10):
         allure.attach(query, 'query', attachment_type=allure.attachment_type.TEXT)
-        query = ydb.ScanQuery(query, {})
         settings = ydb.BaseRequestSettings()
         settings = settings.with_timeout(timeout)
         try:
-            it = cls.get_ydb_driver().table_client.scan_query(query, settings=settings)
+            session = ydb.query.QueryClientSync(cls.get_ydb_driver()).session().create()
+            it = session.execute(query, settings=settings)
             result = next(it)
-            return result.result_set.rows[0][0]
+            return result.rows[0][0]
         except BaseException:
             LOGGER.error("Cannot connect to YDB")
             raise
@@ -363,53 +366,54 @@ def _check_node(n: YdbCluster.Node):
         errors = []
         warnings = []
         try:
-            nodes = cls.get_cluster_nodes(db_only=True)
-            expected_nodes_count = cls.get_dyn_nodes_count()
-            nodes_count = len(nodes)
-            if expected_nodes_count:
-                LOGGER.debug(f'Expected nodes count: {expected_nodes_count}')
-                if nodes_count < expected_nodes_count:
-                    errors.append(f"{expected_nodes_count - nodes_count} nodes from {expected_nodes_count} don't alive")
-            ok_node_count = 0
-            node_errors = []
-            for n in nodes:
-                error = _check_node(n)
-                if error:
-                    node_errors.append(error)
-                else:
-                    ok_node_count += 1
-            if ok_node_count < nodes_count:
-                errors.append(f'Only {ok_node_count} from {nodes_count} dynnodes are ok: {",".join(node_errors)}')
-            paths_to_balance = []
-            if isinstance(balanced_paths, str):
-                paths_to_balance += cls.get_tables(balanced_paths)
-            elif isinstance(balanced_paths, list):
-                for path in balanced_paths:
-                    paths_to_balance += cls.get_tables(path)
-            for p in paths_to_balance:
-                table_nodes = cls.get_cluster_nodes(p)
-                min = None
-                max = None
+            if cls.ydb_mon_port != 0:
+                nodes = cls.get_cluster_nodes(db_only=True)
+                expected_nodes_count = cls.get_dyn_nodes_count()
+                nodes_count = len(nodes)
                 if expected_nodes_count:
-                    if len(table_nodes) < expected_nodes_count:
-                        min = 0
-                for tn in table_nodes:
-                    tablet_count = 0
-                    for tablet in tn.tablets:
-                        if tablet.count > 0 and tablet.state != "Green":
-                            warnings.append(f'Node {tn.host}: {tablet.count} tablets of type {tablet.type} in {tablet.state} state')
-                        if tablet.type in {"ColumnShard", "DataShard"}:
-                            tablet_count += tablet.count
-                    if tablet_count > 0:
-                        if min is None or tablet_count < min:
-                            min = tablet_count
-                        if max is None or tablet_count > max:
-                            max = tablet_count
-                if min is None or max is None:
-                    warnings.append(f'Table {p} has no tablets')
-                elif max - min > 1:
-                    warnings.append(f'Table {p} is not balanced: {min}-{max} shards.')
-                LOGGER.info(f'Table {p} balance: {min}-{max} shards.')
+                    LOGGER.debug(f'Expected nodes count: {expected_nodes_count}')
+                    if nodes_count < expected_nodes_count:
+                        errors.append(f"{expected_nodes_count - nodes_count} nodes from {expected_nodes_count} don't alive")
+                ok_node_count = 0
+                node_errors = []
+                for n in nodes:
+                    error = _check_node(n)
+                    if error:
+                        node_errors.append(error)
+                    else:
+                        ok_node_count += 1
+                if ok_node_count < nodes_count:
+                    errors.append(f'Only {ok_node_count} from {nodes_count} dynnodes are ok: {",".join(node_errors)}')
+                paths_to_balance = []
+                if isinstance(balanced_paths, str):
+                    paths_to_balance += cls.get_tables(balanced_paths)
+                elif isinstance(balanced_paths, list):
+                    for path in balanced_paths:
+                        paths_to_balance += cls.get_tables(path)
+                for p in paths_to_balance:
+                    table_nodes = cls.get_cluster_nodes(p)
+                    min = None
+                    max = None
+                    if expected_nodes_count:
+                        if len(table_nodes) < expected_nodes_count:
+                            min = 0
+                    for tn in table_nodes:
+                        tablet_count = 0
+                        for tablet in tn.tablets:
+                            if tablet.count > 0 and tablet.state != "Green":
+                                warnings.append(f'Node {tn.host}: {tablet.count} tablets of type {tablet.type} in {tablet.state} state')
+                            if tablet.type in {"ColumnShard", "DataShard"}:
+                                tablet_count += tablet.count
+                        if tablet_count > 0:
+                            if min is None or tablet_count < min:
+                                min = tablet_count
+                            if max is None or tablet_count > max:
+                                max = tablet_count
+                    if min is None or max is None:
+                        warnings.append(f'Table {p} has no tablets')
+                    elif max - min > 1:
+                        warnings.append(f'Table {p} is not balanced: {min}-{max} shards.')
+                    LOGGER.info(f'Table {p} balance: {min}-{max} shards.')
 
             cls.execute_single_result_query("select 1", timeout)
         except BaseException as ex: