kaqing 2.0.114__tar.gz → 2.0.177__tar.gz

{kaqing-2.0.114 → kaqing-2.0.177}/PKG-INFO

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

kaqing-2.0.177/adam/__init__.py

@@ -0,0 +1 @@
+from .version import __version__, __release__

{kaqing-2.0.114 → kaqing-2.0.177}/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, tabulize, 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)
+                                tabulize(lines, header=header, separator='\t')
                         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)

{kaqing-2.0.114 → kaqing-2.0.177}/adam/apps.py

@@ -12,7 +12,8 @@ from adam.utils_k8s.services import Services
 from adam.utils import copy_config_file
 
 class AppAction:
-    def __init__(self, name: str, payload: str = None, args: dict[str, str] = None, help: str = None):
+    def __init__(self, typ: str, name: str, payload: str = None, args: dict[str, str] = None, help: str = None):
+        self.typ = typ
         self.name = name
         self.payload = payload
         self.args = args
@@ -33,6 +34,19 @@ class AppAction:
 
         return args
 
+    def __str__(self):
+        line = None
+
+        args = ','.join(self.arguments())
+        if args:
+            line = f'{self.typ}.{self.name},{args}'
+        else:
+            line = f'{self.typ}.{self.name},'
+        if self.help:
+            line = f'{line},{self.help}'
+
+        return line
+
 class AppType:
     all_types: list['AppType'] = []
 
@@ -62,12 +76,12 @@ class AppType:
                         aa: AppAction = None
 
                         if isinstance(v, str):
-                            aa = AppAction(k, help=v)
+                            aa = AppAction(typ, k, help=v)
                         elif isinstance(v, dict):
                             args = {k: v1 for k, v1 in v.items()}
-                            aa = AppAction(k, payload=v['payload'], args=args, help=v['help'])
+                            aa = AppAction(typ, k, payload=v['payload'], args=args, help=v['help'])
                         else:
-                            aa = AppAction(k)
+                            aa = AppAction(typ, k)
 
                         app_actions.append(aa)
 

{kaqing-2.0.114 → kaqing-2.0.177}/adam/batch.py

@@ -3,7 +3,7 @@ import click
 from adam.commands.audit.audit import Audit, AuditCommandHelper
 from adam.commands.bash.bash import Bash
 from adam.commands.check import Check, CheckCommandHelper
-from adam.commands.cp import ClipboardCopy, CopyCommandHelper
+from adam.commands.clipboard_copy import ClipboardCopy, CopyCommandHelper
 from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterCommandHelper, ClusterOrPodCommandHelper, PodCommandHelper
 from adam.commands.cql.cqlsh import CqlCommandHelper, Cqlsh
@@ -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

kaqing-2.0.177/adam/checks/check_utils.py

@@ -0,0 +1,65 @@
+from adam.checks.check import Check
+from adam.checks.check_context import CheckContext
+from adam.checks.check_result import CheckResult
+from adam.checks.compactionstats import CompactionStats
+from adam.checks.cpu import Cpu
+from adam.checks.disk import Disk
+from adam.checks.gossip import Gossip
+from adam.checks.issue import Issue
+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.secrets import Secrets
+from adam.utils_k8s.statefulsets import StatefulSets
+from adam.utils import parallelize, log2
+
+def all_checks() -> list[Check]:
+    return [CompactionStats(), Cpu(), Gossip(), Memory(), Disk(), Status()]
+
+def checks_from_csv(check_str: str):
+    checks: list[Check] = []
+
+    checks_by_name = {c.name(): c for c in all_checks()}
+
+    if check_str:
+        for check_name in check_str.strip(' ').split(','):
+            if check_name in checks_by_name:
+                checks.append(checks_by_name[check_name])
+            else:
+                log2(f'Invalid check name: {check_name}.')
+
+                return None
+
+    return checks
+
+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()
+
+    sts_ns: list[tuple[str, str]] = StatefulSets.list_sts_name_and_ns()
+
+    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))
+
+    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_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_out))
+        if check_results.details:
+            results = results | {check_results.name: check_results.details}
+        if check_results.issues:
+            issues.extend(check_results.issues)
+
+    return CheckResult(None, results, issues)

{kaqing-2.0.114 → kaqing-2.0.177}/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')

kaqing-2.0.177/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'

{kaqing-2.0.114 → kaqing-2.0.177}/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'))

{kaqing-2.0.114 → kaqing-2.0.177}/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

{kaqing-2.0.114 → kaqing-2.0.177}/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}%'

kaqing-2.0.177/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})"

{kaqing-2.0.114 → kaqing-2.0.177}/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'

kaqing-2.0.177/adam/commands/__init__.py

@@ -0,0 +1,22 @@
+from collections.abc import Callable
+from adam.commands.command import ExtractAllOptionsHandler, ExtractOptionsHandler, ExtractSequenceOptionsHandler, ExtractTrailingOptionsHandler, ValidateArgCountHandler
+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)
+
+def validate_args(args: list[str], state: ReplState, count_at_least: int = 1, name: str = None, msg: Callable[[], None] = None, default: str = None):
+    return ValidateArgCountHandler(args, state, count_at_least=count_at_least, name=name, msg=msg, default=default)

kaqing-2.0.177/adam/commands/alter_tables.py

@@ -0,0 +1,66 @@
+from adam.commands import extract_options, validate_args
+from adam.commands.command import Command
+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, log_exc
+
+class AlterTables(Command):
+    COMMAND = 'alter tables with'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(AlterTables, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def required(self):
+        return RequiredState.CLUSTER
+
+    def command(self):
+        return AlterTables.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, state)
+
+        with self.validate(args, state) as (args, state):
+            with extract_options(args, '--include-reaper') as (args, include_reaper):
+                with validate_args(args, state, name='gc grace in seconds') as arg_str:
+                    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])
+                                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
+        return {}
+
+    def help(self, _: ReplState) -> str:
+        return f'{AlterTables.COMMAND} <param = value> [--include-reaper] \t alter schema on all tables'

{kaqing-2.0.114 → kaqing-2.0.177}/adam/commands/audit/audit.py

@@ -1,5 +1,6 @@
 import click
 
+from adam.commands import validate_args
 from adam.commands.audit.audit_repair_tables import AuditRepairTables
 from adam.commands.audit.audit_run import AuditRun
 from adam.commands.audit.show_last10 import ShowLast10
@@ -7,14 +8,13 @@ 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.config import Config
+from adam.commands.intermediate_command import IntermediateCommand
 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,22 @@ 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':
+                with validate_args(args, state, default='select * from audit order by ts desc limit 10') as sql:
+                    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 +61,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 +78,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)

{kaqing-2.0.114 → kaqing-2.0.177}/adam/commands/audit/audit_repair_tables.py

@@ -1,6 +1,6 @@
-import concurrent
 import time
 
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.config import Config
 from adam.repl_state import ReplState
@@ -31,25 +31,20 @@ 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):
+            with validate_args(args, state, default=Config().get('audit.athena.repair-partition-tables', 'audit').split(',')) as tables:
+                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 +60,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"