kaqing 2.0.14__py3-none-any.whl → 2.0.145__py3-none-any.whl

adam/apps.py

@@ -5,10 +5,10 @@ from typing import cast
 import yaml
 
 from adam.config import Config
-from adam.k8s_utils.ingresses import Ingresses
+from adam.utils_k8s.ingresses import Ingresses
 
 from . import __version__
-from adam.k8s_utils.services import Services
+from adam.utils_k8s.services import Services
 from adam.utils import copy_config_file
 
 class AppAction:

adam/batch.py

@@ -1,11 +1,12 @@
 import click
 
-from adam.commands.bash import Bash
+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.command import Command
 from adam.commands.command_helpers import ClusterCommandHelper, ClusterOrPodCommandHelper, PodCommandHelper
-from adam.commands.cqlsh import CqlCommandHelper, Cqlsh
+from adam.commands.cql.cqlsh import CqlCommandHelper, Cqlsh
 from adam.commands.deploy.deploy import Deploy, DeployCommandHelper
 from adam.commands.deploy.undeploy import Undeploy, UndeployCommandHelper
 from adam.commands.issues import Issues
@@ -23,11 +24,20 @@ from adam.commands.restart import Restart
 from adam.commands.rollout import RollOut
 from adam.commands.show.show import Show, ShowCommandHelper
 from adam.commands.watch import Watch
-from adam.k8s_utils.kube_context import KubeContext
+from adam.utils_k8s.kube_context import KubeContext
 from adam.repl import enter_repl
 from adam.repl_state import ReplState
 from adam.cli_group import cli
 
+@cli.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True), cls=AuditCommandHelper, help='Run audit functions.')
+@click.option('--kubeconfig', '-k', required=False, metavar='path', help='path to kubeconfig file')
+@click.option('--config', default='params.yaml', metavar='path', help='path to kaqing parameters file')
+@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)
+
+
 @cli.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True), cls=ClusterOrPodCommandHelper, help='Run a single bash command.')
 @click.option('--kubeconfig', '-k', required=False, metavar='path', help='path to kubeconfig file')
 @click.option('--config', default='params.yaml', metavar='path', help='path to kaqing parameters file')

adam/checks/check_utils.py

@@ -13,10 +13,10 @@ 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.k8s_utils.cassandra_nodes import CassandraNodes
-from adam.k8s_utils.kube_context import KubeContext
-from adam.k8s_utils.secrets import Secrets
-from adam.k8s_utils.statefulsets import StatefulSets
+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
 
 def all_checks() -> list[Check]:

adam/checks/compactionstats.py

@@ -5,7 +5,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.k8s_utils.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
 
 class CompactionStats(Check):
     def name(self):

adam/checks/cpu.py

@@ -5,8 +5,8 @@ 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.k8s_utils.cassandra_nodes import CassandraNodes
-from adam.k8s_utils.custom_resources import CustomResources
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.custom_resources import CustomResources
 
 class Cpu(Check):
     def name(self):

adam/checks/disk.py

@@ -6,7 +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.k8s_utils.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
 
 class Disk(Check):
     def name(self):

adam/checks/gossip.py

@@ -4,7 +4,7 @@ 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.k8s_utils.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
 
 class Gossip(Check):
     def name(self):

adam/checks/memory.py

@@ -2,9 +2,9 @@ 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.k8s_utils.cassandra_nodes import CassandraNodes
-from adam.k8s_utils.custom_resources import CustomResources
-from adam.k8s_utils.pods import Pods
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.custom_resources import CustomResources
+from adam.utils_k8s.pods import Pods
 
 class Memory(Check):
     def name(self):

adam/checks/status.py

@@ -4,7 +4,7 @@ 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.k8s_utils.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
 
 class Status(Check):
     def name(self):

adam/commands/alter_tables.py

@@ -0,0 +1,81 @@
+from adam.commands.command import Command
+from adam.commands.cql.cql_utils import cassandra_tables as get_tables, run_cql
+from adam.config import Config
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+
+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)
+
+        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:
+                            # 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
+
+    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'

adam/commands/app.py

@@ -22,14 +22,14 @@ class App(Command):
         return App.COMMAND
 
     def required(self):
-        return RequiredState.CLUSTER_OR_POD
+        return RequiredState.APP_APP
 
     def run(self, cmd: str, state: ReplState):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
         state, args = self.apply_state(args, state)
-        if not self.validate_state(state, app_required=RequiredState.APP_APP):
+        if not self.validate_state(state):
             return state
 
         args, forced = Command.extract_options(args, '--force')
@@ -64,4 +64,4 @@ class App(Command):
         return {}
 
     def help(self, _: ReplState):
-        return f"{App.COMMAND} [--force]\t post app API query with type and function"
+        return f"<AppType>.<AppAction> <args> [--force]\t post app action; check with 'show app actions' command"

adam/commands/app_ping.py

@@ -18,14 +18,14 @@ class AppPing(Command):
         return AppPing.COMMAND
 
     def required(self):
-        return RequiredState.CLUSTER_OR_POD
+        return RequiredState.APP_APP
 
     def run(self, cmd: str, state: ReplState):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
         state, args = self.apply_state(args, state)
-        if not self.validate_state(state, app_required=RequiredState.APP_APP):
+        if not self.validate_state(state):
             return state
 
         _, forced = Command.extract_options(args, '--force')

adam/commands/audit/audit.py

@@ -0,0 +1,86 @@
+import click
+
+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
+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.repl_state import ReplState
+from adam.sql.sql_completer import SqlCompleter, SqlVariant
+from adam.utils import log2
+from adam.utils_athena import Athena
+from adam.utils_audits import Audits
+
+class Audit(Command):
+    COMMAND = 'audit'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(Audit, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+        self.schema_read = False
+
+    def command(self):
+        return Audit.COMMAND
+
+    def required(self):
+        return ReplState.L
+
+    def run(self, cmd: str, state: ReplState):
+        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)
+
+        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)
+
+        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...')
+                self.schema_read = True
+                # warm up the caches first time when l: drive is accessed
+                Athena.table_names()
+                Athena.column_names()
+                Athena.column_names(partition_cols_only=True)
+
+            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=SqlVariant.ATHENA
+            ).completions_for_nesting()
+
+        return {}
+
+    def cmd_list():
+        return [AuditRepairTables(), AuditRun(), ShowLast10(), ShowSlow10(), ShowTop10()]
+
+    def help(self, _: ReplState):
+        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):
+        Command.intermediate_help(super().get_help(ctx), Audit.COMMAND, Audit.cmd_list(), show_cluster_help=False)

adam/commands/audit/audit_repair_tables.py

@@ -0,0 +1,77 @@
+import concurrent
+import time
+
+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 Athena
+from adam.utils_audits import AuditMeta, Audits
+
+class AuditRepairTables(Command):
+    COMMAND = 'audit repair'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(AuditRepairTables, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+        self.auto_repaired = False
+
+    def command(self):
+        return AuditRepairTables.COMMAND
+
+    def required(self):
+        return ReplState.L
+
+    def run(self, cmd: str, state: ReplState):
+        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
+
+        meta = Audits.get_meta()
+        self.repair(tables, meta)
+
+        return state
+
+    def completion(self, state: ReplState):
+        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,)
+
+            return super().completion(state)
+
+        return {}
+
+    def auto_repair(self, hours: int):
+        self.auto_repaired = True
+
+        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.')
+
+    def repair(self, tables: list[str], meta: AuditMeta, show_sql = False):
+        with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
+            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,)
+
+    def help(self, _: ReplState):
+        return f"{AuditRepairTables.COMMAND} \t run MSCK REPAIR command for new partition discovery"

adam/commands/audit/audit_run.py

@@ -0,0 +1,58 @@
+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 Athena
+from adam.utils_audits import AuditMeta, Audits
+
+class AuditRun(Command):
+    COMMAND = 'audit run'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(AuditRun, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+        self.auto_repaired = False
+
+    def command(self):
+        return AuditRun.COMMAND
+
+    def required(self):
+        return ReplState.L
+
+    def run(self, cmd: str, state: ReplState):
+        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
+
+        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:
+            log2(f'Added {len(clusters)} new clusters.')
+            tables = Config().get('audit.athena.repair-cluster-tables', 'cluster').split(',')
+            with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
+                for table in tables:
+                    executor.submit(Athena.run_query, f'MSCK REPAIR TABLE {table}', None,)
+        else:
+            log2(f'No new clusters were found.')
+
+        return state
+
+    def completion(self, state: ReplState):
+        if state.device == ReplState.L:
+            return super().completion(state)
+
+        return {}
+
+    def help(self, _: ReplState):
+        return f"{AuditRun.COMMAND} \t run"

adam/commands/audit/show_last10.py

@@ -0,0 +1,51 @@
+from adam.commands.audit.utils_show_top10 import extract_limit_and_duration
+from adam.commands.command import Command
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_athena import Athena
+from adam.utils_audits import Audits
+
+class ShowLast10(Command):
+    COMMAND = 'show last'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowLast10, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowLast10.COMMAND
+
+    def required(self):
+        return ReplState.L
+
+    def run(self, cmd: str, state: ReplState):
+        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
+
+        limit, date_condition = extract_limit_and_duration(args)
+
+        query = '\n    '.join([
+            "SELECT * FROM audit",
+            f"WHERE drive <> 'z' and ({date_condition})",
+            f"ORDER BY ts DESC LIMIT {limit};"])
+        log2(query)
+        log2()
+        Athena.run_query(query)
+
+        return state
+
+    def completion(self, _: ReplState):
+
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{ShowLast10.COMMAND} [limit]\t show last <limit> audit lines'

adam/commands/audit/show_slow10.py

@@ -0,0 +1,50 @@
+from adam.commands.audit.utils_show_top10 import extract_limit_and_duration
+from adam.commands.command import Command
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_athena import Athena
+from adam.utils_audits import Audits
+
+class ShowSlow10(Command):
+    COMMAND = 'show slow'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowSlow10, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowSlow10.COMMAND
+
+    def required(self):
+        return ReplState.L
+
+    def run(self, cmd: str, state: ReplState):
+        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
+
+        limit, date_condition = extract_limit_and_duration(args)
+
+        query = '\n    '.join([
+            "SELECT * FROM audit",
+            f"WHERE drive <> 'z' and ({date_condition})",
+            f"ORDER BY CAST(duration AS REAL) DESC LIMIT {limit};"])
+        log2(query)
+        log2()
+        Athena.run_query(query)
+
+        return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{ShowSlow10.COMMAND} [limit]\t show slow <limit> audit lines'

adam/commands/audit/show_top10.py

@@ -0,0 +1,48 @@
+from adam.commands.audit.utils_show_top10 import extract_limit_and_duration
+from adam.commands.command import Command
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_athena import Athena
+
+class ShowTop10(Command):
+    COMMAND = 'show top'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowTop10, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowTop10.COMMAND
+
+    def required(self):
+        return ReplState.L
+
+    def run(self, cmd: str, state: ReplState):
+        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
+
+        limit, date_condition = extract_limit_and_duration(args)
+        query = '\n    '.join([
+            "SELECT min(c) AS cluster, line, COUNT(*) AS cnt, avg(CAST(duration AS REAL)) AS duration",
+            f"FROM audit WHERE drive <> 'z' and ({date_condition})",
+            f"GROUP BY line ORDER BY cnt DESC LIMIT {limit};"])
+        log2(query)
+        log2()
+        Athena.run_query(query)
+
+        return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{ShowTop10.COMMAND} [limit]\t show top <limit> audit lines'