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

adam/utils_k8s/pods.py

@@ -1,24 +1,27 @@
 from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
 import sys
 import time
-from typing import TypeVar, cast
+from typing import TypeVar
 from kubernetes import client
 from kubernetes.stream import stream
-from kubernetes.stream.ws_client import ERROR_CHANNEL
+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 elapsed_time, log2
+from adam.utils import ParallelMapHandler, log2, debug, log_exc
 from .kube_context import KubeContext
 
+from websocket._core import WebSocket
+
 T = TypeVar('T')
 _TEST_POD_EXEC_OUTS: PodExecResult = None
 
 # utility collection on pods; methods are all static
 class Pods:
+    _TEST_POD_CLOSE_SOCKET: bool = False
+
     def set_test_pod_exec_outs(outs: PodExecResult):
         global _TEST_POD_EXEC_OUTS
         _TEST_POD_EXEC_OUTS = outs
@@ -26,11 +29,9 @@ class Pods:
         return _TEST_POD_EXEC_OUTS
 
     def delete(pod_name: str, namespace: str, grace_period_seconds: int = None):
-        try:
+        with log_exc(lambda e: "Exception when calling CoreV1Api->delete_namespaced_pod: %s\n" % e):
             v1 = client.CoreV1Api()
             v1.delete_namespaced_pod(pod_name, namespace, grace_period_seconds=grace_period_seconds)
-        except Exception as e:
-            log2("Exception when calling CoreV1Api->delete_namespaced_pod: %s\n" % e)
 
     def delete_with_selector(namespace: str, label_selector: str, grace_period_seconds: int = None):
         v1 = client.CoreV1Api()
@@ -39,69 +40,20 @@ class Pods:
         for i in ret.items:
             v1.delete_namespaced_pod(name=i.metadata.name, namespace=namespace, grace_period_seconds=grace_period_seconds)
 
-    def on_pods(pods: list[str],
-                namespace: str,
-                body: Callable[[ThreadPoolExecutor, str, str, bool], T],
-                post: Callable[[T], T] = None,
-                action: str = 'action',
-                max_workers=0,
-                show_out=True,
-                on_any = False,
-                background = False) -> list[T]:
-        show_out = KubeContext.show_out(show_out)
-
+    def parallelize(collection: list, max_workers: int = 0, samples = sys.maxsize, msg: str = None, action: str = 'action'):
         if not max_workers:
             max_workers = Config().action_workers(action, 0)
-        if not on_any and max_workers > 0:
-            # if parallel, node sampling is suppressed
-            if KubeContext.show_parallelism():
-                log2(f'Executing on all nodes from statefulset in parallel...')
-            start_time = time.time()
-            try:
-                with ThreadPoolExecutor(max_workers=max_workers) as executor:
-                    # disable stdout from the pod_exec, then show the output in a for loop
-                    futures = [body(executor, pod, namespace, show_out) for pod in pods]
-                    if len(futures) == 0:
-                        return cast(list[T], [])
-
-                rs = [future.result() for future in as_completed(futures)]
-                if post:
-                    rs = [post(r, show_out=show_out) for r in rs]
-
-                return rs
-            finally:
-                if KubeContext.show_parallelism():
-                    log2(f"Parallel {action} elapsed time: {elapsed_time(start_time)} with {max_workers} workers")
-        else:
-            results: list[T] = []
-
-            samples = 1 if on_any else Config().action_node_samples(action, sys.maxsize)
-            l = min(len(pods), samples)
-            adj = 'all'
-            if l < len(pods):
-                adj = f'{l} sample'
-            if show_out:
-                log2(f'Executing on {adj} nodes from statefulset...')
-            for pod_name in pods:
-                try:
-                    # disable stdout from the pod_exec, then show the output in a for loop
-                    result = body(None, pod_name, namespace, False)
-                    if post:
-                        result = post(result, show_out=show_out)
-                    results.append(result)
-                    if result:
-                        l -= 1
-                        if not l:
-                            break
-                except Exception as e:
-                    log2(e)
-
-            return results
+        if samples == sys.maxsize:
+            samples = Config().action_node_samples(action, sys.maxsize)
+
+        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,
-             interaction: Callable[[any, list[str]], any] = None):
+             log_file = None,
+             interaction: Callable[[any, list[str]], any] = None,
+             env_prefix: str = None):
         if _TEST_POD_EXEC_OUTS:
             return _TEST_POD_EXEC_OUTS
 
@@ -109,9 +61,11 @@ class Pods:
 
         api = client.CoreV1Api()
 
-        log_file = None
         tty = True
         exec_command = [shell, '-c', command]
+        if env_prefix:
+            exec_command = [shell, '-c', f'{env_prefix} {command}']
+
         if background or command.endswith(' &'):
             # should be false for starting a background process
             tty = False
@@ -122,15 +76,17 @@ class Pods:
                 if command.startswith('nodetool '):
                     cmd_name = f".{'_'.join(command.split(' ')[5:])}"
 
-                log_file = f'/tmp/qing-{datetime.now().strftime("%d%H%M%S")}{cmd_name}.log'
+                if not log_file:
+                    log_file = f'{log_prefix()}-{datetime.now().strftime("%d%H%M%S")}{cmd_name}.log'
                 command = f"nohup {command} > {log_file} 2>&1 &"
+                if env_prefix:
+                    command = f'{env_prefix} {command}'
                 exec_command = [shell, '-c', command]
 
         k_command = f'kubectl exec {pod_name} -c {container} -n {namespace} -- {shell} -c "{command}"'
-        if show_out:
-            print(k_command)
+        debug(k_command)
 
-        resp = stream(
+        resp: WSClient = stream(
             api.connect_get_namespaced_pod_exec,
             pod_name,
             namespace,
@@ -143,6 +99,7 @@ class Pods:
             _preload_content=False,
         )
 
+        s: WebSocket = resp.sock
         stdout = []
         stderr = []
         error_output = None
@@ -161,11 +118,9 @@ class Pods:
                     stderr.append(frag)
                     if show_out: print(frag, end="")
 
-            try:
+            with log_exc():
                 # get the exit code from server
                 error_output = resp.read_channel(ERROR_CHANNEL)
-            except Exception as e:
-                pass
         except Exception as e:
             if throw_err:
                 raise e
@@ -173,6 +128,9 @@ class Pods:
                 log2(e)
         finally:
             resp.close()
+            if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
+                with log_exc():
+                    s.sock.close()
 
         return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output, pod=pod_name, log_file=log_file)
 
@@ -187,24 +145,26 @@ class Pods:
             command=["cat", file_path],
             stderr=True, stdin=False,
             stdout=True, tty=False,
-            _preload_content=False # Important for streaming
+            _preload_content=False, # Important for streaming
         )
 
+        s: WebSocket = resp.sock
         try:
             while resp.is_open():
                 resp.update(timeout=1)
                 if resp.peek_stdout():
                     yield resp.read_stdout()
 
-            try:
+            with log_exc():
                 # get the exit code from server
                 error_output = resp.read_channel(ERROR_CHANNEL)
-            except Exception as e:
-                pass
         except Exception as e:
             raise e
         finally:
             resp.close()
+            if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
+                with log_exc():
+                    s.sock.close()
 
     def get_container(namespace: str, pod_name: str, container_name: str):
         pod = Pods.get(namespace, pod_name)
@@ -314,4 +274,7 @@ class Pods:
             log2(' Timed Out')
 
     def completed(namespace: str, pod_name: str):
-        return Pods.get(namespace, pod_name).status.phase in ['Succeeded', 'Failed']
+        return Pods.get(namespace, pod_name).status.phase in ['Succeeded', 'Failed']
+
+def log_prefix():
+    return Config().get('log-prefix', '/tmp/qing')

adam/utils_k8s/secrets.py

@@ -6,13 +6,13 @@ from kubernetes import client
 from kubernetes.client import V1Secret
 
 from adam.config import Config
-from adam.utils import log2
+from adam.utils import log2, wait_log
 
 # utility collection on secrets; methods are all static
 class Secrets:
     @functools.lru_cache()
     def list_secrets(namespace: str = None, name_pattern: str = None):
-        Config().wait_log('Inspecting Cassandra Instances...')
+        wait_log('Inspecting Cassandra Instances...')
 
         secrets_names = []
 
@@ -39,14 +39,14 @@ class Secrets:
 
         return secrets_names
 
-    def get_user_pass(ss_name: str, namespace: str, secret_path: str = 'cql.secret'):
+    def get_user_pass(sts_or_pod_name: str, namespace: str, secret_path: str = 'cql.secret'):
         # cs-d0767a536f-cs-d0767a536f-default-sts ->
         # cs-d0767a536f-superuser
         # cs-d0767a536f-reaper-ui
         user = 'superuser'
         if secret_path == 'reaper.secret':
             user = 'reaper-ui'
-        groups = re.match(Config().get(f'{secret_path}.cluster-regex', r'(.*?-.*?)-.*'), ss_name)
+        groups = re.match(Config().get(f'{secret_path}.cluster-regex', r'(.*?-.*?)-.*'), sts_or_pod_name)
         secret_name = Config().get(f'{secret_path}.name', '{cluster}-' + user).replace('{cluster}', groups[1], 1)
 
         secret = Secrets.get_data(namespace, secret_name)

adam/utils_k8s/service_accounts.py

@@ -1,6 +1,7 @@
 from kubernetes import client, config
 
 from adam.config import Config
+from adam.utils import debug
 
 # utility collection on service accounts; methods are all static
 class ServiceAccounts:
@@ -37,7 +38,7 @@ class ServiceAccounts:
             namespace=namespace,
             body=service_account
         )
-        Config().debug(f"Service Account '{api_response.metadata.name}' created in namespace '{namespace}'.")
+        debug(f"Service Account '{api_response.metadata.name}' created in namespace '{namespace}'.")
 
     def delete_service_account(namespace: str, label_selector: str) -> list:
         refs = []
@@ -45,7 +46,7 @@ class ServiceAccounts:
         v1 = client.CoreV1Api()
         sas = v1.list_namespaced_service_account(namespace=namespace, label_selector=label_selector).items
         for sa in sas:
-            Config().debug(f'delete {sa.metadata.name}')
+            debug(f'delete {sa.metadata.name}')
             v1.delete_namespaced_service_account(name=sa.metadata.name, namespace=namespace)
             refs.append(sa)
 
@@ -102,7 +103,7 @@ class ServiceAccounts:
         v1_rbac = client.RbacAuthorizationV1Api()
         cluster_role_bindings = v1_rbac.list_namespaced_role_binding(namespace=namespace, label_selector=label_selector).items
         for binding in cluster_role_bindings:
-            Config().debug(f'delete {binding.metadata.name}')
+            debug(f'delete {binding.metadata.name}')
             v1_rbac.delete_namespaced_role_binding(name=binding.metadata.name, namespace=namespace)
             refs.append(binding)
 
@@ -162,7 +163,7 @@ class ServiceAccounts:
         v1_rbac = client.RbacAuthorizationV1Api()
         cluster_role_bindings = v1_rbac.list_cluster_role_binding(label_selector=label_selector).items
         for binding in cluster_role_bindings:
-            Config().debug(f'delete {binding.metadata.name}')
+            debug(f'delete {binding.metadata.name}')
             v1_rbac.delete_cluster_role_binding(binding.metadata.name)
             refs.append(binding)
 

adam/utils_k8s/services.py

@@ -2,7 +2,7 @@ from typing import List
 from kubernetes import client
 
 from adam.config import Config
-from adam.utils import log2
+from adam.utils import debug, log2
 
 from .kube_context import KubeContext
 
@@ -71,7 +71,7 @@ class Services:
                     namespace=namespace,
                     body=delete_options
                 )
-                Config().debug(f"200 Service '{name}' in namespace '{namespace}' deleted successfully.")
+                debug(f"200 Service '{name}' in namespace '{namespace}' deleted successfully.")
         except client.ApiException as e:
             log2(f"Error deleting Service '{name}': {e}")
         except Exception as e:

adam/utils_k8s/statefulsets.py

@@ -1,12 +1,9 @@
-from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 import functools
 import re
 from typing import List, TypeVar, cast
 from kubernetes import client
 
-from .pods import Pods
 from .kube_context import KubeContext
 from adam.utils import log2
 
@@ -58,15 +55,6 @@ class StatefulSets:
 
         return statefulset_pods
 
-    def on_cluster(statefulset: str,
-                namespace: str,
-                body: Callable[[ThreadPoolExecutor, str, str, bool], T],
-                post: Callable[[T], T] = None,
-                action: str = 'action', max_workers=0, show_out=True, on_any = False, background = False) -> list[T]:
-        pods = StatefulSets.pod_names(statefulset, namespace)
-
-        return Pods.on_pods(pods, namespace, body, post=post, action=action, max_workers=max_workers, show_out=show_out, on_any=on_any, background=background)
-
     @functools.lru_cache()
     def pod_names(ss: str, ns: str):
         return [pod.metadata.name for pod in StatefulSets.pods(ss, ns)]
@@ -92,6 +80,7 @@ class StatefulSets:
 
         return restarted, False
 
+    @functools.lru_cache()
     def get_datacenter(sts: str, ns: str) -> str:
         v1 = client.AppsV1Api()
         namespace = ns

adam/utils_net.py

@@ -18,7 +18,7 @@ def get_my_host():
     return MY_HOST
 
 def get_ip_from_hostname(hostname):
-  try:
-    return socket.gethostbyname(hostname)
-  except socket.gaierror:
-    return None
+    try:
+        return socket.gethostbyname(hostname)
+    except socket.gaierror:
+        return None

adam/utils_repl/automata_completer.py

@@ -0,0 +1,48 @@
+from typing import Generic, Iterable, TypeVar
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion, WordCompleter
+from prompt_toolkit.document import Document
+
+from adam.utils_repl.state_machine import StateMachine, State
+
+__all__ = [
+    "AutomataCompleter",
+]
+
+T = TypeVar('T')
+
+class AutomataCompleter(Completer, Generic[T]):
+    def __init__(self,
+                 state_machine: StateMachine,
+                 first_term: str = '',
+                 debug = False):
+        super().__init__()
+        self.machine = state_machine
+        self.first_term = first_term
+        self.debug = debug
+
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        text = document.text_before_cursor.lstrip()
+        state = ''
+        if self.first_term:
+            text = f'{self.first_term} {text}'
+
+        completer: Completer = None
+        state: State = self.machine.traverse_tokens(self.tokens(text), State(state))
+        if self.debug:
+            print('\n  =>', state.state if isinstance(state, State) else '')
+
+        if state.state in self.machine.suggestions:
+            if completer := self.suggestions_completer(state, self.machine.suggestions[state.state].strip(' ')):
+                for c in completer.get_completions(document, complete_event):
+                    yield c
+
+    def tokens(self, text: str) -> list[T]:
+        return text.split(' ')
+
+    def suggestions_completer(self, _: State, suggestions: str) -> list[str]:
+        if not suggestions:
+            return None
+
+        return WordCompleter(suggestions.split(','))

adam/utils_repl/repl_completer.py

@@ -0,0 +1,46 @@
+import re
+from typing import Iterable, TypeVar
+from prompt_toolkit.completion import CompleteEvent, Completion, NestedCompleter, WordCompleter
+from prompt_toolkit.document import Document
+
+__all__ = [
+    "ReplCompleter",
+]
+
+T = TypeVar('T')
+
+class ReplCompleter(NestedCompleter):
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        # Split document.
+        text = document.text_before_cursor.lstrip()
+        stripped_len = len(document.text_before_cursor) - len(text)
+
+        # If there is a space, check for the first term, and use a
+        # subcompleter.
+        if " " in text:
+            first_term = text.split()[0]
+            completer = self.options.get(first_term)
+
+            # If we have a sub completer, use this for the completions.
+            if completer is not None:
+                remaining_text = text[len(first_term) :].lstrip()
+                move_cursor = len(text) - len(remaining_text) + stripped_len
+
+                new_document = Document(
+                    remaining_text,
+                    cursor_position=document.cursor_position - move_cursor,
+                )
+
+                for c in completer.get_completions(new_document, complete_event):
+                    yield c
+
+        # No space in the input: behave exactly like `WordCompleter`.
+        else:
+            completer = WordCompleter(
+                # Allow dot in the middle or a word
+                list(self.options.keys()), ignore_case=self.ignore_case, pattern=re.compile(r"([a-zA-Z0-9_\.\@\&]+|[^a-zA-Z0-9_\.\@\&\s]+)")
+            )
+            for c in completer.get_completions(document, complete_event):
+                yield c

adam/utils_repl/state_machine.py

@@ -0,0 +1,173 @@
+from abc import abstractmethod
+from typing import Generic, TypeVar
+
+from adam.utils import log_exc
+
+__all__ = [
+    'State',
+    'StateMachine',
+]
+
+T = TypeVar('T')
+
+class State:
+    def __init__(self, state: str, comeback_token: str = None, comeback_state: str = None):
+        self.state = state
+        self.comeback_token = comeback_token
+        self.comeback_state = comeback_state
+        self.context: dict[str, str] = {}
+
+    def __str__(self):
+        return f'{self.state if self.state else None} comeback[{self.comeback_token} {self.comeback_state}]'
+
+class StateMachine(Generic[T]):
+    @abstractmethod
+    def spec(self) -> str:
+        return None
+
+    @abstractmethod
+    def keywords(self) -> list[str]:
+        return None
+
+    def expandable_names(self):
+        return []
+
+    def incomplete_name_transition_condition(self, from_s: str, token: str, to_s: str, suggestions: str) -> list[str]:
+        if not suggestions:
+            return None
+
+        tokens = [token]
+        if '|' in token:
+            tokens = token.split('|')
+
+        if 'name' not in tokens:
+            return None
+
+        return tokens
+
+    def __init__(self, indent=0, push_level = 0, debug = False):
+        self.states: dict[str, State] = {}
+        self.suggestions: dict[str, str] = {}
+
+        self.indent = indent
+        self.push_level = push_level
+        self.comebacks: dict[int, str] = {}
+        self.debug = debug
+
+        from_ss_to_add = []
+        from_ss = ['']
+        words: str = None
+        for l in self.spec():
+            t_and_w = l.split('^')
+            if len(t_and_w) > 1:
+                words = t_and_w[1].strip()
+            else:
+                words = None
+
+            tks = t_and_w[0].strip(' ').split('>')
+            if not l.startswith('-'):
+                if words:
+                    self.suggestions[tks[0].strip(' ')] = words
+
+                if len(tks) == 1:
+                    from_ss_to_add.append(tks[0].strip(' '))
+                    continue
+
+                from_ss = []
+                from_ss.extend(from_ss_to_add)
+                from_ss_to_add = []
+                from_ss.append(tks[0].strip(' '))
+
+            self.add_transitions(from_ss, tks, words)
+
+    def add_transitions(self, from_ss: list[str], tks: list[str], words: str):
+        token = tks[1].strip(' ')
+        if len(tks) > 2:
+            to_s = tks[2].strip(' ')
+            for from_s in from_ss:
+                self.add_whitespace_transition(from_s, to_s)
+                self.add_transition(from_s, token, to_s)
+                self.add_incomplete_name_transition(from_s, token, to_s, words)
+        elif '<' in tks[0]:
+            from_and_token = tks[0].split('<')
+            if len(from_and_token) > 1:
+                for from_s in from_ss:
+                    self.add_comeback_transition(from_s, from_and_token[1], tks[1].strip(' '))
+
+    def add_whitespace_transition(self, from_s: str, to_s: str):
+        if self.witespace_transition_condition(from_s, to_s):
+            if self.debug:
+                print(f'{from_s[:-1]} > _ = {to_s}')
+            self.states[f'{from_s[:-1]} > _'] = State(from_s)
+
+    def witespace_transition_condition(self, from_s: str, to_s: str):
+        return from_s.endswith('_')
+
+    def add_incomplete_name_transition(self, from_s: str, token: str, to_s: str, words: str):
+        if tokens := self.incomplete_name_transition_condition(from_s, token, to_s, words):
+            self.suggestions[to_s] = words
+            for token in tokens:
+                if self.debug:
+                    print(f'{to_s} > {token} = {to_s}')
+                self.states[f'{to_s} > {token}'] = State(to_s)
+
+    def add_transition(self, from_s: str, token: str, to_s: str):
+        tokens = [token]
+        if '|' in token:
+            tokens = token.split('|')
+
+        for t in tokens:
+            if t == '_or_':
+                t = '||'
+            elif t == 'pipe':
+                t = '|'
+            elif t == '_rdr0_':
+                t = '<'
+            elif t == '_rdr1_':
+                t = '>'
+            elif t == '_rdr2_':
+                t = '2>'
+
+            if self.debug:
+                print(f'{from_s} > {t} = {to_s}')
+            self.states[f'{from_s} > {t}'] = State(to_s)
+
+    def add_comeback_transition(self, from_s: str, token: str, to_s: str):
+        key = f'{from_s} > ('
+        orig = self.states[key]
+        if not orig:
+            raise Exception(f'from state not found for {key}')
+
+        orig.comeback_token = token
+        orig.comeback_state = to_s
+        if self.debug:
+            print(f'{from_s} > ) = {to_s}')
+        self.states[key] = orig
+
+    def traverse_tokens(self, tokens: list[str], state: State = State('')):
+        for token in tokens[:-1]:
+            if not token:
+                continue
+
+            if self.debug:
+                print(f'{token} ', end='')
+
+            last_name = None
+
+            if (t := token.lower()) in self.keywords():
+                token = t
+            elif token in ['*', ',', '.']:
+                pass
+            else:
+                last_name = token
+                token = 'word'
+
+            with log_exc():
+                context = state.context
+                state = self.states[f'{state.state} > {token}']
+                state.context = context
+
+                if last_name:
+                    state.context['last_name'] = last_name
+
+        return state