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

adam/repl_commands.py

@@ -1,6 +1,9 @@
 from adam.commands.alter_tables import AlterTables
 from adam.commands.app import App
 from adam.commands.app_ping import AppPing
+from adam.commands.audit.audit import Audit
+from adam.commands.audit.audit_repair_tables import AuditRepairTables
+from adam.commands.audit.audit_run import AuditRun
 from adam.commands.deploy.code_start import CodeStart
 from adam.commands.deploy.code_stop import CodeStop
 from adam.commands.deploy.deploy import Deploy
@@ -11,11 +14,6 @@ from adam.commands.deploy.undeploy import Undeploy
 from adam.commands.deploy.undeploy_frontend import UndeployFrontend
 from adam.commands.deploy.undeploy_pg_agent import UndeployPgAgent
 from adam.commands.deploy.undeploy_pod import UndeployPod
-from adam.commands.describe.describe import Describe
-from adam.commands.describe.describe_keyspace import DescribeKeyspace
-from adam.commands.describe.describe_keyspaces import DescribeKeyspaces
-from adam.commands.describe.describe_table import DescribeTable
-from adam.commands.describe.describe_tables import DescribeTables
 from adam.commands.shell import Shell
 from adam.commands.show.show_app_queues import ShowAppQueues
 from adam.commands.cp import ClipboardCopy
@@ -24,7 +22,7 @@ from adam.commands.cd import Cd
 from adam.commands.check import Check
 from adam.commands.command import Command
 from adam.commands.cql.cqlsh import Cqlsh
-from adam.commands.devices import DeviceApp, DeviceCass, DevicePostgres
+from adam.commands.devices import DeviceApp, DeviceAuditLog, DeviceCass, DevicePostgres
 from adam.commands.exit import Exit
 from adam.commands.medusa.medusa import Medusa
 from adam.commands.param_get import GetParam
@@ -46,6 +44,7 @@ from adam.commands.show.show_app_id import ShowAppId
 from adam.commands.show.show_cassandra_status import ShowCassandraStatus
 from adam.commands.show.show_cassandra_version import ShowCassandraVersion
 from adam.commands.show.show_commands import ShowKubectlCommands
+from adam.commands.show.show_host import ShowHost
 from adam.commands.show.show_login import ShowLogin
 from adam.commands.show.show_params import ShowParams
 from adam.commands.show.show_processes import ShowProcesses
@@ -56,10 +55,10 @@ from adam.commands.watch import Watch
 
 class ReplCommands:
     def repl_cmd_list() -> list[Command]:
-        cmds: list[Command] = ReplCommands.navigation() + ReplCommands.cassandra_check() + ReplCommands.cassandra_ops() + \
-            ReplCommands.tools() + ReplCommands.app() + ReplCommands.exit()
+        cmds: list[Command] = ReplCommands.navigation() + ReplCommands.cassandra_ops() + ReplCommands.postgres_ops() + \
+            ReplCommands.app_ops() + ReplCommands.audit_ops() + ReplCommands.tools() + ReplCommands.exit()
 
-        intermediate_cmds: list[Command] = [App(), Reaper(), Repair(), Deploy(), Describe(), Show(), Undeploy()]
+        intermediate_cmds: list[Command] = [App(), Audit(), Reaper(), Repair(), Deploy(), Show(), Undeploy()]
         ic = [c.command() for c in intermediate_cmds]
         # 1. dedup commands
         deduped = []
@@ -76,22 +75,24 @@ class ReplCommands:
         return deduped
 
     def navigation() -> list[Command]:
-        return [Ls(), PreviewTable(), DeviceApp(), DevicePostgres(), DeviceCass(), Cd(), Pwd(), ClipboardCopy(),
-                GetParam(), SetParam(), ShowParams(), ShowKubectlCommands(), ShowLogin(), ShowAdam()]
-
-    def cassandra_check() -> list[Command]:
-        return Describe.cmd_list() + [ShowCassandraStatus(),
-                ShowCassandraVersion(), ShowRepairs(), ShowStorage(), ShowProcesses(), Check(), Issues(), NodeTool(), Report()]
+        return [Ls(), PreviewTable(), DeviceApp(), DevicePostgres(), DeviceCass(), DeviceAuditLog(), Cd(), Pwd(), ClipboardCopy(),
+                GetParam(), SetParam(), ShowParams(), ShowKubectlCommands(), ShowLogin(), ShowAdam(), ShowHost()]
 
     def cassandra_ops() -> list[Command]:
-        return [AlterTables()] + Medusa.cmd_list() + [Restart(), RollOut(), Watch()] + Reaper.cmd_list() + Repair.cmd_list()
+        return [Cqlsh(), ShowCassandraStatus(), ShowCassandraVersion(), ShowRepairs(), ShowStorage(), ShowProcesses(), Check(), Issues(), NodeTool(), Report()] + \
+               [AlterTables()] + Medusa.cmd_list() + [Restart(), RollOut(), Watch()] + Reaper.cmd_list() + Repair.cmd_list()
 
-    def tools() -> list[Command]:
-        return [Cqlsh(), Postgres(), Bash(), Shell(), CodeStart(), CodeStop(), DeployFrontend(), UndeployFrontend(),
-                DeployPod(), UndeployPod(), DeployPgAgent(), UndeployPgAgent()]
+    def postgres_ops() -> list[Command]:
+        return [Postgres(), DeployPgAgent(), UndeployPgAgent()]
 
-    def app() -> list[Command]:
+    def app_ops() -> list[Command]:
         return [ShowAppActions(), ShowAppId(), ShowAppQueues(), AppPing(), App()]
 
+    def audit_ops() -> list[Command]:
+        return [Audit()] + Audit.cmd_list()
+
+    def tools() -> list[Command]:
+        return [Bash(), Shell(), CodeStart(), CodeStop(), DeployFrontend(), UndeployFrontend(), DeployPod(), UndeployPod()]
+
     def exit() -> list[Command]:
         return [Exit()]

adam/repl_state.py

@@ -1,11 +1,14 @@
 import copy
 from enum import Enum
-
-from adam.commands.postgres.postgres_session import PostgresSession
-from adam.k8s_utils.cassandra_clusters import CassandraClusters
-from adam.k8s_utils.cassandra_nodes import CassandraNodes
-from adam.k8s_utils.kube_context import KubeContext
-from adam.k8s_utils.secrets import Secrets
+import re
+import traceback
+from typing import Callable
+
+from adam.commands.postgres.postgres_context import PostgresContext
+from adam.utils_k8s.cassandra_clusters import CassandraClusters
+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 import display_help, log2, random_alphanumeric
 
 class BashSession:
@@ -47,8 +50,12 @@ class RequiredState(Enum):
 class ReplState:
     A = 'a'
     C = 'c'
+    L = 'l'
     P = 'p'
 
+    ANY = [A, C, L, P]
+    NON_L = [A, C, P]
+
     def __init__(self, device: str = None,
                  sts: str = None, pod: str = None, namespace: str = None, ns_sts: str = None,
                  pg_path: str = None,
@@ -82,12 +89,12 @@ class ReplState:
     def __hash__(self):
         return hash((self.sts, self.pod))
 
-    def apply_args(self, args: list[str], cmd: list[str] = None, resolve_pg = True) -> tuple['ReplState', list[str]]:
+    def apply_args(self, args: list[str], cmd: list[str] = None, resolve_pg = True, args_to_check = 6) -> tuple['ReplState', list[str]]:
         state = self
 
         new_args = []
         for index, arg in enumerate(args):
-            if index < 6:
+            if index < args_to_check:
                 state = copy.copy(state)
 
                 s, n = KubeContext.is_sts_name(arg)
@@ -128,8 +135,32 @@ class ReplState:
             if index < 6:
                 state = copy.copy(state)
 
-                if arg in [f'{ReplState.A}:', f'{ReplState.C}:', f'{ReplState.P}:']:
-                    state.device = arg.strip(':')
+                groups = re.match(r'^([a|c|l|p]):(.*)$', arg)
+                if groups:
+                    if groups[1] == 'p':
+                        state.device = 'p'
+                        state.pg_path = groups[2]
+                    elif groups[1] == 'c':
+                        state.device = 'c'
+                        if path := groups[2]:
+                            p_and_ns = path.split('@')
+                            sts_and_pod = p_and_ns[0].split('/')
+                            state.sts = sts_and_pod[0]
+                            if len(sts_and_pod) > 1:
+                                state.pod = sts_and_pod[1]
+                            if len(p_and_ns) > 1:
+                                state.namespace = p_and_ns[1]
+                            elif ns := KubeContext.in_cluster_namespace():
+                                state.namespace = ns
+                    elif groups[1] == 'l':
+                        state.device = 'l'
+                    else:
+                        state.device = 'a'
+                        if path := groups[2]:
+                            env_and_app = path.split('/')
+                            state.app_env = env_and_app[0]
+                            if len(env_and_app) > 1:
+                                state.app_app = env_and_app[1]
                 else:
                     new_args.append(arg)
             else:
@@ -140,72 +171,169 @@ class ReplState:
 
         return (state, new_args)
 
-    def validate(self, required: RequiredState = None, pg_required: RequiredState = None, app_required: RequiredState = None):
-        if not pg_required and not app_required:
-            if required == RequiredState.CLUSTER:
-                if not self.namespace or not self.sts:
+    def validate(self, required: list[RequiredState] = [], show_err = True):
+        if not required:
+            return True
+
+        def default_err():
+            if self.in_repl:
+                log2(f'Not a valid command on {self.device}: drive.')
+            else:
+                log2('* on a wrong device.')
+                log2()
+                display_help()
+
+        if type(required) is not list:
+            valid, err = self._validate(required)
+            if valid:
+                return True
+
+            if show_err:
+                if err:
+                    err()
+                else:
+                    default_err()
+
+            return False
+
+        devices = [r for r in required if r in [ReplState.L, ReplState.A, ReplState.C, ReplState.P]]
+        non_devices = [r for r in required if r not in [ReplState.L, ReplState.A, ReplState.C, ReplState.P]]
+
+        first_error: Callable = None
+        for r in non_devices:
+            valid, err = self._validate(r)
+            if valid:
+                return True
+
+            if not first_error:
+                first_error = err
+
+        if devices:
+            valid, err = self._validate_device(devices)
+            if valid:
+                return True
+
+            if not first_error:
+                first_error = err
+
+        if show_err and first_error:
+            if first_error:
+                first_error()
+            else:
+                default_err()
+
+        return False
+
+    def _validate(self, required: RequiredState):
+        if required == RequiredState.CLUSTER:
+            if self.device != ReplState.C:
+                return (False, None)
+
+            if not self.namespace or not self.sts:
+                def error():
                     if self.in_repl:
                         log2('cd to a Cassandra cluster first.')
                     else:
                         log2('* Cassandra cluster is missing.')
                         log2()
                         display_help()
+                return (False, error)
 
-                    return False
-            elif required == RequiredState.POD:
-                if not self.namespace or not self.pod:
+        elif required == RequiredState.POD:
+            if self.device != ReplState.C:
+                return (False, None)
+
+            if not self.namespace or not self.pod:
+                def error():
                     if self.in_repl:
                         log2('cd to a pod first.')
                     else:
                         log2('* Pod is missing.')
                         log2()
                         display_help()
+                return (False, error)
+
+        elif required == RequiredState.CLUSTER_OR_POD:
+            if self.device != ReplState.C:
+                return (False, None)
 
-                    return False
-            elif required == RequiredState.CLUSTER_OR_POD:
-                if not self.namespace or not self.sts and not self.pod:
+            if not self.namespace or not self.sts and not self.pod:
+                def error():
                     if self.in_repl:
                         log2('cd to a Cassandra cluster first.')
                     else:
                         log2('* Cassandra cluster or pod is missing.')
                         log2()
                         display_help()
+                return (False, error)
+
+        elif required == RequiredState.NAMESPACE:
+            if self.device != ReplState.C:
+                return (False, None)
 
-                    return False
-            elif required == RequiredState.NAMESPACE:
-                if not self.namespace:
+            if not self.namespace:
+                def error():
                     if self.in_repl:
                         log2('Namespace is required.')
                     else:
                         log2('* namespace is missing.')
                         log2()
                         display_help()
+                return (False, error)
 
-                    return False
+        elif required == RequiredState.PG_DATABASE:
+            if self.device != ReplState.P:
+                return (False, None)
 
-        if pg_required == RequiredState.PG_DATABASE:
-            pg = PostgresSession(self.namespace, self.pg_path)
+            pg: PostgresContext = PostgresContext.apply(self.namespace, self.pg_path)
             if not pg.db:
+                def error():
+                    if self.in_repl:
+                        log2('cd to a database first.')
+                    else:
+                        log2('* database is missing.')
+                        log2()
+                        display_help()
+                return (False, error)
+
+        elif required == RequiredState.APP_APP:
+            if self.device != ReplState.A:
+                return (False, None)
+
+            if not self.app_app:
+                def error():
+                    if self.in_repl:
+                        log2('cd to an app first.')
+                    else:
+                        log2('* app is missing.')
+                        log2()
+                        display_help()
+                return (False, error)
+
+        elif required in [ReplState.L, ReplState.A, ReplState.C, ReplState.P] and self.device != required:
+            def error():
                 if self.in_repl:
-                    log2('cd to a database first.')
+                    log2(f'Switch to {required}: first.')
                 else:
-                    log2('* database is missing.')
+                    log2('* on a wrong device.')
                     log2()
                     display_help()
+            return (False, error)
 
-                return False
-
-        if app_required == RequiredState.APP_APP and not self.app_app:
-            if self.in_repl:
-                log2('cd to an app first.')
-            else:
-                log2('* app is missing.')
-                log2()
-                display_help()
+        return (True, None)
 
-            return False
+    def _validate_device(self, devices: list[RequiredState]):
+        if self.device not in devices:
+            def error():
+                if self.in_repl:
+                    log2(f'Not a valid command on {self.device}: drive.')
+                else:
+                    log2('* on a wrong device.')
+                    log2()
+                    display_help()
+            return (False, error)
 
-        return True
+        return (True, None)
 
     def user_pass(self, secret_path = 'cql.secret'):
         return Secrets.get_user_pass(self.pod if self.pod else self.sts, self.namespace, secret_path=secret_path)

adam/sql/sql_completer.py

@@ -1,53 +1,93 @@
-from typing import Dict, Iterable, Optional
-from prompt_toolkit.completion import CompleteEvent, Completer, Completion, NestedCompleter, WordCompleter
-from prompt_toolkit.document import Document
-
-from adam.sql.any_completer import AnyCompleter as any
-from adam.sql.sql_utils import safe_terms
-from adam.sql.table_name_completer import TableNameCompleter
-
-class SqlCompleter(NestedCompleter):
-    def __init__(
-        self, options: Dict[str, Optional[Completer]], ignore_case: bool = True
-    ) -> None:
-        super().__init__(options, ignore_case)
-
-    def get_completions(
-        self, document: Document, complete_event: CompleteEvent
-    ) -> Iterable[Completion]:
-        text = document.text_before_cursor.lstrip()
-        stripped_len = len(document.text_before_cursor) - len(text)
-
-        terms, has_space = safe_terms(text)
-        if has_space:
-            first_term = terms[0]
-            completer = self.options.get(first_term)
-
-            if completer is not None:
-                remaining_text = text[len(first_term) :].lstrip()
-                move_cursor = len(text) - len(remaining_text) + stripped_len
-
-                new_document = Document(
-                    remaining_text,
-                    cursor_position=document.cursor_position - move_cursor,
-                )
-
-                for c in completer.get_completions(new_document, complete_event):
-                    yield c
+from typing import Callable
+
+import sqlparse
+from sqlparse.sql import Statement, Token
+
+from adam.sql.term_completer import TermCompleter
+from adam.utils_repl.automata_completer import AutomataCompleter
+from adam.sql.sql_state_machine import AthenaStateMachine, CqlStateMachine, SqlStateMachine
+from adam.utils_repl.state_machine import State
+
+__all__ = [
+    "SqlCompleter",
+]
+
+def default_columns(_: list[str]):
+    return 'id,x.,y.,z.'.split(',')
+
+class SqlCompleter(AutomataCompleter[Token]):
+    def tokens(self, text: str) -> list[Token]:
+        tokens = []
+
+        stmts = sqlparse.parse(text)
+        if not stmts:
+            tokens = []
+        else:
+            statement: Statement = stmts[0]
+            tokens = statement.tokens
+
+        return tokens
+
+    def __init__(self,
+                 tables: Callable[[], list[str]],
+                 dml: str = None,
+                 columns: Callable[[list[str]], list[str]] = default_columns,
+                 partition_columns: Callable[[list[str]], list[str]] = lambda x: [],
+                 table_props: Callable[[], dict[str,list[str]]] = lambda: [],
+                 variant = 'sql',
+                 debug = False):
+        machine = SqlStateMachine(debug=debug)
+        if variant == 'cql':
+            machine = CqlStateMachine(debug=debug)
+        elif variant == 'athena':
+            machine = AthenaStateMachine(debug=debug)
+        super().__init__(machine, dml, debug)
+
+        self.tables = tables
+        self.columns = columns
+        self.partition_columns = partition_columns
+        self.table_props = table_props
+        self.variant = variant
+        self.debug = debug
+
+    def suggestions_completer(self, state: State, suggestions: str) -> list[str]:
+        if not suggestions:
+            return None
+
+        terms = []
+        for suggestion in suggestions.split(','):
+            terms.extend(self._terms(state, suggestion))
+
+        return TermCompleter(terms)
+
+    def _terms(self, state: State, word: str) -> list[str]:
+        terms = []
+
+        if word == 'tables':
+            terms.extend(self.tables())
+        elif word == '`tables`':
+            terms.append('tables')
+        elif word == 'columns':
+            terms.extend(self.columns([]))
+        elif word == 'partition-columns':
+            terms.extend(self.partition_columns([]))
+        elif word == 'table-props':
+            terms.extend(self.table_props().keys())
+        elif word == 'table-prop-values':
+            if 'last_name' in state.context and state.context['last_name']:
+                terms.extend(self.table_props()[state.context['last_name']])
+        elif word == 'single':
+            terms.append("'")
+        elif word == 'comma':
+            terms.append(",")
         else:
-            completer = WordCompleter(
-                list(self.options.keys()), ignore_case=self.ignore_case
-            )
-            for c in completer.get_completions(document, complete_event):
-                yield c
+            terms.append(word)
+
+        return terms
 
-    def completions(table: TableNameCompleter):
+    def completions_for_nesting(self):
         return {
-            'delete': {'from': table.nested({'where': any('id').nested({'=': any("'id'")})})},
-            'insert': {'into': table.nested({'values(': None})},
-            'select': any('*').nested({'from': table.nested({
-                'limit': any('1'),
-                'where': any('id').nested({'=': any("'id'").nested({'limit': any('1')})})
-            })}),
-            'update': table.nested({'set': {'column': {'=': None}}}),
+            word : SqlCompleter(self.tables, word, columns=self.columns, partition_columns=self.partition_columns,
+                                table_props=self.table_props, variant=self.variant)
+            for word in self.machine.suggestions[''].strip(' ').split(',')
         }