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

adam/commands/__init__.py

@@ -1,4 +1,5 @@
-from adam.commands.command import ExtractAllOptionsHandler, ExtractOptionsHandler, ExtractSequenceOptionsHandler, ExtractTrailingOptionsHandler
+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
 
@@ -15,4 +16,7 @@ def extract_all_options(args: list[str], trailing = None, sequence = None, optio
     return ExtractAllOptionsHandler(args, trailing = trailing, sequence = sequence, options = options)
 
 def extract_sequence(args: list[str], sequence: list[str]):
-    return ExtractSequenceOptionsHandler(args, sequence = sequence)
+    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):
+    return ValidateArgCountHandler(args, state, count_at_least=count_at_least, name=name, msg=msg)

adam/commands/alter_tables.py

@@ -1,4 +1,4 @@
-from adam.commands import extract_options
+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
@@ -29,45 +29,34 @@ class AlterTables(Command):
 
         with self.validate(args, state) as (args, state):
             with extract_options(args, '--include-reaper') as (args, include_reaper):
-                if not args:
-                    if state.in_repl:
-                        log2('Please enter gc grace in seconds. e.g. alter gc-grace-seconds 3600')
-                    else:
-                        log2('* gc grace second is missing.')
-                        log2()
-                        Command.display_help()
-
-                    return 'missing-arg'
-
-                arg_str = ' '.join(args)
-
-                excludes = [e.strip(' \r\n') for e in Config().get(
-                    'cql.alter-tables.excludes',
-                    'system_auth,system_traces,reaper_db,system_distributed,system_views,system,system_schema,system_virtual_schema').split(',')]
-                batching = Config().get('cql.alter-tables.batching', True)
-                tables = get_tables(state, on_any=True)
-                for k, v in tables.items():
-                    if k not in excludes or k == 'reaper_db' and include_reaper:
-                        if batching:
-                            # alter table <table_name> with GC_GRACE_SECONDS = <timeout>;
-                            cql = ';\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 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):
-                                    # 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)
+                                        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}.')
+                            log2(f'{len(v)} tables altered in {k}.')
 
-                # do not continue to cql route
-                return state
+                    # do not continue to cql route
+                    return state
 
     def completion(self, _: ReplState) -> dict[str, any]:
         # auto completion is taken care of by sql completer

adam/commands/cat.py

@@ -1,9 +1,8 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.devices.devices import Devices
 from adam.repl_state import ReplState, RequiredState
 from adam.utils import log2
-from adam.utils_k8s.app_pods import AppPods
-from adam.utils_k8s.statefulsets import StatefulSets
 
 class Cat(Command):
     COMMAND = 'cat'
@@ -28,17 +27,8 @@ class Cat(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if len(args) < 1:
-                if state.in_repl:
-                    log2('File is required.')
-                    log2()
-                else:
-                    log2('* File is missing.')
-                    Command.display_help()
-
-                return 'command-missing'
-
-            return Devices.device(state).cat(cmd, state)
+            with validate_args(args, state, name='file'):
+                return Devices.device(state).cat(cmd, state)
 
     def completion(self, state: ReplState):
         if super().completion(state):

adam/commands/cd.py

@@ -1,3 +1,4 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.devices.device import Device
 from adam.commands.devices.devices import Devices
@@ -26,15 +27,12 @@ class Cd(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state, apply=False) as (args, state):
-            if len(args) < 2:
-                return state
-
-            device: Device = Devices.device(state)
-            arg = args[1]
-            for dir in arg.split('/'):
-                device.cd(dir, state)
+            with validate_args(args, state, name='directory') as arg_str:
+                device: Device = Devices.device(state)
+                for dir in arg_str.split('/'):
+                    device.cd(dir, state)
 
-            return state
+                return state
 
     def completion(self, state: ReplState):
         return Devices.device(state).cd_completion(Cd.COMMAND, state, default = {})

adam/commands/check.py

@@ -2,7 +2,7 @@ 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
+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
@@ -34,25 +34,17 @@ class Check(Issues):
 
         with self.validate(args, state) as (args, state):
             with extract_options(args, ['-s', '--show']) as (args, 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()
+                with validate_args(args, state, name='check name', msg=lambda:log(lines_to_tabular([check.help() for check in all_checks()], separator=':'))) as arg:
+                    checks = checks_from_csv(args[0])
+                    if not checks:
+                        return 'invalid check name'
 
-                    return 'arg missing'
+                    results = run_checks(state.sts, state.namespace, state.pod, checks=checks, show_out=show_out)
 
-                checks = checks_from_csv(args[0])
-                if not checks:
-                    return 'invalid check name'
+                    issues = CheckResult.collect_issues(results)
+                    IssuesUtils.show_issues(issues, in_repl=state.in_repl)
 
-                results = run_checks(state.sts, state.namespace, state.pod, checks=checks, show_out=show_out)
-
-                issues = CheckResult.collect_issues(results)
-                IssuesUtils.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(): {'-s': None} for check in all_checks()}}

adam/commands/command.py

@@ -1,11 +1,12 @@
 from abc import abstractmethod
+from collections.abc import Callable
 import copy
 import subprocess
 import sys
 from typing import Union
 
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import is_lambda
+from adam.utils import is_lambda, log2
 
 repl_cmds: list['Command'] = []
 
@@ -45,7 +46,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)
@@ -148,11 +149,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
@@ -164,15 +165,48 @@ 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, count_at_least: int = 1, name: str = None, msg: Callable[[], None] = None):
+        self.args = args
+        self.state = state
+        self.count_at_least = count_at_least
+        self.name = name
+        self.msg = msg
+
+    def __enter__(self) -> tuple[list[str], ReplState]:
+        if len(self.args) < self.count_at_least:
+            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()
+
+        return ' '.join(self.args)
+
+    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

adam/commands/cp.py

@@ -1,7 +1,8 @@
 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.repl_state import ReplState, RequiredState
@@ -30,37 +31,28 @@ class ClipboardCopy(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if len(args) < 1:
-                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('* 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'
+            def display_keys():
+                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=','))
+
+            with validate_args(args, state, name='key', msg=display_keys) as key:
+                if not key in CliCommands.values(state):
+                    if state.in_repl:
+                        display_keys()
+                    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()})

adam/commands/cql/cql_completions.py

@@ -1,6 +1,5 @@
-import time
 from adam.commands.cql.utils_cql import cassandra_keyspaces, cassandra_table_names
-from adam.commands.export.exporter import Exporter
+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
@@ -23,8 +22,8 @@ def cql_completions(state: ReplState) -> dict[str, any]:
             'GC_GRACE_SECONDS': ps
         },
         'export-dbs': lambda: ExportDatabases.database_names(),
-        'export-sessions': lambda: Exporter.export_session_names(state.sts, state.pod, state.namespace),
-        'export-sessions-incomplete': lambda: Exporter.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import'),
+        'export-sessions': lambda: ExportSessions.export_session_names(state.sts, state.pod, state.namespace),
+        'export-sessions-incomplete': lambda: ExportSessions.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import'),
     }
 
     return SqlCompleter(

adam/commands/cql/utils_cql.py

@@ -242,8 +242,6 @@ class CqlShHandler:
             for index, arg in enumerate(args):
                 if arg.startswith('--'):
                     opts.append(arg)
-                # elif index == len(args) -1 and arg == '&':
-                #     background = True
                 elif arg != '-e':
                     cqls.append(arg)
             if not cqls:

adam/commands/devices/device_app.py

@@ -1,7 +1,7 @@
 from adam.apps import Apps
 from adam.commands import app
 from adam.commands.bash.bash_completer import BashCompleter
-from adam.commands.command import Command, InvalidState
+from adam.commands.command import Command, InvalidStateException
 from adam.commands.devices.device import Device
 from adam.config import Config
 from adam.repl_state import ReplState
@@ -78,7 +78,7 @@ class DeviceApp(Command, Device):
         if state.app_pod:
             return self.bash(state, state, cmd.split(' '))
 
-        return InvalidState()
+        raise InvalidStateException()
 
     def cat_completion(self, cmd: str, state: ReplState, default: dict = {}):
         if state.app_app:

adam/commands/devices/device_cass.py

@@ -1,5 +1,5 @@
 from adam.commands.bash.bash_completer import BashCompleter
-from adam.commands.command import Command, InvalidState
+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
@@ -62,7 +62,7 @@ class DeviceCass(Command, Device):
         if state.pod:
             return self.bash(state, state, cmd.split(' '))
 
-        return InvalidState()
+        raise InvalidStateException()
 
     def cat_completion(self, cmd: str, state: ReplState, default: dict = {}):
         if state.sts:

adam/commands/devices/device_export.py

@@ -1,6 +1,6 @@
 from adam.commands.command import Command
 from adam.commands.devices.device import Device
-from adam.commands.export.export_databases import ExportDatabases
+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 lines_to_tabular, log, log2, wait_log
@@ -36,16 +36,14 @@ class DeviceExport(Command, Device):
 
     def ls(self, cmd: str, state: ReplState):
         if state.export_session:
-            self.show_export_tables(state.export_session)
+            log(lines_to_tabular(ExportDatabases.table_names(state.export_session), 'NAME', separator=','))
         else:
-            self.show_export_databases()
+            ExportDatabases.show_databases()
 
-    def show_export_databases(self, importer: str = None):
-        lines = [f'{k}\t{v}' for k, v in ExportDatabases.database_names_with_keyspace_cnt(importer).items()]
-        log(lines_to_tabular(lines, 'NAME\tKEYSPACES', separator='\t'))
-
-    def show_export_tables(self, export_session: str):
-        log(lines_to_tabular(ExportDatabases.table_names(export_session), 'NAME', separator=','))
+    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 ['', '..']:

adam/commands/devices/device_postgres.py

@@ -1,5 +1,5 @@
 from adam.commands.bash.bash_completer import BashCompleter
-from adam.commands.command import Command, InvalidState
+from adam.commands.command import Command, InvalidStateException
 from adam.commands.devices.device import Device
 from adam.commands.postgres.postgres_context import PostgresContext
 from adam.commands.postgres.utils_postgres import pg_database_names, pg_table_names, postgres
@@ -50,7 +50,7 @@ class DevicePostgres(Command, Device):
         if state.pod:
             return self.bash(state, state, cmd.split(' '))
 
-        return InvalidState()
+        raise InvalidStateException()
 
     def show_pg_hosts(self, state: ReplState):
         if state.namespace:

adam/commands/export/clean_up_all_export_sessions.py

@@ -1,5 +1,5 @@
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_sessions import ExportSessions
 from adam.repl_state import ReplState, RequiredState
 
 class CleanUpAllExportSessions(Command):
@@ -25,8 +25,8 @@ class CleanUpAllExportSessions(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if Exporter.clean_up_all_sessions(state.sts, state.pod, state.namespace):
-                Exporter.clear_export_session_cache()
+            if ExportSessions.clean_up_all_sessions(state.sts, state.pod, state.namespace):
+                ExportSessions.clear_export_session_cache()
 
             return state
 

adam/commands/export/clean_up_export_sessions.py

@@ -1,4 +1,6 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
+from adam.commands.export.export_sessions import ExportSessions
 from adam.commands.export.exporter import Exporter
 from adam.repl_state import ReplState, RequiredState
 from adam.utils import log, log2
@@ -26,23 +28,14 @@ class CleanUpExportSessions(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not args:
-                if state.in_repl:
-                    log2('Specify export session name.')
-                else:
-                    log2('* Session name is missing.')
+            with validate_args(args, state, name='export session') as arg_str:
+                sessions = [arg.strip(' ') for arg in arg_str.split(',')]
+                csv_cnt, log_cnt = ExportSessions.clean_up_sessions(state.sts, state.pod, state.namespace, sessions)
+                log(f'Removed {csv_cnt} csv and {log_cnt} log files.')
 
-                    Command.display_help()
+                ExportSessions.clear_export_session_cache()
 
-                return 'command-missing'
-
-            sessions = [arg.strip(' ') for arg in ' '.join(args).split(',')]
-            csv_cnt, log_cnt = Exporter.clean_up_sessions(state.sts, state.pod, state.namespace, sessions)
-            log(f'Removed {csv_cnt} csv and {log_cnt} log files.')
-
-            Exporter.clear_export_session_cache()
-
-            return state
+                return state
 
     def completion(self, _: ReplState):
         return {}

adam/commands/export/drop_export_database.py

@@ -1,9 +1,7 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
-from adam.utils import log2
-from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class DropExportDatabase(Command):
     COMMAND = 'drop export database'
@@ -28,25 +26,11 @@ class DropExportDatabase(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not len(args):
-                if state.in_repl:
-                    log2('Database name is required.')
-                    log2()
-                else:
-                    log2('* Database name is missing.')
-                    Command.display_help()
+            with validate_args(args, state, name='database name') as db:
+                with export_db(state) as dbs:
+                    dbs.drop(args[0])
 
-                return 'command-missing'
-
-            Exporter.drop_databases(state.sts, state.pod, state.namespace, args[0])
-
-            SQLite.clear_cache()
-            Athena.clear_cache()
-
-            if state.export_session == args[0]:
-                state.export_session = None
-
-            return state
+                return state
 
     def completion(self, _: ReplState):
         return {}

adam/commands/export/drop_export_databases.py

@@ -1,8 +1,6 @@
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
-from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class DropExportDatabases(Command):
     COMMAND = 'drop all export databases'
@@ -27,12 +25,8 @@ class DropExportDatabases(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            Exporter.drop_databases(state.sts, state.pod, state.namespace)
-
-            SQLite.clear_cache()
-            Athena.clear_cache()
-
-            state.export_session = None
+            with export_db(state) as dbs:
+                dbs.drop_all()
 
             return state
 

adam/commands/export/export.py

@@ -2,7 +2,7 @@ from adam.commands import extract_options
 from adam.commands.command import Command
 from adam.commands.cql.utils_cql import cassandra_keyspaces, cassandra_table_names
 from adam.commands.export.export_databases import ExportDatabases
-from adam.commands.export.export_handlers import export
+from adam.commands.export.exporter import export
 from adam.repl_state import ReplState, RequiredState
 from adam.sql.sql_completer import SqlCompleter, SqlVariant
 from adam.utils import log