kaqing 2.0.184__py3-none-any.whl → 2.0.214__py3-none-any.whl

adam/sso/cred_cache.py

@@ -2,8 +2,7 @@ import os
 from pathlib import Path
 from dotenv import load_dotenv
 
-from adam.config import Config
-from adam.utils import debug, log_exc
+from adam.utils import creating_dir, debug, log_exc
 from adam.utils_k8s.kube_context import KubeContext
 
 class CredCache:
@@ -15,7 +14,7 @@ class CredCache:
 
     def __init__(self):
         if not hasattr(self, 'env_f'):
-            self.dir = f'{Path.home()}/.kaqing'
+            self.dir = creating_dir(f'{Path.home()}/.kaqing')
             self.env_f = f'{self.dir}/.credentials'
             # immutable - cannot reload with different file content
             load_dotenv(dotenv_path=self.env_f)
@@ -44,8 +43,6 @@ class CredCache:
             updated.append(f'IDP_PASSWORD={password}')
 
         if updated:
-            if not os.path.exists(self.env_f):
-                os.makedirs(self.dir, exist_ok=True)
             with open(self.env_f, 'w') as file:
                 file.write('\n'.join(updated))
 

adam/utils.py

@@ -1,3 +1,4 @@
+from abc import ABC
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import redirect_stdout
 import copy
@@ -12,13 +13,14 @@ import random
 import string
 import threading
 import traceback
-from typing import Callable, Iterator, TypeVar, Union
+from typing import Callable, Iterator, TextIO, TypeVar, Union
 from dateutil import parser
 import subprocess
 import sys
 import time
 import click
 import yaml
+from prompt_toolkit.completion import Completer
 
 from . import __version__
 
@@ -26,10 +28,28 @@ T = TypeVar('T')
 
 log_state = threading.local()
 
-class LogConfig:
-    is_debug = lambda: False
-    is_debug_timing = lambda: False
-    is_display_help = True
+class ConfigReadable:
+    def is_debug() -> bool:
+        pass
+
+    def get(self, key: str, default: T) -> T:
+        pass
+
+class ConfigHolder:
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ConfigHolder, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self):
+        if not hasattr(self, 'config'):
+            # set by Config
+            self.config: 'ConfigReadable' = None
+            # only for testing
+            self.is_display_help = True
+            # set by ReplSession
+            self.append_command_history = lambda entry: None
 
 NO_SORT = 0
 SORT = 1
@@ -109,14 +129,24 @@ def log(s = None):
 
     return True
 
-def log2(s = None, nl = True):
+def log2(s = None, nl = True, file: str = None):
     if not loggable():
         return False
 
     if s:
-        click.echo(s, err=True, nl=nl)
+        if file:
+            with open(file, 'at') as f:
+                f.write(s)
+                if nl:
+                    f.write('\n')
+        else:
+            click.echo(s, err=True, nl=nl)
     else:
-        print(file=sys.stderr)
+        if file:
+            with open(file, 'at') as f:
+                f.write('\n')
+        else:
+            print(file=sys.stderr)
 
     return True
 
@@ -160,7 +190,11 @@ def deep_merge_dicts(dict1, dict2):
             merged_dict[key] = deep_merge_dicts(merged_dict[key], value)
         elif key not in merged_dict or value:
             # Otherwise, overwrite or add the value from dict2
-            merged_dict[key] = value
+            if key in merged_dict and isinstance(merged_dict[key], Completer):
+                pass
+                # print('SEAN completer found, ignoring', key, value)
+            else:
+                merged_dict[key] = value
     return merged_dict
 
 def deep_sort_dict(d):
@@ -198,7 +232,7 @@ def get_deep_keys(d, current_path=""):
     return keys
 
 def display_help(replace_arg = False):
-    if not LogConfig.is_display_help:
+    if not ConfigHolder().is_display_help:
         return
 
     args = copy.copy(sys.argv)
@@ -250,28 +284,10 @@ def json_to_csv(json_data: list[dict[any, any]], delimiter: str = ','):
     else:
         return None
 
-def log_to_file(config: dict[any, any]):
-    with log_exc():
-        base = f"/kaqing/logs"
-        os.makedirs(base, exist_ok=True)
-
-        now = datetime.now()
-        timestamp_str = now.strftime("%Y%m%d-%H%M%S")
-        filename = f"{base}/login.{timestamp_str}.txt"
-        with open(filename, 'w') as f:
-            if isinstance(config, dict):
-                try:
-                    json.dump(config, f, indent=4)
-                except:
-                    f.write(config)
-            else:
-                    f.write(config)
-
 def copy_config_file(rel_path: str, module: str, suffix: str = '.yaml', show_out = True):
-    dir = f'{Path.home()}/.kaqing'
+    dir = creating_dir(f'{Path.home()}/.kaqing')
     path = f'{dir}/{rel_path}'
     if not os.path.exists(path):
-        os.makedirs(dir, exist_ok=True)
         module = importlib.import_module(module)
         with open(path, 'w') as f:
             yaml.dump(module.config(), f, default_flow_style=False)
@@ -287,11 +303,15 @@ def is_lambda(func):
     return callable(func) and hasattr(func, '__name__') and func.__name__ == '<lambda>'
 
 def debug(s = None):
-    if LogConfig.is_debug():
+    if ConfigHolder().config.is_debug():
         log2(f'DEBUG {s}')
 
+def debug_complete(s = None):
+    CommandLog.log(f'DEBUG {s}', config=ConfigHolder().config.get('debugs.complete', 'off'))
+
 def debug_trace():
-    if LogConfig.is_debug():
+    if ConfigHolder().config.is_debug():
+    # if LogConfig.is_debug():
         log2(traceback.format_exc())
 
 def in_docker() -> bool:
@@ -309,34 +329,42 @@ def in_docker() -> bool:
     return False
 
 class Ing:
-    def __init__(self, msg: str, suppress_log=False):
+    def __init__(self, msg: str, suppress_log=False, job_log: str = None, condition = True):
         self.msg = msg
         self.suppress_log = suppress_log
+        self.job_log = job_log
+        self.condition = condition
 
     def __enter__(self):
+        if not self.condition:
+            return None
+
         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)
+                if not self.suppress_log and not ConfigHolder().config.is_debug():
+                    log2(f'{self.msg}...', nl=False, file=self.job_log)
 
             return None
         finally:
             log_state.ing_cnt += 1
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        if not self.condition:
+            return False
+
         log_state.ing_cnt -= 1
         if not log_state.ing_cnt:
-            if not self.suppress_log and not LogConfig.is_debug():
-                log2(' OK')
+            if not self.suppress_log and not ConfigHolder().config.is_debug():
+                log2(' OK', file=self.job_log)
 
         return False
 
-def ing(msg: str, body: Callable[[], None]=None, suppress_log=False):
+def ing(msg: str, body: Callable[[], None]=None, suppress_log=False, job_log: str = None, condition = True):
     if not body:
-        return Ing(msg, suppress_log=suppress_log)
+        return Ing(msg, suppress_log=suppress_log, job_log=job_log, condition=condition)
 
     r = None
 
@@ -350,7 +378,7 @@ def ing(msg: str, body: Callable[[], None]=None, suppress_log=False):
     return r
 
 def loggable():
-    return LogConfig.is_debug() or not hasattr(log_state, 'ing_cnt') or not log_state.ing_cnt
+    return ConfigHolder().config and ConfigHolder().config.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):
@@ -378,7 +406,7 @@ class LogTiming:
         self.s0 = s0
 
     def __enter__(self):
-        if not LogConfig.is_debug_timing():
+        if (config := ConfigHolder().config.get('debugs.timings', 'off')) not in ['on', 'file']:
             return
 
         if not hasattr(log_state, 'timings'):
@@ -390,7 +418,7 @@ class LogTiming:
             self.s0 = time.time()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if not LogConfig.is_debug_timing():
+        if (config := ConfigHolder().config.get('debugs.timings', 'off')) not in ['on', 'file']:
             return False
 
         child = log_state.timings
@@ -401,7 +429,9 @@ class LogTiming:
         log_state.timings = self.me
 
         if not self.me.depth:
-            log2(self.me.tree())
+            # log timings finally
+            CommandLog.log(self.me.tree(), config)
+
             log_state.timings = TimingNode(0)
 
         return False
@@ -410,7 +440,7 @@ 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 not ConfigHolder().config.get('debugs.timings', False):
         if body:
             return body()
 
@@ -430,7 +460,9 @@ def log_timing(msg: str, body: Callable[[], None]=None, s0: time.time = None):
 
 def timing_log_line(depth: int, msg: str, s0: time.time):
     elapsed = time.time() - s0
-    prefix = '[timings] '
+    offloaded = '-' if threading.current_thread().name.startswith('offload') or threading.current_thread().name.startswith('async') else '+'
+    prefix = f'[{offloaded} timings] '
+
     if depth:
         if elapsed > 0.01:
             prefix = ('  ' * (depth-1)) + '* '
@@ -526,7 +558,7 @@ class LogTrace:
             elif self.err_msg is not False and self.err_msg:
                 log2(self.err_msg)
 
-            if self.err_msg is not False and LogConfig.is_debug():
+            if self.err_msg is not False and ConfigHolder().config.is_debug():
                 traceback.print_exception(exc_type, exc_val, exc_tb, file=sys.stderr)
 
         # swallow exception
@@ -567,37 +599,44 @@ class ParallelService:
         else:
             return iterator
 
+thread_pools: dict[str, ThreadPoolExecutor] = {}
+thread_pool_lock = threading.Lock()
+
 class ParallelMapHandler:
-    def __init__(self, collection: list, workers: int, samples: int = sys.maxsize, msg: str = None, collect = True):
+    def __init__(self, collection: list, workers: int, samples: int = sys.maxsize, msg: str = None, collect = True, name = None):
         self.collection = collection
         self.workers = workers
         self.executor = None
         self.samples = samples
         self.msg = msg
         if msg and msg.startswith('d`'):
-            if LogConfig.is_debug():
+            if ConfigHolder().config.is_debug():
                 self.msg = msg.replace('d`', '', 1)
             else:
                 self.msg = None
         self.collect = collect
+        self.name = name
 
         self.begin = []
         self.end = []
         self.start_time = None
 
     def __enter__(self):
+        self.start_time = None
+
         self.calc_msgs()
 
         if self.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.workers)
+            self.executor = self.pool()
+            # self.executor = ThreadPoolExecutor(max_workers=self.workers)
             self.executor.__enter__()
 
         return ParallelService(self)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.executor:
+        if not self.name and self.executor:
             self.executor.__exit__(exc_type, exc_val, exc_tb)
 
         if self.end:
@@ -605,6 +644,15 @@ class ParallelMapHandler:
 
         return False
 
+    def pool(self, thread_name_prefix: str = None):
+        if not self.name:
+            return ThreadPoolExecutor(max_workers=self.workers)
+
+        if self.name not in thread_pools:
+            thread_pools[self.name] = ThreadPoolExecutor(max_workers=self.workers, thread_name_prefix=thread_name_prefix)
+
+        return thread_pools[self.name]
+
     def size(self):
         if not self.collection:
             return 0
@@ -615,25 +663,28 @@ class ParallelMapHandler:
         if not self.msg:
             return
 
+        self.begin = []
+        self.end = []
         size = self.size()
         offloaded = False
         serially = False
         sampling = False
         if size == 0:
             offloaded = True
-            self.msg = self.msg.replace('{size}', '1')
+            msg = self.msg.replace('{size}', '1')
         elif self.workers > 1 and size > 1 and self.samples == sys.maxsize:
-            self.msg = self.msg.replace('{size}', f'{size}')
+            msg = self.msg.replace('{size}', f'{size}')
         elif self.samples < sys.maxsize:
             sampling = True
+            samples = self.samples
             if self.samples > size:
-                self.samples = size
-            self.msg = self.msg.replace('{size}', f'{self.samples}/{size} sample')
+                samples = size
+            msg = self.msg.replace('{size}', f'{samples}/{size} sample')
         else:
             serially = True
-            self.msg = self.msg.replace('{size}', f'{size}')
+            msg = self.msg.replace('{size}', f'{size}')
 
-        for token in self.msg.split(' '):
+        for token in msg.split(' '):
             if '|' in token:
                 self.begin.append(token.split('|')[0])
                 if not sampling and not serially and not offloaded:
@@ -650,8 +701,19 @@ class ParallelMapHandler:
         else:
             log2(f'{" ".join(self.begin)} with {self.workers} workers...')
 
-def parallelize(collection: list, workers: int = 0, samples = sys.maxsize, msg: str = None, collect = True):
-    return ParallelMapHandler(collection, workers, samples = samples, msg = msg, collect = collect)
+# parallelizers: dict[str, ParallelMapHandler] = {}
+# parallelizer_lock = threading.Lock()
+
+def parallelize(collection: list, workers: int = 0, samples = sys.maxsize, msg: str = None, collect = True, name = None):
+    return ParallelMapHandler(collection, workers, samples = samples, msg = msg, collect = collect, name = name)
+    # if not name:
+    #     return ParallelMapHandler(collection, workers, samples = samples, msg = msg, collect = collect)
+
+    # with parallelizer_lock:
+    #     if name not in parallelizers:
+    #         parallelizers[name] = ParallelMapHandler(collection, workers, samples = samples, msg = msg, collect = collect, name = name)
+
+    # return parallelizers[name]
 
 class OffloadService:
     def __init__(self, handler: 'OffloadHandler'):
@@ -670,22 +732,24 @@ class OffloadService:
             return future
 
 class OffloadHandler(ParallelMapHandler):
-    def __init__(self, max_workers: int, msg: str = None):
-        super().__init__(None, max_workers, msg=msg, collect=False )
+    def __init__(self, max_workers: int, msg: str = None, name: str = None):
+        super().__init__(None, max_workers, msg=msg, collect=False, name=f'offload-{name}')
 
     def __enter__(self):
+        self.start_time = None
         self.calc_msgs()
 
-        if self.workers > 1 and (not self.size() or self.size()) and self.samples == sys.maxsize:
+        if self.workers > 1:
             self.start_time = time.time()
 
-            self.executor = ThreadPoolExecutor(max_workers=self.workers)
+            self.executor = self.pool(thread_name_prefix='offload')
+            # self.executor = ThreadPoolExecutor(max_workers=self.workers, thread_name_prefix='offload')
             self.executor.__enter__()
 
         return OffloadService(self)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.executor:
+        if not self.name and self.executor:
             self.executor.__exit__(exc_type, exc_val, exc_tb)
 
         if self.end:
@@ -693,38 +757,33 @@ class OffloadHandler(ParallelMapHandler):
 
         return False
 
-    def size(self):
-        if not self.collection:
-            return 0
-
-        return len(self.collection)
-
     def calc_msgs(self):
         if not self.msg:
             return
 
+        self.begin = []
+        self.end = []
         size = self.size()
-            # return
 
         offloaded = False
         serially = False
         sampling = False
         if size == 0:
             offloaded = True
-            self.msg = self.msg.replace('{size}', '1')
+            msg = self.msg.replace('{size}', '1')
         elif self.workers > 1 and size > 1 and self.samples == sys.maxsize:
-            self.msg = self.msg.replace('{size}', f'{size}')
+            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')
+            samples = self.samples
+            if samples > size:
+                samples = size
+            msg = self.msg.replace('{size}', f'{samples}/{size} sample')
         else:
             serially = True
-            self.msg = self.msg.replace('{size}', f'{size}')
-            # return
+            msg = self.msg.replace('{size}', f'{size}')
 
-        for token in self.msg.split(' '):
+        for token in msg.split(' '):
             if '|' in token:
                 self.begin.append(token.split('|')[0])
                 if not sampling and not serially and not offloaded:
@@ -741,5 +800,83 @@ class OffloadHandler(ParallelMapHandler):
         else:
             log2(f'{" ".join(self.begin)} with {self.workers} workers...')
 
-def offload(max_workers: int = 3, msg: str = None):
-    return OffloadHandler(max_workers, msg = msg)
+def offload(max_workers: int = 3, msg: str = None, name: str = None):
+    return OffloadHandler(max_workers, msg = msg, name = name)
+
+def kaqing_log_file_name(suffix = 'log'):
+    return f"{log_dir()}/{datetime.now().strftime('%d%H%M%S')}.{suffix}"
+
+def log_dir():
+    return creating_dir(ConfigHolder().config.get('log-dir', '/tmp/qing-db/q/logs'))
+
+class LogFileHandler:
+    def __init__(self, suffix = 'log', condition=True):
+        self.suffix = suffix
+        self.condition = condition
+
+    def __enter__(self):
+        self.f = None
+        if self.condition:
+            self.f = open(kaqing_log_file_name(suffix=self.suffix), 'w')
+            self.f.__enter__()
+
+        return self.f
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.f:
+            self.f.__exit__(exc_type, exc_val, exc_tb)
+
+            if ConfigHolder().append_command_history:
+                ConfigHolder().append_command_history(f':cat {self.f.name}')
+
+        return False
+
+def kaqing_log_file(suffix = 'log', condition=True):
+    return LogFileHandler(suffix = suffix, condition=condition)
+
+class CommandLog:
+    log_file = None
+
+    def log(line: str, config: str = 'off'):
+        if config == 'file':
+            if not CommandLog.log_file:
+                try:
+                    CommandLog.log_file = open(kaqing_log_file_name(suffix='cmd.log'), 'w')
+                except:
+                    pass
+
+            try:
+                CommandLog.log_file.write(line + '\n')
+            except:
+                pass
+        elif config == 'on':
+            log2(line)
+
+    def close_log_file():
+        if CommandLog.log_file:
+            try:
+                CommandLog.log_file.close()
+            except:
+                pass
+
+            if ConfigHolder().append_command_history:
+                ConfigHolder().append_command_history(f':cat {CommandLog.log_file.name}')
+
+            CommandLog.log_file = None
+
+class ExecResult(ABC):
+    def exit_code(self) -> int:
+        pass
+
+    def cat_log_file_cmd(self) -> str:
+        pass
+
+_dirs_created = set()
+
+def creating_dir(dir):
+    if dir not in _dirs_created:
+        _dirs_created.add(dir)
+        if not os.path.exists(dir):
+            os.makedirs(dir, exist_ok=True)
+
+    return dir

adam/utils_k8s/app_clusters.py

@@ -2,7 +2,7 @@ import sys
 from typing import TypeVar
 
 from adam.utils_k8s.app_pods import AppPods
-from adam.pod_exec_result import PodExecResult
+from adam.utils_k8s.pod_exec_result import PodExecResult
 from adam.utils import log, log2
 from adam.utils_k8s.pods import Pods
 from .kube_context import KubeContext
@@ -11,8 +11,15 @@ T = TypeVar('T')
 
 # utility collection on app clusters; methods are all static
 class AppClusters:
-    def exec(pods: list[str], namespace: str, command: str, action: str = 'action',
-             max_workers=0, show_out=True, on_any = False, shell = '/bin/sh', backgrounded = False) -> list[PodExecResult]:
+    def exec(pods: list[str],
+             namespace: str,
+             command: str,
+             action: str = 'action',
+             max_workers=0,
+             show_out=True,
+             on_any = False,
+             shell = '/bin/sh',
+             backgrounded = False) -> list[PodExecResult]:
         samples = 1 if on_any else sys.maxsize
         msg = 'd`Running|Ran ' + action + ' command onto {size} pods'
         with Pods.parallelize(pods, max_workers, samples, msg, action=action) as exec:
@@ -23,6 +30,6 @@ class AppClusters:
                     if result.stdout:
                         log(result.stdout)
                     if result.stderr:
-                        log2(result.stderr, file=sys.stderr)
+                        log2(result.stderr)
 
             return results

adam/utils_k8s/app_pods.py

@@ -4,7 +4,7 @@ from kubernetes import client
 
 from adam.config import Config
 from adam.utils_k8s.pods import Pods
-from adam.pod_exec_result import PodExecResult
+from adam.utils_k8s.pod_exec_result import PodExecResult
 from adam.repl_session import ReplSession
 
 # utility collection on app pods; methods are all static
@@ -25,12 +25,17 @@ class AppPods:
 
         return v1.list_namespaced_pod(namespace, label_selector=label_selector).items
 
-    def exec(pod_name: str, namespace: str, command: str, show_out = True, throw_err = False, shell = '/bin/sh', backgrounded = False) -> PodExecResult:
+    def exec(pod_name: str,
+             namespace: str,
+             command: str,
+             show_out = True,
+             throw_err = False,
+             shell = '/bin/sh',
+             backgrounded = False) -> PodExecResult:
         container = Config().get('app.container-name', 'c3-server')
         r = Pods.exec(pod_name, container, namespace, command, show_out = show_out, throw_err = throw_err, shell = shell, backgrounded = backgrounded)
 
-        if r and Config().get('repl.history.push-cat-remote-log-file', True):
-            if r.log_file and ReplSession().prompt_session:
-                ReplSession().prompt_session.history.append_string(f'@{r.pod} cat {r.log_file}')
+        if r and r.log_file:
+            ReplSession().append_history(f':cat {r.log_file}')
 
         return r

adam/utils_k8s/cassandra_clusters.py

@@ -3,7 +3,7 @@ from typing import TypeVar
 
 from adam.config import Config
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
-from adam.pod_exec_result import PodExecResult
+from adam.utils_k8s.pod_exec_result import PodExecResult
 from adam.utils import log, log2
 from adam.utils_k8s.pods import Pods
 from adam.utils_k8s.statefulsets import StatefulSets
@@ -21,20 +21,24 @@ class CassandraClusters:
              on_any = False,
              shell = '/bin/sh',
              backgrounded = False,
-             log_file = None) -> list[PodExecResult]:
+             log_file = None,
+             history=True) -> list[PodExecResult]:
 
         pods = StatefulSets.pod_names(sts, namespace)
         samples = 1 if on_any else sys.maxsize
+        if (backgrounded or command.endswith(' &')) and not log_file:
+            log_file = Pods.log_file(command)
+
         msg = 'd`Running|Ran ' + action + ' command onto {size} pods'
         with Pods.parallelize(pods, max_workers, samples, msg, action=action) as exec:
-            results: list[PodExecResult] = exec.map(lambda pod: CassandraNodes.exec(pod, namespace, command, False, False, shell, backgrounded, log_file))
+            results: list[PodExecResult] = exec.map(lambda pod: CassandraNodes.exec(pod, namespace, command, False, False, shell, backgrounded, log_file, history))
             for result in results:
                 if show_out and not Config().is_debug():
                     log(result.command)
                     if result.stdout:
                         log(result.stdout)
                     if result.stderr:
-                        log2(result.stderr, file=sys.stderr)
+                        log2(result.stderr)
 
             return results