kaqing 2.0.95__py3-none-any.whl → 2.0.115__py3-none-any.whl

adam/utils_athena.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 import functools
 import time
 import boto3
@@ -5,88 +6,121 @@ import boto3
 from adam.config import Config
 from adam.utils import lines_to_tabular, log, log2
 
-@functools.lru_cache()
-def audit_table_names():
-   region_name = Config().get('audit.athena.region', 'us-west-2')
-   database_name = Config().get('audit.athena.database', 'audit')
-   catalog_name = Config().get('audit.athena.catalog', 'AwsDataCatalog')
+# no state utility class
+class Athena:
+   @functools.lru_cache()
+   def database_names(prefix: str = 'export_'):
+      # this function is called only from export currently
+      Config().wait_log(f'Inspecting export database schema...')
 
-   athena_client = boto3.client('athena', region_name=region_name)
-   paginator = athena_client.get_paginator('list_table_metadata')
+      query = f"SELECT schema_name FROM information_schema.schemata WHERE schema_name <> 'information_schema'"
+      if prefix:
+         query = f"{query} AND schema_name like '{prefix}%'"
 
-   table_names = []
-   for page in paginator.paginate(CatalogName=catalog_name, DatabaseName=database_name):
-      for table_metadata in page.get('TableMetadataList', []):
-         table_names.append(table_metadata['Name'])
+      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
+
+      return []
 
-   return table_names
+   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 audit_column_names(tables: list[str] = [], database: str = None):
-   if not database:
-      database = Config().get('audit.athena.database', 'audit')
+   @functools.lru_cache()
+   def table_names(database: str = 'audit', function: str = 'audit'):
+      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')
 
-   if not tables:
-      tables = Config().get('audit.athena.tables', 'audit').split(',')
+      athena_client = boto3.client('athena', region_name=region_name)
+      paginator = athena_client.get_paginator('list_table_metadata')
 
-   table_names = "'" + "','".join([table.strip() for table in tables]) + "'"
+      table_names = []
+      for page in paginator.paginate(CatalogName=catalog_name, DatabaseName=database_name):
+         for table_metadata in page.get('TableMetadataList', []):
+            table_names.append(table_metadata['Name'])
 
-   query = f"select column_name from information_schema.columns where table_name in ({table_names}) and table_schema = '{database}'"
-   _, _, rs = audit_query(query)
-   if rs:
-      return [row['Data'][0].get('VarCharValue') for row in rs[1:]]
+      return table_names
 
-   return []
+   @functools.lru_cache()
+   def column_names(tables: list[str] = [], database: str = None, function: str = 'audit', partition_cols_only = False):
+      if not database:
+         database = Config().get(f'{function}.athena.database', 'audit')
 
-def run_audit_query(sql: str, database: str = None):
-   state, reason, rs = audit_query(sql, database)
+      if not tables:
+         tables = Config().get(f'{function}.athena.tables', 'audit').split(',')
 
-   if state == 'SUCCEEDED':
+      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:
-         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') for col in row['Data']]
-               lines.append('\t'.join(row_data))
-
-         log(lines_to_tabular(lines, header='\t'.join(columns), separator='\t'))
-   else:
-      log2(f"Query failed or was cancelled. State: {state}")
-      log2(f"Reason: {reason}")
-
-def audit_query(sql: str, database: str = None) -> tuple[str, str, list]:
-   athena_client = boto3.client('athena')
-
-   if not database:
-      database = Config().get('audit.athena.database', 'audit')
-
-   s3_output_location = Config().get('audit.athena.output', 's3://s3.ops--audit/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'), [])
+         return [row['Data'][0].get('VarCharValue') for row in rs[1:]]
+
+      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'))
+      else:
+         log2(f"Query failed or was cancelled. State: {state}")
+         log2(f"Reason: {reason}")
+
+   def query(sql: str, database: str = None, function: str = 'audit') -> tuple[str, str, list]:
+      athena_client = boto3.client('athena')
+
+      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,104 @@
+from datetime import datetime
+import functools
+import getpass
+import time
+import boto3
+import requests
+
+from adam.config import Config
+from adam.utils import lines_to_tabular, log, 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_export.py

@@ -0,0 +1,42 @@
+import functools
+import boto3
+
+from adam.config import Config
+
+# no state utility class
+class Exports:
+   @functools.lru_cache()
+   def table_names(db: str):
+      region_name = Config().get('audit.athena.region', 'us-west-2')
+      database_name = Config().get('audit.athena.database', db)
+      catalog_name = Config().get('audit.athena.catalog', 'AwsDataCatalog')
+
+      athena_client = boto3.client('athena', region_name=region_name)
+      paginator = athena_client.get_paginator('list_table_metadata')
+
+      table_names = []
+      for page in paginator.paginate(CatalogName=catalog_name, DatabaseName=database_name):
+         for table_metadata in page.get('TableMetadataList', []):
+            table_names.append(table_metadata['Name'])
+
+      return table_names
+
+   @functools.lru_cache()
+   def column_names(tables: list[str] = [], database: str = None, partition_cols_only = False):
+      if not database:
+         database = Config().get('audit.athena.database', 'audit')
+
+      if not tables:
+         tables = Config().get('audit.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 = Exports.audit_query(query)
+      if rs:
+         return [row['Data'][0].get('VarCharValue') for row in rs[1:]]
+
+      return []

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

@@ -6,7 +6,6 @@ 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')
@@ -14,12 +13,12 @@ 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, on_any = False, shell = '/bin/sh') -> list[PodExecResult]:
+             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(CassandraNodes.exec, pod, namespace, command, False, False, shell)
+                return executor.submit(CassandraNodes.exec, pod, namespace, command, False, False, shell, background)
 
-            return CassandraNodes.exec(pod, namespace, command, show_out=show_out)
+            return CassandraNodes.exec(pod, namespace, command, show_out=show_out, background=background)
 
         def post(result, show_out: bool):
             if KubeContext.show_out(show_out):
@@ -31,4 +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, on_any=on_any)
+        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/utils_k8s/cassandra_nodes.py

@@ -6,12 +6,12 @@ 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, shell = '/bin/sh') -> PodExecResult:
-        r = Pods.exec(pod_name, "cassandra", namespace, command, show_out = show_out, throw_err = throw_err, shell = shell)
+    def exec(pod_name: str, namespace: str, command: str, show_out = True, throw_err = False, shell = '/bin/sh', background = False) -> PodExecResult:
+        r = Pods.exec(pod_name, "cassandra", 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:
-                ReplSession().prompt_session.history.append_string(f'bash {r.pod} cat {r.log_file}')
+            if r.log_file and ReplSession().prompt_session:
+                ReplSession().prompt_session.history.append_string(f'@{r.pod} cat {r.log_file}')
 
         return r
 

adam/utils_k8s/pods.py

@@ -11,7 +11,7 @@ from kubernetes.stream.ws_client import ERROR_CHANNEL
 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, random_alphanumeric
+from adam.utils import elapsed_time, log2
 from .kube_context import KubeContext
 
 T = TypeVar('T')
@@ -28,7 +28,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)
 
@@ -43,7 +43,11 @@ 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, on_any = False) -> 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:
@@ -94,7 +98,9 @@ class Pods:
 
             return results
 
-    def exec(pod_name: str, container: str, namespace: str, command: str, show_out = True, throw_err = False, shell = '/bin/sh',
+    def exec(pod_name: str, container: str, namespace: str, command: str,
+             show_out = True, throw_err = False, shell = '/bin/sh',
+             background = False,
              interaction: Callable[[any, list[str]], any] = None):
         if _TEST_POD_EXEC_OUTS:
             return _TEST_POD_EXEC_OUTS
@@ -106,8 +112,8 @@ class Pods:
         log_file = None
         tty = True
         exec_command = [shell, '-c', command]
-        if command.endswith(' &'):
-            # should be false for starting a backgroud process
+        if background or command.endswith(' &'):
+            # should be false for starting a background process
             tty = False
 
             if Config().get('repl.background-process.auto-nohup', True):
@@ -170,6 +176,36 @@ class Pods:
 
         return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output, pod=pod_name, log_file=log_file)
 
+    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
+        )
+
+        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()
+
     def get_container(namespace: str, pod_name: str, container_name: str):
         pod = Pods.get(namespace, pod_name)
         if not pod:

adam/utils_k8s/statefulsets.py

@@ -62,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, on_any = False) -> 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, on_any=on_any)
+        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/version.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-__version__ = "2.0.95"  #: the working version
+__version__ = "2.0.115"  #: the working version
 __release__ = "1.0.0"  #: the release version

{kaqing-2.0.95.dist-info → kaqing-2.0.115.dist-info}/METADATA

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