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

adam/utils_athena.py

@@ -1,23 +1,24 @@
+from collections.abc import Callable
 import functools
 import time
 import boto3
 import botocore
 
 from adam.config import Config
-from adam.utils import lines_to_tabular, log, log2
+from adam.utils import lines_to_tabular, log, log2, log_exc, wait_log
 
 # no state utility class
 class Athena:
    @functools.lru_cache()
    def database_names(like: str = None):
       # this function is called only from export currently
-      Config().wait_log(f'Inspecting export database schema...')
+      wait_log(f'Inspecting export database schema...')
 
       query = f"SELECT schema_name FROM information_schema.schemata WHERE schema_name <> 'information_schema'"
       if like:
          query = f"{query} AND schema_name like '{like}'"
 
-      try:
+      with log_exc():
          state, reason, rs = Athena.query(query)
          if rs:
             names = []
@@ -26,8 +27,6 @@ class Athena:
                   names.append(row_data[0])
 
             return names
-      except:
-         pass
 
       return []
 
@@ -64,7 +63,7 @@ class Athena:
 
    @functools.lru_cache()
    def column_names(tables: list[str] = [], database: str = None, function: str = 'audit', partition_cols_only = False):
-      try:
+      with log_exc():
          if not database:
             database = Config().get(f'{function}.athena.database', 'audit')
 
@@ -80,15 +79,13 @@ class Athena:
          _, _, rs = Athena.query(query)
          if rs:
             return [row['Data'][0].get('VarCharValue') for row in rs[1:]]
-      except:
-         # aws credentials not found
-         pass
 
       return []
 
-   def run_query(sql: str, database: str = None):
+   def run_query(sql: str, database: str = None, output: Callable[[str], str] = None):
       state, reason, rs = Athena.query(sql, database)
 
+      log_file = None
       if state == 'SUCCEEDED':
          if rs:
             column_info = rs[0]['Data']
@@ -98,14 +95,18 @@ class Athena:
                   row_data = [col.get('VarCharValue') if col else '' for col in row['Data']]
                   lines.append('\t'.join(row_data))
 
-            log(lines_to_tabular(lines, header='\t'.join(columns), separator='\t'))
+            out = lines_to_tabular(lines, header='\t'.join(columns), separator='\t')
+            if output:
+               log_file = output(out)
+            else:
+               log(out)
 
-            return len(lines)
+            return len(lines), log_file
       else:
          log2(f"Query failed or was cancelled. State: {state}")
          log2(f"Reason: {reason}")
 
-      return 0
+      return 0, log_file
 
    def query(sql: str, database: str = None, function: str = 'audit') -> tuple[str, str, list]:
       region_name = Config().get(f'{function}.athena.region', 'us-west-2')

adam/utils_audits.py

@@ -4,7 +4,8 @@ import time
 import requests
 
 from adam.config import Config
-from adam.utils import log2
+from adam.utils import OffloadHandler, debug, log2, log_exc, offload
+from adam.utils_athena import Athena
 from adam.utils_net import get_my_host
 
 class AuditMeta:
@@ -32,7 +33,7 @@ class Audits:
       try:
          response = requests.post(audit_endpoint, json=payload, timeout=Config().get("audit.timeout", 10))
          if response.status_code in [200, 201]:
-               Config().debug(response.text)
+               debug(response.text)
          else:
                log2(f"Error: {response.status_code} {response.text}")
       except requests.exceptions.Timeout as e:
@@ -42,15 +43,13 @@ class Audits:
       checked_in = 0.0
       cluster_last_checked = 0.0
 
-      state, _, rs = Audits.audit_query(f'select partitions_last_checked, clusters_last_checked from meta')
+      state, _, rs = Athena.query(f'select partitions_last_checked, clusters_last_checked from meta')
       if state == 'SUCCEEDED':
          if len(rs) > 1:
-            try:
+            with log_exc():
                row = rs[1]['Data']
                checked_in = float(row[0]['VarCharValue'])
                cluster_last_checked = float(row[1]['VarCharValue'])
-            except:
-               pass
 
       return AuditMeta(checked_in, cluster_last_checked)
 
@@ -67,7 +66,7 @@ class Audits:
       try:
          response = requests.post(audit_endpoint, json=payload, timeout=Config().get("audit.timeout", 10))
          if response.status_code in [200, 201]:
-               Config().debug(response.text)
+               debug(response.text)
          else:
                log2(f"Error: {response.status_code} {response.text}")
       except requests.exceptions.Timeout as e:
@@ -84,13 +83,11 @@ class Audits:
          f'(select distinct c as name from audit where {Audits.date_from(dt_object)}) as c2',
          'on c1.name = c2.name where c1.name is null'])
       log2(query)
-      state, _, rs = Audits.audit_query(query)
+      state, _, rs = Athena.query(query)
       if state == 'SUCCEEDED':
          if len(rs) > 1:
-               try:
+               with log_exc():
                   return [r['Data'][0]['VarCharValue'] for r in rs[1:]]
-               except:
-                  pass
 
       return []
 
@@ -99,4 +96,7 @@ class Audits:
         m = dt_object.strftime("%m")
         d = dt_object.strftime("%d")
 
-        return f"y = '{y}' and m = '{m}' and d >= '{d}' or y = '{y}' and m > '{m}' or y > '{y}'"
+        return f"y = '{y}' and m = '{m}' and d >= '{d}' or y = '{y}' and m > '{m}' or y > '{y}'"
+
+   def offload() -> OffloadHandler:
+       return offload(max_workers=Config().get('audit.workers', 3))

adam/utils_issues.py

@@ -0,0 +1,32 @@
+from adam.checks.check_result import CheckResult
+from adam.checks.issue import Issue
+from adam.repl_session import ReplSession
+from adam.utils import lines_to_tabular, log, log2
+
+class IssuesUtils:
+    def show(check_results: list[CheckResult], in_repl = False):
+        IssuesUtils.show_issues(CheckResult.collect_issues(check_results), in_repl=in_repl)
+
+    def show_issues(issues: list[Issue], in_repl = False):
+        if not issues:
+            log2('No issues found.')
+        else:
+            suggested = 0
+            log2(f'* {len(issues)} issues found.')
+            lines = []
+            for i, issue in enumerate(issues, start=1):
+                lines.append(f"{i}||{issue.category}||{issue.desc}")
+                lines.append(f"||statefulset||{issue.statefulset}@{issue.namespace}")
+                lines.append(f"||pod||{issue.pod}@{issue.namespace}")
+                if issue.details:
+                    lines.append(f"||details||{issue.details}")
+
+                if issue.suggestion:
+                    lines.append(f'||suggestion||{issue.suggestion}')
+                    if in_repl:
+                        ReplSession().prompt_session.history.append_string(issue.suggestion)
+                        suggested += 1
+            log(lines_to_tabular(lines, separator='||'))
+            if suggested:
+                log2()
+                log2(f'* {suggested} suggested commands are added to history. Press <Up> arrow to access them.')

adam/utils_k8s/app_clusters.py

@@ -1,10 +1,9 @@
-from concurrent.futures import ThreadPoolExecutor
 import sys
 from typing import TypeVar
 
 from adam.utils_k8s.app_pods import AppPods
 from adam.pod_exec_result import PodExecResult
-from adam.utils import log2
+from adam.utils import log, log2
 from adam.utils_k8s.pods import Pods
 from .kube_context import KubeContext
 
@@ -14,20 +13,16 @@ T = TypeVar('T')
 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', background = False) -> list[PodExecResult]:
-        def body(executor: ThreadPoolExecutor, pod: str, namespace: str, show_out: bool):
-            if executor:
-                return executor.submit(AppPods.exec, pod, namespace, command, False, False, shell, background)
+        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:
+            results: list[PodExecResult] = exec.map(lambda pod: AppPods.exec(pod, namespace, command, False, False, shell, background))
+            for result in results:
+                if KubeContext.show_out(show_out):
+                    log(result.command)
+                    if result.stdout:
+                        log(result.stdout)
+                    if result.stderr:
+                        log2(result.stderr, file=sys.stderr)
 
-            return AppPods.exec(pod, namespace, command, show_out=show_out, background=background)
-
-        def post(result, show_out: bool):
-            if KubeContext.show_out(show_out):
-                print(result.command)
-                if result.stdout:
-                    print(result.stdout)
-                if result.stderr:
-                    log2(result.stderr, file=sys.stderr)
-
-            return result
-
-        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)
+            return results

adam/utils_k8s/app_pods.py

@@ -1,3 +1,4 @@
+import functools
 from typing import List
 from kubernetes import client
 
@@ -8,6 +9,7 @@ from adam.repl_session import ReplSession
 
 # utility collection on app pods; methods are all static
 class AppPods:
+    @functools.lru_cache()
     def pod_names(namespace: str, env: str, app: str):
         return [pod.metadata.name for pod in AppPods.app_pods(namespace, env, app)]
 

adam/utils_k8s/cassandra_clusters.py

@@ -1,33 +1,36 @@
-from concurrent.futures import ThreadPoolExecutor
 import sys
 from typing import TypeVar
 
+from adam.config import Config
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.pod_exec_result import PodExecResult
-from adam.utils import log2
-from .statefulsets import StatefulSets
-from .kube_context import KubeContext
+from adam.utils import log, log2
+from adam.utils_k8s.pods import Pods
+from adam.utils_k8s.statefulsets import StatefulSets
 
 T = TypeVar('T')
 
 # utility collection on cassandra clusters; methods are all static
 class CassandraClusters:
-    def exec(statefulset: str, namespace: str, command: str, action: str = 'action',
+    def exec(sts: str, namespace: str, command: str, action: str = 'action',
              max_workers=0, show_out=True, on_any = False, shell = '/bin/sh', background = False, log_file = None) -> list[PodExecResult]:
-        def body(executor: ThreadPoolExecutor, pod: str, namespace: str, show_out: bool):
-            if executor:
-                return executor.submit(CassandraNodes.exec, pod, namespace, command, False, False, shell, background, log_file)
 
-            return CassandraNodes.exec(pod, namespace, command, show_out=show_out, background=background, log_file=log_file)
+        pods = StatefulSets.pod_names(sts, namespace)
+        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:
+            results: list[PodExecResult] = exec.map(lambda pod: CassandraNodes.exec(pod, namespace, command, False, False, shell, background, log_file))
+            for result in results:
+                if show_out and not Config().is_debug():
+                    log(result.command)
+                    if result.stdout:
+                        log(result.stdout)
+                    if result.stderr:
+                        log2(result.stderr, file=sys.stderr)
 
-        def post(result, show_out: bool):
-            if KubeContext.show_out(show_out):
-                print(result.command)
-                if result.stdout:
-                    print(result.stdout)
-                if result.stderr:
-                    log2(result.stderr, file=sys.stderr)
+            return results
 
-            return result
+    def pod_names_by_host_id(sts: str, ns: str):
+        pods = StatefulSets.pods(sts, ns)
 
-        return StatefulSets.on_cluster(statefulset, namespace, body, post=post, action=action, max_workers=max_workers, show_out=show_out, on_any=on_any, background=background)
+        return {CassandraNodes.get_host_id(pod.metadata.name, ns): pod.metadata.name for pod in pods}

adam/utils_k8s/custom_resources.py

@@ -1,10 +1,10 @@
+import functools
 import re
-import time
 from kubernetes import client
 
 from adam.config import Config
 from .kube_context import KubeContext
-from adam.utils import elapsed_time, lines_to_tabular, log2
+from adam.utils import log2, log_exc
 
 
 # utility collection; methods are all static
@@ -19,14 +19,12 @@ class CustomResources:
         strip = Config().get('app.strip', '0')
 
         v1 = client.CustomObjectsApi()
-        try:
+        with log_exc():
             c3cassandras = v1.list_cluster_custom_object(group=group, version=v, plural=plural)
             for c in c3cassandras.items():
                 if c[0] == 'items':
                     for item in c[1]:
                         app_ids_by_ss[f"{item['metadata']['name']}@{item['metadata']['namespace']}"] = item['metadata']['labels'][label].strip(strip)
-        except Exception:
-            pass
 
         return app_ids_by_ss
 
@@ -121,11 +119,10 @@ class CustomResources:
         body = bkspecs
         pretty = 'true'
 
-        try:
+        with log_exc(lambda e: "Exception when calling create_medusa_backupjob.create_namespaced_custom_object: %s\n" % e):
             api_instance.create_namespaced_custom_object(group, version, namespace, plural, body, pretty=pretty)
             log2(f"create_medusa_backupjob: created Full Backup {bkname}: {api_instance}")
-        except Exception as e:
-            log2("Exception when calling create_medusa_backupjob.create_namespaced_custom_object: %s\n" % e)
+
         return None
 
     def create_medusa_restorejob(restorejobname: str, bkname: str, dc: str, ns: str):
@@ -154,11 +151,10 @@ class CustomResources:
         body = rtspecs
         pretty = 'true'
 
-        try:
+        with log_exc(lambda e: "Exception when calling create_medusa_restorejob.create_namespaced_custom_object: %s\n" % e):
             api_instance.create_namespaced_custom_object(group, version, namespace, plural, body, pretty=pretty)
             log2(f"create_medusa_restorejob: created Restore Job {restorejobname}: {api_instance}")
-        except Exception as e:
-            log2("Exception when calling create_medusa_restorejob.create_namespaced_custom_object: %s\n" % e)
+
         return None
 
     def medusa_show_backup_names(dc: str, ns: str) -> list[dict]:
@@ -171,6 +167,10 @@ class CustomResources:
 
         return None
 
+    def clear_caches():
+        CustomResources.medusa_show_backupjobs.cache_clear()
+
+    @functools.lru_cache()
     def medusa_show_backupjobs(dc: str, ns: str) -> list[dict]:
         api_instance = client.CustomObjectsApi()
         group = 'medusa.k8ssandra.io'
@@ -180,11 +180,10 @@ class CustomResources:
         pretty = 'true'
         label_selector = 'cassandra.datastax.com/datacenter=' + dc
 
-        try:
+        with log_exc(lambda e: "Exception when calling medusa_show_backupjobs.list_namespaced_custom_object: %s\n" % e):
             api_response = api_instance.list_namespaced_custom_object(group, version, namespace, plural, pretty=pretty, label_selector=label_selector)
             return api_response['items']
-        except Exception as e:
-            log2("Exception when calling medusa_show_backupjobs.list_namespaced_custom_object: %s\n" % e)
+
         return None
 
     def medusa_show_restorejobs(dc: str, ns: str):
@@ -196,11 +195,11 @@ class CustomResources:
         pretty = 'true'
         label_selector = 'cassandra.datastax.com/datacenter=' + dc
         rtlist = []
-        try:
+
+        with log_exc(lambda e: "Exception when calling medusa_show_restorejobs.list_namespaced_custom_object: %s\n" % e):
             api_response = api_instance.list_namespaced_custom_object(group, version, namespace, plural, pretty=pretty, label_selector=label_selector)
             for x in api_response['items']:
                 rtlist.append(f"{x['metadata']['name']}\t{x['metadata']['creationTimestamp']}\t{x['status'].get('finishTime', '')}")
             return rtlist
-        except Exception as e:
-            log2("Exception when calling medusa_show_restorejobs.list_namespaced_custom_object: %s\n" % e)
+
         return None

adam/utils_k8s/ingresses.py

@@ -1,7 +1,7 @@
 from kubernetes import client
 
 from adam.config import Config
-from adam.utils import log2
+from adam.utils import debug, log2
 
 # utility collection on ingresses; methods are all static
 class Ingresses:
@@ -63,7 +63,7 @@ class Ingresses:
                 log2(f"200 Ingress '{name}' in namespace '{namespace}' deleted successfully.")
             else:
                 api.delete_namespaced_ingress(name=name, namespace=namespace)
-                Config().debug(f"200 Ingress '{name}' in namespace '{namespace}' deleted successfully.")
+                debug(f"200 Ingress '{name}' in namespace '{namespace}' deleted successfully.")
         except client.ApiException as e:
             log2(f"Error deleting Ingress: {e}")
 

adam/utils_k8s/jobs.py

@@ -1,7 +1,7 @@
 from kubernetes import client
 from time import sleep
 from .pods import Pods
-from adam.utils import log2
+from adam.utils import log2, log_exc
 
 # utility collection on jobs; methods are all static
 class Jobs:
@@ -20,11 +20,10 @@ class Jobs:
             metadata=client.V1ObjectMeta(name=job_name),
             spec=spec)
 
-        try:
+        with log_exc(lambda e: "Exception when calling BatchV1Apii->create_namespaced_job: %s\n" % e):
             client.BatchV1Api().create_namespaced_job(body=job, namespace=namespace)
             log2(f"Job {job_name} created in {namespace}")
-        except Exception as e:
-            log2("Exception when calling BatchV1Apii->create_namespaced_job: %s\n" % e)
+
         return
 
     def get_job_pods(job_name: str, namespace: str):
@@ -32,7 +31,7 @@ class Jobs:
         return pods
 
     def delete(job_name: str, namespace: str, wait=True):
-        try:
+        with log_exc(lambda e: "Exception when calling BatchV1Apii->delete_namespaced_job: %s\n" % e):
             client.BatchV1Api().delete_namespaced_job(name=job_name, namespace=namespace, propagation_policy='Background')
             if wait:
                 while True:
@@ -41,14 +40,11 @@ class Jobs:
                         return
                     sleep(5)
             log2(f"Job {job_name} in {namespace} deleted.")
-        except Exception as e:
-            log2("Exception when calling BatchV1Apii->delete_namespaced_job: %s\n" % e)
+
         return
 
     def get_logs(job_name: str, namespace: str):
         v1 = client.CoreV1Api()
-        try:
+        with log_exc(lambda e: "Exception when calling CorV1Apii->list_namespaced_pod, cannot find job pod: %s\n" % e):
             pod_name = Jobs.get_job_pods(job_name, namespace).items[0].metadata.name
-            log2(v1.read_namespaced_pod_log(name=pod_name, namespace=namespace))
-        except Exception as e:
-            log2("Exception when calling CorV1Apii->list_namespaced_pod, cannot find job pod: %s\n" % e)
+            log2(v1.read_namespaced_pod_log(name=pod_name, namespace=namespace))

adam/utils_k8s/k8s.py

@@ -0,0 +1,87 @@
+from collections.abc import Callable
+import re
+import portforward
+
+from adam.commands.command import InvalidState
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_k8s.kube_context import KubeContext
+
+class PortForwardHandler:
+    connections: dict[str, int] = {}
+
+    def __init__(self, state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
+        self.state = state
+        self.local_port = local_port
+        self.svc_or_pod = svc_or_pod
+        self.target_port = target_port
+        self.forward_connection = None
+        self.pod = None
+
+    def __enter__(self) -> tuple[str, str]:
+        state = self.state
+
+        if not self.svc_or_pod:
+            log2('No service or pod found.')
+
+            raise InvalidState(state)
+
+        if KubeContext.in_cluster():
+            svc_name = self.svc_or_pod(True)
+            if not svc_name:
+                log2('No service found.')
+
+                raise InvalidState(state)
+
+            # cs-a526330d23-cs-a526330d23-default-sts-0 ->
+            # curl http://cs-a526330d23-cs-a526330d23-reaper-service.stgawsscpsr.svc.cluster.local:8080
+            groups = re.match(r'^(.*?-.*?-.*?-.*?-).*', state.sts)
+            if groups:
+                svc = f'{groups[1]}{svc_name}.{state.namespace}.svc.cluster.local:{self.target_port}'
+                return (svc, svc)
+            else:
+                raise InvalidState(state)
+        else:
+            pod = self.svc_or_pod(False)
+            if not pod:
+                log2('No pod found.')
+
+                raise InvalidState(state)
+
+            self.pod = pod
+            self.forward_connection = portforward.forward(state.namespace, pod, self.local_port, self.target_port)
+            if self.inc_connection_cnt() == 1:
+                self.forward_connection.__enter__()
+
+            return (f'localhost:{self.local_port}', f'{pod}:{self.target_port}')
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.forward_connection:
+            if not self.dec_connection_cnt():
+                return self.forward_connection.__exit__(exc_type, exc_val, exc_tb)
+
+        return False
+
+    def inc_connection_cnt(self):
+        id = self.connection_id(self.pod)
+        if id not in PortForwardHandler.connections:
+            PortForwardHandler.connections[id] = 1
+        else:
+            PortForwardHandler.connections[id] += 1
+
+        return PortForwardHandler.connections[id]
+
+    def dec_connection_cnt(self):
+        id = self.connection_id(self.pod)
+        if id not in PortForwardHandler.connections:
+            PortForwardHandler.connections[id] = 0
+        elif PortForwardHandler.connections[id] > 0:
+            PortForwardHandler.connections[id] -= 1
+
+        return PortForwardHandler.connections[id]
+
+    def connection_id(self, pod: str):
+        return f'{self.local_port}:{pod}:{self.target_port}'
+
+def port_forwarding(state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
+    return PortForwardHandler(state, local_port, svc_or_pod, target_port)