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

adam/utils_k8s/pods.py

@@ -1,9 +1,8 @@
 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 typing import TypeVar
 from kubernetes import client
 from kubernetes.stream import stream
 from kubernetes.stream.ws_client import ERROR_CHANNEL, WSClient
@@ -11,7 +10,7 @@ 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 ParallelMapHandler, log2, debug, log_exc
 from .kube_context import KubeContext
 
 from websocket._core import WebSocket
@@ -30,11 +29,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,64 +40,13 @@ 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
+        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',
@@ -138,8 +84,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 +118,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 +129,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 +155,17 @@ 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 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,15 +55,6 @@ class StatefulSets:
 
         return statefulset_pods
 
-    def on_cluster(statefulset: str,
-                namespace: str,
-                body: Callable[[ThreadPoolExecutor, str, str, bool], T],
-                post: Callable[[T], T] = None,
-                action: str = 'action', max_workers=0, show_out=True, on_any = False, background = False) -> list[T]:
-        pods = StatefulSets.pod_names(statefulset, namespace)
-
-        return Pods.on_pods(pods, namespace, body, post=post, action=action, max_workers=max_workers, show_out=show_out, on_any=on_any, background=background)
-
     @functools.lru_cache()
     def pod_names(ss: str, ns: str):
         return [pod.metadata.name for pod in StatefulSets.pods(ss, ns)]
@@ -92,6 +80,7 @@ class StatefulSets:
 
         return restarted, False
 
+    @functools.lru_cache()
     def get_datacenter(sts: str, ns: str) -> str:
         v1 = client.AppsV1Api()
         namespace = ns

adam/utils_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,63 @@ import sqlite3
 import pandas
 
 from adam.config import Config
-from adam.utils import lines_to_tabular, log
+from adam.utils import lines_to_tabular, 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:
+        if self.keyspace:
+            self.conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{self.database}_{self.keyspace}.db')
+        else:
+            self.conn = SQLite.connect(self.database)
+
+        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 +84,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 +97,41 @@ class SQLite:
          if conn:
                conn.close()
 
+    def connect(database: str):
+        os.makedirs(SQLite.local_db_dir(), exist_ok=True)
+
+        conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{database}.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 = lines_to_tabular(lines, header='\t'.join(df.columns.tolist()), separator='\t')
+        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.145"  #: the working version
+__version__ = "2.0.172"  #: the working version
 __release__ = "1.0.0"  #: the release version

{kaqing-2.0.145.dist-info → kaqing-2.0.172.dist-info}/METADATA

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