kaqing 2.0.85__py3-none-any.whl → 2.0.87__py3-none-any.whl

adam/commands/audit/audit.py

@@ -2,10 +2,11 @@ import click
 
 from adam.commands.audit.audit_repair_tables import AuditRepairTables
 from adam.commands.command import Command
+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_table_names, run_audit_query
+from adam.utils_athena import audit_column_names, audit_table_names, run_audit_query
 
 class Audit(Command):
     COMMAND = 'audit'
@@ -18,6 +19,7 @@ class Audit(Command):
 
     def __init__(self, successor: Command=None):
         super().__init__(successor)
+        self.schema_read = False
 
     def command(self):
         return Audit.COMMAND
@@ -45,7 +47,16 @@ class Audit(Command):
 
     def completion(self, state: ReplState):
         if state.device == ReplState.L:
-            return super().completion(state) | SqlCompleter.completions(lambda: audit_table_names())
+            if not self.schema_read:
+                Config().wait_log(f'Inspecting audit database schema...')
+                self.schema_read = True
+                audit_column_names()
+
+            def columns(_):
+                return audit_column_names()
+
+            return super().completion(state) | SqlCompleter.completions(lambda: audit_table_names(), columns=columns) | {
+                'desc': {table: None for table in audit_table_names()}}
 
         return {}
 
@@ -53,7 +64,7 @@ class Audit(Command):
         return [AuditRepairTables()]
 
     def help(self, _: ReplState):
-        return f"{Audit.COMMAND} \t run AWS Athena query on 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/help.py

@@ -23,11 +23,12 @@ class Help(Command):
             return super().run(cmd, state)
 
         def section(cmds : list[ReplCommands]):
-            return [f'  {c.help(state)}' for c in cmds if c.help(state)]
+            sorted_cmds = sorted(cmds, key=lambda cmd: cmd.command())
+            return [f'  {c.help(state)}' for c in sorted_cmds if c.help(state)]
 
         lines = []
         lines.append('NAVIGATION')
-        lines.append('  a: | c: | p:\t switch to another operational device: App, Cassandra or Postgres')
+        lines.append('  a: | c: | l: | p:\t switch to another operational device: App, Cassandra, Audit or Postgres')
         lines.extend(section(ReplCommands.navigation()))
         lines.append('CHECK CASSANDRA')
         lines.extend(section(ReplCommands.cassandra_check()))

adam/embedded_params.py

@@ -1,2 +1,2 @@
 def config():
-    return {'app': {'console-endpoint': 'https://{host}/{env}/{app}/static/console/index.html', 'cr': {'cluster-regex': '(.*?-.*?)-.*', 'group': 'ops.c3.ai', 'v': 'v2', 'plural': 'c3cassandras'}, 'label': 'c3__app_id-0', 'login': {'admin-group': '{host}/C3.ClusterAdmin', 'ingress': '{app_id}-k8singr-appleader-001', 'timeout': 5, 'session-check-url': 'https://{host}/{env}/{app}/api/8/C3/userSessionToken', 'cache-creds': True, 'cache-username': True, 'url': 'https://{host}/{env}/{app}', 'another': "You're logged in to {has}. However, for this app, you need to log in to {need}.", 'token-server-url': 'http://localhost:{port}', 'password-max-length': 128}, 'strip': '0'}, 'audit': {'endpoint': 'https://4psvtaxlcb.execute-api.us-west-2.amazonaws.com/prod/', 'workers': 3, 'timeout': 10, 'athena': {'region': 'us-west-2', 'catalog': 'AwsDataCatalog', 'database': 'audit', 'output': 's3://s3.ops--audit/ddl/results'}}, 'bash': {'workers': 32}, 'cassandra': {'service-name': 'all-pods-service'}, 'cql': {'workers': 32, 'samples': 3, 'secret': {'cluster-regex': '(.*?-.*?)-.*', 'name': '{cluster}-superuser', 'password-item': 'password'}, 'alter-tables': {'excludes': 'system_auth,system_traces,reaper_db,system_distributed,system_views,system,system_schema,system_virtual_schema', 'gc-grace-periods': '3600,86400,864000,7776000', 'batching': True}}, 'checks': {'compactions-threshold': 250, 'cpu-busy-threshold': 98.0, 'cpu-threshold': 0.0, 'cassandra-data-path': '/c3/cassandra', 'root-disk-threshold': 50, 'cassandra-disk-threshold': 50, 'snapshot-size-cmd': "ls /c3/cassandra/data/data/*/*/snapshots | grep snapshots | sed 's/:$//g' | xargs -I {} du -sk {} | awk '{print $1}' | awk '{s+=$1} END {print s}'", 'snapshot-size-threshold': '40G', 'table-sizes-cmd': "ls -Al /c3/cassandra/data/data/ | awk '{print $9}' | sed 's/\\^r//g' | xargs -I {} du -sk /c3/cassandra/data/data/{}"}, 'get-host-id': {'workers': 32}, 'idps': {'ad': {'email-pattern': '.*@c3.ai', 'uri': 'https://login.microsoftonline.com/53ad779a-93e7-485c-ba20-ac8290d7252b/oauth2/v2.0/authorize?response_type=id_token&response_mode=form_post&client_id=00ff94a8-6b0a-4715-98e0-95490012d818&scope=openid+email+profile&redirect_uri=https%3A%2F%2Fplat.c3ci.cloud%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY', 'jwks-uri': 'https://login.microsoftonline.com/common/discovery/keys', 'contact': 'Please contact ted.tran@c3.ai.', 'whitelist-file': '/kaqing/members'}, 'okta': {'default': True, 'email-pattern': '.*@c3iot.com', 'uri': 'https://c3energy.okta.com/oauth2/v1/authorize?response_type=id_token&response_mode=form_post&client_id={client_id}&scope=openid+email+profile+groups&redirect_uri=https%3A%2F%2F{host}%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY', 'jwks-uri': 'https://c3energy.okta.com/oauth2/v1/keys'}}, 'issues': {'workers': 32}, 'logs': {'path': '/c3/cassandra/logs/system.log'}, 'medusa': {'restore-auto-complete': False}, 'nodetool': {'workers': 32, 'samples': 3, 'commands_in_line': 40}, 'pg': {'name-pattern': '^{namespace}.*-k8spg-.*', 'excludes': '.helm., -admin-secret', 'agent': {'name': 'ops-pg-agent', 'just-in-time': False, 'timeout': 86400, 'image': 'seanahnsf/kaqing'}, 'default-db': 'postgres', 'default-schema': 'postgres', 'secret': {'endpoint-key': 'postgres-db-endpoint', 'port-key': 'postgres-db-port', 'username-key': 'postgres-admin-username', 'password-key': 'postgres-admin-password'}}, 'pod': {'name': 'ops', 'image': 'seanahnsf/kaqing-cloud', 'sa': {'name': 'ops', 'proto': 'c3', 'additional-cluster-roles': 'c3aiops-k8ssandra-operator'}, 'label-selector': 'run=ops'}, 'preview': {'rows': 10}, 'processes': {'columns': 'pod,cpu,mem', 'header': 'POD_NAME,CPU,MEM/LIMIT'}, 'reaper': {'service-name': 'reaper-service', 'port-forward': {'timeout': 86400, 'local-port': 9001}, 'abort-runs-batch': 10, 'show-runs-batch': 100, 'pod': {'cluster-regex': '(.*?-.*?-.*?-.*?)-.*', 'label-selector': 'k8ssandra.io/reaper={cluster}-reaper'}, 'secret': {'cluster-regex': '(.*?-.*?)-.*', 'name': '{cluster}-reaper-ui', 'password-item': 'password'}}, 'repair': {'log-path': '/home/cassrepair/logs/', 'image': 'ci-registry.c3iot.io/cloudops/cassrepair:2.0.14', 'secret': 'ciregistryc3iotio', 'env': {'interval': 24, 'timeout': 60, 'pr': False, 'runs': 1}}, 'repl': {'start-drive': 'a', 'auto-enter-app': 'c3/c3', 'auto-enter-only-cluster': True}, 'status': {'columns': 'status,address,load,tokens,owns,host_id,gossip,compactions', 'header': '--,Address,Load,Tokens,Owns,Host ID,GOSSIP,COMPACTIONS'}, 'storage': {'columns': 'pod,volume_root,volume_cassandra,snapshots,data,compactions', 'header': 'POD_NAME,VOLUME /,VOLUME CASS,SNAPSHOTS,DATA,COMPACTIONS'}, 'watch': {'auto': 'rollout', 'timeout': 3600, 'interval': 10}, 'debug': False, 'debugs': {'timings': False, 'exit-on-error': False, 'show-parallelism': False}}
+    return {'app': {'console-endpoint': 'https://{host}/{env}/{app}/static/console/index.html', 'cr': {'cluster-regex': '(.*?-.*?)-.*', 'group': 'ops.c3.ai', 'v': 'v2', 'plural': 'c3cassandras'}, 'label': 'c3__app_id-0', 'login': {'admin-group': '{host}/C3.ClusterAdmin', 'ingress': '{app_id}-k8singr-appleader-001', 'timeout': 5, 'session-check-url': 'https://{host}/{env}/{app}/api/8/C3/userSessionToken', 'cache-creds': True, 'cache-username': True, 'url': 'https://{host}/{env}/{app}', 'another': "You're logged in to {has}. However, for this app, you need to log in to {need}.", 'token-server-url': 'http://localhost:{port}', 'password-max-length': 128}, 'strip': '0'}, 'audit': {'endpoint': 'https://4psvtaxlcb.execute-api.us-west-2.amazonaws.com/prod/', 'workers': 3, 'timeout': 10, 'log-audit-queries': False, 'athena': {'region': 'us-west-2', 'catalog': 'AwsDataCatalog', 'database': 'audit', 'tables': 'audit', 'output': 's3://s3.ops--audit/ddl/results'}}, 'bash': {'workers': 32}, 'cassandra': {'service-name': 'all-pods-service'}, 'cql': {'workers': 32, 'samples': 3, 'secret': {'cluster-regex': '(.*?-.*?)-.*', 'name': '{cluster}-superuser', 'password-item': 'password'}, 'alter-tables': {'excludes': 'system_auth,system_traces,reaper_db,system_distributed,system_views,system,system_schema,system_virtual_schema', 'gc-grace-periods': '3600,86400,864000,7776000', 'batching': True}}, 'checks': {'compactions-threshold': 250, 'cpu-busy-threshold': 98.0, 'cpu-threshold': 0.0, 'cassandra-data-path': '/c3/cassandra', 'root-disk-threshold': 50, 'cassandra-disk-threshold': 50, 'snapshot-size-cmd': "ls /c3/cassandra/data/data/*/*/snapshots | grep snapshots | sed 's/:$//g' | xargs -I {} du -sk {} | awk '{print $1}' | awk '{s+=$1} END {print s}'", 'snapshot-size-threshold': '40G', 'table-sizes-cmd': "ls -Al /c3/cassandra/data/data/ | awk '{print $9}' | sed 's/\\^r//g' | xargs -I {} du -sk /c3/cassandra/data/data/{}"}, 'get-host-id': {'workers': 32}, 'idps': {'ad': {'email-pattern': '.*@c3.ai', 'uri': 'https://login.microsoftonline.com/53ad779a-93e7-485c-ba20-ac8290d7252b/oauth2/v2.0/authorize?response_type=id_token&response_mode=form_post&client_id=00ff94a8-6b0a-4715-98e0-95490012d818&scope=openid+email+profile&redirect_uri=https%3A%2F%2Fplat.c3ci.cloud%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY', 'jwks-uri': 'https://login.microsoftonline.com/common/discovery/keys', 'contact': 'Please contact ted.tran@c3.ai.', 'whitelist-file': '/kaqing/members'}, 'okta': {'default': True, 'email-pattern': '.*@c3iot.com', 'uri': 'https://c3energy.okta.com/oauth2/v1/authorize?response_type=id_token&response_mode=form_post&client_id={client_id}&scope=openid+email+profile+groups&redirect_uri=https%3A%2F%2F{host}%2Fc3%2Fc3%2Foidc%2Flogin&nonce={nonce}&state=EMPTY', 'jwks-uri': 'https://c3energy.okta.com/oauth2/v1/keys'}}, 'issues': {'workers': 32}, 'logs': {'path': '/c3/cassandra/logs/system.log'}, 'medusa': {'restore-auto-complete': False}, 'nodetool': {'workers': 32, 'samples': 3, 'commands_in_line': 40}, 'pg': {'name-pattern': '^{namespace}.*-k8spg-.*', 'excludes': '.helm., -admin-secret', 'agent': {'name': 'ops-pg-agent', 'just-in-time': False, 'timeout': 86400, 'image': 'seanahnsf/kaqing'}, 'default-db': 'postgres', 'default-schema': 'postgres', 'secret': {'endpoint-key': 'postgres-db-endpoint', 'port-key': 'postgres-db-port', 'username-key': 'postgres-admin-username', 'password-key': 'postgres-admin-password'}}, 'pod': {'name': 'ops', 'image': 'seanahnsf/kaqing-cloud', 'sa': {'name': 'ops', 'proto': 'c3', 'additional-cluster-roles': 'c3aiops-k8ssandra-operator'}, 'label-selector': 'run=ops'}, 'preview': {'rows': 10}, 'processes': {'columns': 'pod,cpu,mem', 'header': 'POD_NAME,CPU,MEM/LIMIT'}, 'reaper': {'service-name': 'reaper-service', 'port-forward': {'timeout': 86400, 'local-port': 9001}, 'abort-runs-batch': 10, 'show-runs-batch': 100, 'pod': {'cluster-regex': '(.*?-.*?-.*?-.*?)-.*', 'label-selector': 'k8ssandra.io/reaper={cluster}-reaper'}, 'secret': {'cluster-regex': '(.*?-.*?)-.*', 'name': '{cluster}-reaper-ui', 'password-item': 'password'}}, 'repair': {'log-path': '/home/cassrepair/logs/', 'image': 'ci-registry.c3iot.io/cloudops/cassrepair:2.0.14', 'secret': 'ciregistryc3iotio', 'env': {'interval': 24, 'timeout': 60, 'pr': False, 'runs': 1}}, 'repl': {'start-drive': 'a', 'auto-enter-app': 'c3/c3', 'auto-enter-only-cluster': True}, 'status': {'columns': 'status,address,load,tokens,owns,host_id,gossip,compactions', 'header': '--,Address,Load,Tokens,Owns,Host ID,GOSSIP,COMPACTIONS'}, 'storage': {'columns': 'pod,volume_root,volume_cassandra,snapshots,data,compactions', 'header': 'POD_NAME,VOLUME /,VOLUME CASS,SNAPSHOTS,DATA,COMPACTIONS'}, 'watch': {'auto': 'rollout', 'timeout': 3600, 'interval': 10}, 'debug': False, 'debugs': {'timings': False, 'exit-on-error': False, 'show-parallelism': False}}

adam/repl.py

@@ -175,7 +175,7 @@ def enter_repl(state: ReplState):
                     log2(f'Timing command {cmd}: {time.time() - s0:.2f}')
 
                 # offload audit logging
-                if cmd:
+                if cmd and (state.device != ReplState.L or Config().get('audit.log-audit-queries', False)):
                     executor.submit(audit_log, cmd, state)
 
 def audit_log(cmd: str, state: ReplState):

adam/sql/sql_completer.py

@@ -13,11 +13,15 @@ __all__ = [
 
 DML_COMPLETER = TermCompleter(['select', 'insert', 'delete', 'update'])
 
+def default_columns(tables: list[str]):
+    return 'id,x.,y.,z.'.split(',')
+
 class SqlCompleter(Completer):
-    def __init__(self, tables: Callable[[], list[str]], dml: str = None, debug = False):
+    def __init__(self, tables: Callable[[], list[str]], dml: str = None, columns: Callable[[list[str]], list[str]] = default_columns, debug = False):
         super().__init__()
         self.dml = dml
         self.tables = tables
+        self.columns = columns
         self.debug = debug
         self.machine = StateMachine(debug=self.debug)
 
@@ -49,6 +53,8 @@ class SqlCompleter(Completer):
                 for word in self.machine.suggestions[state.to_s].strip(' ').split(','):
                     if word == 'tables':
                         terms.extend(self.tables())
+                    elif word == 'columns':
+                        terms.extend(self.columns([]))
                     elif word == 'single':
                         terms.append("'")
                     elif word == 'comma':
@@ -63,10 +69,10 @@ class SqlCompleter(Completer):
             for c in completer.get_completions(document, complete_event):
                 yield c
 
-    def completions(table_names: Callable[[], list[str]]):
+    def completions(table_names: Callable[[], list[str]], columns: Callable[[list[str]], list[str]] = default_columns):
         return {
-            'delete': SqlCompleter(table_names, 'delete'),
-            'insert': SqlCompleter(table_names, 'insert'),
-            'select': SqlCompleter(table_names, 'select'),
-            'update': SqlCompleter(table_names, 'update'),
+            'delete': SqlCompleter(table_names, 'delete', columns=columns),
+            'insert': SqlCompleter(table_names, 'insert', columns=columns),
+            'select': SqlCompleter(table_names, 'select', columns=columns),
+            'update': SqlCompleter(table_names, 'update', columns=columns),
         }

adam/sql/state_machine.py

@@ -141,10 +141,12 @@ SPEC = [
     'select_from_x                   > ,                > select_from_x_comma_',
     'select_from_sq_                 > as               > select_from_x_as                        ^ as',
     'select_from_x_comma_            > name             > select_from_x                           ^ tables',
-    'select_from_x_                                                                               ^ as,where,inner join,left outer join,right outer join,full outer join,group by,limit',
-    'select_from_x_as_x_             > ,                > select_from_x_comma_                    ^ where,inner join,left outer join,right outer join,full outer join,group by,limit',
+    'select_from_x_                                                                               ^ as,where,inner join,left outer join,right outer join,full outer join,group by,order by,limit',
+    'select_from_x_as_x_             > ,                > select_from_x_comma_                    ^ where,inner join,left outer join,right outer join,full outer join,group by,order by,limit',
     '-                               > as               > select_from_x_as',
     '-                               > where            > select_where',
+    '-                               > order            > select_order',
+    '-                               > order by         > select_order_by',
     '-                               > limit            > select_where_sc_limit',
     '-                               > group            > select_group',
     '-                               > group by         > select_group_by',
@@ -161,7 +163,7 @@ SPEC = [
     'select_from_x_as_               > name             > select_from_x_as_x                      ^ x,y,z',
     'select_from_x_as_x              > ,                > select_from_x_as_x_comma_',
     'select_from_x_as_x_comma_       > name             > select_from_x                           ^ tables',
-    'select_where_                   > name             > select_where_a                          ^ id,x.,y.,z.',
+    'select_where_                   > name             > select_where_a                          ^ columns',
     'select_where_a                  > comparison       > select_where_a_op                       ^ =,<,<=,>,>=,<>,like,not,in',
     'select_where_a_                 > comparison       > select_where_a_op                       ^ =,<,<=,>,>=,<>,like,not,in',
     '-                               > not              > select_where_a_not',
@@ -177,15 +179,27 @@ SPEC = [
     'select_where_a_in_lp_a_comma_   > name|single|num  > select_where_a_in_lp_a                  ^ single',
     'select_where_a_not_op           > name|single|num  > select_where_sc                         ^ single',
     'select_where_a_op               > name|single|num  > select_where_sc                         ^ single',
-    'select_where_sc_                > and|or           > select_where                            ^ and,or,group by,limit',
+    'select_where_sc_                > and|or           > select_where                            ^ and,or,order by,group by,limit',
     '-                               > group            > select_group',
     '-                               > group by         > select_group_by',
+    '-                               > order            > select_order',
+    '-                               > order by         > select_order_by',
     '-                               > limit            > select_where_sc_limit',
     'select_group_                   > by               > select_group_by                         ^ by',
-    'select_group_by_                > name             > select_group_by_a                       ^ id,x.,y.,z.',
+    'select_group_by_                > name             > select_group_by_a                       ^ columns',
     'select_group_by_a               > ,                > select_group_by_a_comma_',
-    'select_group_by_a_comma_        > name             > select_group_by_a                       ^ id,x.,y.,z.',
-    'select_group_by_a_              > limit            > select_where_sc_limit                   ^ limit',
+    'select_group_by_a_comma_        > name             > select_group_by_a                       ^ columns',
+    'select_group_by_a_              > limit            > select_where_sc_limit                   ^ limit,order by',
+    '-                               > order            > select_order',
+    '-                               > order by         > select_order_by',
+    'select_order_                   > by               > select_order_by                         ^ by',
+    'select_order_by_                > name             > select_order_by_a                       ^ columns',
+    'select_order_by_a               > ,                > select_order_by_a_comma_',
+    'select_order_by_a_comma_        > name             > select_order_by_a                       ^ columns',
+    'select_order_by_a_              > desc|asc         > select_order_by_a_desc                  ^ desc,asc,limit',
+    '-                               > limit            > select_where_sc_limit',
+    'select_order_by_a_desc          > ,                > select_order_by_a_comma_',
+    'select_order_by_a_desc_         > limit            > select_where_sc_limit                   ^ limit',
     'select_where_sc_limit_          > num              > select_where_sc_limit_num               ^ 1',
     'select_where_sc_limit_num_rp__  > as               > select_from_x_as                        ^ as',
     'select_where_x_inner_           > join             > select_join',
@@ -203,9 +217,9 @@ SPEC = [
     '-                               > on               > select_x_join_y_on                      ^ as,on',
     'select_x_join_y_as_             > name             > select_x_join_y_as_y                    ^ x,y,z',
     'select_x_join_y_as_y_           > on               > select_x_join_y_on                      ^ on',
-    'select_x_join_y_on_             > name             > select_x_join_y_on_a                    ^ id,x.,y.,z.',
+    'select_x_join_y_on_             > name             > select_x_join_y_on_a                    ^ columns',
     'select_x_join_y_on_a            > comparison       > select_x_join_y_on_a_op                 ^ =',
-    'select_x_join_y_on_a_op         > name             > select_x_join_y_on_a_op_b               ^ id,x.,y.,z.',
+    'select_x_join_y_on_a_op         > name             > select_x_join_y_on_a_op_b               ^ columns',
     'select_x_join_y_on_a_op_b       > _                > select_from_x_as_x_',
 
 
@@ -264,7 +278,7 @@ SPEC = [
 
 KEYWORDS = [
     'select', 'from', 'as', 'not', 'in', 'where',
-    'and', 'or', 'group', 'by', 'group by', 'limit',
+    'and', 'or', 'group', 'by', 'group by', 'order', 'order by', 'limit', 'asc', 'desc',
     'inner join', 'on', 'left', 'right', 'full', 'outer', 'left outer join',
     'left join', 'right outer join', 'right join', 'full join', 'full outer join',
     'insert', 'into', 'values',

adam/utils_athena.py

@@ -21,16 +21,52 @@ def audit_table_names():
 
    return table_names
 
-def run_audit_query(sql: str):
+@functools.lru_cache()
+def audit_column_names(tables: list[str] = [], database: str = None):
+   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}'"
+   _, _, 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') 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')
 
-   database_name = Config().get('audit.athena.database', 'audit')
+   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_name
+            'Database': database
       },
       ResultConfiguration={
             'OutputLocation': s3_output_location
@@ -49,14 +85,8 @@ def run_audit_query(sql: str):
    if state == 'SUCCEEDED':
       results_response = athena_client.get_query_results(QueryExecutionId=query_execution_id)
       if results_response['ResultSet']['Rows']:
-         column_info = results_response['ResultSet']['Rows'][0]['Data']
-         columns = [col.get('VarCharValue') for col in column_info]
-         lines = []
-         for row in results_response['ResultSet']['Rows'][1:]:
-               row_data = [col.get('VarCharValue') for col in row['Data']]
-               lines.append('\t'.join(row_data))
+         return (state, None, results_response['ResultSet']['Rows'])
 
-         log(lines_to_tabular(lines, header='\t'.join(columns), separator='\t'))
+      return (state, None, [])
    else:
-      log2(f"Query failed or was cancelled. State: {state}")
-      log2(f"Reason: {query_status['QueryExecution']['Status'].get('StateChangeReason')}")
+      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.85"  #: the working version
+__version__ = "2.0.87"  #: the working version
 __release__ = "1.0.0"  #: the release version

{kaqing-2.0.85.dist-info → kaqing-2.0.87.dist-info}/METADATA

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

{kaqing-2.0.85.dist-info → kaqing-2.0.87.dist-info}/RECORD

@@ -6,17 +6,17 @@ adam/cli.py,sha256=03pIZdomAu7IL-GSP6Eun_PKwwISShRAmfx6eVRPGC0,458
 adam/cli_group.py,sha256=W3zy1BghCtVcEXizq8fBH-93ZRVVwgAyGPzy0sHno1Y,593
 adam/config.py,sha256=hDxd86MbKVVKoHhTT077L3YRKYcr0i6xVADLlnoE684,2818
 adam/embedded_apps.py,sha256=lKPx63mKzJbNmwz0rgL4gF76M9fDGxraYTtNAIGnZ_s,419
-adam/embedded_params.py,sha256=TGZ6j51c8jMjzDbzn4LU01fKZD9BfCfubQFfRNsk3ss,4831
+adam/embedded_params.py,sha256=FRG1eCyTHX5ynyt3Q8gc3EF2FeelBhW_mh94SSc6PXE,4878
 adam/log.py,sha256=gg5DK52wLPc9cjykeh0WFHyAk1qI3HEpGaAK8W2dzXY,1146
 adam/pod_exec_result.py,sha256=nq0xnCNOpUGBSijGF0H-YNrwBc9vUQs4DkvLMIFS5LQ,951
-adam/repl.py,sha256=gZlakdkKYreK6BQ2NvmTjvIDGVzTkNrmRbP4ZE8fEEI,9323
+adam/repl.py,sha256=xDG3l76O4LXyWhobmyvOd-Txs9ND-pQrTB0I9uhMU1I,9407
 adam/repl_commands.py,sha256=1AFE8AJoyQdJd0tlcwtMnwMjV5G7jL9MDkxJZxdo150,4963
 adam/repl_session.py,sha256=uIogcvWBh7wd8QQ-p_JgLsyJ8YJgINw5vOd6JIsd7Vo,472
 adam/repl_state.py,sha256=dXyGlWXcSsfCjrYwMhU44PVn_oThSZ8dhJ5HCVE8-qU,8743
 adam/utils.py,sha256=sbsNZP3qGJtb6fXCa4dDXHry5ay9ev583cCZIQzy07s,7382
-adam/utils_athena.py,sha256=_E-Ik-rgpIhDhpwlllH3qlyI3P_8Jcw9b9qSj2Ql2UM,2319
+adam/utils_athena.py,sha256=tU6Arg4g7eKV6ei9SLgakOJJqxKgSCsII-7a68OI7_g,3199
 adam/utils_net.py,sha256=65fhBnWMCkhGtyHqz95qcHaCo35q-WX1RBkkXG8dKpI,416
-adam/version.py,sha256=SfFdSx99xjv6dmyMa6T5bmSm3TNRAiaI4-KWzbffs6Q,139
+adam/version.py,sha256=3Mhc1yrzD17NIXSGACmAcbjGjH1pg8z7tUhhzRQwwgQ,139
 adam/checks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 adam/checks/check.py,sha256=Qopr3huYcMu2bzQgb99dEUYjFzkjKHRI76S6KA9b9Rk,702
 adam/checks/check_context.py,sha256=FEHkQ32jY1EDopQ2uYWqy9v7aEEX1orLpJWhopwAlh4,402
@@ -62,7 +62,7 @@ adam/commands/commands_utils.py,sha256=ShUcxtDSd9B3NM0GDj3NBvKdmjCGY8qXgeUJpzNF6
 adam/commands/cp.py,sha256=dyQViRDPNqsKRkxPb7WyEVIBNw7YB6IfYa2q3VtfzyA,3107
 adam/commands/devices.py,sha256=ygpvjYa_NcezM_hctGOAAuBEHiDJtLtJNlOcLGvFJAc,3184
 adam/commands/exit.py,sha256=5MWUAmzYBlsrp0CoiTDB13SUkX9Ya18UlGeOIPia6TA,798
-adam/commands/help.py,sha256=Ey3R1X8w_CMhdADI0t8dSQ28euhDHheJm7NermiGni4,1645
+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=T-O9DYXmWEm4G1I5SM6MwyeRwq2aT-WMqNW0XA2MWmo,1165
@@ -79,7 +79,7 @@ adam/commands/rollout.py,sha256=52_4ijna3v-8Oug12et43DRHFDNhiN34p6xLTQmhdbQ,2959
 adam/commands/shell.py,sha256=wY_PIx7Lt6vuxhFArlfxdEnBbrouCJ3yNHhFn17DEqw,848
 adam/commands/watch.py,sha256=mmBFpB8T1V7zrNs5b2YNyDDztMym_ILPDdkrbdAXTas,2438
 adam/commands/audit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-adam/commands/audit/audit.py,sha256=xLXwDSYYTuVr_O5gDo87Wa4qSR71zjnj1Q6vN8dF_dE,1864
+adam/commands/audit/audit.py,sha256=uXyzGRRW4mfg4AAjxMe6yrTrC3UNQnZh9_8n7-Y3EVE,2324
 adam/commands/audit/audit_repair_tables.py,sha256=5Glsy29md3KH_4_smuVBeu6c4zygICPK8Jiv-fNRC38,1186
 adam/commands/cql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 adam/commands/cql/cql_completions.py,sha256=dXe51NTWEJis76587IWSn9Av-cjC0J6KMaxlBKfF4wM,411
@@ -170,8 +170,8 @@ adam/k8s_utils/services.py,sha256=EOJJGACVbbRvu5T3rMKqIJqgYic1_MSJ17EA0TJ6UOk,31
 adam/k8s_utils/statefulsets.py,sha256=5g7KxGRHgEewT8rnZneDTaJDylUf-dHH2edWJEoorr8,4667
 adam/k8s_utils/volumes.py,sha256=RIBmlOSWM3V3QVXLCFT0owVOyh4rGG1ETp521a-6ndo,1137
 adam/sql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-adam/sql/sql_completer.py,sha256=R5boj_bdHlzeJkGEy7XQDz_qtE-VF1yjg8pXn5zYGJA,2510
-adam/sql/state_machine.py,sha256=6wOCVE86S6pPJhLvJudbYhBbKDqOmpcY2lAEK75pbm4,23200
+adam/sql/sql_completer.py,sha256=zH2I-WLc7aZmdIg0PlHPVLdxemI8oN8YTCnG7qxcYqA,2907
+adam/sql/state_machine.py,sha256=BLX0myqei4vn-Ny1AS1Rg-su7Gahfp85sHmE8q8_4kg,24501
 adam/sql/term_completer.py,sha256=bNnHAVf9NZl52xS_BQpikbOK39gDBJADnT9gSvG0iqI,2539
 adam/sso/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 adam/sso/authenticator.py,sha256=BCm16L9zf5aLU47-sTCnudn2zLPwd8M2wwRminJfsqw,615
@@ -183,8 +183,8 @@ adam/sso/idp.py,sha256=fvcwUw_URTgsO6ySaqTIw0zQT2qRO1IPSGhf6rPtybo,5804
 adam/sso/idp_login.py,sha256=QAtCUeDTVWliJy40RK_oac8Vgybr13xH8wzeBoxPaa8,1754
 adam/sso/idp_session.py,sha256=9BUHNRf70u4rVKrVY1HKPOEmOviXvkjam8WJxmXSKIM,1735
 adam/sso/sso_config.py,sha256=5N8WZgIJQBtHUy585XLRWKjpU87_v6QluyNK9E27D5s,2459
-kaqing-2.0.85.dist-info/METADATA,sha256=H9R7Qb1ZTmYMDrzCFBJziK45RLM3IDkx82kFaDk1kpE,132
-kaqing-2.0.85.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-kaqing-2.0.85.dist-info/entry_points.txt,sha256=SkzhuQJUWsXOzHeZ5TgQ2c3_g53UGK23zzJU_JTZOZI,39
-kaqing-2.0.85.dist-info/top_level.txt,sha256=8_2PZkwBb-xDcnc8a2rAbQeJhXKXskc7zTP7pSPa1fw,5
-kaqing-2.0.85.dist-info/RECORD,,
+kaqing-2.0.87.dist-info/METADATA,sha256=u4J5o2OrBxTqMu35ZbFvIYMcfQSfpnn4TCtjfUR7M-Y,132
+kaqing-2.0.87.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+kaqing-2.0.87.dist-info/entry_points.txt,sha256=SkzhuQJUWsXOzHeZ5TgQ2c3_g53UGK23zzJU_JTZOZI,39
+kaqing-2.0.87.dist-info/top_level.txt,sha256=8_2PZkwBb-xDcnc8a2rAbQeJhXKXskc7zTP7pSPa1fw,5
+kaqing-2.0.87.dist-info/RECORD,,