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

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/utils_k8s/config_maps.py

@@ -0,0 +1,34 @@
+from kubernetes import client
+
+# utility collection on config maps; methods are all static
+class ConfigMaps:
+    def create(name: str, namespace: str, data: dict[str, str], labels: dict[str, str] = {}):
+        v1 = client.CoreV1Api()
+
+        metadata = client.V1ObjectMeta(
+            name=name,
+            namespace=namespace,
+            labels=labels
+        )
+
+        configmap = client.V1ConfigMap(
+            api_version="v1",
+            kind="ConfigMap",
+            metadata=metadata,
+            data=data
+        )
+
+        try:
+            api_response = v1.create_namespaced_config_map(body=configmap, namespace=namespace)
+            # print(f"ConfigMap '{name}' created successfully in namespace '{namespace}'.")
+            # print(api_response)
+        except client.ApiException as e:
+            # print(f"Error creating ConfigMap: {e}")
+            raise e
+
+    def delete_with_selector(namespace: str, label_selector: str):
+        v1 = client.CoreV1Api()
+
+        ret = v1.list_namespaced_config_map(namespace=namespace, label_selector=label_selector)
+        for i in ret.items:
+            v1.delete_namespaced_config_map(name=i.metadata.name, namespace=namespace)

adam/utils_k8s/deployment.py

@@ -0,0 +1,56 @@
+from kubernetes import client
+
+from adam.utils_k8s.pods import Pods
+from adam.utils_k8s.volumes import ConfigMapMount
+
+# utility collection on deployments; methods are all static
+class Deployments:
+    def delete_with_selector(namespace: str, label_selector: str, grace_period_seconds: int = None):
+        v1 = client.AppsV1Api()
+
+        ret = v1.list_namespaced_deployment(namespace=namespace, label_selector=label_selector)
+        for i in ret.items:
+            v1.delete_namespaced_deployment(name=i.metadata.name, namespace=namespace, grace_period_seconds=grace_period_seconds)
+
+    def create_deployment_spec(name: str, image: str, image_pull_secret: str,
+                               envs: list, container_security_context: client.V1SecurityContext,
+                               volume_name: str, pvc_name:str, mount_path:str,
+                               command: list[str]=None, sa_name=None, labels: dict[str, str] = {},
+                               config_map_mount: ConfigMapMount = None):
+        return client.V1DeploymentSpec(
+            replicas=1,
+            selector=client.V1LabelSelector(match_labels=labels),
+            template=client.V1PodTemplateSpec(
+                metadata=client.V1ObjectMeta(labels=labels),
+                spec=Pods.create_pod_spec(name, image, image_pull_secret, envs, container_security_context,
+                                          volume_name, pvc_name, mount_path, command=command, sa_name=sa_name,
+                                          restart_policy="Always", config_map_mount=config_map_mount),
+            ),
+        )
+
+    def create(namespace: str, deployment_name: str, image: str,
+               command: list[str] = None,
+               secret: str = None,
+               env: dict[str, any] = {},
+               container_security_context: client.V1SecurityContext = None,
+               labels: dict[str, str] = {},
+               volume_name: str = None,
+               pvc_name: str = None,
+               mount_path: str = None,
+               sa_name=None,
+               config_map_mount: ConfigMapMount = None):
+        v1 = client.AppsV1Api()
+        envs = []
+        for k, v in env.items():
+            envs.append(client.V1EnvVar(name=str(k), value=str(v)))
+        deployment = Deployments.create_deployment_spec(deployment_name, image, secret, envs,
+                                                        container_security_context, volume_name, pvc_name,
+                                                        mount_path, command=command, sa_name=sa_name, labels=labels,
+                                                        config_map_mount=config_map_mount)
+        return v1.create_namespaced_deployment(
+            namespace=namespace,
+            body=client.V1Deployment(spec=deployment, metadata=client.V1ObjectMeta(
+                name=deployment_name,
+                labels=labels
+            ))
+        )

adam/{k8s_utils → utils_k8s}/jobs.py

@@ -12,7 +12,7 @@ class Jobs:
             envs.append(client.V1EnvVar(name=k.upper(), value=str(v)))
         for k, v in env_from.items():
             envs.append(client.V1EnvVar(name=k.upper(), value_from=client.V1EnvVarSource(secret_key_ref=client.V1SecretKeySelector(key=k, name=v))))
-        template = Pods.create_pod_spec(job_name, image, image_pull_secret, envs, volume_name, pvc_name, mount_path, command)
+        template = Pods.create_pod_spec(job_name, image, image_pull_secret, envs, None, volume_name, pvc_name, mount_path, command)
         spec = client.V1JobSpec(template=client.V1PodTemplateSpec(spec=template), backoff_limit=1, ttl_seconds_after_finished=300)
         job = client.V1Job(
             api_version="batch/v1",

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)