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

adam/utils_k8s/k8s.py

@@ -0,0 +1,87 @@
+from collections.abc import Callable
+import re
+import portforward
+
+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
+
+class PortForwardHandler:
+    connections: dict[str, int] = {}
+
+    def __init__(self, state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
+        self.state = state
+        self.local_port = local_port
+        self.svc_or_pod = svc_or_pod
+        self.target_port = target_port
+        self.forward_connection = None
+        self.pod = None
+
+    def __enter__(self) -> tuple[str, str]:
+        state = self.state
+
+        if not self.svc_or_pod:
+            log2('No service or pod found.')
+
+            raise InvalidStateException(state)
+
+        if KubeContext.in_cluster():
+            svc_name = self.svc_or_pod(True)
+            if not svc_name:
+                log2('No service found.')
+
+                raise InvalidStateException(state)
+
+            # cs-a526330d23-cs-a526330d23-default-sts-0 ->
+            # curl http://cs-a526330d23-cs-a526330d23-reaper-service.stgawsscpsr.svc.cluster.local:8080
+            groups = re.match(r'^(.*?-.*?-.*?-.*?-).*', state.sts)
+            if groups:
+                svc = f'{groups[1]}{svc_name}.{state.namespace}.svc.cluster.local:{self.target_port}'
+                return (svc, svc)
+            else:
+                raise InvalidStateException(state)
+        else:
+            pod = self.svc_or_pod(False)
+            if not pod:
+                log2('No pod found.')
+
+                raise InvalidStateException(state)
+
+            self.pod = pod
+            self.forward_connection = portforward.forward(state.namespace, pod, self.local_port, self.target_port)
+            if self.inc_connection_cnt() == 1:
+                self.forward_connection.__enter__()
+
+            return (f'localhost:{self.local_port}', f'{pod}:{self.target_port}')
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.forward_connection:
+            if not self.dec_connection_cnt():
+                return self.forward_connection.__exit__(exc_type, exc_val, exc_tb)
+
+        return False
+
+    def inc_connection_cnt(self):
+        id = self.connection_id(self.pod)
+        if id not in PortForwardHandler.connections:
+            PortForwardHandler.connections[id] = 1
+        else:
+            PortForwardHandler.connections[id] += 1
+
+        return PortForwardHandler.connections[id]
+
+    def dec_connection_cnt(self):
+        id = self.connection_id(self.pod)
+        if id not in PortForwardHandler.connections:
+            PortForwardHandler.connections[id] = 0
+        elif PortForwardHandler.connections[id] > 0:
+            PortForwardHandler.connections[id] -= 1
+
+        return PortForwardHandler.connections[id]
+
+    def connection_id(self, pod: str):
+        return f'{self.local_port}:{pod}:{self.target_port}'
+
+def port_forwarding(state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
+    return PortForwardHandler(state, local_port, svc_or_pod, target_port)

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,24 +1,29 @@
 from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
+import os
 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 GeneratorStream, ParallelMapHandler, log2, debug, log_exc
+from adam.utils_local import local_tmp_dir
 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 +31,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 +42,25 @@ 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
-
-    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):
+        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',
+             backgrounded = False,
+             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,10 +68,12 @@ class Pods:
 
         api = client.CoreV1Api()
 
-        log_file = None
         tty = True
         exec_command = [shell, '-c', command]
-        if background or command.endswith(' &'):
+        if env_prefix:
+            exec_command = [shell, '-c', f'{env_prefix} {command}']
+
+        if backgrounded or command.endswith(' &'):
             # should be false for starting a background process
             tty = False
 
@@ -122,15 +83,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 +106,7 @@ class Pods:
             _preload_content=False,
         )
 
+        s: WebSocket = resp.sock
         stdout = []
         stderr = []
         error_output = None
@@ -161,11 +125,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,9 +135,55 @@ 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)
 
+    def read_file(pod_name: str, container: str, namespace: str, file_path: str):
+        v1 = client.CoreV1Api()
+
+        resp = stream(
+            v1.connect_get_namespaced_pod_exec,
+            name=pod_name,
+            namespace=namespace,
+            container=container,
+            command=["cat", file_path],
+            stderr=True, stdin=False,
+            stdout=True, tty=False,
+            _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()
+
+            with log_exc():
+                # get the exit code from server
+                error_output = resp.read_channel(ERROR_CHANNEL)
+        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 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:
@@ -284,4 +292,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_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_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