kaqing 2.0.171__py3-none-any.whl → 2.0.204__py3-none-any.whl

adam/utils_k8s/pods.py

@@ -1,5 +1,8 @@
 from collections.abc import Callable
 from datetime import datetime
+import os
+import re
+import subprocess
 import sys
 import time
 from typing import TypeVar
@@ -8,9 +11,11 @@ from kubernetes.stream import stream
 from kubernetes.stream.ws_client import ERROR_CHANNEL, WSClient
 
 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 ParallelMapHandler, log2
+from adam.utils_k8s.pod_exec_result import PodExecResult
+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
@@ -29,11 +34,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()
@@ -50,9 +53,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):
@@ -61,6 +69,27 @@ class Pods:
 
         show_out = KubeContext.show_out(show_out)
 
+        if backgrounded or command.endswith(' &'):
+            command = command.strip(' &')
+
+            log_all_file = None
+            log_pod_file = None
+            if log_file:
+                log_pod_file = Pods.log_file_from_template(log_file, pod_name=pod_name)
+                if (a := Pods.log_file_from_template(log_file, pod_name='all')) != log_file:
+                    log_all_file = a
+            else:
+                log_pod_file = Pods.log_file(command, pod_name=pod_name)
+
+            command = command.replace('"', '\\"')
+            cmd = f'nohup kubectl exec {pod_name} -c {container} -- {shell} -c "{command} &" > {log_pod_file} 2>&1 &'
+            if log_all_file:
+                cmd = f'{cmd} >> {log_all_file}'
+
+            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)
+
         api = client.CoreV1Api()
 
         tty = True
@@ -68,26 +97,26 @@ class Pods:
         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
+        # if backgrounded or command.endswith(' &'):
+        #     print('!!!!SEAN backgrounded, but no via-sh!!!!!')
+        #     # 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 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 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 Config().is_debug():
-            log2(k_command)
+        debug(k_command)
 
         resp: WSClient = stream(
             api.connect_get_namespaced_pod_exec,
@@ -121,11 +150,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
@@ -134,13 +161,38 @@ class Pods:
         finally:
             resp.close()
             if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
-                try:
+                with log_exc():
                     s.sock.close()
-                except:
-                    pass
 
         return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output, pod=pod_name, log_file=log_file)
 
+    def log_file(command: str, pod_name: str = None, dt: datetime = None):
+        cmd_name = ''
+        if command.startswith('nodetool '):
+            command = command.strip(' &')
+            cmd_name = f".{'_'.join(command.split(' ')[5:])}"
+
+        pod_suffix = '{pod}'
+        if pod_name:
+            pod_suffix = pod_name
+            if groups := re.match(r'.*-(.*)', pod_name):
+                pod_suffix = f'-{groups[1]}'
+
+        if not dt:
+            dt = datetime.now()
+
+        return f'{log_prefix()}-{dt.strftime("%d%H%M%S")}{cmd_name}{pod_suffix}.log'
+
+    def log_file_from_template(log_file: str, pod_name: str):
+        pod_suffix = pod_name
+        if pod_name and (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()
 
@@ -162,20 +214,30 @@ class Pods:
                 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:
-                try:
+                with log_exc():
                     s.sock.close()
-                except:
-                    pass
+
+    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)
+
+        ReplSession().append_history(f':sh cat {to_path}')
+
+        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,42 @@
+import subprocess
+from adam.config import Config
+from adam.utils import ExecResult, debug
+
+def local_tmp_dir():
+    return Config().get('local-tmp-dir', '/tmp/qing-db')
+
+class LocalExecResult(ExecResult):
+    def __init__(self, stdout: str, stderr: str, command: str = None, code = 0, log_file: str = None):
+        self.stdout: str = stdout
+        self.stderr: str = stderr
+        self.command: str = command
+        self.code = code
+        self.pod = 'local'
+        self.log_file = log_file
+
+    def exit_code(self) -> int:
+        return self.code
+
+    def __str__(self):
+        return f'{"OK" if self.exit_code() == 0 else self.exit_code()} {self.command}'
+
+    def __audit_extra__(self):
+        return self.log_file if self.log_file else None
+
+def local_exec(cmd: list[str], shell=False, show_out=False):
+    stdout = ''
+    stderr = ''
+    returncode = 0
+
+    try:
+        if show_out:
+            debug(' '.join(cmd))
+
+        r = subprocess.run(cmd, capture_output=True, text=True, shell=shell)
+        stdout = r.stdout
+        stderr = r.stderr
+        returncode = r.returncode
+    except FileNotFoundError as e:
+        pass
+
+    return LocalExecResult(stdout, stderr, ' '.join(cmd), returncode)

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

adam/utils_repl/repl_completer.py

@@ -1,14 +1,57 @@
+import copy
 import re
-from typing import Iterable, TypeVar
+import traceback
+from typing import Iterable, TypeVar, cast
 from prompt_toolkit.completion import CompleteEvent, Completion, NestedCompleter, WordCompleter
 from prompt_toolkit.document import Document
 
+from adam.utils import debug_complete, log2
+from adam.utils_repl.appendable_completer import AppendableCompleter
+
+import nest_asyncio
+nest_asyncio.apply()
+
+import asyncio
+
 __all__ = [
     "ReplCompleter",
 ]
 
 T = TypeVar('T')
 
+def merge_completions(dict1, dict2):
+    if isinstance(dict1, dict):
+        target = dict1.copy()
+    else:
+        target = copy.copy(dict1)
+
+    try:
+        for key, value in dict2.items():
+            if key in target:
+                debug_complete(f'[{key}] {type(dict2)} is being merged to {type(target[key])} completions')
+                if isinstance(value, dict):
+                    if isinstance(target[key], dict):
+                        target[key] = merge_completions(target[key], value)
+                    elif isinstance(target[key], AppendableCompleter):
+                        cast(AppendableCompleter, target[key]).append_completions(key, value)
+                    elif isinstance(target[key], NestedCompleter):
+                        cast(NestedCompleter, target[key]).options = merge_completions(cast(NestedCompleter, target[key]).options, value)
+                elif isinstance(value, AppendableCompleter):
+                    if isinstance(target[key], dict):
+                        cast(AppendableCompleter, value).append_completions(key, target[key])
+                        target[key] = value
+                    else:
+                        log2(f'* {key} of {type(value)} is overriding existing {type(target[key])} completions')
+                else:
+                    target[key] = value
+            else:
+                target[key] = value
+
+    except Exception as e:
+        traceback.print_exc()
+
+    return target
+
 class ReplCompleter(NestedCompleter):
     def get_completions(
         self, document: Document, complete_event: CompleteEvent
@@ -43,4 +86,4 @@ class ReplCompleter(NestedCompleter):
                 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
+                yield c

adam/utils_repl/state_machine.py

@@ -1,6 +1,8 @@
 from abc import abstractmethod
 from typing import Generic, TypeVar
 
+from adam.utils import log_exc
+
 __all__ = [
     'State',
     'StateMachine',
@@ -160,14 +162,12 @@ class StateMachine(Generic[T]):
                 last_name = token
                 token = 'word'
 
-            try:
+            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
-            except:
-                pass
 
         return state

adam/utils_sqlite.py

@@ -1,3 +1,5 @@
+from collections.abc import Callable
+
 import functools
 import glob
 import os
@@ -5,11 +7,12 @@ import sqlite3
 import pandas
 
 from adam.config import Config
-from adam.utils import lines_to_tabular, log, wait_log
+from adam.utils import tabulize, log, wait_log
 
 class CursorHandler:
     def __init__(self, conn: sqlite3.Connection):
         self.conn = conn
+        self.cursor = None
 
     def __enter__(self):
         self.cursor = self.conn.cursor()
@@ -22,6 +25,28 @@ class CursorHandler:
 
         return False
 
+class SQLiteConnectionHandler:
+    def __init__(self, database: str, keyspace = None):
+        self.database = database
+        self.keyspace = keyspace
+        self.conn = None
+
+    def __enter__(self) -> sqlite3.Connection:
+        self.conn = SQLite.connect(self.database, self.keyspace)
+
+        self.conn.__enter__()
+
+        return self.conn
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.conn:
+            self.conn.__exit__(exc_type, exc_val, exc_tb)
+
+        return False
+
+def sqlite(database: str, keyspace = None):
+    return SQLiteConnectionHandler(database, keyspace=keyspace)
+
 # no state utility class
 class SQLite:
     def cursor(conn: sqlite3.Connection):
@@ -33,23 +58,6 @@ class SQLite:
     def keyspace(database: str):
         return '_'.join(database.replace(".db", "").split('_')[1:])
 
-    def connect(session: str):
-        os.makedirs(SQLite.local_db_dir(), exist_ok=True)
-
-        conn = None
-
-        try:
-            conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{session}_root.db')
-            with SQLite.cursor(conn) as cursor:
-                for d in SQLite.database_names(session):
-                    if d != f'{session}_root.db':
-                        q = f"ATTACH DATABASE '{SQLite.local_db_dir()}/{d}' AS {SQLite.keyspace(d)};"
-                        cursor.execute(q)
-        finally:
-            pass
-
-        return conn
-
     @functools.lru_cache()
     def database_names(prefix: str = None):
         wait_log('Inspecting export databases...')
@@ -86,24 +94,44 @@ class SQLite:
          if conn:
                conn.close()
 
+    def connect(database: str, keyspace: str = None):
+        os.makedirs(SQLite.local_db_dir(), exist_ok=True)
+
+        if keyspace:
+            return sqlite3.connect(f'{SQLite.local_db_dir()}/{database}_{keyspace}.db')
+        else:
+            conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{database}_root.db')
+            with SQLite.cursor(conn) as cursor:
+                for d in SQLite.database_names(database):
+                    if d != f'{database}.db':
+                        q = f"ATTACH DATABASE '{SQLite.local_db_dir()}/{d}' AS {SQLite.keyspace(d)}"
+                        cursor.execute(q)
+
+            return conn
+
     @functools.lru_cache()
     def column_names(tables: list[str] = [], database: str = None, function: str = 'audit', partition_cols_only = False):
         pass
 
-    def run_query(query: str, database: str = None, conn_passed = None):
-        conn = None
-        try:
-            if not conn_passed:
-                conn = SQLite.connect(database)
+    def run_query(query: str, database: str = None, output: Callable[[str], str] = None) -> tuple[int, str]:
+        with sqlite(database) as conn:
+            return SQLite.run_query_with_conn(conn, query, output=output)
+
+    def run_query_with_conn(conn, query: str, output: Callable[[str], str] = None) -> tuple[int, str]:
+        log_file = None
 
-            df = SQLite.query(conn_passed if conn_passed else conn, query)
-            lines = ['\t'.join(map(str, line)) for line in df.values.tolist()]
-            log(lines_to_tabular(lines, header='\t'.join(df.columns.tolist()), separator='\t'))
+        df = SQLite.query(conn, query)
+        lines = ['\t'.join(map(str, line)) for line in df.values.tolist()]
+        out = tabulize(lines, header='\t'.join(df.columns.tolist()), separator='\t', to=0)
+        if output:
+            log_file = output(out)
+        else:
+            log(out)
 
-            return len(lines)
-        finally:
-            if conn:
-                conn.close()
+        return len(lines), log_file
 
     def query(conn, sql: str) -> tuple[str, str, list]:
         return pandas.read_sql_query(sql, conn)
+
+def log_prefix():
+    return Config().get('export.log-prefix', '/tmp/qing')

adam/version.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-__version__ = "2.0.171"  #: the working version
+__version__ = "2.0.204"  #: the working version
 __release__ = "1.0.0"  #: the release version

{kaqing-2.0.171.dist-info → kaqing-2.0.204.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaqing
-Version: 2.0.171
+Version: 2.0.204
 Summary: UNKNOWN
 Home-page: UNKNOWN
 License: UNKNOWN