kaqing 2.0.172__py3-none-any.whl → 2.0.186__py3-none-any.whl

adam/utils.py

@@ -20,19 +20,33 @@ import time
 import click
 import yaml
 
+from prompt_toolkit.completion import Completer
+
 from . import __version__
 
+T = TypeVar('T')
+
 log_state = threading.local()
 
 class LogConfig:
     is_debug = lambda: False
     is_debug_timing = lambda: False
+    is_debug_complete = 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)
+NO_SORT = 0
+SORT = 1
+REVERSE_SORT = -1
+
+def tabulize(lines: list[T], fn: Callable[..., T] = None, header: str = None, dashed_line = False, separator = ' ', to: int = 1, sorted: int = NO_SORT):
+    if fn:
+        lines = list(map(fn, lines))
+
+    if sorted == SORT:
+        lines.sort()
+    elif sorted == REVERSE_SORT:
+        lines.sort(reverse=True)
 
-def lines_to_tabular(lines: list[str], header: str = None, dashed_line = False, separator = ' '):
     maxes = []
     nls = []
 
@@ -63,7 +77,14 @@ def lines_to_tabular(lines: list[str], header: str = None, dashed_line = False,
     for line in lines:
         format_line(line)
 
-    return '\n'.join(nls)
+    table = '\n'.join(nls)
+
+    if to == 1:
+        log(table)
+    elif to == 2:
+        log2(table)
+
+    return table
 
 def convert_seconds(total_seconds_float):
     total_seconds_int = int(total_seconds_float)  # Convert float to integer seconds
@@ -117,7 +138,7 @@ def duration(start_time: float, end_time: float = None):
     d = convert_seconds(end_time - start_time)
     t = []
     if d:
-        t.append(f'{d}h')
+        t.append(f'{d[0]}h')
     if t or d[1]:
         t.append(f'{d[1]}m')
     t.append(f'{d[2]}s')
@@ -142,7 +163,10 @@ 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):
+                print('SEAN completer found, ignoring', key, value)
+            else:
+                merged_dict[key] = value
     return merged_dict
 
 def deep_sort_dict(d):
@@ -272,6 +296,10 @@ def debug(s = None):
     if LogConfig.is_debug():
         log2(f'DEBUG {s}')
 
+def debug_complete(s = None):
+    if LogConfig.is_debug_complete():
+        log2(f'DEBUG {s}')
+
 def debug_trace():
     if LogConfig.is_debug():
         log2(traceback.format_exc())
@@ -411,8 +439,11 @@ def log_timing(msg: str, body: Callable[[], None]=None, s0: time.time = None):
     return r
 
 def timing_log_line(depth: int, msg: str, s0: time.time):
+    # print('SEAN log timing', msg, threading.current_thread().name)
     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)) + '* '
@@ -432,6 +463,66 @@ def wait_log(msg: str):
 def clear_wait_log_flag():
     WaitLog.wait_log_flag = False
 
+def bytes_generator_from_file(file_path, chunk_size=4096):
+    with open(file_path, 'rb') as f:
+        while True:
+            chunk = f.read(chunk_size)
+            if not chunk:
+                break
+            yield chunk
+
+class GeneratorStream(io.RawIOBase):
+    def __init__(self, generator):
+        self._generator = generator
+        self._buffer = b''  # Buffer to store leftover bytes from generator yields
+
+    def readable(self):
+        return True
+
+    def _read_from_generator(self):
+        try:
+            chunk = next(self._generator)
+            if isinstance(chunk, str):
+                chunk = chunk.encode('utf-8')  # Encode if generator yields strings
+            self._buffer += chunk
+        except StopIteration:
+            pass  # Generator exhausted
+
+    def readinto(self, b):
+        # Fill the buffer if necessary
+        while len(self._buffer) < len(b):
+            old_buffer_len = len(self._buffer)
+            self._read_from_generator()
+            if len(self._buffer) == old_buffer_len:  # Generator exhausted and buffer empty
+                break
+
+        bytes_to_read = min(len(b), len(self._buffer))
+        b[:bytes_to_read] = self._buffer[:bytes_to_read]
+        self._buffer = self._buffer[bytes_to_read:]
+        return bytes_to_read
+
+    def read(self, size=-1):
+        if size == -1:  # Read all remaining data
+            while True:
+                old_buffer_len = len(self._buffer)
+                self._read_from_generator()
+                if len(self._buffer) == old_buffer_len:
+                    break
+            data = self._buffer
+            self._buffer = b''
+            return data
+        else:
+            # Ensure enough data in buffer
+            while len(self._buffer) < size:
+                old_buffer_len = len(self._buffer)
+                self._read_from_generator()
+                if len(self._buffer) == old_buffer_len:
+                    break
+
+            data = self._buffer[:size]
+            self._buffer = self._buffer[size:]
+            return data
+
 class LogTrace:
     def __init__(self, err_msg: Union[str, callable, bool] = None):
         self.err_msg = err_msg
@@ -457,8 +548,6 @@ class LogTrace:
 def log_exc(err_msg: Union[str, callable, bool] = None):
     return LogTrace(err_msg=err_msg)
 
-T = TypeVar('T')
-
 class ParallelService:
     def __init__(self, handler: 'ParallelMapHandler'):
         self.handler = handler
@@ -492,9 +581,9 @@ class ParallelService:
             return iterator
 
 class ParallelMapHandler:
-    def __init__(self, collection: list, max_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):
         self.collection = collection
-        self.max_workers = max_workers
+        self.workers = workers
         self.executor = None
         self.samples = samples
         self.msg = msg
@@ -512,10 +601,10 @@ class ParallelMapHandler:
     def __enter__(self):
         self.calc_msgs()
 
-        if self.max_workers > 1 and (not self.size() or self.size()) and self.samples == sys.maxsize:
+        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.max_workers)
+            self.executor = ThreadPoolExecutor(max_workers=self.workers)
             self.executor.__enter__()
 
         return ParallelService(self)
@@ -540,15 +629,13 @@ class ParallelMapHandler:
             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:
+        elif self.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
@@ -558,7 +645,6 @@ class ParallelMapHandler:
         else:
             serially = True
             self.msg = self.msg.replace('{size}', f'{size}')
-            # return
 
         for token in self.msg.split(' '):
             if '|' in token:
@@ -575,7 +661,10 @@ class ParallelMapHandler:
         elif sampling or serially:
             log2(f'{" ".join(self.begin)} serially...')
         else:
-            log2(f'{" ".join(self.begin)} with {self.max_workers} workers...')
+            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)
 
 class OffloadService:
     def __init__(self, handler: 'OffloadHandler'):
@@ -600,10 +689,11 @@ class OffloadHandler(ParallelMapHandler):
     def __enter__(self):
         self.calc_msgs()
 
-        if self.max_workers > 1 and (not self.size() or self.size()) and self.samples == sys.maxsize:
+        if self.workers > 1:
+        # 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.max_workers)
+            self.executor = ThreadPoolExecutor(max_workers=self.workers, thread_name_prefix='offload')
             self.executor.__enter__()
 
         return OffloadService(self)
@@ -617,11 +707,11 @@ class OffloadHandler(ParallelMapHandler):
 
         return False
 
-    def size(self):
-        if not self.collection:
-            return 0
+    # def size(self):
+    #     if not self.collection:
+    #         return 0
 
-        return len(self.collection)
+    #     return len(self.collection)
 
     def calc_msgs(self):
         if not self.msg:
@@ -636,7 +726,7 @@ class OffloadHandler(ParallelMapHandler):
         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:
+        elif self.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
@@ -663,10 +753,7 @@ class OffloadHandler(ParallelMapHandler):
         elif sampling or serially:
             log2(f'{" ".join(self.begin)} serially...')
         else:
-            log2(f'{" ".join(self.begin)} with {self.max_workers} workers...')
-
-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)
+            log2(f'{" ".join(self.begin)} with {self.workers} workers...')
 
 def offload(max_workers: int = 3, msg: str = None):
     return OffloadHandler(max_workers, msg = msg)

adam/utils_athena.py

@@ -5,7 +5,7 @@ import boto3
 import botocore
 
 from adam.config import Config
-from adam.utils import lines_to_tabular, log, log2, log_exc, wait_log
+from adam.utils import tabulize, log, log2, log_exc, wait_log
 
 # no state utility class
 class Athena:
@@ -90,18 +90,17 @@ class Athena:
          if rs:
             column_info = rs[0]['Data']
             columns = [col.get('VarCharValue') for col in column_info]
-            lines = []
-            for row in rs[1:]:
-                  row_data = [col.get('VarCharValue') if col else '' for col in row['Data']]
-                  lines.append('\t'.join(row_data))
-
-            out = lines_to_tabular(lines, header='\t'.join(columns), separator='\t')
+            out = tabulize(rs[1:],
+                           lambda r: '\t'.join(col.get('VarCharValue') if col else '' for col in r['Data']),
+                           header='\t'.join(columns),
+                           separator='\t',
+                           to=0)
             if output:
                log_file = output(out)
             else:
                log(out)
 
-            return len(lines), log_file
+            return len(rs)-1, log_file
       else:
          log2(f"Query failed or was cancelled. State: {state}")
          log2(f"Reason: {reason}")

adam/utils_issues.py

@@ -1,7 +1,7 @@
 from adam.checks.check_result import CheckResult
 from adam.checks.issue import Issue
 from adam.repl_session import ReplSession
-from adam.utils import lines_to_tabular, log, log2
+from adam.utils import tabulize, log2
 
 class IssuesUtils:
     def show(check_results: list[CheckResult], in_repl = False):
@@ -26,7 +26,7 @@ class IssuesUtils:
                     if in_repl:
                         ReplSession().prompt_session.history.append_string(issue.suggestion)
                         suggested += 1
-            log(lines_to_tabular(lines, separator='||'))
+            tabulize(lines, separator='||', to=2)
             if suggested:
                 log2()
                 log2(f'* {suggested} suggested commands are added to history. Press <Up> arrow to access them.')

adam/utils_k8s/app_clusters.py

@@ -12,11 +12,11 @@ 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', background = False) -> list[PodExecResult]:
+             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:
-            results: list[PodExecResult] = exec.map(lambda pod: AppPods.exec(pod, namespace, command, False, False, shell, background))
+            results: list[PodExecResult] = exec.map(lambda pod: AppPods.exec(pod, namespace, command, False, False, shell, backgrounded))
             for result in results:
                 if KubeContext.show_out(show_out):
                     log(result.command)

adam/utils_k8s/app_pods.py

@@ -11,6 +11,9 @@ from adam.repl_session import ReplSession
 class AppPods:
     @functools.lru_cache()
     def pod_names(namespace: str, env: str, app: str):
+        if not env or not app:
+            return []
+
         return [pod.metadata.name for pod in AppPods.app_pods(namespace, env, app)]
 
     def app_pods(namespace: str, env: str, app: str) -> List[client.V1Pod]:
@@ -22,9 +25,9 @@ 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', background = 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, background = background)
+        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:

adam/utils_k8s/cassandra_clusters.py

@@ -12,14 +12,22 @@ T = TypeVar('T')
 
 # utility collection on cassandra clusters; methods are all static
 class CassandraClusters:
-    def exec(sts: str, namespace: str, command: str, action: str = 'action',
-             max_workers=0, show_out=True, on_any = False, shell = '/bin/sh', background = False, log_file = None) -> list[PodExecResult]:
+    def exec(sts: str,
+             namespace: str,
+             command: str,
+             action: str = 'action',
+             max_workers=0,
+             show_out=True,
+             on_any = False,
+             shell = '/bin/sh',
+             backgrounded = False,
+             log_file = None) -> list[PodExecResult]:
 
         pods = StatefulSets.pod_names(sts, namespace)
         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:
-            results: list[PodExecResult] = exec.map(lambda pod: CassandraNodes.exec(pod, namespace, command, False, False, shell, background, log_file))
+            results: list[PodExecResult] = exec.map(lambda pod: CassandraNodes.exec(pod, namespace, command, False, False, shell, backgrounded, log_file))
             for result in results:
                 if show_out and not Config().is_debug():
                     log(result.command)

adam/utils_k8s/cassandra_nodes.py

@@ -6,8 +6,8 @@ from adam.repl_session import ReplSession
 
 # utility collection on cassandra nodes; methods are all static
 class CassandraNodes:
-    def exec(pod_name: str, namespace: str, command: str, show_out = True, throw_err = False, shell = '/bin/sh', background = False, log_file = None) -> PodExecResult:
-        r = Pods.exec(pod_name, "cassandra", namespace, command, show_out = show_out, throw_err = throw_err, shell = shell, background = background, log_file=log_file)
+    def exec(pod_name: str, namespace: str, command: str, show_out = True, throw_err = False, shell = '/bin/sh', backgrounded = False, log_file = None) -> PodExecResult:
+        r = Pods.exec(pod_name, "cassandra", namespace, command, show_out = show_out, throw_err = throw_err, shell = shell, backgrounded = backgrounded, log_file=log_file)
 
         if r and Config().get('repl.history.push-cat-remote-log-file', True):
             if r.log_file and ReplSession().prompt_session:

adam/utils_k8s/k8s.py

@@ -1,8 +1,9 @@
 from collections.abc import Callable
+import inspect
 import re
 import portforward
 
-from adam.commands.command import InvalidState
+from adam.commands.command import InvalidStateException
 from adam.repl_state import ReplState
 from adam.utils import log2
 from adam.utils_k8s.kube_context import KubeContext
@@ -24,14 +25,14 @@ class PortForwardHandler:
         if not self.svc_or_pod:
             log2('No service or pod found.')
 
-            raise InvalidState(state)
+            raise InvalidStateException(state)
 
         if KubeContext.in_cluster():
             svc_name = self.svc_or_pod(True)
             if not svc_name:
                 log2('No service found.')
 
-                raise InvalidState(state)
+                raise InvalidStateException(state)
 
             # cs-a526330d23-cs-a526330d23-default-sts-0 ->
             # curl http://cs-a526330d23-cs-a526330d23-reaper-service.stgawsscpsr.svc.cluster.local:8080
@@ -40,15 +41,23 @@ class PortForwardHandler:
                 svc = f'{groups[1]}{svc_name}.{state.namespace}.svc.cluster.local:{self.target_port}'
                 return (svc, svc)
             else:
-                raise InvalidState(state)
+                raise InvalidStateException(state)
         else:
             pod = self.svc_or_pod(False)
             if not pod:
                 log2('No pod found.')
 
-                raise InvalidState(state)
+                raise InvalidStateException(state)
 
             self.pod = pod
+
+            # pf = portforward.forward(state.namespace, pod, self.local_port + 1, self.target_port, log_level=portforward.LogLevel.DEBUG)
+            # print(inspect.getsource(pf.__enter__))
+            # print('test portforward START', state.namespace, pod, self.local_port + 1, self.target_port, pf.__enter__)
+            # with pf:
+            #     print('test portforward BODY')
+            # print('test portforward OK')
+
             self.forward_connection = portforward.forward(state.namespace, pod, self.local_port, self.target_port)
             if self.inc_connection_cnt() == 1:
                 self.forward_connection.__enter__()

adam/utils_k8s/kube_context.py

@@ -3,7 +3,7 @@ import re
 from kubernetes import config as kconfig
 
 from adam.config import Config
-from adam.utils import idp_token_from_env, lines_to_tabular, log2
+from adam.utils import idp_token_from_env, log2, tabulize
 
 class KubeContext:
     _in_cluster = False
@@ -56,7 +56,7 @@ class KubeContext:
             log2('Use -v <key>=<value> format.')
             log2()
             lines = [f'{key}\t{Config().get(key, None)}' for key in Config().keys()]
-            log2(lines_to_tabular(lines, separator='\t'))
+            tabulize(lines, separator='\t', to=2)
 
         for p in param_ovrs:
             tokens = p.split('=')

adam/utils_k8s/pods.py

@@ -1,5 +1,6 @@
 from collections.abc import Callable
 from datetime import datetime
+import os
 import sys
 import time
 from typing import TypeVar
@@ -10,7 +11,8 @@ from kubernetes.stream.ws_client import ERROR_CHANNEL, WSClient
 from adam.config import Config
 from adam.utils_k8s.volumes import ConfigMapMount
 from adam.pod_exec_result import PodExecResult
-from adam.utils import ParallelMapHandler, log2, debug, log_exc
+from adam.utils import GeneratorStream, ParallelMapHandler, log2, debug, log_exc
+from adam.utils_local import local_tmp_dir
 from .kube_context import KubeContext
 
 from websocket._core import WebSocket
@@ -48,9 +50,14 @@ class Pods:
 
         return ParallelMapHandler(collection, max_workers, samples = samples, msg = msg)
 
-    def exec(pod_name: str, container: str, namespace: str, command: str,
-             show_out = True, throw_err = False, shell = '/bin/sh',
-             background = False,
+    def exec(pod_name: str,
+             container: str,
+             namespace: str,
+             command: str,
+             show_out = True,
+             throw_err = False,
+             shell = '/bin/sh',
+             backgrounded = False,
              log_file = None,
              interaction: Callable[[any, list[str]], any] = None,
              env_prefix: str = None):
@@ -66,7 +73,7 @@ class Pods:
         if env_prefix:
             exec_command = [shell, '-c', f'{env_prefix} {command}']
 
-        if background or command.endswith(' &'):
+        if backgrounded or command.endswith(' &'):
             # should be false for starting a background process
             tty = False
 
@@ -166,6 +173,17 @@ class Pods:
                 with log_exc():
                     s.sock.close()
 
+    def download_file(pod_name: str, container: str, namespace: str, from_path: str, to_path: str = None):
+        if not to_path:
+            to_path = f'{local_tmp_dir()}/{os.path.basename(from_path)}'
+
+        bytes = Pods.read_file(pod_name, container, namespace, from_path)
+        with open(to_path, 'wb') as f:
+            for item in GeneratorStream(bytes):
+                f.write(item)
+
+        return to_path
+
     def get_container(namespace: str, pod_name: str, container_name: str):
         pod = Pods.get(namespace, pod_name)
         if not pod:

adam/utils_k8s/statefulsets.py

@@ -56,8 +56,11 @@ class StatefulSets:
         return statefulset_pods
 
     @functools.lru_cache()
-    def pod_names(ss: str, ns: str):
-        return [pod.metadata.name for pod in StatefulSets.pods(ss, ns)]
+    def pod_names(sts: str, ns: str):
+        if not sts:
+            return []
+
+        return [pod.metadata.name for pod in StatefulSets.pods(sts, ns)]
 
     def restarted_at(ss: str, ns: str):
         # returns timestamp and if being rolled out

adam/utils_local.py

@@ -0,0 +1,4 @@
+from adam.config import Config
+
+def local_tmp_dir():
+    return Config().get('local-tmp-dir', '/tmp/qing-db')

adam/utils_repl/appendable_completer.py

@@ -0,0 +1,6 @@
+from abc import abstractmethod
+
+class AppendableCompleter:
+    @abstractmethod
+    def append_completions(self, key: str, value: dict[str, any]):
+        pass