kaqing 2.0.145__py3-none-any.whl → 2.0.189__py3-none-any.whl

adam/utils_k8s/pods.py

@@ -1,9 +1,9 @@
 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, WSClient
@@ -11,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 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
@@ -30,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()
@@ -43,68 +42,22 @@ 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,
+        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):
@@ -120,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
 
@@ -138,8 +91,7 @@ class Pods:
                 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: WSClient = stream(
             api.connect_get_namespaced_pod_exec,
@@ -173,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
@@ -186,10 +136,8 @@ 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)
 
@@ -214,20 +162,28 @@ 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)
+
+        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/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,18 +55,12 @@ 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)]
+    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
@@ -92,6 +83,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_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,58 @@
+from concurrent.futures import ThreadPoolExecutor
+import copy
+import inspect
 import re
-from typing import Iterable, TypeVar
-from prompt_toolkit.completion import CompleteEvent, Completion, NestedCompleter, WordCompleter
+import threading
+import traceback
+from typing import Iterable, TypeVar, cast
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion, NestedCompleter, WordCompleter
 from prompt_toolkit.document import Document
 
+from adam.utils import debug, debug_complete, log2, log_timing
+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(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
@@ -42,5 +86,87 @@ class ReplCompleter(NestedCompleter):
                 # 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
+
+lock = threading.Lock()
+in_queue = set()
+
+def preload(action: callable, log_key: str = None):
+    with lock:
+        if not LazyNestedCompleter.loop:
+            LazyNestedCompleter.loop = asyncio.new_event_loop()
+            LazyNestedCompleter.async_exec = ThreadPoolExecutor(max_workers=6, thread_name_prefix='async')
+            LazyNestedCompleter.loop.set_default_executor(LazyNestedCompleter.async_exec)
+
+        # some lib does not handle asyncio loop properly, as sync exec submit does not work, use another async loop
+        async def a():
+            try:
+                arg_needed = len(action.__code__.co_varnames)
+
+                if log_key:
+                    with log_timing(log_key):
+                        r = action(None) if arg_needed else action()
+                else:
+                    r = action(None) if arg_needed else action()
+                if inspect.isawaitable(r):
+                    await r
+
+                in_queue.remove(log_key)
+            except Exception as e:
+                log2('preloading error', e, inspect.getsourcelines(action)[0][0])
+                traceback.print_exc()
+
+        if log_key not in in_queue:
+            in_queue.add(log_key)
+            LazyNestedCompleter.async_exec.submit(lambda: LazyNestedCompleter.loop.run_until_complete(a()))
+
+class LazyNestedCompleter(NestedCompleter):
+    loop: asyncio.AbstractEventLoop = None
+    async_exec: ThreadPoolExecutor = None
+
+    def __init__(self, name: str, options_lambda: callable, ignore_case: bool = True, auto: str = 'lazy') -> None:
+        super().__init__(None, ignore_case)
+        self.options_lambda = options_lambda
+        if auto == 'lazy':
+            preload(options_lambda, log_key=name)
+
+    def __repr__(self) -> str:
+        return "LazyNestedCompleter(%r, ignore_case=%r)" % (self.options, self.ignore_case)
+
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        if not self.options:
+            self.options = self.options_lambda()
+
+        # 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(
+                list(self.options.keys()), ignore_case=self.ignore_case
+            )
             for c in completer.get_completions(document, complete_event):
                 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,41 +7,60 @@ import sqlite3
 import pandas
 
 from adam.config import Config
-from adam.utils import lines_to_tabular, 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()
+
+        return self.cursor
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.cursor:
+            self.cursor.close()
+
+        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):
+        return CursorHandler(conn)
+
     def local_db_dir():
         return Config().get('export.sqlite.local-db-dir', '/tmp/qing-db')
 
     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')
-            cursor = None
-            try:
-                cursor = conn.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:
-                if cursor:
-                    cursor.close()
-        finally:
-            pass
-
-        return conn
-
     @functools.lru_cache()
     def database_names(prefix: str = None):
-        Config().wait_log('Inspecting export databases...')
+        wait_log('Inspecting export databases...')
 
         pattern = f'{SQLite.local_db_dir()}/s*.db'
         if prefix:
@@ -60,15 +81,10 @@ class SQLite:
       tables = []
       try:
          conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{ts_prefix}_{keyspace}.db')
-         cursor = None
-         try:
-            cursor = conn.cursor()
+         with SQLite.cursor(conn) as cursor:
             cursor.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;")
 
             tables = [row[0] for row in cursor.fetchall() if row[0] != "sqlite_sequence"]
-         finally:
-               if cursor:
-                  cursor.close()
 
          return tables
       except sqlite3.Error as e:
@@ -78,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')