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

adam/__init__.py

@@ -1,3 +1 @@
-import click
-
 from .version import __version__, __release__

adam/app_session.py

@@ -1,7 +1,6 @@
 import json
 import threading
 import time
-import traceback
 import requests
 from urllib.parse import urlparse
 
@@ -9,7 +8,7 @@ from adam.log import Log
 from adam.sso.idp import Idp
 from adam.sso.idp_login import IdpLogin
 from adam.config import Config
-from adam.utils import json_to_csv, lines_to_tabular, log, log2
+from adam.utils import debug, debug_trace, json_to_csv, lines_to_tabular, log, log2, log_exc
 from adam.apps import Apps
 
 class AppLogin:
@@ -62,11 +61,9 @@ class AppSession:
                     if r.status_code >= 200 and r.status_code < 300 or r.status_code == 400:
                         try:
                             js = r.json()
-                            try:
+                            with log_exc(js):
                                 header, lines = json_to_csv(js, delimiter='\t')
                                 log(lines_to_tabular(lines, header=header, separator='\t'))
-                            except:
-                                log(js)
                         except:
                             if urlparse(r.url).hostname != urlparse(uri).hostname and not retried:
                                 app_login = app_session.login(idp_uri=app_login.idp_uri, forced=forced, use_token_from_env=False, use_cached_creds=False)
@@ -76,7 +73,7 @@ class AppSession:
 
                             if r.text:
                                 log2(f'{r.status_code} {r.url} Failed parsing the results.')
-                                Config().debug(r.text)
+                                debug(r.text)
                     else:
                         log2(r.status_code)
                         log2(r.text)
@@ -115,7 +112,7 @@ class AppSession:
             try:
                 # oidc/login may hang
                 timeout = Config().get('app.login.timeout', 5)
-                Config().debug(f'-> {idp_login.app_login_url}')
+                debug(f'-> {idp_login.app_login_url}')
                 session.post(idp_login.app_login_url, headers=headers, data=form_data, timeout=timeout)
             except Exception:
                 pass
@@ -133,7 +130,7 @@ class AppSession:
                 check_uri = Config().get('app.login.session-check-url', 'https://{host}/{env}/{app}/api/8/C3/userSessionToken')
                 check_uri = check_uri.replace('{host}', self.host).replace('{env}', self.env).replace('{app}', 'c3')
                 r = session.get(check_uri)
-                Config().debug(f'{r.status_code} {check_uri}')
+                debug(f'{r.status_code} {check_uri}')
 
                 res_text = r.text
                 js = json.loads(res_text)
@@ -142,10 +139,10 @@ class AppSession:
                     break
 
                 app_access_token = js['signedToken']
-                Config().debug(f'{r.text}')
+                debug(f'{r.text}')
 
                 self.app_login = AppLogin(session, app_access_token, idp_uri)
-            except Exception:
+            except Exception as e:
                 try:
                     need = urlparse(r.url).hostname
                     if idp_login.idp_uri:
@@ -158,7 +155,7 @@ class AppSession:
                         log2(f"Invalid username/password.")
                     break
                 finally:
-                    Config().debug(traceback.format_exc())
+                    debug_trace()
 
                     if 'res_text' in locals():
                         Log.log_to_file(res_text)

adam/batch.py

@@ -35,7 +35,7 @@ from adam.cli_group import cli
 @click.option('--param', '-v', multiple=True, metavar='<key>=<value>', help='parameter override')
 @click.argument('extra_args', nargs=-1, metavar='repair', type=click.UNPROCESSED)
 def audit(kubeconfig: str, config: str, param: list[str], extra_args):
-    run_command(Audit(), kubeconfig, config, param, None, None, None, extra_args)
+    run_command(Audit(), kubeconfig, config, param, None, None, None, extra_args, device=ReplState.L)
 
 
 @cli.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True), cls=ClusterOrPodCommandHelper, help='Run a single bash command.')
@@ -281,14 +281,14 @@ def watch(kubeconfig: str, config: str, param: list[str], cluster: str, namespac
     run_command(Watch(), kubeconfig, config, param, cluster, namespace, None, extra_args)
 
 
-def run_command(cmd: Command, kubeconfig: str, config: str, params: list[str], cluster:str, namespace: str, pod: str, extra_args):
+def run_command(cmd: Command, kubeconfig: str, config: str, params: list[str], cluster:str, namespace: str, pod: str, extra_args, device=ReplState.C):
     is_user_entry = False
 
     KubeContext.init_config(kubeconfig, is_user_entry=is_user_entry)
     if not KubeContext.init_params(config, params, is_user_entry=is_user_entry):
         return
 
-    state = ReplState(ns_sts=cluster, pod=pod, namespace=namespace)
+    state = ReplState(device=device, ns_sts=cluster, pod=pod, namespace=namespace)
     if cmd.command() == 'pg' and not extra_args:
         state, _ = state.apply_args(extra_args)
         state.device = ReplState.P

adam/checks/check_utils.py

@@ -1,7 +1,3 @@
-from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor, as_completed
-import time
-
 from adam.checks.check import Check
 from adam.checks.check_context import CheckContext
 from adam.checks.check_result import CheckResult
@@ -14,10 +10,9 @@ from adam.checks.memory import Memory
 from adam.checks.status import Status
 from adam.config import Config
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
-from adam.utils_k8s.kube_context import KubeContext
 from adam.utils_k8s.secrets import Secrets
 from adam.utils_k8s.statefulsets import StatefulSets
-from adam.utils import elapsed_time, log2
+from adam.utils import parallelize, log2
 
 def all_checks() -> list[Check]:
     return [CompactionStats(), Cpu(), Gossip(), Memory(), Disk(), Status()]
@@ -38,57 +33,30 @@ def checks_from_csv(check_str: str):
 
     return checks
 
-def run_checks(cluster: str = None, namespace: str = None, pod: str = None, checks: list[Check] = None, show_output=True):
+def run_checks(cluster: str = None, namespace: str = None, pod: str = None, checks: list[Check] = None, show_out=True):
     if not checks:
         checks = all_checks()
 
-    sss: list[tuple[str, str]] = StatefulSets.list_sts_name_and_ns()
-
-    action = 'issues'
-    crs: list[CheckResult] = []
-
-    def on_clusters(f: Callable[[any, list[str]], any]):
-        for ss, ns in sss:
-            if (not cluster or cluster == ss) and (not namespace or namespace == ns):
-                pods = StatefulSets.pods(ss, ns)
-                for pod_name in [pod.metadata.name for pod in pods]:
-                    if not pod or pod == pod_name:
-                        f(ss, ns, pod_name, show_output)
-
-    max_workers = Config().action_workers(action, 30)
-    if max_workers < 2:
-        def serial(ss, ns, pod_name, show_output):
-            if not pod or pod == pod_name:
-                crs.append(run_checks_on_pod(checks, ss[0], ns, pod_name, show_output))
-
-        on_clusters(serial)
-    else:
-        if KubeContext.show_parallelism():
-            log2(f'Executing on all nodes from statefulset with {max_workers} workers...')
-        start_time = time.time()
-        try:
-            futures = []
-            with ThreadPoolExecutor(max_workers=max_workers) as executor:
-                def submit(ss, ns, pod_name, show_output):
-                    f = executor.submit(run_checks_on_pod, checks, ss, ns, pod_name, show_output,)
-                    if f: futures.append(f)
-
-                on_clusters(submit)
+    sts_ns: list[tuple[str, str]] = StatefulSets.list_sts_name_and_ns()
 
-            crs = [future.result() for future in as_completed(futures)]
-        finally:
-            if KubeContext.show_parallelism():
-                log2(f"Parallel {action} elapsed time: {elapsed_time(start_time)} with {max_workers} workers")
+    sts_ns_pods: list[tuple[str, str, str]] = []
+    for sts, ns in sts_ns:
+        if (not cluster or cluster == sts) and (not namespace or namespace == ns):
+            pods = StatefulSets.pods(sts, ns)
+            for pod_name in [pod.metadata.name for pod in pods]:
+                if not pod or pod == pod_name:
+                    sts_ns_pods.append((sts, ns, pod_name))
 
-    return crs
+    with parallelize(sts_ns_pods, Config().action_workers('issues', 30), msg='d`Running|Ran checks on {size} pods') as exec:
+        return exec.map(lambda sts_ns_pod: run_checks_on_pod(checks, sts_ns_pod[0], sts_ns_pod[1], sts_ns_pod[2], show_out))
 
-def run_checks_on_pod(checks: list[Check], cluster: str = None, namespace: str = None, pod: str = None, show_output=True):
+def run_checks_on_pod(checks: list[Check], cluster: str = None, namespace: str = None, pod: str = None, show_out=True):
     host_id = CassandraNodes.get_host_id(pod, namespace)
     user, pw = Secrets.get_user_pass(pod, namespace)
     results = {}
     issues: list[Issue] = []
     for c in checks:
-        check_results = c.check(CheckContext(cluster, host_id, pod, namespace, user, pw, show_output=show_output))
+        check_results = c.check(CheckContext(cluster, host_id, pod, namespace, user, pw, show_output=show_out))
         if check_results.details:
             results = results | {check_results.name: check_results.details}
         if check_results.issues:

adam/checks/cpu.py

@@ -7,6 +7,7 @@ from adam.checks.issue import Issue
 from adam.config import Config
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.custom_resources import CustomResources
+from adam.utils_k8s.pods import Pods
 
 class Cpu(Check):
     def name(self):
@@ -20,10 +21,15 @@ class Cpu(Check):
             'namespace': ctx.namespace,
             'statefulset': ctx.statefulset,
             'cpu': 'Unknown',
-            'idle': 'Unknown'
+            'idle': 'Unknown',
+            'limit': 'NA'
         }
 
         try:
+            container = Pods.get_container(ctx.namespace, ctx.pod, container_name='cassandra')
+            if container.resources.limits and "cpu" in container.resources.limits:
+                details['limit'] = container.resources.limits["cpu"]
+
             idle = 'Unknown'
             result = CassandraNodes.exec(ctx.pod, ctx.namespace, "mpstat 5 2 | grep Average | awk '{print $NF}'", show_out=ctx.show_output)
             lines = result.stdout.strip(' \r\n').split('\n')

adam/checks/cpu_metrics.py

@@ -0,0 +1,52 @@
+from kubernetes.utils import parse_quantity
+
+from adam.checks.check import Check
+from adam.checks.check_context import CheckContext
+from adam.checks.check_result import CheckResult
+from adam.checks.issue import Issue
+from adam.config import Config
+from adam.utils_k8s.custom_resources import CustomResources
+from adam.utils_k8s.pods import Pods
+
+class CpuMetrics(Check):
+    def name(self):
+        return 'cpu-metrics'
+
+    def check(self, ctx: CheckContext) -> CheckResult:
+        issues: list[Issue] = []
+
+        details = {
+            'name': ctx.pod,
+            'namespace': ctx.namespace,
+            'statefulset': ctx.statefulset,
+            'cpu': 'Unknown',
+            'limit': 'NA'
+        }
+
+        try:
+            container = Pods.get_container(ctx.namespace, ctx.pod, container_name='cassandra')
+            if container.resources.limits and "cpu" in container.resources.limits:
+                details['limit'] = container.resources.limits["cpu"]
+
+            metrics = CustomResources.get_metrics(ctx.namespace, ctx.pod, container_name='cassandra')
+            usage = 'Unknown'
+            if metrics:
+                usage = details['cpu'] = metrics["usage"]["cpu"]
+
+            cpu_threshold = Config().get('checks.cpu-threshold', 0.0)
+            if  cpu_threshold != 0.0 and usage != "Unknown" and parse_quantity(usage) > cpu_threshold:
+                issues.append(Issue(
+                    statefulset=ctx.statefulset,
+                    namespace=ctx.namespace,
+                    pod=ctx.pod,
+                    category='cpu',
+                    desc=f'CPU is too busy: {usage}',
+                    suggestion=f"qing restart {ctx.pod}@{ctx.namespace}"
+                ))
+        except Exception as e:
+            issues.append(self.issue_from_err(sts_name=ctx.statefulset, ns=ctx.namespace, pod_name=ctx.pod, exception=e))
+
+        return CheckResult(self.name(), details, issues)
+
+    def help(self):
+        return f'{CpuMetrics().name()}: check cpu busy percentage with metrics'

adam/checks/disk.py

@@ -6,6 +6,7 @@ from adam.checks.check_context import CheckContext
 from adam.checks.check_result import CheckResult
 from adam.checks.issue import Issue
 from adam.config import Config
+from adam.utils import log_exc
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 
 class Disk(Check):
@@ -87,10 +88,8 @@ class Disk(Check):
 
         ss_size = 0.0
         if ss_out:
-            try:
+            with log_exc():
                 ss_size = round(float(ss_out.strip(' \r\n')) / 1024 / 1024, 2)
-            except:
-                pass
 
         def parse_du_out(l: str, default: str = None):
             groups = re.match(r'^(\S+)\s+(\S+)$', l.strip('\r'))

adam/columns/columns.py

@@ -1,6 +1,7 @@
 from adam.columns.column import Column
 from adam.columns.compactions import Compactions
 from adam.columns.cpu import Cpu
+from adam.columns.cpu_metrics import CpuMetrics
 from adam.columns.dir_data import DataDir
 from adam.columns.dir_snapshots import SnapshotsDir
 from adam.columns.gossip import Gossip
@@ -23,7 +24,7 @@ class Columns:
     COLUMNS_BY_NAME = None
 
     def all_columns():
-        return [Compactions(), Cpu(), DataDir(), SnapshotsDir(), Gossip(), HostId(), Memory(),
+        return [Compactions(), Cpu(), CpuMetrics(), DataDir(), SnapshotsDir(), Gossip(), HostId(), Memory(),
                 NodeAddress(), NodeLoad(), NodeOwns(), NodeStatus(),NodeTokens(), PodName(), CassandraVolume(), RootVolume()]
 
     def columns_by_name():
@@ -38,6 +39,7 @@ class Columns:
             name = name.strip(' ')
             if not name in Columns.COLUMNS_BY_NAME:
                 return None
+
             cols.append(Columns.COLUMNS_BY_NAME[name]())
 
         return cols

adam/columns/cpu.py

@@ -1,3 +1,5 @@
+from kubernetes.utils.quantity import parse_quantity
+
 from adam.checks.check_result import CheckResult
 from adam.checks.cpu import Cpu as CpuCheck
 from adam.columns.column import Column
@@ -14,4 +16,4 @@ class Cpu(Column):
         cpu = r.details[CpuCheck().name()]
         busy = 100.0 - float(cpu['idle'])
 
-        return f'{round(busy)}%'
+        return f'{round(busy)}%/{parse_quantity(cpu["limit"]) * 100}%'

adam/columns/cpu_metrics.py

@@ -0,0 +1,22 @@
+from kubernetes.utils.quantity import parse_quantity
+
+from adam.checks.check_result import CheckResult
+from adam.checks.cpu_metrics import CpuMetrics as CpuCheck
+from adam.columns.column import Column
+
+class CpuMetrics(Column):
+    def name(self):
+        return 'cpu-metrics'
+
+    def checks(self):
+        return [CpuCheck()]
+
+    def pod_value(self, check_results: list[CheckResult], pod_name: str):
+        r = self.result_by_pod(check_results, pod_name)
+        cpu = r.details[CpuCheck().name()]
+
+        cpu_decimal = parse_quantity(cpu['cpu'])
+        cpu_limit = parse_quantity(cpu['limit'])
+        business = cpu_decimal * 100 / cpu_limit
+
+        return f"{business:.2f}%({cpu_decimal}/{cpu_limit})"

adam/columns/memory.py

@@ -3,6 +3,7 @@ from kubernetes.utils import parse_quantity
 from adam.checks.check_result import CheckResult
 from adam.checks.memory import Memory as MemoryCheck
 from adam.columns.column import Column
+from adam.utils import log_exc
 
 class Memory(Column):
     def name(self):
@@ -18,7 +19,5 @@ class Memory(Column):
         return f"{Memory.to_g(mem['used'])}/{Memory.to_g(mem['limit'])}"
 
     def to_g(v: str):
-        try:
-            return f'{round(parse_quantity(v) / 1024 / 1024 / 1024, 2)}G'
-        except:
-            return v
+        with log_exc():
+            return f'{round(parse_quantity(v) / 1024 / 1024 / 1024, 2)}G'

adam/commands/__init__.py

@@ -0,0 +1,18 @@
+from adam.commands.command import ExtractAllOptionsHandler, ExtractOptionsHandler, ExtractSequenceOptionsHandler, ExtractTrailingOptionsHandler
+from adam.repl_state import ReplState
+from adam.commands.app.utils_app import AppHandler
+
+def app(state: ReplState) -> AppHandler:
+    return AppHandler(state)
+
+def extract_options(args: list[str], options: list[str]):
+    return ExtractOptionsHandler(args, options = options)
+
+def extract_trailing_options(args: list[str], trailing: list[str]):
+    return ExtractTrailingOptionsHandler(args, trailing = trailing)
+
+def extract_all_options(args: list[str], trailing = None, sequence = None, options = None):
+    return ExtractAllOptionsHandler(args, trailing = trailing, sequence = sequence, options = options)
+
+def extract_sequence(args: list[str], sequence: list[str]):
+    return ExtractSequenceOptionsHandler(args, sequence = sequence)

adam/commands/alter_tables.py

@@ -1,8 +1,9 @@
+from adam.commands import extract_options
 from adam.commands.command import Command
-from adam.commands.cql.cql_utils import tables as get_tables, run_cql
+from adam.commands.cql.utils_cql import cassandra, cassandra_tables as get_tables
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
+from adam.utils import log2, log_exc
 
 class AlterTables(Command):
     COMMAND = 'alter tables with'
@@ -26,52 +27,47 @@ class AlterTables(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return state
-
-        if not args:
-            if state.in_repl:
-                log2('Please enter gc grace in seconds. e.g. alter gc-grace-seconds 3600')
-            else:
-                log2('* gc grace second is missing.')
-                log2()
-                Command.display_help()
-
-            return 'missing-arg'
-
-        args, include_reaper = Command.extract_options(args, '--include-reaper')
-        arg_str = ' '.join(args)
-
-        excludes = [e.strip(' \r\n') for e in Config().get(
-            'cql.alter-tables.excludes',
-            'system_auth,system_traces,reaper_db,system_distributed,system_views,system,system_schema,system_virtual_schema').split(',')]
-        batching = Config().get('cql.alter-tables.batching', True)
-        tables = get_tables(state, on_any=True)
-        for k, v in tables.items():
-            if k not in excludes or k == 'reaper_db' and include_reaper:
-                if batching:
-                    # alter table <table_name> with GC_GRACE_SECONDS = <timeout>;
-                    cql = ';\n'.join([f'alter table {k}.{t} with {arg_str}' for t in v])
-                    try:
-                        run_cql(state, cql, [], show_out=Config().is_debug(), on_any=True)
-                    except Exception as e:
-                        log2(e)
-                        continue
-                else:
-                    for t in v:
-                        try:
+        with self.validate(args, state) as (args, state):
+            with extract_options(args, '--include-reaper') as (args, include_reaper):
+                if not args:
+                    if state.in_repl:
+                        log2('Please enter gc grace in seconds. e.g. alter gc-grace-seconds 3600')
+                    else:
+                        log2('* gc grace second is missing.')
+                        log2()
+                        Command.display_help()
+
+                    return 'missing-arg'
+
+                arg_str = ' '.join(args)
+
+                excludes = [e.strip(' \r\n') for e in Config().get(
+                    'cql.alter-tables.excludes',
+                    'system_auth,system_traces,reaper_db,system_distributed,system_views,system,system_schema,system_virtual_schema').split(',')]
+                batching = Config().get('cql.alter-tables.batching', True)
+                tables = get_tables(state, on_any=True)
+                for k, v in tables.items():
+                    if k not in excludes or k == 'reaper_db' and include_reaper:
+                        if batching:
                             # alter table <table_name> with GC_GRACE_SECONDS = <timeout>;
-                            cql = f'alter table {k}.{t} with {arg_str}'
-                            run_cql(state, cql, [], show_out=Config().is_debug(), on_any=True)
-                        except Exception as e:
-                            log2(e)
-                            continue
-
-                log2(f'{len(v)} tables altered in {k}.')
-
-        # do not continue to cql route
-        return state
+                            cql = ';\n'.join([f'alter table {k}.{t} with {arg_str}' for t in v])
+                            with log_exc(True):
+                                with cassandra(state) as pods:
+                                    pods.cql(cql, show_out=Config().is_debug(), show_query=not Config().is_debug(), on_any=True)
+                                continue
+                        else:
+                            for t in v:
+                                with log_exc(True):
+                                    # alter table <table_name> with GC_GRACE_SECONDS = <timeout>;
+                                    cql = f'alter table {k}.{t} with {arg_str}'
+                                    with cassandra(state) as pods:
+                                        pods.cql(show_out=Config().is_debug(), show_query=not Config().is_debug(), on_any=True)
+                                    continue
+
+                        log2(f'{len(v)} tables altered in {k}.')
+
+                # do not continue to cql route
+                return state
 
     def completion(self, _: ReplState) -> dict[str, any]:
         # auto completion is taken care of by sql completer

adam/commands/audit/audit.py

@@ -7,14 +7,14 @@ from adam.commands.audit.show_slow10 import ShowSlow10
 from adam.commands.audit.show_top10 import ShowTop10
 from adam.commands.audit.utils_show_top10 import show_top10_completions_for_nesting
 from adam.commands.command import Command
+from adam.commands.intermediate_command import IntermediateCommand
 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.sql.sql_completer import SqlCompleter, SqlVariant
+from adam.utils import log2, wait_log
 from adam.utils_athena import Athena
-from adam.utils_audits import Audits
 
-class Audit(Command):
+class Audit(IntermediateCommand):
     COMMAND = 'audit'
 
     # the singleton pattern
@@ -37,29 +37,26 @@ class Audit(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return state
-
-        r = None
-        if len(args) > 0:
-            r = super().intermediate_run(cmd, state, args, Audit.cmd_list(), display_help=False)
+        with self.validate(args, state) as (args, state):
+            r = None
+            if len(args) > 0:
+                r = self.intermediate_run(cmd, state, args, self.cmd_list(), display_help=False)
 
-        if not r or isinstance(r, str) and r == 'command-missing':
-            sql = 'select * from audit order by ts desc limit 10'
-            if args:
-                sql = ' '.join(args)
-            else:
-                log2(sql)
+            if not r or isinstance(r, str) and r == 'command-missing':
+                sql = 'select * from audit order by ts desc limit 10'
+                if args:
+                    sql = ' '.join(args)
+                else:
+                    log2(sql)
 
-            Athena.run_query(sql)
+                Athena.run_query(sql)
 
-        return state
+            return state
 
     def completion(self, state: ReplState):
         if state.device == ReplState.L:
             if not self.schema_read:
-                Config().wait_log(f'Inspecting audit database schema...')
+                wait_log(f'Inspecting audit database schema...')
                 self.schema_read = True
                 # warm up the caches first time when l: drive is accessed
                 Athena.table_names()
@@ -68,14 +65,16 @@ class Audit(Command):
 
             return super().completion(state) | show_top10_completions_for_nesting() | SqlCompleter(
                 lambda: Athena.table_names(),
-                columns=lambda table: Athena.column_names(),
-                partition_columns=lambda table: Athena.column_names(partition_cols_only=True),
-                variant='athena'
+                expandables={
+                    'columns': lambda table: Athena.column_names(),
+                    'partition_columns': lambda table: Athena.column_names(partition_cols_only=True)
+                },
+                variant=SqlVariant.ATHENA
             ).completions_for_nesting()
 
         return {}
 
-    def cmd_list():
+    def cmd_list(self):
         return [AuditRepairTables(), AuditRun(), ShowLast10(), ShowSlow10(), ShowTop10()]
 
     def help(self, _: ReplState):
@@ -83,4 +82,4 @@ class Audit(Command):
 
 class AuditCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):
-        Command.intermediate_help(super().get_help(ctx), Audit.COMMAND, Audit.cmd_list(), show_cluster_help=False)
+        IntermediateCommand.intermediate_help(super().get_help(ctx), Audit.COMMAND, Audit().cmd_list(), show_cluster_help=False)

adam/commands/audit/audit_repair_tables.py

@@ -1,4 +1,3 @@
-import concurrent
 import time
 
 from adam.commands.command import Command
@@ -31,25 +30,23 @@ class AuditRepairTables(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return state
-
-        tables = Config().get('audit.athena.repair-partition-tables', 'audit').split(',')
-        if args:
-            tables = args
+        with self.validate(args, state) as (args, state):
+            tables = Config().get('audit.athena.repair-partition-tables', 'audit').split(',')
+            if args:
+                tables = args
 
-        meta = Audits.get_meta()
-        self.repair(tables, meta)
+            meta = Audits.get_meta()
+            self.repair(tables, meta)
 
-        return state
+            return state
 
     def completion(self, state: ReplState):
+        # trigger auto repair if on l: drive
         if state.device == ReplState.L:
             if not self.auto_repaired:
                 if hours := Config().get('audit.athena.auto-repair.elapsed_hours', 12):
-                    with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
-                        executor.submit(self.auto_repair, hours,)
+                    with Audits.offload() as exec:
+                        exec.submit(lambda: self.auto_repair(hours))
 
             return super().completion(state)
 
@@ -65,13 +62,13 @@ class AuditRepairTables(Command):
             log2(f'Audit tables have been auto-repaired.')
 
     def repair(self, tables: list[str], meta: AuditMeta, show_sql = False):
-        with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
+        with Audits.offload() as exec:
             for table in tables:
                 if show_sql:
                     log(f'MSCK REPAIR TABLE {table}')
 
-                executor.submit(Athena.run_query, f'MSCK REPAIR TABLE {table}', None,)
-            executor.submit(Audits.put_meta, Audits.PARTITIONS_ADDED, meta,)
+                exec.submit(Athena.query, f'MSCK REPAIR TABLE {table}', None,)
+            exec.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"
+        return f"{AuditRepairTables.COMMAND}\t run MSCK REPAIR command for new partition discovery"