From 02133f93b5d3bf81712e2d546dbea473f1c7cbb4 Mon Sep 17 00:00:00 2001 From: Kriukov Aleksandr Date: Fri, 27 Jun 2025 16:07:37 +0300 Subject: [PATCH 1/6] Add iam token argument --- ydb/apps/dstool/lib/common.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py index fae1480f1588..c3bdac4521f0 100644 --- a/ydb/apps/dstool/lib/common.py +++ b/ydb/apps/dstool/lib/common.py @@ -152,14 +152,21 @@ def make_url(self, endpoint, path, params): location = endpoint.host_with_port return urllib.parse.urlunsplit((endpoint.protocol, location, path, urllib.parse.urlencode(params), '')) - def parse_token(self, token_file): + def parse_token(self, token_file, iam_token_file): if token_file: self.token = token_file.readline().rstrip('\r\n') token_file.close() + if iam_token_file: + self.token = iam_token_file.readline().rstrip('\r\n') + self.token_type = 'Bearer' + iam_token_file.close() if self.token is None: self.token = os.getenv('YDB_TOKEN') if self.token is not None: self.token = self.token.strip() + if self.token is None: + self.token = os.getenv('IAM_TOKEN') + self.token_type = 'Bearer' if self.token is None: try: path = os.path.expanduser(os.path.join('~', '.ydb', 'token')) @@ -199,7 +206,7 @@ def apply_args(self, args, with_localhost=True): if 'http' not in protocols and 'https' in protocols: self.mon_protocol = 'https' - self.parse_token(args.token_file) + self.parse_token(args.token_file, args.iam_token_file) self.domain = 1 self.verbose = args.verbose or args.debug self.debug = args.debug @@ -218,7 +225,9 @@ def add_host_access_options(self, parser, with_endpoint=True): g.add_argument('--endpoint', '-e', metavar='[PROTOCOL://]HOST[:PORT]', type=str, required=True, action='append', help=ConnectionParams.ENDPOINT_HELP) g.add_argument('--grpc-port', type=int, default=2135, metavar='PORT', help='GRPC port to use for procedure invocation') g.add_argument('--mon-port', type=int, default=8765, metavar='PORT', help='HTTP monitoring port for viewer JSON access') - g.add_argument('--token-file', type=FileType(encoding='ascii'), metavar='PATH', help='Path to token file') + token_group = g.add_mutually_exclusive_group() + token_group.add_argument('--token-file', type=FileType(encoding='ascii'), metavar='PATH', help='Path to token file') + token_group.add_argument('--iam-token-file', type=FileType(encoding='ascii'), metavar='PATH', help='Path to IAM token file') g.add_argument('--ca-file', metavar='PATH', dest='cafile', type=str, help='File containing PEM encoded root certificates for SSL/TLS connections. ' 'If this parameter is empty, the default roots will be used.') g.add_argument('--http-timeout', type=int, default=5, help='Timeout for blocking socket I/O operations during HTTP(s) queries') From 89e5efbb90624267fbade64a321b975fffa218c1 Mon Sep 17 00:00:00 2001 From: Kriukov Aleksandr Date: Fri, 27 Jun 2025 16:34:05 +0300 Subject: [PATCH 2/6] improve token parsing --- ydb/apps/dstool/lib/common.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py index c3bdac4521f0..8fd56c0b4e79 100644 --- a/ydb/apps/dstool/lib/common.py +++ b/ydb/apps/dstool/lib/common.py @@ -153,13 +153,11 @@ def make_url(self, endpoint, path, params): return urllib.parse.urlunsplit((endpoint.protocol, location, path, urllib.parse.urlencode(params), '')) def parse_token(self, token_file, iam_token_file): - if token_file: - self.token = token_file.readline().rstrip('\r\n') - token_file.close() + if token_file or iam_token_file: + self.token = self.read_stripped_file_and_close(token_file) if iam_token_file: - self.token = iam_token_file.readline().rstrip('\r\n') + self.token = self.read_stripped_file_and_close(iam_token_file) self.token_type = 'Bearer' - iam_token_file.close() if self.token is None: self.token = os.getenv('YDB_TOKEN') if self.token is not None: @@ -180,6 +178,13 @@ def parse_token(self, token_file, iam_token_file): else: self.token_type = 'OAuth' + def read_stripped_file_and_close(self, token_file): + if token_file is None: + return None + result = token_file.readline().rstrip('\r\n') + token_file.close() + return result + def apply_args(self, args, with_localhost=True): self.args = args self.grpc_port = args.grpc_port From 680decc6015685239d8f34ba48762b410c86be57 Mon Sep 17 00:00:00 2001 From: Kriukov Aleksandr Date: Fri, 27 Jun 2025 16:37:05 +0300 Subject: [PATCH 3/6] fix --- ydb/apps/dstool/lib/common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py index 8fd56c0b4e79..a9f97e82b73d 100644 --- a/ydb/apps/dstool/lib/common.py +++ b/ydb/apps/dstool/lib/common.py @@ -153,9 +153,9 @@ def make_url(self, endpoint, path, params): return urllib.parse.urlunsplit((endpoint.protocol, location, path, urllib.parse.urlencode(params), '')) def parse_token(self, token_file, iam_token_file): - if token_file or iam_token_file: + if token_file: self.token = self.read_stripped_file_and_close(token_file) - if iam_token_file: + elif iam_token_file: self.token = self.read_stripped_file_and_close(iam_token_file) self.token_type = 'Bearer' if self.token is None: From 19aef9f6c222e61f117a4a81471d86cd9a944d6e Mon Sep 17 00:00:00 2001 From: Kriukov Aleksandr Date: Fri, 27 Jun 2025 16:58:56 +0300 Subject: [PATCH 4/6] rewrite parse token method --- ydb/apps/dstool/lib/common.py | 71 ++++++++++++++++++++++------------- 1 file changed, 44 insertions(+), 27 deletions(-) diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py index a9f97e82b73d..2335a9951b9a 100644 --- a/ydb/apps/dstool/lib/common.py +++ b/ydb/apps/dstool/lib/common.py @@ -152,38 +152,55 @@ def make_url(self, endpoint, path, params): location = endpoint.host_with_port return urllib.parse.urlunsplit((endpoint.protocol, location, path, urllib.parse.urlencode(params), '')) - def parse_token(self, token_file, iam_token_file): + def parse_token(self, token_file, iam_token_file=None): if token_file: - self.token = self.read_stripped_file_and_close(token_file) - elif iam_token_file: - self.token = self.read_stripped_file_and_close(iam_token_file) - self.token_type = 'Bearer' - if self.token is None: - self.token = os.getenv('YDB_TOKEN') + self.token_type, self.token = self.read_token_from_file(token_file, 'OAuth') + token_file.close() + return + + if iam_token_file: + self.token_type, self.token = self.read_token_from_file(iam_token_file, 'Bearer') + iam_token_file.close() + return + + token_value = os.getenv('YDB_TOKEN') + if token_value is not None: + self.token_type, self.token = self.parse_token_value(token_value, 'OAuth') + return + + token_value = os.getenv('IAM_TOKEN') + if token_value is not None: + self.token_type, self.token = self.parse_token_value(token_value, 'Bearer') + return + + default_token_paths = [ + ('OAuth', os.path.expanduser(os.path.join('~', '.ydb', 'token'))), + ('Bearer', os.path.expanduser(os.path.join('~', '.ydb', 'iam_token'))), + ] + for token_type, token_file_path in default_token_paths: + self.token_type, self.token = self.read_token_file(token_file_path, token_type) if self.token is not None: - self.token = self.token.strip() - if self.token is None: - self.token = os.getenv('IAM_TOKEN') - self.token_type = 'Bearer' - if self.token is None: - try: - path = os.path.expanduser(os.path.join('~', '.ydb', 'token')) - with open(path) as f: - self.token = f.readline().strip('\r\n') - except Exception: - pass - - if self.token is not None and len(self.token.split(' ')) == 2: - self.token_type, self.token = self.token.split(' ') - else: - self.token_type = 'OAuth' + return - def read_stripped_file_and_close(self, token_file): + def read_token_from_file(self, token_file, default_token_type): if token_file is None: + return default_token_type, None + token_value = token_file.readline().rstrip('\r\n') + return self.parse_token_value(token_value, default_token_type) + + def read_token_file(self, token_file_path, default_token_type): + if token_file_path is None: + return default_token_type, None + return self.read_token_from_file_and_close(open(token_file_path, 'r'), default_token_type) + + def parse_token_value(self, token_value, default_token_type): + if token_value is None: return None - result = token_file.readline().rstrip('\r\n') - token_file.close() - return result + splitted = token_value.strip().split(' ') + if len(splitted) == 2: + return splitted + else: + return default_token_type, token_value def apply_args(self, args, with_localhost=True): self.args = args From d40e6d68db63cd877f3f7ca214c508cc360bb0da Mon Sep 17 00:00:00 2001 From: Kriukov Aleksandr Date: Fri, 27 Jun 2025 17:02:12 +0300 Subject: [PATCH 5/6] fix bug in parse_token_value --- ydb/apps/dstool/lib/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py index 2335a9951b9a..d4edaca64249 100644 --- a/ydb/apps/dstool/lib/common.py +++ b/ydb/apps/dstool/lib/common.py @@ -195,7 +195,7 @@ def read_token_file(self, token_file_path, default_token_type): def parse_token_value(self, token_value, default_token_type): if token_value is None: - return None + return default_token_type, None splitted = token_value.strip().split(' ') if len(splitted) == 2: return splitted From db493fea265498ff67d8c4adfb8315056161fc50 Mon Sep 17 00:00:00 2001 From: Kriukov Aleksandr Date: Fri, 27 Jun 2025 17:41:07 +0300 Subject: [PATCH 6/6] fix missed try except block --- ydb/apps/dstool/lib/common.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ydb/apps/dstool/lib/common.py b/ydb/apps/dstool/lib/common.py index d4edaca64249..8e26c94d4009 100644 --- a/ydb/apps/dstool/lib/common.py +++ b/ydb/apps/dstool/lib/common.py @@ -191,7 +191,10 @@ def read_token_from_file(self, token_file, default_token_type): def read_token_file(self, token_file_path, default_token_type): if token_file_path is None: return default_token_type, None - return self.read_token_from_file_and_close(open(token_file_path, 'r'), default_token_type) + try: + return self.read_token_from_file_and_close(open(token_file_path, 'r'), default_token_type) + except Exception: + return default_token_type, None def parse_token_value(self, token_value, default_token_type): if token_value is None: