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

adam/sql/term_completer.py

@@ -0,0 +1,76 @@
+from typing import Callable, Iterable, List, Mapping, Optional, Pattern, Union
+
+from prompt_toolkit.completion import CompleteEvent, Completion, WordCompleter
+from prompt_toolkit.document import Document
+from prompt_toolkit.formatted_text import AnyFormattedText
+
+__all__ = [
+    "TermCompleter",
+]
+
+class TermCompleter(WordCompleter):
+    def __init__(
+        self,
+        words: Union[List[str], Callable[[], List[str]]],
+        ignore_case: bool = False,
+        display_dict: Optional[Mapping[str, AnyFormattedText]] = None,
+        meta_dict: Optional[Mapping[str, AnyFormattedText]] = None,
+        WORD: bool = False,
+        sentence: bool = False,
+        match_middle: bool = False,
+        pattern: Optional[Pattern[str]] = None,
+    ) -> None:
+        super().__init__(words, ignore_case, display_dict, meta_dict, WORD, sentence, match_middle, pattern)
+
+    def __str__(self):
+        return ','.join(self.words)
+
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        # Get list of words.
+        words = self.words
+        if callable(words):
+            words = words()
+
+        # Get word/text before cursor.
+        if self.sentence:
+            word_before_cursor = document.text_before_cursor
+        else:
+            word_before_cursor = document.get_word_before_cursor(
+                WORD=self.WORD, pattern=self.pattern
+            )
+
+        if self.ignore_case:
+            word_before_cursor = word_before_cursor.lower()
+
+        def word_matches(word: str) -> bool:
+            """True when the word before the cursor matches."""
+            if self.ignore_case:
+                word = word.lower()
+
+            if self.match_middle:
+                return word_before_cursor in word
+            else:
+                return word.startswith(word_before_cursor)
+
+        for a in words:
+            if word_before_cursor in ['(', ',', '=', 'in', "',"]:
+                display = self.display_dict.get(a, a)
+                display_meta = self.meta_dict.get(a, "")
+                yield Completion(
+                    a,
+                    0,
+                    display=display,
+                    display_meta=display_meta,
+                )
+
+            if word_matches(a):
+                display = self.display_dict.get(a, a)
+                display_meta = self.meta_dict.get(a, "")
+                yield Completion(
+                    a,
+                    -len(word_before_cursor),
+                    display=display,
+                    display_meta=display_meta,
+                )

adam/sso/authn_ad.py

@@ -137,7 +137,7 @@ class AdAuthenticator(Authenticator):
         return []
 
     def parse_id_token(self, id_token: str) -> IdToken:
-        jwks_url = Config().get('idps.ad.jwks-uri', 'https://login.microsoftonline.com/common/discovery/keys')
+        jwks_url = Config().get('idps.ad.jwks-uri', '')
         try:
             jwks_client = jwt.PyJWKClient(jwks_url, cache_jwk_set=True, lifespan=360)
             signing_key = jwks_client.get_signing_key_from_jwt(id_token)

adam/sso/cred_cache.py

@@ -4,7 +4,7 @@ import traceback
 from dotenv import load_dotenv
 
 from adam.config import Config
-from adam.k8s_utils.kube_context import KubeContext
+from adam.utils_k8s.kube_context import KubeContext
 
 class CredCache:
     # the singleton pattern

adam/sso/idp.py

@@ -9,7 +9,7 @@ import requests
 from kubernetes import config
 import yaml
 
-from adam.k8s_utils.secrets import Secrets
+from adam.utils_k8s.secrets import Secrets
 
 from .cred_cache import CredCache
 from .idp_session import IdpSession

adam/utils.py

@@ -9,7 +9,8 @@ import os
 from pathlib import Path
 import random
 import string
-from typing import cast
+import threading
+from typing import Callable
 from dateutil import parser
 import subprocess
 import sys
@@ -19,6 +20,8 @@ import yaml
 
 from . import __version__
 
+is_debug_holder = [lambda: False]
+
 def to_tabular(lines: str, header: str = None, dashed_line = False):
     return lines_to_tabular(lines.split('\n'), header, dashed_line)
 
@@ -70,6 +73,9 @@ def epoch(timestamp_string: str):
     return parser.parse(timestamp_string).timestamp()
 
 def log(s = None):
+    if not loggable():
+        return
+
     # want to print empty line for False or empty collection
     if s == None:
         print()
@@ -77,6 +83,9 @@ def log(s = None):
         click.echo(s)
 
 def log2(s = None, nl = True):
+    if not loggable():
+        return
+
     if s:
         click.echo(s, err=True, nl=nl)
     else:
@@ -125,6 +134,19 @@ def deep_merge_dicts(dict1, dict2):
             merged_dict[key] = value
     return merged_dict
 
+def deep_sort_dict(d):
+    """
+    Recursively sorts a dictionary by its keys, and any nested lists by their elements.
+    """
+    if not isinstance(d, (dict, list)):
+        return d
+
+    if isinstance(d, dict):
+        return {k: deep_sort_dict(d[k]) for k in sorted(d)}
+
+    if isinstance(d, list):
+        return sorted([deep_sort_dict(item) for item in d])
+
 def get_deep_keys(d, current_path=""):
     """
     Recursively collects all combined keys (paths) from a deep dictionary.
@@ -228,4 +250,63 @@ def copy_config_file(rel_path: str, module: str, suffix: str = '.yaml', show_out
     return path
 
 def idp_token_from_env():
-    return os.getenv('IDP_TOKEN')
+    return os.getenv('IDP_TOKEN')
+
+def is_lambda(func):
+    return callable(func) and hasattr(func, '__name__') and func.__name__ == '<lambda>'
+
+class Ing:
+    state = threading.local()
+
+    def __init__(self, msg: str, suppress_log=False):
+        self.msg = msg
+        self.suppress_log = suppress_log
+        self.nested = False
+
+    def __enter__(self):
+        if not hasattr(Ing.state, 'ing_cnt'):
+            Ing.state.ing_cnt = 0
+
+        try:
+            if not Ing.state.ing_cnt:
+                if not self.suppress_log and not is_debug_holder[0]():
+                    log2(f'{self.msg}...', nl=False)
+
+            return None
+        finally:
+            Ing.state.ing_cnt += 1
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        Ing.state.ing_cnt -= 1
+        if not Ing.state.ing_cnt:
+            if not self.suppress_log and not is_debug_holder[0]():
+                log2(' OK')
+
+        return False
+
+def ing(msg: str, body: Callable[[], None]=None, suppress_log=False):
+    if not body:
+        return Ing(msg, suppress_log=suppress_log)
+
+    r = None
+
+    if not hasattr(Ing.state, 'ing_cnt'):
+        Ing.state.ing_cnt = 0
+
+    if not Ing.state.ing_cnt:
+        if not suppress_log and not is_debug_holder[0]():
+            log2(f'{msg}...', nl=False)
+
+    Ing.state.ing_cnt += 1
+    try:
+        r = body()
+    finally:
+        Ing.state.ing_cnt -= 1
+        if not Ing.state.ing_cnt:
+            if not suppress_log and not is_debug_holder[0]():
+                log2(' OK')
+
+    return r
+
+def loggable():
+    return is_debug_holder[0]() or not hasattr(Ing.state, 'ing_cnt') or not Ing.state.ing_cnt

adam/utils_athena.py

@@ -0,0 +1,145 @@
+import functools
+import time
+import boto3
+import botocore
+
+from adam.config import Config
+from adam.utils import lines_to_tabular, log, log2
+
+# 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...')
+
+      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:
+         state, reason, rs = Athena.query(query)
+         if rs:
+            names = []
+            for row in rs[1:]:
+                  row_data = [col.get('VarCharValue') if col else '' for col in row['Data']]
+                  names.append(row_data[0])
+
+            return names
+      except:
+         pass
+
+      return []
+
+   def clear_cache(cache: str = None):
+      if not cache or cache == 'databases':
+         Athena.database_names.cache_clear()
+      if not cache or cache == 'tables':
+         Athena.table_names.cache_clear()
+      if not cache or cache == 'columns':
+         Athena.column_names.cache_clear()
+
+   @functools.lru_cache()
+   def table_names(database: str = 'audit', function: str = 'audit'):
+      table_names = []
+      try:
+         region_name = Config().get(f'{function}.athena.region', 'us-west-2')
+         database_name = Config().get(f'{function}.athena.database', database)
+         catalog_name = Config().get(f'{function}.athena.catalog', 'AwsDataCatalog')
+
+         athena_client = boto3.client('athena', region_name=region_name)
+         paginator = athena_client.get_paginator('list_table_metadata')
+
+         for page in paginator.paginate(CatalogName=catalog_name, DatabaseName=database_name):
+            for table_metadata in page.get('TableMetadataList', []):
+               table_names.append(table_metadata['Name'])
+      except botocore.exceptions.NoCredentialsError as e:
+         # aws credentials not found
+         if function == 'audit':
+            log2(f'Please configure AWS credentials to Audit Log Database.')
+      except:
+         pass
+
+      return table_names
+
+   @functools.lru_cache()
+   def column_names(tables: list[str] = [], database: str = None, function: str = 'audit', partition_cols_only = False):
+      try:
+         if not database:
+            database = Config().get(f'{function}.athena.database', 'audit')
+
+         if not tables:
+            tables = Config().get(f'{function}.athena.tables', 'audit').split(',')
+
+         table_names = "'" + "','".join([table.strip() for table in tables]) + "'"
+
+         query = f"select column_name from information_schema.columns where table_name in ({table_names}) and table_schema = '{database}'"
+         if partition_cols_only:
+            query = f"{query} and extra_info = 'partition key'"
+
+         _, _, 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):
+      state, reason, rs = Athena.query(sql, database)
+
+      if state == 'SUCCEEDED':
+         if rs:
+            column_info = rs[0]['Data']
+            columns = [col.get('VarCharValue') for col in column_info]
+            lines = []
+            for row in rs[1:]:
+                  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'))
+
+            return len(lines)
+      else:
+         log2(f"Query failed or was cancelled. State: {state}")
+         log2(f"Reason: {reason}")
+
+      return 0
+
+   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')
+      athena_client = boto3.client('athena', region_name=region_name)
+
+      if not database:
+         database = Config().get(f'{function}.athena.database', 'audit')
+
+      s3_output_location = Config().get(f'{function}.athena.output', f's3://s3.ops--{function}/ddl/results')
+
+      response = athena_client.start_query_execution(
+         QueryString=sql,
+         QueryExecutionContext={
+               'Database': database
+         },
+         ResultConfiguration={
+               'OutputLocation': s3_output_location
+         }
+      )
+
+      query_execution_id = response['QueryExecutionId']
+
+      while True:
+         query_status = athena_client.get_query_execution(QueryExecutionId=query_execution_id)
+         state = query_status['QueryExecution']['Status']['State']
+         if state in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
+               break
+         time.sleep(1)
+
+      if state == 'SUCCEEDED':
+         results_response = athena_client.get_query_results(QueryExecutionId=query_execution_id)
+         if results_response['ResultSet']['Rows']:
+            return (state, None, results_response['ResultSet']['Rows'])
+
+         return (state, None, [])
+      else:
+         return (state, query_status['QueryExecution']['Status'].get('StateChangeReason'), [])

adam/utils_audits.py

@@ -0,0 +1,102 @@
+from datetime import datetime
+import getpass
+import time
+import requests
+
+from adam.config import Config
+from adam.utils import log2
+from adam.utils_net import get_my_host
+
+class AuditMeta:
+   def __init__(self, partitions_last_checked: float, cluster_last_checked: float):
+      self.partitions_last_checked = partitions_last_checked
+      self.cluster_last_checked = cluster_last_checked
+
+# no state utility class
+class Audits:
+   PARTITIONS_ADDED = 'partitions-added'
+   ADD_CLUSTERS = 'add-clusters'
+
+   def log(cmd: str, cluster = 'NA', drive: str = 'NA', duration: float = 0.0, audit_extra = None):
+      payload = {
+         'cluster': cluster if cluster else 'NA',
+         'ts': time.time(),
+         'host': get_my_host(),
+         'user': getpass.getuser(),
+         'line': cmd.replace('"', '""').replace('\n', ' '),
+         'drive': drive,
+         'duration': duration,
+         'audit_extra': audit_extra if audit_extra else '',
+      }
+      audit_endpoint = Config().get("audit.endpoint", "https://4psvtaxlcb.execute-api.us-west-2.amazonaws.com/prod/")
+      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)
+         else:
+               log2(f"Error: {response.status_code} {response.text}")
+      except requests.exceptions.Timeout as e:
+         log2(f"Timeout occurred: {e}")
+
+   def get_meta() -> AuditMeta:
+      checked_in = 0.0
+      cluster_last_checked = 0.0
+
+      state, _, rs = Audits.audit_query(f'select partitions_last_checked, clusters_last_checked from meta')
+      if state == 'SUCCEEDED':
+         if len(rs) > 1:
+            try:
+               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)
+
+   def put_meta(action: str, meta: AuditMeta, clusters: list[str] = None):
+      payload = {
+         'action': action,
+         'partitions-last-checked': meta.partitions_last_checked,
+         'clusters-last-checked': meta.cluster_last_checked
+      }
+      if clusters:
+         payload['clusters'] = clusters
+
+      audit_endpoint = Config().get("audit.endpoint", "https://4psvtaxlcb.execute-api.us-west-2.amazonaws.com/prod/")
+      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)
+         else:
+               log2(f"Error: {response.status_code} {response.text}")
+      except requests.exceptions.Timeout as e:
+         log2(f"Timeout occurred: {e}")
+
+   def find_new_clusters(cluster_last_checked: float) -> list[str]:
+      dt_object = datetime.fromtimestamp(cluster_last_checked)
+
+      # select distinct c2.name from cluster as c1 right outer join
+      #     (select distinct c as name from audit where y = '1969' and m = '12' and d >= '31' or y = '1969' and m > '12' or y > '1969') as c2
+      #     on c1.name = c2.name where c1.name is null
+      query = '\n    '.join([
+         'select distinct c2.name from cluster as c1 right outer join',
+         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)
+      if state == 'SUCCEEDED':
+         if len(rs) > 1:
+               try:
+                  return [r['Data'][0]['VarCharValue'] for r in rs[1:]]
+               except:
+                  pass
+
+      return []
+
+   def date_from(dt_object: datetime):
+        y = dt_object.strftime("%Y")
+        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}'"

adam/utils_k8s/app_clusters.py

@@ -0,0 +1,33 @@
+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_k8s.pods import Pods
+from .kube_context import KubeContext
+
+T = TypeVar('T')
+
+# utility collection on app clusters; methods are all static
+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)
+
+            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)

adam/utils_k8s/app_pods.py

@@ -0,0 +1,31 @@
+from typing import List
+from kubernetes import client
+
+from adam.config import Config
+from adam.utils_k8s.pods import Pods
+from adam.pod_exec_result import PodExecResult
+from adam.repl_session import ReplSession
+
+# utility collection on app pods; methods are all static
+class AppPods:
+    def pod_names(namespace: str, env: str, app: str):
+        return [pod.metadata.name for pod in AppPods.app_pods(namespace, env, app)]
+
+    def app_pods(namespace: str, env: str, app: str) -> List[client.V1Pod]:
+        v1 = client.CoreV1Api()
+
+        env_key = Config().get('app.env', 'c3__env-0')
+        app_key = Config().get('app.app', 'c3__app-0')
+        label_selector = f'applicationGroup=c3,{env_key}=0{env}0,{app_key}=0{app}0'
+
+        return v1.list_namespaced_pod(namespace, label_selector=label_selector).items
+
+    def exec(pod_name: str, namespace: str, command: str, show_out = True, throw_err = False, shell = '/bin/sh', background = False) -> PodExecResult:
+        container = Config().get('app.container-name', 'c3-server')
+        r = Pods.exec(pod_name, container, namespace, command, show_out = show_out, throw_err = throw_err, shell = shell, background = background)
+
+        if r and Config().get('repl.history.push-cat-remote-log-file', True):
+            if r.log_file and ReplSession().prompt_session:
+                ReplSession().prompt_session.history.append_string(f'@{r.pod} cat {r.log_file}')
+
+        return r

adam/{k8s_utils → utils_k8s}/cassandra_clusters.py

@@ -1,25 +1,24 @@
-from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
 import sys
 from typing import TypeVar
 
-from adam.k8s_utils.cassandra_nodes import CassandraNodes
+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 .pods import Pods
 from .kube_context import KubeContext
 
 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', max_workers=0, show_out=True) -> list[PodExecResult]:
+    def exec(statefulset: 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,)
+                return executor.submit(CassandraNodes.exec, pod, namespace, command, False, False, shell, background, log_file)
 
-            return CassandraNodes.exec(pod, namespace, command, show_out=show_out)
+            return CassandraNodes.exec(pod, namespace, command, show_out=show_out, background=background, log_file=log_file)
 
         def post(result, show_out: bool):
             if KubeContext.show_out(show_out):
@@ -31,18 +30,4 @@ class CassandraClusters:
 
             return result
 
-        return StatefulSets.on_cluster(statefulset, namespace, body, post=post, action=action, max_workers=max_workers, show_out=show_out)
-
-    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) -> 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)
-
-    def pod_names_by_host_id(ss: str, ns: str):
-        pods = StatefulSets.pods(ss, ns)
-
-        return {CassandraNodes.get_host_id(pod.metadata.name, ns): pod.metadata.name for pod in pods}
+        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)

adam/{k8s_utils → utils_k8s}/cassandra_nodes.py

@@ -1,18 +1,25 @@
 from adam.config import Config
-from adam.k8s_utils.pods import Pods
-from adam.k8s_utils.secrets import Secrets
+from adam.utils_k8s.pods import Pods
+from adam.utils_k8s.secrets import Secrets
 from adam.pod_exec_result import PodExecResult
+from adam.repl_session import ReplSession
 
 # utility collection on cassandra nodes; methods are all static
 class CassandraNodes:
-    def exec(pod_name: str, namespace: str, command: str, show_out = True, throw_err = False) -> PodExecResult:
-        return Pods.exec(pod_name, "cassandra", namespace, command, show_out, throw_err)
+    def exec(pod_name: str, namespace: str, command: str, show_out = True, throw_err = False, shell = '/bin/sh', background = False, log_file = None) -> PodExecResult:
+        r = Pods.exec(pod_name, "cassandra", namespace, command, show_out = show_out, throw_err = throw_err, shell = shell, background = background, log_file=log_file)
+
+        if r and Config().get('repl.history.push-cat-remote-log-file', True):
+            if r.log_file and ReplSession().prompt_session:
+                ReplSession().prompt_session.history.append_string(f'@{r.pod} cat {r.log_file}')
+
+        return r
 
     def get_host_id(pod_name: str, ns: str):
         try:
             user, pw = Secrets.get_user_pass(pod_name, ns)
             command = f'echo "SELECT host_id FROM system.local; exit" | cqlsh --no-color -u {user} -p {pw}'
-            result: PodExecResult = CassandraNodes.exec(pod_name, ns, command, show_out=Config().get('debug.trace', False))
+            result: PodExecResult = CassandraNodes.exec(pod_name, ns, command, show_out=Config().is_debug())
             next = False
             for line in result.stdout.splitlines():
                 if next:

adam/{k8s_utils → utils_k8s}/deployment.py

@@ -1,7 +1,7 @@
 from kubernetes import client
 
-from adam.k8s_utils.pods import Pods
-from adam.k8s_utils.volumes import ConfigMapMount
+from adam.utils_k8s.pods import Pods
+from adam.utils_k8s.volumes import ConfigMapMount
 
 # utility collection on deployments; methods are all static
 class Deployments:

adam/{k8s_utils → utils_k8s}/kube_context.py

@@ -105,4 +105,4 @@ class KubeContext:
         return s or Config().is_debug()
 
     def show_parallelism():
-        return Config().get('debug.show-parallelism', False)
+        return Config().get('debugs.show-parallelism', False)