kaqing 2.0.104__py3-none-any.whl → 2.0.105__py3-none-any.whl

adam/commands/audit/audit.py

@@ -7,7 +7,7 @@ from adam.config import Config
 from adam.repl_state import ReplState
 from adam.sql.sql_completer import SqlCompleter
 from adam.utils import log2
-from adam.utils_athena import audit_column_names, audit_table_names, run_audit_query
+from adam.utils_audits import Audits
 
 class Audit(Command):
     COMMAND = 'audit'
@@ -42,7 +42,7 @@ class Audit(Command):
             else:
                 log2(sql)
 
-            run_audit_query(sql)
+            Audits.run_audit_query(sql)
 
         return state
 
@@ -52,13 +52,13 @@ class Audit(Command):
                 Config().wait_log(f'Inspecting audit database schema...')
                 self.schema_read = True
                 # warm up the caches first time when l: drive is accessed
-                audit_column_names()
-                audit_column_names(partition_cols_only=True)
+                Audits.audit_column_names()
+                Audits.audit_column_names(partition_cols_only=True)
 
             return super().completion(state) | SqlCompleter(
-                lambda: audit_table_names(),
-                columns=lambda table: audit_column_names(),
-                partition_columns=lambda table: audit_column_names(partition_cols_only=True),
+                lambda: Audits.audit_table_names(),
+                columns=lambda table: Audits.audit_column_names(),
+                partition_columns=lambda table: Audits.audit_column_names(partition_cols_only=True),
                 variant='athena'
             ).completions_for_nesting()
 
@@ -68,7 +68,7 @@ class Audit(Command):
         return [AuditRepairTables(), AuditRun()]
 
     def help(self, _: ReplState):
-        return f'[{Audit.COMMAND}] <sql-statements>\t run SQL queries on Authena audit database'
+        return f'[{Audit.COMMAND}] [<sql-statements>]\t run SQL queries on Authena audit database'
 
 class AuditCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):

adam/commands/audit/audit_repair_tables.py

@@ -1,12 +1,11 @@
 import concurrent
 import time
-import requests
 
 from adam.commands.command import Command
 from adam.config import Config
 from adam.repl_state import ReplState
 from adam.utils import log, log2
-from adam.utils_athena import AuditMeta, audit_query, get_meta, put_meta, run_audit_query
+from adam.utils_audits import AuditMeta, Audits
 
 class AuditRepairTables(Command):
     COMMAND = 'audit repair'
@@ -34,7 +33,7 @@ class AuditRepairTables(Command):
         if args:
             tables = args
 
-        meta = get_meta()
+        meta = Audits.get_meta()
         self.repair(tables, meta)
 
         return state
@@ -53,8 +52,8 @@ class AuditRepairTables(Command):
     def auto_repair(self, hours: int):
         self.auto_repaired = True
 
-        meta = get_meta()
-        if meta.checked_in + hours * 60 * 60 < time.time():
+        meta: AuditMeta = Audits.get_meta()
+        if meta.partitions_last_checked + hours * 60 * 60 < time.time():
             tables = Config().get('audit.athena.repair-partition-tables', 'audit').split(',')
             self.repair(tables, meta, show_sql=True)
             log2(f'Audit tables have been auto-repaired.')
@@ -65,8 +64,8 @@ class AuditRepairTables(Command):
                 if show_sql:
                     log(f'MSCK REPAIR TABLE {table}')
 
-                executor.submit(run_audit_query, f'MSCK REPAIR TABLE {table}', None,)
-            executor.submit(put_meta, 'check-in', meta,)
+                executor.submit(Audits.run_audit_query, f'MSCK REPAIR TABLE {table}', None,)
+            executor.submit(Audits.put_meta, Audits.PARTITIONS_ADDED, meta,)
 
     def help(self, _: ReplState):
         return f"{AuditRepairTables.COMMAND} \t run MSCK REPAIR command for new partition discovery"

adam/commands/audit/audit_run.py

@@ -1,8 +1,9 @@
+import concurrent
 from adam.commands.command import Command
 from adam.config import Config
 from adam.repl_state import ReplState
 from adam.utils import log2
-from adam.utils_athena import AuditMeta, find_new_clusters, get_meta, put_meta, run_audit_query
+from adam.utils_audits import AuditMeta, Audits
 
 class AuditRun(Command):
     COMMAND = 'audit run'
@@ -26,14 +27,15 @@ class AuditRun(Command):
 
         state, args = self.apply_state(args, state)
 
-        meta: AuditMeta = get_meta()
-        clusters = find_new_clusters(meta.cluster_last_checked)
+        meta: AuditMeta = Audits.get_meta()
+        clusters = Audits.find_new_clusters(meta.cluster_last_checked)
+        Audits.put_meta(Audits.ADD_CLUSTERS, meta, clusters=clusters)
         if clusters:
-            put_meta('add-clusters', meta, clusters=clusters)
             log2(f'Added {len(clusters)} new clusters.')
             tables = Config().get('audit.athena.repair-cluster-tables', 'cluster').split(',')
-            for table in tables:
-                run_audit_query(f'MSCK REPAIR TABLE {table}')
+            with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
+                for table in tables:
+                    Audits.run_audit_query(f'MSCK REPAIR TABLE {table}')
         else:
             log2(f'No new clusters were found.')
 

adam/commands/ls.py

@@ -14,7 +14,7 @@ from adam.pod_exec_result import PodExecResult
 from adam.repl_state import ReplState
 from adam.utils import lines_to_tabular, log, log2
 from adam.apps import Apps
-from adam.utils_athena import audit_table_names
+from adam.utils_audits import Audits
 
 class Ls(Command):
     COMMAND = 'ls'
@@ -131,7 +131,7 @@ class Ls(Command):
         log(lines_to_tabular(pg_table_names(pg.namespace, pg.path()), 'NAME', separator=','))
 
     def show_audit_log_tables(self):
-        log(lines_to_tabular(audit_table_names(), 'NAME', separator=','))
+        log(lines_to_tabular(Audits.audit_table_names(), 'NAME', separator=','))
 
     def completion(self, state: ReplState):
         if state.pod:

adam/commands/preview_table.py

@@ -4,7 +4,7 @@ from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
 from adam.utils import lines_to_tabular, log, log2
-from adam.utils_athena import audit_table_names, run_audit_query
+from adam.utils_audits import Audits
 
 class PreviewTable(Command):
     COMMAND = 'preview'
@@ -43,7 +43,7 @@ class PreviewTable(Command):
                     lines = [db["name"] for db in pg.tables() if db["schema"] == PostgresContext.default_schema()]
                     log(lines_to_tabular(lines, separator=','))
                 elif state.device == ReplState.L:
-                    log(lines_to_tabular(audit_table_names(), separator=','))
+                    log(lines_to_tabular(Audits.audit_table_names(), separator=','))
                 else:
                     log(lines_to_tabular(cassandra_table_names(state), separator=','))
 
@@ -66,7 +66,7 @@ class PreviewTable(Command):
         if state.device == ReplState.P:
             PostgresContext.apply(state.namespace, state.pg_path).run_sql(f'select * from {table} limit {rows}')
         elif state.device == ReplState.L:
-            run_audit_query(f'select * from {table} limit {rows}')
+            Audits.run_audit_query(f'select * from {table} limit {rows}')
         else:
             run_cql(state, f'select * from {table} limit {rows}', show_out=True, use_single_quotes=True, on_any=True)
 

adam/repl_commands.py

@@ -58,7 +58,7 @@ class ReplCommands:
         cmds: list[Command] = ReplCommands.navigation() + ReplCommands.cassandra_check() + ReplCommands.cassandra_ops() + \
             ReplCommands.tools() + ReplCommands.app() + ReplCommands.exit()
 
-        intermediate_cmds: list[Command] = [App(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
+        intermediate_cmds: list[Command] = [App(), Audit(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
         ic = [c.command() for c in intermediate_cmds]
         # 1. dedup commands
         deduped = []
@@ -86,7 +86,7 @@ class ReplCommands:
 
     def tools() -> list[Command]:
         return [Cqlsh(), Postgres(), Bash(), Shell(), CodeStart(), CodeStop(), DeployFrontend(), UndeployFrontend(),
-                DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), AuditRepairTables(), AuditRun(), Audit()]
+                DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent(), Audit()] + Audit.cmd_list()
 
     def app() -> list[Command]:
         return [ShowAppActions(), ShowAppId(), ShowAppQueues(), AppPing(), App()]

adam/utils_audits.py

@@ -0,0 +1,167 @@
+from datetime import datetime
+import functools
+import time
+import boto3
+import requests
+
+from adam.config import Config
+from adam.utils import lines_to_tabular, log, log2
+
+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 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 find_new_clusters(cluster_last_checked: float) -> list[str]:
+      dt_object = datetime.fromtimestamp(cluster_last_checked)
+
+      y = dt_object.strftime("%Y")
+      m = dt_object.strftime("%m")
+      d = dt_object.strftime("%d")
+      # 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
+      where = f"y = '{y}' and m = '{m}' and d >= '{d}' or y = '{y}' and m > '{m}' or y > '{y}'"
+      query = '\n    '.join([
+         'select distinct c2.name from cluster as c1 right outer join',
+         f'(select distinct c as name from audit where {where}) 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 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}")
+
+   @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')
+
+      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 audit_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 = Audits.audit_query(query)
+      if rs:
+         return [row['Data'][0].get('VarCharValue') for row in rs[1:]]
+
+      return []
+
+   def run_audit_query(sql: str, database: str = None):
+      state, reason, rs = Audits.audit_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 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'), [])

adam/version.py

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

{kaqing-2.0.104.dist-info → kaqing-2.0.105.dist-info}/METADATA

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

{kaqing-2.0.104.dist-info → kaqing-2.0.105.dist-info}/RECORD

@@ -10,13 +10,13 @@ adam/embedded_params.py,sha256=YeyY3QTLBtjiMsBt16A7z8lvJ_hlVtAPJR6RbjStIwI,5060
 adam/log.py,sha256=gg5DK52wLPc9cjykeh0WFHyAk1qI3HEpGaAK8W2dzXY,1146
 adam/pod_exec_result.py,sha256=WBXJSvxzXp9TfsfXeHtIvgz8GvfMAAcH5M03GISLqzw,1046
 adam/repl.py,sha256=ZfYEZ6s1Hhw-aIFZk358bb7KCAJ-aTDxSPJvfd2ciHQ,10904
-adam/repl_commands.py,sha256=eaxwvI1z2BnuRY4HTNqpzr0L7DhCl9j3FzVvVmi5SuI,4651
+adam/repl_commands.py,sha256=iKcw208Iqs4jxJYX2pCT4woGijpcg0IM58tVlhZQNo8,4646
 adam/repl_session.py,sha256=uIogcvWBh7wd8QQ-p_JgLsyJ8YJgINw5vOd6JIsd7Vo,472
 adam/repl_state.py,sha256=GuDq3C42hMX3pmjLlMbARJ8nGiBw8-gNT-xhkKQu-a4,8797
 adam/utils.py,sha256=sbsNZP3qGJtb6fXCa4dDXHry5ay9ev583cCZIQzy07s,7382
-adam/utils_athena.py,sha256=hNvwd3F1qqs-3rGoK42aq0b5XdrHuRXGQtBRU4SBIqY,5840
+adam/utils_audits.py,sha256=kPw6B0lw2q4PmYdk_k16MX_V-ONnrhJX3T-cMTFYZdM,6376
 adam/utils_net.py,sha256=65fhBnWMCkhGtyHqz95qcHaCo35q-WX1RBkkXG8dKpI,416
-adam/version.py,sha256=LKihFK6mt-4FST7GwFS-_WfQG7UfoEZB6WZ6szHFz5Q,140
+adam/version.py,sha256=W2f85EjavUMXa9WADSC-bmOjXgpifTdBQ8mFdwXyUIA,140
 adam/checks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 adam/checks/check.py,sha256=Qopr3huYcMu2bzQgb99dEUYjFzkjKHRI76S6KA9b9Rk,702
 adam/checks/check_context.py,sha256=FEHkQ32jY1EDopQ2uYWqy9v7aEEX1orLpJWhopwAlh4,402
@@ -66,12 +66,12 @@ adam/commands/help.py,sha256=4IzR4p8UiXr00o1TaymHWm8957EWbWRyuvhrJzZzvc0,1734
 adam/commands/issues.py,sha256=VS-PC7e-2lywsa-lbmoUX8IY77OPGzFudwbw1g8XmQc,2599
 adam/commands/login.py,sha256=bj95WWIF7mJDJhnyS9T8xvaZUGL37dj7GlH8TgmODbk,1877
 adam/commands/logs.py,sha256=GBVztFlCQfd4jfMtqydPjWS9xsB5mV4Aj4ohSQFm6i0,1165
-adam/commands/ls.py,sha256=cNs_pl9UZpG43NrOVc3EBAbzgHye1I0E-vN72tIG6ok,5809
+adam/commands/ls.py,sha256=49bS2V6jL_HAu6Y9kJzf8e3v4krxn_wNXUf47_R3Tkw,5805
 adam/commands/nodetool.py,sha256=k4Gr63slw0BGinbbj3nHue2GHf55cjAjTdt636UKjb8,2363
 adam/commands/nodetool_commands.py,sha256=5IgWC3rmeDD1cgwqQjiiWzi-wJpJ3n_8pAzz_9phXuk,2635
 adam/commands/param_get.py,sha256=kPAAppK2T0tEFRnSIVFLDPIIGHhgLA7drJhn8TRyvvE,1305
 adam/commands/param_set.py,sha256=QDIuqfU80aWCB16OK49yf7XRaRTWwiLkwMsJuVikq9I,1271
-adam/commands/preview_table.py,sha256=scGOpKkSc95IHGWfQOrM0cctSgeDmjr2E6BoPp5xk1U,2857
+adam/commands/preview_table.py,sha256=mz9Cstk_UfNNi3rBye4kZ4M0mtjeyfRCUtqP5zifvDc,2843
 adam/commands/pwd.py,sha256=AvM1gMgxBfIupevtYkGmNWGyEcsIx7-S8L9_v_DPEgk,2469
 adam/commands/report.py,sha256=Ky45LIzSlB_X4V12JZWjU3SA2u4_FKRencRTq7psOWU,1944
 adam/commands/restart.py,sha256=SAxWHvglTckQJ0tJe5t-HWsVerbreNMM-7Nb9PAqno4,2044
@@ -79,9 +79,9 @@ adam/commands/rollout.py,sha256=Db9P4Owd3aPcRLIGhwyEElBNm_2Ke54KbiXyVKmztcE,2959
 adam/commands/shell.py,sha256=wY_PIx7Lt6vuxhFArlfxdEnBbrouCJ3yNHhFn17DEqw,848
 adam/commands/watch.py,sha256=fU2LGll-Igl08HpUQALOnh8l3s3AMGFX26NCLhqbfcw,2438
 adam/commands/audit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-adam/commands/audit/audit.py,sha256=MBCvQn6ZwytZPC47W2aV6iVUw1a9_JHMaogsejf5mgw,2588
-adam/commands/audit/audit_repair_tables.py,sha256=2lyAbnni1EY01k9fGNckBvhALFLLNndnsRS_wvIt2wQ,2542
-adam/commands/audit/audit_run.py,sha256=2Z9VqsAejToWlSYd0JEy1FuT7_ywqTj2D7gZgBkUCKY,1578
+adam/commands/audit/audit.py,sha256=DcyolsPy9Qzan5x3lwomkmkwpAeAa9PHj8Y1Q2bwI70,2584
+adam/commands/audit/audit_repair_tables.py,sha256=zNyhOK_-dbeK-7rH-b_HNeH5Qwrk_ccqakWKujVP7Lk,2549
+adam/commands/audit/audit_run.py,sha256=Bfzv73aaYvHtRhr_mtzeH1muJtexhPfoUPENAiGKCgI,1699
 adam/commands/cql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 adam/commands/cql/cql_completions.py,sha256=29TJmNHg2_EwGyBrgh04dtZl0BDxGEnLXf_hbPQF9no,660
 adam/commands/cql/cql_utils.py,sha256=EWswF4JEj4EI_lIa1tqhNV-0sasf81SK64aNf8JPqZo,4103
@@ -176,8 +176,8 @@ adam/utils_k8s/service_accounts.py,sha256=v2oQSqCrNvt2uRnKlNwR3fjtpUG7oF5nqgzEB7
 adam/utils_k8s/services.py,sha256=EOJJGACVbbRvu5T3rMKqIJqgYic1_MSJ17EA0TJ6UOk,3156
 adam/utils_k8s/statefulsets.py,sha256=0J_cYRqH96PCcq3tdsRrs4Q4ewv5dT_FMBR0HGAJ3d8,4710
 adam/utils_k8s/volumes.py,sha256=RIBmlOSWM3V3QVXLCFT0owVOyh4rGG1ETp521a-6ndo,1137
-kaqing-2.0.104.dist-info/METADATA,sha256=8r0K3oowsddpx2JIoV0koNZpv9hiASyAa3XvYNFvseE,133
-kaqing-2.0.104.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-kaqing-2.0.104.dist-info/entry_points.txt,sha256=SkzhuQJUWsXOzHeZ5TgQ2c3_g53UGK23zzJU_JTZOZI,39
-kaqing-2.0.104.dist-info/top_level.txt,sha256=8_2PZkwBb-xDcnc8a2rAbQeJhXKXskc7zTP7pSPa1fw,5
-kaqing-2.0.104.dist-info/RECORD,,
+kaqing-2.0.105.dist-info/METADATA,sha256=-XxUVOCHMejB3MGXatLONCSDP3GKrOQE2MBoEXNp_U4,133
+kaqing-2.0.105.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+kaqing-2.0.105.dist-info/entry_points.txt,sha256=SkzhuQJUWsXOzHeZ5TgQ2c3_g53UGK23zzJU_JTZOZI,39
+kaqing-2.0.105.dist-info/top_level.txt,sha256=8_2PZkwBb-xDcnc8a2rAbQeJhXKXskc7zTP7pSPa1fw,5
+kaqing-2.0.105.dist-info/RECORD,,

adam/utils_athena.py

@@ -1,162 +0,0 @@
-from datetime import datetime
-import functools
-import time
-import boto3
-import requests
-
-from adam.config import Config
-from adam.utils import lines_to_tabular, log, log2
-
-class AuditMeta:
-   def __init__(self, checked_in: float, cluster_last_checked: float):
-      self.checked_in = checked_in
-      self.cluster_last_checked = cluster_last_checked
-
-def get_meta() -> AuditMeta:
-   checked_in = 0.0
-   cluster_last_checked = 0.0
-
-   state, _, rs = 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 find_new_clusters(cluster_last_checked: float) -> list[str]:
-   dt_object = datetime.fromtimestamp(cluster_last_checked)
-
-   y = dt_object.strftime("%Y")
-   m = dt_object.strftime("%m")
-   d = dt_object.strftime("%d")
-   # 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
-   where = f"y = '{y}' and m = '{m}' and d >= '{d}' or y = '{y}' and m > '{m}' or y > '{y}'"
-   query = '\n    '.join([
-      'select distinct c2.name from cluster as c1 right outer join',
-      f'(select distinct c as name from audit where {where}) as c2',
-      'on c1.name = c2.name where c1.name is null'])
-   log2(query)
-   state, _, rs = 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 put_meta(action: str, meta: AuditMeta, clusters: list[str] = None):
-   payload = {
-      'action': action,
-      'partitions-last-checked': meta.checked_in,
-      '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}")
-
-@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')
-
-   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 audit_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 = audit_query(query)
-   if rs:
-      return [row['Data'][0].get('VarCharValue') for row in rs[1:]]
-
-   return []
-
-def run_audit_query(sql: str, database: str = None):
-   state, reason, rs = audit_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 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'), [])