kaqing 2.0.110__py3-none-any.whl → 2.0.214__py3-none-any.whl

adam/commands/devices/device_auit_log.py

@@ -0,0 +1,49 @@
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+from adam.repl_state import ReplState
+from adam.utils import tabulize, log
+from adam.utils_athena import Athena
+
+class DeviceAuditLog(Command, Device):
+    COMMAND = f'{ReplState.L}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DeviceAuditLog, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DeviceAuditLog.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.L
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DeviceAuditLog.COMMAND}\t move to Audit Log Operations device'
+
+    def ls(self, cmd: str, _: ReplState):
+        tabulize(Athena.table_names(), header='NAME', separator=',')
+
+    def pwd(self, _: ReplState):
+        return '\t'.join([f'{ReplState.L}:>', '/'])
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        return True, chain.run(f'audit {cmd}', state)
+
+    def show_tables(self, _: ReplState):
+        tabulize(Athena.table_names(), separator=',')
+
+    def show_table_preview(self, _: ReplState, table: str, rows: int):
+        Athena.run_query(f'select * from {table} limit {rows}')

adam/commands/devices/device_cass.py

@@ -0,0 +1,179 @@
+from adam.commands.bash.bash_completer import BashCompleter
+from adam.commands.command import Command, InvalidStateException
+from adam.commands.commands_utils import show_pods, show_rollout
+from adam.commands.cql.utils_cql import cassandra, cassandra_table_names
+from adam.commands.devices.device import Device
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import tabulize, log, log2, wait_log
+from adam.utils_k8s.cassandra_clusters import CassandraClusters
+from adam.utils_k8s.custom_resources import CustomResources
+from adam.utils_k8s.kube_context import KubeContext
+from adam.utils_k8s.statefulsets import StatefulSets
+
+class DeviceCass(Command, Device):
+    COMMAND = f'{ReplState.C}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DeviceCass, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DeviceCass.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.C
+
+        return state
+
+    def pod(self, state: ReplState) -> str:
+        return state.pod
+
+    def pod_names(self, state: ReplState) -> list[str]:
+        return StatefulSets.pod_names(state.sts, state.namespace)
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DeviceCass.COMMAND}\t move to Cassandra Operations device'
+
+    def default_container(self, _: ReplState) -> str:
+        return 'cassandra'
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.pod:
+            return self.bash(state, state, cmd.split(' '))
+        elif state.sts and state.namespace:
+            show_pods(StatefulSets.pods(state.sts, state.namespace), state.namespace, show_namespace=not KubeContext.in_cluster_namespace())
+            show_rollout(state.sts, state.namespace)
+        else:
+            self.show_statefulsets()
+
+    def ls_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.sts:
+            return super().completion(state) | {f'@{p}': {'ls': None} for p in self.pod_names(state)}
+        else:
+            return {cmd: {n: None for n in StatefulSets.list_sts_names()}}
+
+    def show_statefulsets(self):
+        ss = StatefulSets.list_sts_names()
+        if len(ss) == 0:
+            log2('No Cassandra clusters found.')
+            return
+
+        app_ids = CustomResources.get_app_ids()
+        list = []
+        for s in ss:
+            cr_name = CustomResources.get_cr_name(s)
+            app_id = 'Unknown'
+            if cr_name in app_ids:
+                app_id = app_ids[cr_name]
+            list.append(f"{s} {app_id}")
+
+        header = 'STATEFULSET_NAME@NAMESPACE APP_ID'
+        if KubeContext.in_cluster_namespace():
+            header = 'STATEFULSET_NAME APP_ID'
+        tabulize(list, header=header)
+
+    def cd(self, dir: str, state: ReplState):
+        if dir == '':
+            state.sts = None
+            state.pod = None
+        elif dir == '..':
+            if state.pod:
+                state.pod = None
+            else:
+                state.sts = None
+        else:
+            if not state.sts:
+                ss_and_ns = dir.split('@')
+                state.sts = ss_and_ns[0]
+                if len(ss_and_ns) > 1:
+                    state.namespace = ss_and_ns[1]
+            elif not state.pod:
+                p, _ = KubeContext.is_pod_name(dir)
+                if p:
+                    state.pod = p
+                else:
+                    names = CassandraClusters.pod_names_by_host_id(state.sts, state.namespace)
+                    if dir in names:
+                        state.pod = names[dir]
+                    else:
+                        state.pod = dir
+
+    def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.pod:
+            return {cmd: {'..': None}}
+        elif state.sts:
+            return {cmd: {'..': None} | {p: None for p in self.pod_names(state)}}
+        else:
+            return {cmd: {p: None for p in StatefulSets.list_sts_names()}}
+
+    def pwd(self, state: ReplState):
+        words = []
+
+        if state.sts:
+            words.append(f'sts/{state.sts}')
+        if state.pod:
+            words.append(f'pod/{state.pod}')
+
+        return '\t'.join([f'{ReplState.C}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        if state.sts:
+            return True, chain.run(f'cql {cmd}', state)
+
+        return False, None
+
+    def enter(self, state: ReplState):
+        auto_enter = Config().get('repl.c.auto-enter', 'cluster')
+        if auto_enter and auto_enter in ['cluster', 'first-pod']:
+            sts = StatefulSets.list_sts_name_and_ns()
+            if not sts:
+                log2("No Cassandra clusters found.")
+            elif not state.sts and len(sts) == 1:
+                cluster = sts[0]
+                state.sts = cluster[0]
+                state.namespace = cluster[1]
+                if auto_enter == 'first-pod':
+                    state.pod = f'{state.sts}-0'
+                if KubeContext().in_cluster_namespace:
+                    wait_log(f'Moving to the only Cassandra cluster: {state.sts}...')
+                else:
+                    wait_log(f'Moving to the only Cassandra cluster: {state.sts}@{state.namespace}...')
+
+    def show_tables(self, state: ReplState):
+        tabulize(cassandra_table_names(state), separator=',')
+
+    def show_table_preview(self, state: ReplState, table: str, rows: int):
+        with cassandra(state) as pods:
+            pods.cql(f'select * from {table} limit {rows}', show_out=True, use_single_quotes=True, on_any=True)
+
+    def bash_target_changed(self, s0: ReplState, s1: ReplState):
+        return s0.sts != s1.sts or s0.pod != s1.pod
+
+    def exec_no_dir(self, command: str, state: ReplState):
+        with cassandra(state) as pods:
+            return pods.exec(command, action='bash', show_out=True, shell='bash')
+
+    def exec_with_dir(self, command: str, session_just_created: bool, state: ReplState):
+        with cassandra(state) as pods:
+            return pods.exec(command, action='bash', show_out=not session_just_created, shell='bash')
+
+    def bash_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        completions = {cmd: BashCompleter(lambda: [])}
+
+        if state.sts and state.namespace:
+            completions |= {f'@{p}': {cmd: BashCompleter(lambda: [])} for p in self.pod_names(state)}
+
+        return completions
+

adam/commands/devices/device_export.py

@@ -0,0 +1,87 @@
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+from adam.commands.export.export_databases import ExportDatabases, export_db
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import tabulize, log2, wait_log
+
+class DeviceExport(Command, Device):
+    COMMAND = f'{ReplState.X}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DeviceExport, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DeviceExport.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.X
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DeviceExport.COMMAND}\t move to Export Database Operations device'
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.export_session:
+            tabulize(ExportDatabases.table_names(state.export_session), header='NAME', separator=',')
+        else:
+            ExportDatabases.show_databases()
+
+    def show_table_preview(self, state: ReplState, table: str, rows: int):
+        if state.export_session:
+            with export_db(state) as dbs:
+                dbs.sql(f'select * from {table} limit {rows}')
+
+    def cd(self, dir: str, state: ReplState):
+        if dir in ['', '..']:
+            state.export_session = None
+        else:
+            state.export_session = dir
+
+    def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.export_session:
+            return {cmd: {'..': None} | {n: None for n in ExportDatabases.database_names()}}
+        else:
+            return {cmd: {n: None for n in ExportDatabases.database_names()}}
+
+    def pwd(self, state: ReplState):
+        words = []
+
+        if state.export_session:
+            words.append(state.export_session)
+
+        return '\t'.join([f'{ReplState.X}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        if cmd.startswith('select '):
+            cmd = f'xelect {cmd[7:]}'
+
+        result = chain.run(cmd, state)
+        if type(result) is ReplState:
+            if state.export_session and not result.export_session:
+                state.export_session = None
+
+        return True, result
+
+    def enter(self, state: ReplState):
+        if auto_enter := Config().get('repl.x.auto-enter', 'no'):
+            if auto_enter == 'latest':
+                wait_log(f'Moving to the latest export database...')
+                if dbs := ExportDatabases.database_names():
+                    state.export_session = sorted(dbs)[-1]
+                else:
+                    log2('No export database found.')
+

adam/commands/devices/device_postgres.py

@@ -0,0 +1,160 @@
+from adam.commands.bash.bash_completer import BashCompleter
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+from adam.commands.postgres.postgres_databases import PostgresDatabases, pg_path
+from adam.commands.postgres.utils_postgres import pg_database_names, pg_table_names, postgres
+from adam.repl_state import ReplState
+from adam.utils import tabulize, log2, wait_log
+
+class DevicePostgres(Command, Device):
+    COMMAND = f'{ReplState.P}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DevicePostgres, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DevicePostgres.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.P
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DevicePostgres.COMMAND}\t move to Postgres Operations device'
+
+    def pod(self, state: ReplState) -> str:
+        pod, _ = PostgresDatabases.pod_and_container(state.namespace)
+        return pod
+
+    def pod_names(self, state: ReplState) -> list[str]:
+        return [self.pod(state)]
+
+    def default_container(self, state: ReplState) -> str:
+        _, container = PostgresDatabases.pod_and_container(state.namespace)
+        return container
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.pod_targetted:
+            self.bash(state, state, ['ls'])
+            return
+
+        with pg_path(state) as (host, database):
+            if database:
+                tabulize(pg_table_names(state), header='NAME', separator=',')
+            elif host:
+                tabulize(pg_database_names(state), header='DATABASE', separator=',')
+            else:
+                self.show_pg_hosts(state)
+
+    def show_pg_hosts(self, state: ReplState):
+        if state.namespace:
+            tabulize(PostgresDatabases.hosts(state),
+                     lambda p: f'{p.host},{p.endpoint()}:{p.port()},{p.username()},{p.password()}',
+                     header='NAME,ENDPOINT,USERNAME,PASSWORD',
+                     separator=',')
+        else:
+            tabulize(PostgresDatabases.hosts(state),
+                     lambda p: f'{p.host},{p.namespace},{p.endpoint()}:{p.port()},{p.username()},{p.password()}',
+                     header='NAME,NAMESPACE,ENDPOINT,USERNAME,PASSWORD',
+                     separator=',')
+
+    def cd(self, dir: str, state: ReplState):
+        if dir == '':
+            state.pg_path = None
+            return
+
+        with pg_path(state) as (host, database):
+            if dir == '..':
+                if database:
+                    database = None
+                else:
+                    host = None
+            else:
+                tks = dir.split('@')
+                if not host:
+                    host = tks[0]
+                else:
+                    database = tks[0]
+
+                if len(tks) > 1:
+                    state.namespace = tks[1]
+
+            if database:
+                state.pg_path = f'{host}/{database}'
+            else:
+                state.pg_path = host
+
+    def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        return {cmd: {d: None for d in self.direct_dirs(cmd, state, default=default)}}
+
+    def direct_dirs(self, cmd: str, state: ReplState, default: dict = {}) -> list[str]:
+        with pg_path(state) as (host, database):
+            if database:
+                return ['..']
+            elif host:
+                return ['..'] + [p for p in pg_database_names(state)]
+            else:
+                return [p for p in PostgresDatabases.host_names(state.namespace)]
+
+    def pwd(self, state: ReplState):
+        words = []
+
+        with pg_path(state) as (host, database):
+            if host:
+                words.append(f'host/{host}')
+            if database:
+                words.append(f'database/{database}')
+
+            return '\t'.join([f'{ReplState.P}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        with pg_path(state) as (_, database):
+            if not database:
+                database = PostgresDatabases.default_db()
+
+            if database:
+                return True, chain.run(f'pg {cmd}', state)
+
+        return False, None
+
+    def enter(self, _: ReplState):
+        wait_log('Inspecting postgres database instances...')
+
+    def show_tables(self, state: ReplState):
+        tabulize(PostgresDatabases.tables(state, default_schema=True), lambda d: d['name'], separator=',')
+
+    def show_table_preview(self, state: ReplState, table: str, rows: int):
+        PostgresDatabases.run_sql(state, f'select * from {table} limit {rows}')
+
+    def bash(self, s0: ReplState, s1: ReplState, args: list[str]):
+        pod, container = PostgresDatabases.pod_and_container(s1.namespace)
+        log2(f'Running on {pod}(container:{container})...')
+
+        return super().bash(s0, s1, args)
+
+    def bash_target_changed(self, s0: ReplState, s1: ReplState):
+        return s0.pg_path != s1.pg_path
+
+    def exec_no_dir(self, command: str, state: ReplState):
+        with postgres(state) as pod:
+            return pod.exec(command, show_out=True)
+
+    def exec_with_dir(self, command: str, session_just_created: bool, state: ReplState):
+        with postgres(state) as pod:
+            return pod.exec(command, show_out=not session_just_created)
+
+    def bash_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        return {cmd: BashCompleter(lambda: [])}

adam/commands/devices/devices.py

@@ -0,0 +1,25 @@
+from adam.commands.devices.device import Device
+from adam.commands.devices.device_app import DeviceApp
+from adam.commands.devices.device_auit_log import DeviceAuditLog
+from adam.commands.devices.device_cass import DeviceCass
+from adam.commands.devices.device_export import DeviceExport
+from adam.commands.devices.device_postgres import DevicePostgres
+from adam.repl_state import ReplState
+
+class Devices:
+    def of(state: ReplState) -> Device:
+        if state.device == ReplState.A:
+            return DeviceApp()
+        elif state.device == ReplState.C:
+            return DeviceCass()
+        elif state.device == ReplState.L:
+            return DeviceAuditLog()
+        elif state.device == ReplState.P:
+            return DevicePostgres()
+        elif state.device == ReplState.X:
+            return DeviceExport()
+
+        return DeviceCass()
+
+    def all():
+        return [DeviceApp(), DeviceCass(), DeviceAuditLog(), DevicePostgres(), DeviceExport()]

adam/commands/{check.py → diag/check.py}

@@ -2,11 +2,13 @@ import click
 
 from adam.checks.check_result import CheckResult
 from adam.checks.check_utils import all_checks, checks_from_csv, run_checks
+from adam.commands import extract_options, validate_args
 from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterOrPodCommandHelper
-from adam.commands.issues import Issues
+from adam.commands.diag.issues import Issues
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log
+from adam.utils import tabulize, log
+from adam.utils_issues import IssuesUtils
 
 class Check(Issues):
     COMMAND = 'check'
@@ -30,36 +32,25 @@ class Check(Issues):
         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
+        with self.validate(args, state) as (args, state):
+            with extract_options(args, ['-s', '--show']) as (args, show_out):
+                with validate_args(args, state, name='check name', msg=lambda: tabulize([check.help() for check in all_checks()], separator=':')) as arg:
+                    checks = checks_from_csv(args[0])
+                    if not checks:
+                        return 'invalid check name'
 
-        args, show = Command.extract_options(args, ['-s', '--show'])
+                    results = run_checks(state.sts, state.namespace, state.pod, checks=checks, show_out=show_out)
 
-        if not args:
-            if state.in_repl:
-                log(lines_to_tabular([check.help() for check in all_checks()], separator=':'))
-            else:
-                log('* Check name is missing.')
-                Command.display_help()
-            return 'arg missing'
+                    issues = CheckResult.collect_issues(results)
+                    IssuesUtils.show_issues(issues, in_repl=state.in_repl)
 
-        checks = checks_from_csv(args[0])
-        if not checks:
-            return 'invalid check name'
-
-        results = run_checks(state.sts, state.namespace, state.pod, checks=checks, show_output=show)
-
-        issues = CheckResult.collect_issues(results)
-        Issues.show_issues(issues, in_repl=state.in_repl)
-
-        return issues if issues else 'no issues found'
+                    return issues if issues else 'no issues found'
 
     def completion(self, _: ReplState):
-        return {Check.COMMAND: {check.name(): None for check in all_checks()}}
+        return {Check.COMMAND: {check.name(): {'-s': None} for check in all_checks()}}
 
     def help(self, _: ReplState):
-        return f'{Check.COMMAND} <check-name>\t run a single check'
+        return f'{Check.COMMAND} <check-name> [-s]\t run a single check'
 
 class CheckCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):

adam/commands/diag/generate_report.py

@@ -0,0 +1,52 @@
+import click
+import json
+
+from adam.checks.check_result import CheckResult
+from adam.checks.check_utils import run_checks
+from adam.commands import extract_options
+from adam.commands.command import Command
+from adam.commands.commands_utils import kaqing_log_file
+from adam.repl_state import ReplState
+from adam.utils import log2
+
+class GenerateReport(Command):
+    COMMAND = 'generate report'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(GenerateReport, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return GenerateReport.COMMAND
+
+    def required(self):
+        return ReplState.NON_L
+
+    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, ['-s', '--show']) as (args, show_out):
+                results = run_checks(state.sts, state.namespace, state.pod, show_out=show_out)
+                output = CheckResult.report(results)
+
+                if state.in_repl:
+                    with kaqing_log_file() as json_file:
+                        json.dump(output, json_file, indent=2)
+                        log2(f'Report stored in {json_file.name}.')
+                else:
+                    click.echo(json.dumps(output, indent=2))
+
+                return output
+
+    def completion(self, state: ReplState):
+        return super().completion(state, {'-s': None})
+
+    def help(self, _: ReplState):
+        return f"{GenerateReport.COMMAND} [-s]\t generate report"

adam/commands/diag/issues.py

@@ -0,0 +1,43 @@
+from adam.checks.check_result import CheckResult
+from adam.checks.check_utils import run_checks
+from adam.commands import extract_options
+from adam.commands.command import Command
+from adam.repl_state import ReplState
+from adam.utils_issues import IssuesUtils
+
+class Issues(Command):
+    COMMAND = 'issues'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(Issues, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return Issues.COMMAND
+
+    def required(self):
+        return ReplState.NON_L
+
+    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, ['-s', '--show']) as (args, show_out):
+                results = run_checks(state.sts, state.namespace, state.pod, show_out=show_out)
+
+                issues = CheckResult.collect_issues(results)
+                IssuesUtils.show_issues(issues, in_repl=state.in_repl)
+
+                return issues if issues else 'issues'
+
+    def completion(self, state: ReplState):
+        return super().completion(state, {'-s': None})
+
+    def help(self, _: ReplState):
+        return f'{Issues.COMMAND} [-s]\t find all issues'

adam/commands/exit.py

@@ -23,10 +23,7 @@ class Exit(Command):
         exit()
 
     def completion(self, state: ReplState):
-        if state.pod:
-            return {Exit.COMMAND: None}
-
-        return {}
+        return {Exit.COMMAND: None}
 
     def help(self, _: ReplState):
         return f'{Exit.COMMAND}\t exit kaqing <Ctrl-D>'