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

adam/utils_async_job.py

@@ -0,0 +1,73 @@
+from datetime import datetime
+import re
+import traceback
+
+from adam.utils import log_dir
+
+class AsyncJobs:
+    _last_id: str = None
+    _last_command: str = None
+
+    def log_file(command: str, pod_name: str = None, job_id: str = None, pod_suffix: str = None, err = False, dir: str = None):
+        try:
+            if not job_id:
+                job_id = AsyncJobs.new_id()
+
+            AsyncJobs._last_id = job_id
+            AsyncJobs._last_command = command
+
+            AsyncJobs.write_last_command(job_id, command)
+
+            cmd_name = ''
+            if command.startswith('nodetool '):
+                command = command.strip(' &')
+                cmd_name = f".{'_'.join(command.split(' ')[5:])}"
+
+            if pod_suffix is None:
+                pod_suffix = '{pod}'
+                if pod_name:
+                    pod_suffix = pod_name
+                    if groups := re.match(r'.*-(.*)', pod_name):
+                        pod_suffix = f'-{groups[1]}'
+
+            if not dir:
+                dir = log_dir()
+
+            return f'{dir}/{job_id}{cmd_name}{pod_suffix}.{"err" if err else "log"}'
+        except:
+            traceback.print_exc()
+
+    def new_id(dt: datetime = None):
+        if not dt:
+            dt = datetime.now()
+
+        id = dt.strftime("%d%H%M%S")
+        AsyncJobs._last_id = id
+
+        return id
+
+    def last_command():
+        last_id = AsyncJobs._last_id
+        last_command = AsyncJobs._last_command
+
+        if not last_id or not last_command:
+            l0, l1 = AsyncJobs.read_last_command()
+            if not last_id:
+                last_id = l0
+            if not last_command:
+                last_command = l1
+
+        return last_id, last_command
+
+    def write_last_command(job_id: str, command: str):
+        with open(f'{log_dir()}/last', 'wt') as f:
+            f.write(job_id)
+            f.write('\n')
+            f.write(command)
+
+    def read_last_command():
+        with open(f'{log_dir()}/last', 'rt') as f:
+            job_id = f.readline().strip(' \r\n')
+            command = f.readline().strip(' \r\n')
+
+            return job_id, command

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

@@ -2,9 +2,10 @@ import sys
 from typing import TypeVar
 
 from adam.config import Config
+from adam.utils_async_job import AsyncJobs
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
-from adam.pod_exec_result import PodExecResult
-from adam.utils import log, log2
+from adam.utils_k8s.pod_exec_result import PodExecResult
+from adam.utils import log, log2, log_to_pods, pod_log_dir
 from adam.utils_k8s.pods import Pods
 from adam.utils_k8s.statefulsets import StatefulSets
 
@@ -21,20 +22,31 @@ class CassandraClusters:
              on_any = False,
              shell = '/bin/sh',
              backgrounded = False,
-             log_file = None) -> list[PodExecResult]:
+             log_file = None,
+             history=True,
+             text_color: str = None) -> 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:
+            pod_suffix = None
+            dir = None
+            if log_to_pods():
+                pod_suffix = ''
+                dir = pod_log_dir()
+
+            log_file = AsyncJobs.log_file(command, job_id=AsyncJobs.new_id(), pod_suffix=pod_suffix, dir=dir)
+
         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, text_color))
             for result in results:
                 if show_out and not Config().is_debug():
-                    log(result.command)
+                    log(result.command, text_color=text_color)
                     if result.stdout:
-                        log(result.stdout)
+                        log(result.stdout, text_color=text_color)
                     if result.stderr:
-                        log2(result.stderr, file=sys.stderr)
+                        log2(result.stderr, text_color=text_color)
 
             return results
 

adam/utils_k8s/cassandra_nodes.py

@@ -1,17 +1,27 @@
 from adam.config import Config
 from adam.utils_k8s.pods import Pods
 from adam.utils_k8s.secrets import Secrets
-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 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', 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)
+    def exec(pod_name: str,
+             namespace: str,
+             command: str,
+             show_out = True,
+             throw_err = False,
+             shell = '/bin/sh',
+             backgrounded = False,
+             log_file = None,
+             history = True,
+             text_color: str = None) -> PodExecResult:
+        r: PodExecResult = Pods.exec(pod_name, "cassandra", namespace, command, show_out = show_out, throw_err = throw_err, shell = shell, backgrounded = backgrounded, log_file=log_file, text_color=text_color)
 
-        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 history and r and r.log_file:
+            # entry = f':tail {r.log_file}'
+
+            ReplSession().append_history(r.log_file)
 
         return r
 

adam/utils_k8s/k8s.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+import inspect
 import re
 import portforward
 
@@ -49,6 +50,14 @@ class PortForwardHandler:
                 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

@@ -102,7 +102,4 @@ class KubeContext:
         return name if re.match(r"^(?!pg-).*-k8spg-.*$", name) else None
 
     def show_out(s: bool):
-        return s or Config().is_debug()
-
-    def show_parallelism():
-        return Config().get('debugs.show-parallelism', False)
+        return s or Config().is_debug()

adam/{pod_exec_result.py → utils_k8s/pod_exec_result.py}

@@ -1,8 +1,8 @@
 import yaml
 
-from adam.utils import log_exc
+from adam.utils import ExecResult, log_exc
 
-class PodExecResult:
+class PodExecResult(ExecResult):
     # {
     #   'metadata': {},
     #   'status': 'Failure',
@@ -34,6 +34,12 @@ class PodExecResult:
 
         return code
 
+    def cat_log_file_cmd(self):
+        if self.pod and self.log_file:
+            return f'@{self.pod} cat {self.log_file}'
+
+        return None
+
     def __str__(self):
         return f'{"OK" if self.exit_code() == 0 else self.exit_code()} {self.command}'
 

adam/utils_k8s/pods.py

@@ -1,18 +1,22 @@
 from collections.abc import Callable
-from datetime import datetime
 import os
+import re
+import subprocess
 import sys
 import time
 from typing import TypeVar
 from kubernetes import client
 from kubernetes.stream import stream
 from kubernetes.stream.ws_client import ERROR_CHANNEL, WSClient
+from prompt_toolkit import print_formatted_text, HTML
 
 from adam.config import Config
+from adam.repl_session import ReplSession
 from adam.utils_k8s.volumes import ConfigMapMount
-from adam.pod_exec_result import PodExecResult
-from adam.utils import GeneratorStream, ParallelMapHandler, log2, debug, log_exc
-from adam.utils_local import local_tmp_dir
+from adam.utils_k8s.pod_exec_result import PodExecResult
+from adam.utils import GeneratorStream, ParallelMapHandler, PodLogFile, log2, debug, log_exc, log_to_pods, pod_log_dir
+from adam.utils_local import local_downloads_dir, local_exec
+from adam.utils_async_job import AsyncJobs
 from .kube_context import KubeContext
 
 from websocket._core import WebSocket
@@ -48,7 +52,7 @@ class Pods:
         if samples == sys.maxsize:
             samples = Config().action_node_samples(action, sys.maxsize)
 
-        return ParallelMapHandler(collection, max_workers, samples = samples, msg = msg)
+        return ParallelMapHandler(collection, max_workers, samples = samples, msg = msg, name=action)
 
     def exec(pod_name: str,
              container: str,
@@ -60,12 +64,19 @@ class Pods:
              backgrounded = False,
              log_file = None,
              interaction: Callable[[any, list[str]], any] = None,
-             env_prefix: str = None):
+             env_prefix: str = None,
+             text_color: str = None):
         if _TEST_POD_EXEC_OUTS:
             return _TEST_POD_EXEC_OUTS
 
         show_out = KubeContext.show_out(show_out)
 
+        if backgrounded or command.endswith(' &'):
+            if log_to_pods():
+                return Pods.exec_backgrounded_logging_to_pod(pod_name, container, namespace, command, show_out, shell, log_file, env_prefix, text_color)
+            else:
+                return Pods.exec_backgrounded(pod_name, container, command, show_out, shell, log_file, env_prefix, text_color)
+
         api = client.CoreV1Api()
 
         tty = True
@@ -73,25 +84,28 @@ class Pods:
         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
-
-            if Config().get('repl.background-process.auto-nohup', True):
-                command = command.strip(' &')
-                cmd_name = ''
-                if command.startswith('nodetool '):
-                    cmd_name = f".{'_'.join(command.split(' ')[5:])}"
-
-                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]
+        # if backgrounded or command.endswith(' &'):
+        #     # should be false for starting a background process
+        #     tty = False
+
+        #     if Config().get('repl.background-process.auto-nohup', True):
+        #         command = command.strip(' &')
+        #         cmd_name = ''
+        #         if command.startswith('nodetool '):
+        #             cmd_name = f".{'_'.join(command.split(' ')[5:])}"
+
+        #         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}"'
-        debug(k_command)
+        if Config().is_debug():
+            debug(k_command)
+        elif show_out:
+            log2(k_command, text_color=text_color)
 
         resp: WSClient = stream(
             api.connect_get_namespaced_pod_exec,
@@ -116,14 +130,22 @@ class Pods:
                 if resp.peek_stdout():
                     frag = resp.read_stdout()
                     stdout.append(frag)
-                    if show_out: print(frag, end="")
+                    if show_out:
+                        if text_color:
+                            print_formatted_text(HTML(f'<ansi{text_color}>{frag}</ansi{text_color}>'), end="")
+                        else:
+                            print(frag, end="")
 
                     if interaction:
                         interaction(resp, stdout)
                 if resp.peek_stderr():
                     frag = resp.read_stderr()
                     stderr.append(frag)
-                    if show_out: print(frag, end="")
+                    if show_out:
+                        if text_color:
+                            print_formatted_text(HTML(f'<ansi{text_color}>{frag}</ansi{text_color}>'), end="")
+                        else:
+                            print(frag, end="")
 
             with log_exc():
                 # get the exit code from server
@@ -132,7 +154,7 @@ class Pods:
             if throw_err:
                 raise e
             else:
-                log2(e)
+                log2(e, text_color=text_color)
         finally:
             resp.close()
             if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
@@ -141,6 +163,96 @@ class Pods:
 
         return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output, pod=pod_name, log_file=log_file)
 
+    def exec_backgrounded(pod_name: str,
+             container: str,
+             command: str,
+             show_out = True,
+             shell = '/bin/sh',
+             log_file = None,
+             env_prefix: str = None,
+             text_color: str = None):
+        # nohup kubectl exec cs-a7b13e29bd-cs-a7b13e29bd-default-sts-0 -c cassandra -- /bin/sh -c "nohup nodetool -u cs-a7b13e29bd-superuser -pw ... repair > /tmp/qing-db/q/logs/19230320.repair-0.log 2> /tmp/qing-db/q/logs/19230320.repair-0.err &" &
+
+        command = command.strip(' &')
+
+        log_pod_file = None
+        if log_file:
+            log_pod_file = Pods.log_file_from_template(log_file, pod_name=pod_name)
+        else:
+            log_pod_file = AsyncJobs.log_file(command, pod_name=pod_name)
+
+        if env_prefix:
+            command = f'{env_prefix} {command}'
+
+        log_err_file = log_pod_file.replace('.log', '.err')
+
+        command = command.replace('"', '\\"')
+        # nohup kubectl exec cs-a7b13e29bd-cs-a7b13e29bd-default-sts-0 -c cassandra -- /bin/sh -c "nodetool -u cs-a7b13e29bd-superuser -pw ... repair &"
+        #   > /tmp/qing-db/q/logs/19080002.repair-0.log 2> /tmp/qing-db/q/logs/19080002.repair-0.err &
+        cmd = f'nohup kubectl exec {pod_name} -c {container} -- {shell} -c "{command} &" > {log_pod_file} 2> {log_err_file}'
+        cmd = f'{cmd} &'
+
+        if show_out:
+            log2(cmd, text_color=text_color)
+
+        result = subprocess.run(cmd, capture_output=True, text=True, shell=True)
+        return PodExecResult(result.stdout, result.stderr, cmd, None, pod=pod_name, log_file=log_pod_file)
+
+    def exec_backgrounded_logging_to_pod(pod_name: str,
+                                         container: str,
+                                         namespace: str,
+                                         command: str,
+                                         show_out = True,
+                                         shell = '/bin/sh',
+                                         log_file = None,
+                                         env_prefix: str = None,
+                                         text_color: str = None):
+        command = command.strip(' &')
+
+        log_pod_file = None
+        if log_file:
+            log_pod_file = log_file
+        else:
+            pod_suffix = None
+            dir = None
+            if log_to_pods():
+                pod_suffix = ''
+                dir = pod_log_dir()
+
+            log_pod_file = AsyncJobs.log_file(command, job_id=AsyncJobs.new_id(), pod_suffix=pod_suffix, dir=dir)
+
+        if env_prefix:
+            command = f'{env_prefix} {command}'
+
+        log_err_file = log_pod_file.replace('.log', '.err')
+
+        command = command.replace('"', '\\"')
+        # nohup kubectl exec sts-0 -c cassandra -- /bin/sh -c "nohup nodetool -u cs-superuser -pw ... repair > /tmp/qing-db/q/logs/19230320.repair-0.log 2> /tmp/qing-db/q/logs/19230320.repair-0.err &" &
+        cmd = f'nohup kubectl exec {pod_name} -c {container} -- {shell} -c "nohup {command} > {log_pod_file} 2> {log_err_file} &"'
+        cmd = f'{cmd} &'
+
+        if show_out:
+            log2(cmd, text_color=text_color)
+
+        result = subprocess.run(cmd, capture_output=True, text=True, shell=True)
+        return PodExecResult(result.stdout, result.stderr, cmd, None, pod=pod_name, log_file=PodLogFile(log_pod_file, pod=pod_name))
+
+    def log_file_from_template(log_file: str, pod_name: str):
+        if not log_file:
+            return None
+
+        pod_suffix = pod_name
+        if pod_name:
+            if groups := re.match(r'.*-(.*)', pod_name):
+                pod_suffix = f'-{groups[1]}'
+            elif groups := re.match(r'.*_(.*)', pod_name):
+                pod_suffix = f'-{groups[1]}'
+
+        if not pod_suffix.startswith('-'):
+            pod_suffix = f'-{pod_suffix}'
+
+        return log_file.replace('{pod}', pod_suffix)
+
     def read_file(pod_name: str, container: str, namespace: str, file_path: str):
         v1 = client.CoreV1Api()
 
@@ -175,13 +287,15 @@ class Pods:
 
     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)}'
+            to_path = f'{local_downloads_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)
 
+        ReplSession().append_history(f':cat {to_path}')
+
         return to_path
 
     def get_container(namespace: str, pod_name: str, container_name: str):
@@ -294,5 +408,51 @@ class Pods:
     def completed(namespace: str, pod_name: str):
         return Pods.get(namespace, pod_name).status.phase in ['Succeeded', 'Failed']
 
-def log_prefix():
-    return Config().get('log-prefix', '/tmp/qing')
+    _dirs_created = set()
+
+    def creating_dir(pod_name: str,
+                     container: str,
+                     namespace: str,
+                     dir: str,
+                     show_out = False):
+        key = f'{dir}@{pod_name}'
+        if key not in Pods._dirs_created:
+            Pods._dirs_created.add(key)
+            Pods.exec(pod_name, container, namespace, f'mkdir -p {dir}', show_out=show_out, shell='bash')
+
+        return dir
+
+    def find_files(pod: str, container: str, namespace: str, pattern: str, mmin: int = 0, remote = False):
+        stdout = ''
+        if not remote:
+            # find . -maxdepth 1 -type f -name '*'
+            dir = os.path.dirname(pattern)
+            base = os.path.basename(pattern)
+            cmd = ['find', dir, '-name', base]
+            if mmin:
+                cmd += ['-mmin',  f'-{mmin}']
+
+            cmd += ["-exec", "stat", "-c", "'%n %s'", "{}", "\;"]
+
+            stdout = local_exec(cmd, show_out=Config().is_debug()).stdout
+        else:
+            cmd = f'find {pattern}'
+            if mmin:
+                cmd = f'{cmd} -mmin -{mmin}'
+
+            cmd += " -exec stat -c '%n %s' {} \;"
+
+            stdout = Pods.exec(pod, container, namespace, cmd, show_out=Config().is_debug(), shell='bash').stdout
+
+        log_files: list[PodLogFile] = []
+        for line in stdout.split('\n'):
+            line = line.strip(' \r')
+            if line:
+                tokens = line.split(' ')
+                log_files.append(PodLogFile(tokens[0], pod, size=tokens[1]))
+
+        return log_files
+
+    def file_size(pod: str, container:str, namespace: str, file: str):
+        stdout = Pods.exec(pod, container, namespace, f'stat -c %s {file}', show_out=Config().is_debug(), shell='bash').stdout
+        return stdout.strip('\r\n ')

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