kaqing 2.0.98__py3-none-any.whl → 2.0.171__py3-none-any.whl

adam/sso/authn_ad.py

@@ -8,6 +8,7 @@ from urllib.parse import urlparse, parse_qs
 from adam.log import Log
 from adam.sso.authenticator import Authenticator
 from adam.sso.id_token import IdToken
+from adam.utils import debug
 from .idp_login import IdpLogin
 from adam.config import Config
 
@@ -33,7 +34,7 @@ class AdAuthenticator(Authenticator):
 
         session = requests.Session()
         r = session.get(idp_uri)
-        Config().debug(f'{r.status_code} {idp_uri}')
+        debug(f'{r.status_code} {idp_uri}')
 
         config = self.validate_and_return_config(r)
 
@@ -52,7 +53,7 @@ class AdAuthenticator(Authenticator):
         r = session.post(login_uri, data=body, headers={
             'Content-Type': 'application/x-www-form-urlencoded'
         })
-        Config().debug(f'{r.status_code} {login_uri}')
+        debug(f'{r.status_code} {login_uri}')
 
         config = self.validate_and_return_config(r)
 
@@ -69,7 +70,7 @@ class AdAuthenticator(Authenticator):
         r = session.post(kmsi_uri, data=body, headers={
             'Content-Type': 'application/x-www-form-urlencoded'
         })
-        Config().debug(f'{r.status_code} {kmsi_uri}')
+        debug(f'{r.status_code} {kmsi_uri}')
 
         if (config := self.extract_config_object(r.text)):
             if 'sErrorCode' in config and config['sErrorCode'] == '50058':
@@ -101,7 +102,7 @@ class AdAuthenticator(Authenticator):
 
     def validate_and_return_config(self, r: requests.Response):
         if r.status_code < 200 or r.status_code >= 300:
-            Config().debug(r.text)
+            debug(r.text)
 
             return None
 
@@ -164,6 +165,6 @@ class AdAuthenticator(Authenticator):
                 exp=data['exp'] if 'exp' in data else 0
             )
         except:
-            Config().debug(traceback.format_exc())
+            debug(traceback.format_exc())
 
         return None

adam/sso/authn_okta.py

@@ -8,7 +8,7 @@ from adam.sso.id_token import IdToken
 
 from .idp_login import IdpLogin
 from adam.config import Config
-from adam.utils import log2
+from adam.utils import debug, log2
 
 class OktaException(Exception):
     pass
@@ -49,7 +49,7 @@ class OktaAuthenticator(Authenticator):
 
         session = requests.Session()
         response = session.post(authn_uri, headers=headers, data=json.dumps(payload))
-        Config().debug(f'{response.status_code} {authn_uri}')
+        debug(f'{response.status_code} {authn_uri}')
         auth_response = response.json()
 
         if 'sessionToken' not in auth_response:
@@ -59,7 +59,7 @@ class OktaAuthenticator(Authenticator):
 
         url = f'{idp_uri}&sessionToken={session_token}'
         r = session.get(url)
-        Config().debug(f'{r.status_code} {url}')
+        debug(f'{r.status_code} {url}')
 
         id_token = OktaAuthenticator().extract(r.text, r'.*name=\"id_token\" value=\"(.*?)\".*')
         if not id_token:

adam/sso/cred_cache.py

@@ -4,6 +4,7 @@ import traceback
 from dotenv import load_dotenv
 
 from adam.config import Config
+from adam.utils import debug
 from adam.utils_k8s.kube_context import KubeContext
 
 class CredCache:
@@ -37,7 +38,7 @@ class CredCache:
                 try:
                     file.truncate()
                 except:
-                    Config().debug(traceback.format_exc())
+                    debug(traceback.format_exc())
 
         updated = []
         updated.append(f'IDP_USERNAME={username}')
@@ -56,4 +57,4 @@ class CredCache:
             if password:
                 self.overrides['IDP_PASSWORD'] = password
 
-            Config().debug(f'Cached username: {username}, password: {password}, try load: {self.get_username()}')
+            debug(f'Cached username: {username}, password: {password}, try load: {self.get_username()}')

adam/sso/idp.py

@@ -15,7 +15,7 @@ from .cred_cache import CredCache
 from .idp_session import IdpSession
 from .idp_login import IdpLogin
 from adam.config import Config
-from adam.utils import log, log2
+from adam.utils import debug, log
 
 T = TypeVar('T')
 
@@ -57,7 +57,7 @@ class Idp:
                 default_user: str = None
                 if use_cached_creds:
                     default_user = CredCache().get_username()
-                    Config().debug(f'User read from cache: {default_user}')
+                    debug(f'User read from cache: {default_user}')
 
                 if from_env := os.getenv('USERNAME'):
                     default_user = from_env
@@ -137,7 +137,7 @@ class Idp:
 
                 return IdpLogin(None, None, None, username)
         except:
-            Config().debug(traceback.format_exc())
+            debug(traceback.format_exc())
             pass
 
         return None

adam/utils.py

@@ -1,3 +1,5 @@
+from abc import abstractmethod
+from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import redirect_stdout
 import copy
 import csv
@@ -9,6 +11,8 @@ import os
 from pathlib import Path
 import random
 import string
+import threading
+from typing import Callable, Iterator, TypeVar
 from dateutil import parser
 import subprocess
 import sys
@@ -18,6 +22,13 @@ import yaml
 
 from . import __version__
 
+log_state = threading.local()
+
+class LogConfig:
+    is_debug = lambda: False
+    is_debug_timing = lambda: False
+    is_display_help = True
+
 def to_tabular(lines: str, header: str = None, dashed_line = False):
     return lines_to_tabular(lines.split('\n'), header, dashed_line)
 
@@ -69,18 +80,28 @@ def epoch(timestamp_string: str):
     return parser.parse(timestamp_string).timestamp()
 
 def log(s = None):
+    if not loggable():
+        return False
+
     # want to print empty line for False or empty collection
     if s == None:
         print()
     else:
         click.echo(s)
 
+    return True
+
 def log2(s = None, nl = True):
+    if not loggable():
+        return False
+
     if s:
         click.echo(s, err=True, nl=nl)
     else:
         print(file=sys.stderr)
 
+    return True
+
 def elapsed_time(start_time: float):
     end_time = time.time()
     elapsed_time = end_time - start_time
@@ -95,8 +116,8 @@ def duration(start_time: float, end_time: float = None):
         end_time = time.time()
     d = convert_seconds(end_time - start_time)
     t = []
-    if d[0]:
-        t.append(f'{d[0]}h')
+    if d:
+        t.append(f'{d}h')
     if t or d[1]:
         t.append(f'{d[1]}m')
     t.append(f'{d[2]}s')
@@ -159,6 +180,9 @@ def get_deep_keys(d, current_path=""):
     return keys
 
 def display_help(replace_arg = False):
+    if not LogConfig.is_display_help:
+        return
+
     args = copy.copy(sys.argv)
     if replace_arg:
         args[len(args) - 1] = '--help'
@@ -203,6 +227,7 @@ def json_to_csv(json_data: list[dict[any, any]], delimiter: str = ','):
         with redirect_stdout(body) as f:
             dict_writer = csv.DictWriter(f, keys, delimiter=delimiter)
             dict_writer.writerows(flattened_data)
+
         return header.getvalue().strip('\r\n'), [l.strip('\r') for l in body.getvalue().split('\n')]
     else:
         return None
@@ -240,4 +265,415 @@ def copy_config_file(rel_path: str, module: str, suffix: str = '.yaml', show_out
     return path
 
 def idp_token_from_env():
-    return os.getenv('IDP_TOKEN')
+    return os.getenv('IDP_TOKEN')
+
+def is_lambda(func):
+    return callable(func) and hasattr(func, '__name__') and func.__name__ == '<lambda>'
+
+def debug(s = None):
+    if LogConfig.is_debug():
+        log2(f'DEBUG {s}')
+
+class Ing:
+    def __init__(self, msg: str, suppress_log=False):
+        self.msg = msg
+        self.suppress_log = suppress_log
+
+    def __enter__(self):
+        if not hasattr(log_state, 'ing_cnt'):
+            log_state.ing_cnt = 0
+
+        try:
+            if not log_state.ing_cnt:
+                if not self.suppress_log and not LogConfig.is_debug():
+                    log2(f'{self.msg}...', nl=False)
+
+            return None
+        finally:
+            log_state.ing_cnt += 1
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        log_state.ing_cnt -= 1
+        if not log_state.ing_cnt:
+            if not self.suppress_log and not LogConfig.is_debug():
+                log2(' OK')
+
+        return False
+
+def ing(msg: str, body: Callable[[], None]=None, suppress_log=False):
+    if not body:
+        return Ing(msg, suppress_log=suppress_log)
+
+    r = None
+
+    t = Ing(msg, suppress_log=suppress_log)
+    t.__enter__()
+    try:
+        r = body()
+    finally:
+        t.__exit__(None, None, None)
+
+    return r
+
+def loggable():
+    return LogConfig.is_debug() or not hasattr(log_state, 'ing_cnt') or not log_state.ing_cnt
+
+class TimingNode:
+    def __init__(self, depth: int, s0: time.time = time.time(), line: str = None):
+        self.depth = depth
+        self.s0 = s0
+        self.line = line
+        self.children = []
+
+    def __str__(self):
+        return f'[{self.depth}: {self.line}, children={len(self.children)}]'
+
+    def tree(self):
+        lines = []
+        if self.line:
+            lines.append(self.line)
+
+        for child in self.children:
+            if child.line:
+                lines.append(child.tree())
+        return '\n'.join(lines)
+
+class LogTiming:
+    def __init__(self, msg: str, s0: time.time = None):
+        self.msg = msg
+        self.s0 = s0
+
+    def __enter__(self):
+        if not LogConfig.is_debug_timing():
+            return
+
+        if not hasattr(log_state, 'timings'):
+            log_state.timings = TimingNode(0)
+
+        self.me = log_state.timings
+        log_state.timings = TimingNode(self.me.depth+1)
+        if not self.s0:
+            self.s0 = time.time()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if not LogConfig.is_debug_timing():
+            return False
+
+        child = log_state.timings
+        log_state.timings.line = timing_log_line(self.me.depth, self.msg, self.s0)
+
+        if child and child.line:
+            self.me.children.append(child)
+        log_state.timings = self.me
+
+        if not self.me.depth:
+            log2(self.me.tree())
+            log_state.timings = TimingNode(0)
+
+        return False
+
+def log_timing(msg: str, body: Callable[[], None]=None, s0: time.time = None):
+    if not s0 and not body:
+        return LogTiming(msg, s0=s0)
+
+    if not LogConfig.is_debug_timing():
+        if body:
+            return body()
+
+        return
+
+    r = None
+
+    t = LogTiming(msg, s0=s0)
+    t.__enter__()
+    try:
+        if body:
+            r = body()
+    finally:
+        t.__exit__(None, None, None)
+
+    return r
+
+def timing_log_line(depth: int, msg: str, s0: time.time):
+    elapsed = time.time() - s0
+    prefix = '[timings] '
+    if depth:
+        if elapsed > 0.01:
+            prefix = ('  ' * (depth-1)) + '* '
+        else:
+            prefix = '  ' * depth
+
+    return f'{prefix}{msg}: {elapsed:.2f} sec'
+
+class WaitLog:
+    wait_log_flag = False
+
+def wait_log(msg: str):
+    if not WaitLog.wait_log_flag:
+        log2(msg)
+        WaitLog.wait_log_flag = True
+
+def clear_wait_log_flag():
+    WaitLog.wait_log_flag = False
+
+T = TypeVar('T')
+
+class ParallelService:
+    def __init__(self, handler: 'ParallelMapHandler'):
+        self.handler = handler
+
+    def map(self, fn: Callable[..., T]) -> Iterator[T]:
+        executor = self.handler.executor
+        collection = self.handler.collection
+        collect = self.handler.collect
+        samples_cnt = self.handler.samples
+
+        iterator = None
+        if executor:
+            iterator = executor.map(fn, collection)
+        elif samples_cnt < sys.maxsize:
+            samples = []
+
+            for elem in collection:
+                if not samples_cnt:
+                    break
+
+                samples.append(fn(elem))
+                samples_cnt -= 1
+
+            iterator = iter(samples)
+        else:
+            iterator = map(fn, collection)
+
+        if collect:
+            return list(iterator)
+        else:
+            return iterator
+
+class ParallelMapHandler:
+    def __init__(self, collection: list, max_workers: int, samples: int = sys.maxsize, msg: str = None, collect = True):
+        self.collection = collection
+        self.max_workers = max_workers
+        self.executor = None
+        self.samples = samples
+        self.msg = msg
+        if msg and msg.startswith('d`'):
+            if LogConfig.is_debug():
+                self.msg = msg.replace('d`', '', 1)
+            else:
+                self.msg = None
+        self.collect = collect
+
+        self.begin = []
+        self.end = []
+        self.start_time = None
+
+    def __enter__(self):
+        self.calc_msgs()
+
+        if self.max_workers > 1 and (not self.size() or self.size()) and self.samples == sys.maxsize:
+            self.start_time = time.time()
+
+            self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
+            self.executor.__enter__()
+
+        return ParallelService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.executor:
+            self.executor.__exit__(exc_type, exc_val, exc_tb)
+
+        if self.end:
+            log2(f'{" ".join(self.end)} in {elapsed_time(self.start_time)}.')
+
+        return False
+
+    def size(self):
+        if not self.collection:
+            return 0
+
+        return len(self.collection)
+
+    def calc_msgs(self):
+        if not self.msg:
+            return
+
+        size = self.size()
+            # return
+
+        offloaded = False
+        serially = False
+        sampling = False
+        if size == 0:
+            offloaded = True
+            self.msg = self.msg.replace('{size}', '1')
+        elif self.max_workers > 1 and size > 1 and self.samples == sys.maxsize:
+            self.msg = self.msg.replace('{size}', f'{size}')
+        elif self.samples < sys.maxsize:
+            sampling = True
+            if self.samples > size:
+                self.samples = size
+            self.msg = self.msg.replace('{size}', f'{self.samples}/{size} sample')
+        else:
+            serially = True
+            self.msg = self.msg.replace('{size}', f'{size}')
+            # return
+
+        for token in self.msg.split(' '):
+            if '|' in token:
+                self.begin.append(token.split('|')[0])
+                if not sampling and not serially and not offloaded:
+                    self.end.append(token.split('|')[1])
+            else:
+                self.begin.append(token)
+                if not sampling and not serially and not offloaded:
+                    self.end.append(token)
+
+        if offloaded:
+            log2(f'{" ".join(self.begin)} offloaded...')
+        elif sampling or serially:
+            log2(f'{" ".join(self.begin)} serially...')
+        else:
+            log2(f'{" ".join(self.begin)} with {self.max_workers} workers...')
+
+class OffloadService:
+    def __init__(self, handler: 'OffloadHandler'):
+        self.handler = handler
+
+    def submit(self, fn: Callable[..., T], /, *args, **kwargs) -> Future[T]:
+        executor = self.handler.executor
+
+        if executor:
+            return executor.submit(fn, *args, **kwargs)
+        else:
+            future = Future()
+
+            future.set_result(fn(*args, **kwargs))
+
+            return future
+
+class OffloadHandler(ParallelMapHandler):
+    def __init__(self, max_workers: int, msg: str = None):
+        super().__init__(None, max_workers, msg=msg, collect=False )
+
+    def __enter__(self):
+        self.calc_msgs()
+
+        if self.max_workers > 1 and (not self.size() or self.size()) and self.samples == sys.maxsize:
+            self.start_time = time.time()
+
+            self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
+            self.executor.__enter__()
+
+        return OffloadService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.executor:
+            self.executor.__exit__(exc_type, exc_val, exc_tb)
+
+        if self.end:
+            log2(f'{" ".join(self.end)} in {elapsed_time(self.start_time)}.')
+
+        return False
+
+    def size(self):
+        if not self.collection:
+            return 0
+
+        return len(self.collection)
+
+    def calc_msgs(self):
+        if not self.msg:
+            return
+
+        size = self.size()
+            # return
+
+        offloaded = False
+        serially = False
+        sampling = False
+        if size == 0:
+            offloaded = True
+            self.msg = self.msg.replace('{size}', '1')
+        elif self.max_workers > 1 and size > 1 and self.samples == sys.maxsize:
+            self.msg = self.msg.replace('{size}', f'{size}')
+        elif self.samples < sys.maxsize:
+            sampling = True
+            if self.samples > size:
+                self.samples = size
+            self.msg = self.msg.replace('{size}', f'{self.samples}/{size} sample')
+        else:
+            serially = True
+            self.msg = self.msg.replace('{size}', f'{size}')
+            # return
+
+        for token in self.msg.split(' '):
+            if '|' in token:
+                self.begin.append(token.split('|')[0])
+                if not sampling and not serially and not offloaded:
+                    self.end.append(token.split('|')[1])
+            else:
+                self.begin.append(token)
+                if not sampling and not serially and not offloaded:
+                    self.end.append(token)
+
+        if offloaded:
+            log2(f'{" ".join(self.begin)} offloaded...')
+        elif sampling or serially:
+            log2(f'{" ".join(self.begin)} serially...')
+        else:
+            log2(f'{" ".join(self.begin)} with {self.max_workers} workers...')
+
+# class ParallelMapHandler(ParallelHandler):
+#     def __enter__(self):
+#         self.calc_msgs()
+
+#         if self.max_workers > 1 and (not self.size() or self.size()) and self.samples == sys.maxsize:
+#             self.start_time = time.time()
+
+#             self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
+#             self.executor.__enter__()
+
+#         if self.collection:
+#             return self.map
+#         else:
+#             return self.submit
+
+#     def map(self, fn: Callable[..., T]) -> Iterator[T]:
+#         iterator = None
+#         if self.executor:
+#             iterator = self.executor.map(fn, self.collection)
+#         elif self.samples < sys.maxsize:
+#             samples = []
+
+#             for elem in self.collection:
+#                 if not self.samples:
+#                     break
+
+#                 samples.append(fn(elem))
+#                 self.samples -= 1
+
+#             iterator = iter(samples)
+#         else:
+#             iterator = map(fn, self.collection)
+
+#         if self.collect:
+#             return list(iterator)
+#         else:
+#             return iterator
+
+#     def submit(self, fn: Callable[..., T], /, *args, **kwargs) -> Future[T]:
+#         if self.executor:
+#             return self.executor.submit(fn, *args, **kwargs)
+#         else:
+#             future = Future()
+
+#             future.set_result(fn(*args, **kwargs))
+
+#             return future
+
+def parallelize(collection: list, max_workers: int = 0, samples = sys.maxsize, msg: str = None, collect = True):
+    return ParallelMapHandler(collection, max_workers, samples = samples, msg = msg, collect = collect)
+
+def offload(max_workers: int = 3, msg: str = None):
+    return OffloadHandler(max_workers, msg = msg)