kaqing 2.0.171__py3-none-any.whl → 2.0.203__py3-none-any.whl

adam/commands/{cp.py → clipboard_copy.py}

@@ -1,14 +1,17 @@
+from functools import partial
 import click
 import pyperclip
 
-from adam.commands.command import Command
+from adam.commands import validate_args
+from adam.commands.command import Command, InvalidArgumentsException
 from adam.commands.command_helpers import ClusterOrPodCommandHelper
 from adam.commands.cli_commands import CliCommands
+from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import lines_to_tabular, log, log2
+from adam.utils import tabulize, log, log2
 
 class ClipboardCopy(Command):
-    COMMAND = 'cp'
+    COMMAND = 'cli cp'
 
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
@@ -30,40 +33,37 @@ class ClipboardCopy(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if len(args) < 1:
-                if state.in_repl:
+            def display_keys(missing = False):
+                if missing:
                     log2('Key is required.')
-                    log2()
-                    log2('Keys:')
-                    log2(lines_to_tabular([f'{k},{v}' for k, v in CliCommands.values(state, collapse=True).items()], separator=','))
                 else:
-                    log2('* Key is missing.')
-                    Command.display_help()
-
-                return 'command-missing'
-
-            key = args[0]
-            if not key in CliCommands.values(state):
-                if state.in_repl:
-                    log2('Key is required.')
-                    log2()
-                    log2('Keys:')
-                    log2(lines_to_tabular([f'{k},{v}' for k, v in CliCommands.values(state, collapse=True).items()], separator=','))
-                else:
-                    log2('* Invalid key')
-                    Command.display_help()
-
-                return 'command-invalid'
-
-            value = CliCommands.values(state)[key]
-            pyperclip.copy(value)
-            log2('The following line has been copied to clipboard. Use <Ctrl-V> to use it.')
-            log2(f'  {value}')
-
-            return 'value-copied'
+                    log2('Key is invalid.')
+                log2()
+                tabulize(CliCommands.values(state, collapse=True).items(),
+                         lambda a: f'{a[0]},{a[1]}',
+                         header='KEY,VALUE',
+                         separator=',',
+                         to=2)
+
+            with validate_args(args, state, name='key', msg=partial(display_keys, True)) as key:
+                if not key in CliCommands.values(state):
+                    if state.in_repl:
+                        display_keys(False)
+                    else:
+                        log2('* Invalid key')
+                        Command.display_help()
+
+                    raise InvalidArgumentsException()
+
+                value = CliCommands.values(state)[key]
+                pyperclip.copy(value)
+                log2('The following line has been copied to clipboard. Use <Ctrl-V> to use it.')
+                log2(f'  {value}')
+
+                return state
 
     def completion(self, state: ReplState):
-        return super().completion(state, lambda: {key: None for key in CliCommands.values(state).keys()})
+        return super().completion(state, lambda: {key: None for key in CliCommands.values(state).keys()}, auto_key='cli.cp-auto-complete')
 
     def help(self, _: ReplState):
         return f"{ClipboardCopy.COMMAND} <key>\t copy a value to clipboard for conveninence"
@@ -82,8 +82,6 @@ class CopyCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):
         log(super().get_help(ctx))
         log()
-        log('Keys:')
-
-        log(lines_to_tabular(self.lines(), separator=':'))
+        tabulize(self.lines(), header='KEY:DESC', separator=':')
         log()
         ClusterOrPodCommandHelper.cluter_or_pod_help()

adam/commands/code.py

@@ -4,7 +4,7 @@ from functools import partial
 from adam.commands.command import Command
 from adam.commands.cql.utils_cql import cassandra_table_names, run_cql
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import lines_to_tabular
+from adam.utils import tabulize
 from adam.utils_k8s.statefulsets import StatefulSets
 
 class Code(Command):
@@ -46,7 +46,7 @@ class Code(Command):
                 "cql('select...'): run cql statement"
             ]
 
-            code.interact(local=my_local, banner=lines_to_tabular(lines, header='Variables', separator=':'))
+            code.interact(local=my_local, banner=tabulize(lines, header='Variables', separator=':'))
 
             return state
 

adam/commands/command.py

@@ -1,11 +1,14 @@
 from abc import abstractmethod
+from collections.abc import Callable
 import copy
 import subprocess
 import sys
 from typing import Union
 
+from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import is_lambda
+from adam.sql.lark_completer import LarkCompleter
+from adam.utils import log2
 
 repl_cmds: list['Command'] = []
 
@@ -28,12 +31,38 @@ class Command:
 
         return None
 
-    def completion(self, state: ReplState, leaf: dict[str, any] = None) -> dict[str, any]:
-        if not self.validate_state(state, show_err=False):
+    def completion(self, state: ReplState, leaf: dict[str, any] = None, pods: tuple[list[str], str] = None, auto: str = 'on', auto_key: str = None) -> dict[str, any]:
+        # pods is a tuple of list of pod names and the current pod repl is on
+        if not pods:
+            return self._completion(state, leaf, auto=auto, auto_key=auto_key)
+
+        c = {}
+
+        pod_names = pods[0]
+        pod = pods[1]
+
+        if pod:
+            c |= self._completion(state, leaf, auto=auto, auto_key=auto_key)
+
+        c |= {f'@{p}': self._completion(state.with_pod(p), leaf, to_validate=False, auto=auto, auto_key=auto_key) for p in pod_names if p != pod}
+
+        return c
+
+    def _completion(self, state: ReplState, leaf: dict[str, any] = None, to_validate = True, auto: str = 'on', auto_key: str = None) -> dict[str, any]:
+        if to_validate and not self.validate_state(state, show_err=False):
             return {}
 
-        if is_lambda(leaf):
-            leaf = leaf()
+        if callable(leaf):
+            if auto_key:
+                auto_key = f'auto-complete.{auto_key}'
+                auto = Config().get(auto_key, 'off')
+
+            if auto == 'on':
+                leaf = leaf()
+            elif auto in ['lazy', 'jit']:
+                leaf = LarkCompleter.from_lambda(auto_key, leaf, auto=auto)
+            else:
+                leaf = None
 
         d = leaf
         for t in reversed(self.command().split(' ')):
@@ -45,7 +74,7 @@ class Command:
         return None
 
     def validate(self, args: list[str], state: ReplState, apply = True):
-        return ValidateHandler(self, args, state, apply = apply)
+        return ValidateStateHandler(self, args, state, apply = apply)
 
     def validate_state(self, state: ReplState, show_err = True):
         return state.validate(self.required(), show_err=show_err)
@@ -96,9 +125,10 @@ class Command:
         args.extend(['--help'])
         subprocess.run(args)
 
-    def extract_options(args: list[str], trailing: Union[str, list[str]] = [], options: list[str] = []):
+    def extract_options(args: list[str], trailing: Union[str, list[str]] = [], sequence: list[str] = [], options: list[str] = []):
         found_options: list[str] = []
         found_trailing = None
+        found_sequence = None
 
         if trailing is None:
             trailing = []
@@ -110,16 +140,37 @@ class Command:
         elif isinstance(options, str):
             options = [options]
 
+        if args and trailing:
+            while args and args[-1] in trailing:
+                found_trailing = True
+                args = args[:-1]
+
+        if args and sequence:
+            args, found_sequence = Command.extract_option_sequence(args, sequence)
+
         new_args: list[str] = []
-        for index, arg in enumerate(args):
-            if index == len(args) - 1 and arg in trailing:
-                found_trailing = arg
-            elif arg in options:
-                found_options.append(arg)
-            else:
-                new_args.append(arg)
+        if args:
+            for arg in args:
+                if arg in options:
+                    found_options.append(arg)
+                else:
+                    new_args.append(arg)
+
+        return new_args, found_trailing, found_sequence, found_options
+
+    def extract_option_sequence(args, sequence):
+        new_args = args
+
+        len_sub = len(sequence)
+        len_main = len(args)
+        for i in range(len_main - len_sub + 1):
+            if args[i:i+len_sub] == sequence:
+                new_args = copy.copy(args)
+                del new_args[i:i+len_sub]
+
+                return new_args, True
 
-        return new_args, found_trailing, found_options
+        return new_args, False
 
     def print_chain(cmd: 'Command'):
         print(f'{cmd.command()}', end = '')
@@ -128,11 +179,11 @@ class Command:
             cmd = s
         print()
 
-class InvalidState(Exception):
+class InvalidStateException(Exception):
     def __init__(self, state: ReplState):
         super().__init__(f'Invalid state')
 
-class ValidateHandler:
+class ValidateStateHandler:
     def __init__(self, cmd: Command, args: list[str], state: ReplState, apply = True):
         self.cmd = cmd
         self.args = args
@@ -144,22 +195,75 @@ class ValidateHandler:
         args = self.args
         if self.apply:
             state, args = self.cmd.apply_state(args, state)
+        else:
+            args = args[len(self.cmd.command_tokens()):]
 
         if not self.cmd.validate_state(state):
-            raise InvalidState(state)
+            raise InvalidStateException(state)
 
         return args, state
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         return False
 
+class InvalidArgumentsException(Exception):
+    def __init__(self):
+        super().__init__(f'Invalid arguments')
+
+class ValidateArgCountHandler:
+    def __init__(self, args: list[str], state: ReplState, at_least: int = 1, exactly: int = 1,
+                 name: str = None, msg: Callable[[], None] = None, default: Union[str, list[str]] = None,
+                 separator: str = None):
+        self.args = args
+        self.state = state
+        self.at_least = at_least
+        self.exactly = exactly
+        self.name = name
+        self.msg = msg
+        self.default = default
+        self.separator = separator
+
+    def __enter__(self) -> Union[tuple[list[str], ReplState], tuple[str, ReplState]]:
+        if self.exactly > 0 and len(self.args) != self.exactly or len(self.args) < self.at_least:
+            if self.default:
+                v = self.default
+                if isinstance(v, list):
+                    v = ' '.join(v)
+
+                return v
+
+            if self.state.in_repl:
+                if self.msg:
+                    self.msg()
+                elif self.name:
+                    log2(f'{self.name} is required.')
+            elif self.name:
+                log2(f'* {self.name} is missing.')
+
+                Command.display_help()
+
+            raise InvalidArgumentsException()
+
+        if self.separator == ' ':
+            return self.args
+
+        # join and re-split with separator
+        args_to_return = ' '.join(self.args)
+        if self.separator:
+            args_to_return = [arg.strip(' ') for arg in args_to_return.split(self.separator)]
+
+        return args_to_return
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
 class ExtractOptionsHandler:
     def __init__(self, args: list[str], options: list[str] = None):
         self.args = args
         self.options = options
 
     def __enter__(self) -> tuple[list[str], list[str]]:
-        args, _, options = Command.extract_options(self.args, options=self.options)
+        args, _, _, options = Command.extract_options(self.args, options=self.options)
         return args, options
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -171,20 +275,33 @@ class ExtractTrailingOptionsHandler:
         self.trailing = trailing
 
     def __enter__(self) -> tuple[list[str], list[str]]:
-        args, trailing, _ = Command.extract_options(self.args, trailing=self.trailing)
+        args, trailing, _, _ = Command.extract_options(self.args, trailing=self.trailing)
         return args, trailing
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         return False
 
 class ExtractAllOptionsHandler:
-    def __init__(self, args: list[str], trailing: list[str] = None, options: list[str] = None):
+    def __init__(self, args: list[str], trailing: list[str] = None, sequence: list[str] = None, options: list[str] = None):
         self.args = args
         self.trailing = trailing
+        self.sequence = sequence
         self.options = options
 
     def __enter__(self) -> tuple[list[str], list[str]]:
-        return Command.extract_options(self.args, options=self.options)
+        return Command.extract_options(self.args, trailing=self.trailing, sequence=self.sequence, options=self.options)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+class ExtractSequenceOptionsHandler:
+    def __init__(self, args: list[str], sequence: list[str]):
+        self.args = args
+        self.sequence = sequence
+
+    def __enter__(self) -> tuple[list[str], list[str]]:
+        args, _, sequence, _ = Command.extract_options(self.args, sequence=self.sequence)
+        return args, sequence
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         return False

adam/commands/commands_utils.py

@@ -8,7 +8,7 @@ from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.pods import Pods
 from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState
-from adam.utils import duration, lines_to_tabular, log, log2
+from adam.utils import SORT, duration, kaqing_log_file, tabulize, log2
 
 def show_pods(pods: List[client.V1Pod], ns: str, show_namespace = True, show_host_id = True):
     if len(pods) == 0:
@@ -48,9 +48,7 @@ def show_pods(pods: List[client.V1Pod], ns: str, show_namespace = True, show_hos
             line += f"@{ns}"
         return line + f" {ready}/{pod_cnt} {status}"
 
-    pod_names = [line(pod) for pod in pods]
-
-    log(lines_to_tabular(pod_names, 'HOST_ID POD_NAME READY POD_STATUS' if show_host_id else 'POD_NAME READY POD_STATUS'))
+    tabulize(pods, line, header='HOST_ID POD_NAME READY POD_STATUS' if show_host_id else 'POD_NAME READY POD_STATUS')
 
 def show_rollout(sts: str, ns: str):
     restarted, rollingout = StatefulSets.restarted_at(sts, ns)
@@ -61,18 +59,22 @@ def show_rollout(sts: str, ns: str):
         else:
             log2(f'Cluster has completed rollout {d} ago.')
 
-def show_table(state: ReplState, pods: list[str], cols: str, header: str, show_out=False):
+def show_table(state: ReplState, pods: list[str], cols: str, header: str, show_out=False, backgrounded = False):
     columns = Columns.create_columns(cols)
 
     results = run_checks(cluster=state.sts, pod=state.pod, namespace=state.namespace, checks=collect_checks(columns), show_out=show_out)
 
-    def line(pod_name: str):
-        cells = [c.pod_value(results, pod_name) for c in columns]
-        return ','.join(cells)
+    r = tabulize(pods, lambda p: ','.join([c.pod_value(results, p) for c in columns]), header=header, separator=',', sorted=SORT, to = 0 if backgrounded else 1)
+
+    if backgrounded:
+        r = write_to_kaqing_log_file(r)
+
+    IssuesUtils.show(results, state.in_repl)
 
-    lines = [line(pod) for pod in pods]
-    lines.sort()
+    return r
 
-    log(lines_to_tabular(lines, header, separator=','))
+def write_to_kaqing_log_file(r: str):
+    with kaqing_log_file() as f:
+        f.write(r)
 
-    IssuesUtils.show(results, state.in_repl)
+        return f.name

adam/commands/cql/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 lark completer
+        return {}
+
+    def help(self, _: ReplState) -> str:
+        return f'{AlterTables.COMMAND} <param = value> [--include-reaper] \t alter schema on all tables'

adam/commands/cql/completions_c.py

@@ -0,0 +1,29 @@
+from adam.commands.cql.utils_cql import cassandra_keyspaces, cassandra_table_names
+from adam.commands.export.export_sessions import ExportSessions
+from adam.commands.export.export_databases import ExportDatabases
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.sql.lark_completer import LarkCompleter
+from adam.utils import log_timing
+from adam.utils_k8s.statefulsets import StatefulSets
+
+def completions_c(state: ReplState) -> dict[str, any]:
+    ps = Config().get('cql.alter-tables.gc-grace-periods', '3600,86400,864000,7776000').split(',')
+
+    with log_timing('lark.completions'):
+        return LarkCompleter(
+            expandables = {
+                'tables': lambda x: cassandra_table_names(state),
+                'keyspaces': lambda x: cassandra_keyspaces(state.with_no_pod()),
+                'table-props': lambda x: {
+                    'GC_GRACE_SECONDS': ps
+                },
+                'table-props-value': lambda x: {'GC_GRACE_SECONDS': ps}[x],
+                'export-database-types': lambda x: ['athena', 'sqlite', 'csv'],
+                'export-databases': lambda x: ExportDatabases.database_names(),
+                'export-sessions': lambda x: ExportSessions.export_session_names(state.sts, state.pod, state.namespace),
+                'export-sessions-incomplete': lambda x: ExportSessions.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import'),
+                'hosts': lambda x: [f'@{p}' for p in StatefulSets.pod_names(state.sts, state.namespace)],
+            },
+            variant=ReplState.C
+        ).completions_for_nesting()

adam/commands/cql/cqlsh.py

@@ -3,9 +3,8 @@ import click
 from adam.commands import extract_trailing_options
 from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterOrPodCommandHelper
-from adam.commands.cql.cql_completions import cql_completions
+from adam.commands.cql.completions_c import completions_c
 from adam.commands.cql.utils_cql import cassandra
-from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState, RequiredState
 from adam.utils import log
 
@@ -34,17 +33,14 @@ class Cqlsh(Command):
         with self.validate(args, state) as (args, state):
             with extract_trailing_options(args, '&') as (args, backgrounded):
                 with cassandra(state) as pods:
-                    pods.cql(args, background=backgrounded)
+                    pods.cql(args, backgrounded=backgrounded)
 
     def completion(self, state: ReplState) -> dict[str, any]:
         if state.device != state.C:
-            # conflicts with psql completions
             return {}
 
         if state.sts or state.pod:
-            c = cql_completions(state)
-            return c | \
-                   {f'@{p}': c for p in StatefulSets.pod_names(state.sts, state.namespace)}
+            return completions_c(state)
 
         return {}