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

adam/{k8s_utils → utils_k8s}/pods.py

@@ -1,23 +1,28 @@
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime
 import sys
 import time
 from typing import TypeVar, cast
 from 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.k8s_utils.volumes import ConfigMapMount
+from adam.utils_k8s.volumes import ConfigMapMount
 from adam.pod_exec_result import PodExecResult
 from adam.utils import elapsed_time, log2
 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
@@ -27,7 +32,7 @@ class Pods:
     def delete(pod_name: str, namespace: str, grace_period_seconds: int = None):
         try:
             v1 = client.CoreV1Api()
-            api_response = v1.delete_namespaced_pod(pod_name, namespace, grace_period_seconds=grace_period_seconds)
+            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)
 
@@ -42,12 +47,16 @@ class Pods:
                 namespace: str,
                 body: Callable[[ThreadPoolExecutor, str, str, bool], T],
                 post: Callable[[T], T] = None,
-                action: str = 'action', max_workers=0, show_out=True) -> list[T]:
+                action: str = 'action',
+                max_workers=0,
+                show_out=True,
+                on_any = False,
+                background = False) -> list[T]:
         show_out = KubeContext.show_out(show_out)
 
         if not max_workers:
             max_workers = Config().action_workers(action, 0)
-        if max_workers > 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...')
@@ -70,7 +79,7 @@ class Pods:
         else:
             results: list[T] = []
 
-            samples = Config().action_node_samples(action, sys.maxsize)
+            samples = 1 if on_any else Config().action_node_samples(action, sys.maxsize)
             l = min(len(pods), samples)
             adj = 'all'
             if l < len(pods):
@@ -79,7 +88,8 @@ class Pods:
                 log2(f'Executing on {adj} nodes from statefulset...')
             for pod_name in pods:
                 try:
-                    result = body(None, pod_name, namespace, show_out)
+                    # 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)
@@ -92,7 +102,12 @@ class Pods:
 
             return results
 
-    def exec(pod_name: str, container: str, namespace: str, command: str, show_out = True, throw_err = False, interaction: Callable[[any, list[str]], any] = None):
+    def exec(pod_name: str, container: str, namespace: str, command: str,
+             show_out = True, throw_err = False, shell = '/bin/sh',
+             background = 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
 
@@ -100,12 +115,33 @@ class Pods:
 
         api = client.CoreV1Api()
 
-        exec_command = ["/bin/sh", "-c", command]
-        k_command = f'kubectl exec {pod_name} -c {container} -n {namespace} -- {command}'
+        tty = True
+        exec_command = [shell, '-c', command]
+        if env_prefix:
+            exec_command = [shell, '-c', f'{env_prefix} {command}']
+
+        if background or command.endswith(' &'):
+            # should be false for starting a background process
+            tty = False
+
+            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}"'
         if show_out:
             print(k_command)
 
-        resp = stream(
+        resp: WSClient = stream(
             api.connect_get_namespaced_pod_exec,
             pod_name,
             namespace,
@@ -114,10 +150,11 @@ class Pods:
             stderr=True,
             stdin=True,
             stdout=True,
-            tty=True,
+            tty=tty,
             _preload_content=False,
         )
 
+        s: WebSocket = resp.sock
         stdout = []
         stderr = []
         error_output = None
@@ -139,7 +176,7 @@ class Pods:
             try:
                 # get the exit code from server
                 error_output = resp.read_channel(ERROR_CHANNEL)
-            except Exception:
+            except Exception as e:
                 pass
         except Exception as e:
             if throw_err:
@@ -148,9 +185,49 @@ class Pods:
                 log2(e)
         finally:
             resp.close()
+            if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
+                try:
+                    s.sock.close()
+                except:
+                    pass
+
+        return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output, pod=pod_name, log_file=log_file)
 
-        return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output)
+    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()
+
+            try:
+                # 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:
+                    s.sock.close()
+                except:
+                    pass
     def get_container(namespace: str, pod_name: str, container_name: str):
         pod = Pods.get(namespace, pod_name)
         if not pod:
@@ -233,17 +310,33 @@ class Pods:
             ))
         )
 
-    def wait_for_running(namespace: str, pod_name: str, msg: str=None, label_selector: str = None):
-        msged = False
-
-        while (Pods.get_with_selector(namespace, label_selector) if label_selector else Pods.get(namespace, pod_name)).status.phase != 'Running':
-            if not msged:
-                if not msg:
-                    msg = f'Waiting for the {pod_name} pod to start up...'
-                log2(msg, nl=False)
-                msged = True
-            time.sleep(5)
-        log2(' OK')
+    def wait_for_running(namespace: str, pod_name: str, msg: str = None, label_selector: str = None):
+        cnt = 2
+        while (cnt < 302 and Pods.get_with_selector(namespace, label_selector) if label_selector else Pods.get(namespace, pod_name)).status.phase != 'Running':
+            if not msg:
+                msg = f'Waiting for the {pod_name} pod to start up.'
+
+            max_len = len(msg) + 3
+            mod = cnt % 3
+            padded = ''
+            if mod == 0:
+                padded = f'\r{msg}'.ljust(max_len)
+            elif mod == 1:
+                padded = f'\r{msg}.'.ljust(max_len)
+            else:
+                padded = f'\r{msg}..'.ljust(max_len)
+            log2(padded, nl=False)
+            cnt += 1
+            time.sleep(1)
+
+        log2(f'\r{msg}..'.ljust(max_len), nl=False)
+        if cnt < 302:
+            log2(' OK')
+        else:
+            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/{k8s_utils → utils_k8s}/secrets.py

@@ -1,4 +1,5 @@
 import base64
+import functools
 import re
 from typing import cast
 from kubernetes import client
@@ -9,7 +10,10 @@ from adam.utils import log2
 
 # 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...')
+
         secrets_names = []
 
         v1 = client.CoreV1Api()

adam/{k8s_utils → utils_k8s}/statefulsets.py

@@ -24,11 +24,12 @@ class StatefulSets:
 
         return statefulsets.items
 
+    @functools.lru_cache()
     def list_sts_name_and_ns():
         return [(statefulset.metadata.name, statefulset.metadata.namespace) for statefulset in StatefulSets.list_sts()]
 
-    def list_sts_names(show_namespace = True):
-        if show_namespace:
+    def list_sts_names():
+        if not KubeContext.in_cluster_namespace():
             return [f"{sts}@{ns}" for sts, ns in StatefulSets.list_sts_name_and_ns()]
         else:
             return [f"{sts}" for sts, _ in StatefulSets.list_sts_name_and_ns()]
@@ -61,10 +62,10 @@ class StatefulSets:
                 namespace: str,
                 body: Callable[[ThreadPoolExecutor, str, str, bool], T],
                 post: Callable[[T], T] = None,
-                action: str = 'action', max_workers=0, show_out=True) -> list[T]:
+                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)
+        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):

adam/utils_net.py

@@ -0,0 +1,24 @@
+import socket
+
+MY_HOST = None
+
+def get_my_host():
+    global MY_HOST
+
+    if MY_HOST:
+       return MY_HOST
+
+    MY_HOST = get_ip_from_hostname('host.docker.internal')
+    if not MY_HOST:
+       MY_HOST = socket.gethostname()
+
+    if not MY_HOST:
+       MY_HOST = 'NA'
+
+    return MY_HOST
+
+def get_ip_from_hostname(hostname):
+    try:
+        return socket.gethostbyname(hostname)
+    except socket.gaierror:
+        return None

adam/utils_repl/automata_completer.py

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

adam/utils_repl/repl_completer.py

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

adam/utils_repl/state_machine.py

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

@@ -0,0 +1,101 @@
+import functools
+import glob
+import os
+import sqlite3
+import pandas
+
+from adam.config import Config
+from adam.utils import lines_to_tabular, log
+
+# no state utility class
+class SQLite:
+    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...')
+
+        pattern = f'{SQLite.local_db_dir()}/s*.db'
+        if prefix:
+            pattern = f'{SQLite.local_db_dir()}/{prefix}*'
+        return [os.path.basename(f) for f in glob.glob(pattern)]
+
+    def clear_cache(cache: str = None):
+        SQLite.database_names.cache_clear()
+        SQLite.table_names.cache_clear()
+
+    @functools.lru_cache()
+    def table_names(database: str):
+      tokens = database.replace('.db', '').split('_')
+      ts_prefix = tokens[0]
+      keyspace = '_'.join(tokens[1:])
+
+      conn = None
+      tables = []
+      try:
+         conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{ts_prefix}_{keyspace}.db')
+         cursor = None
+         try:
+            cursor = conn.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:
+         print(f"Error connecting to or querying the database: {e}")
+         return []
+      finally:
+         if conn:
+               conn.close()
+
+    @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)
+
+            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'))
+
+            return len(lines)
+        finally:
+            if conn:
+                conn.close()
+
+    def query(conn, sql: str) -> tuple[str, str, list]:
+        return pandas.read_sql_query(sql, conn)