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

adam/repl_state.py

@@ -1,10 +1,10 @@
-import copy
+from copy import copy
 from enum import Enum
 import re
-import traceback
 from typing import Callable
 
-from adam.commands.postgres.postgres_context import PostgresContext
+from adam.utils_k8s.app_clusters import AppClusters
+from adam.utils_k8s.app_pods import AppPods
 from adam.utils_k8s.cassandra_clusters import CassandraClusters
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.kube_context import KubeContext
@@ -19,25 +19,8 @@ class BashSession:
     def pwd(self, state: 'ReplState'):
         command = f'cat /tmp/.qing-{self.session_id}'
 
-        if state.pod:
-            rs = [CassandraNodes.exec(state.pod, state.namespace, command, show_out=False)]
-        elif state.sts:
-            rs = CassandraClusters.exec(state.sts, state.namespace, command, action='bash', show_out=False)
-
-        dir = None
-        for r in rs:
-            if r.exit_code(): # if fails to read the session file, ignore
-                continue
-
-            dir0 = r.stdout.strip(' \r\n')
-            if dir:
-                if dir != dir0:
-                    log2('Inconsitent working dir found across multiple pods.')
-                    return None
-            else:
-                dir = dir0
-
-        return dir
+        with device(state) as pods:
+            return pods.exec(command, action='bash', show_out=False)
 
 class RequiredState(Enum):
     CLUSTER = 'cluster'
@@ -46,20 +29,22 @@ class RequiredState(Enum):
     NAMESPACE = 'namespace'
     PG_DATABASE = 'pg_database'
     APP_APP = 'app_app'
+    EXPORT_DB = 'export_db'
 
 class ReplState:
     A = 'a'
     C = 'c'
     L = 'l'
     P = 'p'
+    X = 'x'
 
-    ANY = [A, C, L, P]
-    NON_L = [A, C, P]
+    ANY = [A, C, L, P, X]
+    NON_L = [A, C, P, X]
 
     def __init__(self, device: str = None,
                  sts: str = None, pod: str = None, namespace: str = None, ns_sts: str = None,
                  pg_path: str = None,
-                 app_env: str = None, app_app: str = None,
+                 app_env: str = None, app_app: str = None, app_pod: str = None,
                  in_repl = False, bash_session: BashSession = None, remote_dir = None):
         self.namespace = KubeContext.in_cluster_namespace()
 
@@ -69,12 +54,16 @@ class ReplState:
         self.pg_path = pg_path
         self.app_env = app_env
         self.app_app = app_app
+        self.app_pod = app_pod
         if namespace:
             self.namespace = namespace
         self.in_repl = in_repl
         self.bash_session = bash_session
         self.remote_dir = remote_dir
-        # self.wait_log_flag = False
+        self.original_state: ReplState = None
+        self.pod_targetted = False
+
+        self.export_session: str = None
 
         if ns_sts:
             nn = ns_sts.split('@')
@@ -89,13 +78,51 @@ class ReplState:
     def __hash__(self):
         return hash((self.sts, self.pod))
 
+    def __str__(self):
+        msg = ''
+        if self.device == ReplState.P:
+            msg = f'{ReplState.P}:'
+            host, database = self.pg_host_n_database()
+            if database:
+                msg += database
+            elif host:
+                msg += host
+        elif self.device == ReplState.A:
+            msg = f'{ReplState.A}:'
+            if self.app_env:
+                msg += self.app_env
+            if self.app_app:
+                msg += f'/{self.app_app}'
+            if self.app_pod:
+                # azops88-c3-c3-k8sdeploy-appleader-001-79957cf5b6-9k4bw
+                group = re.match(r".*?-.*?-.*?-.*?-(.*?-.*?)-.*", self.app_pod)
+                msg += '/' + group[1]
+        elif self.device == ReplState.L:
+            msg = f'{ReplState.L}:'
+        elif self.device == ReplState.X:
+            msg = f'{ReplState.X}:'
+            if self.export_session:
+                msg += self.export_session
+        else:
+            msg = f'{ReplState.C}:'
+            if self.pod:
+                # cs-d0767a536f-cs-d0767a536f-default-sts-0
+                group = re.match(r".*?-.*?-(.*)", self.pod)
+                msg += group[1]
+            elif self.sts:
+                # cs-d0767a536f-cs-d0767a536f-default-sts
+                group = re.match(r".*?-.*?-(.*)", self.sts)
+                msg += group[1]
+
+        return msg
+
     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 < args_to_check:
-                state = copy.copy(state)
+                state = copy(state)
 
                 s, n = KubeContext.is_sts_name(arg)
                 if s:
@@ -133,9 +160,9 @@ class ReplState:
         new_args = []
         for index, arg in enumerate(args):
             if index < 6:
-                state = copy.copy(state)
+                state = copy(state)
 
-                groups = re.match(r'^([a|c|l|p]):(.*)$', arg)
+                groups = re.match(r'^([a|c|l|p|x]):(.*)$', arg)
                 if groups:
                     if groups[1] == 'p':
                         state.device = 'p'
@@ -154,6 +181,8 @@ class ReplState:
                                 state.namespace = ns
                     elif groups[1] == 'l':
                         state.device = 'l'
+                    elif groups[1] == 'x':
+                        state.device = 'x'
                     else:
                         state.device = 'a'
                         if path := groups[2]:
@@ -196,8 +225,8 @@ class ReplState:
 
             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]]
+        devices = [r for r in required if r in [ReplState.L, ReplState.A, ReplState.C, ReplState.P, ReplState.X]]
+        non_devices = [r for r in required if r not in [ReplState.L, ReplState.A, ReplState.C, ReplState.P, ReplState.X]]
 
         first_error: Callable = None
         for r in non_devices:
@@ -285,8 +314,8 @@ class ReplState:
             if self.device != ReplState.P:
                 return (False, None)
 
-            pg: PostgresContext = PostgresContext.apply(self.namespace, self.pg_path)
-            if not pg.db:
+            _, database = self.pg_host_n_database()
+            if not database:
                 def error():
                     if self.in_repl:
                         log2('cd to a database first.')
@@ -310,7 +339,24 @@ class ReplState:
                         display_help()
                 return (False, error)
 
-        elif required in [ReplState.L, ReplState.A, ReplState.C, ReplState.P] and self.device != required:
+        elif required == RequiredState.EXPORT_DB:
+            if self.device not in [ReplState.C, ReplState.X]:
+                return (False, None)
+
+            if not self.export_session:
+                def error():
+                    if self.in_repl:
+                        if self.device == ReplState.C:
+                            log2("Select an export database first with 'use' command.")
+                        else:
+                            log2('cd to an export database first.')
+                    else:
+                        log2('* export database is missing.')
+                        log2()
+                        display_help()
+                return (False, error)
+
+        elif required in [ReplState.L, ReplState.A, ReplState.C, ReplState.P, ReplState.X] and self.device != required:
             def error():
                 if self.in_repl:
                     log2(f'Switch to {required}: first.')
@@ -339,13 +385,103 @@ class ReplState:
         return Secrets.get_user_pass(self.pod if self.pod else self.sts, self.namespace, secret_path=secret_path)
 
     def enter_bash(self, bash_session: BashSession):
+        self.push()
+
         self.bash_session = bash_session
-        if self.device != ReplState.C:
-            self.device = ReplState.C
-            log2(f'Moved to {ReplState.C}: automatically. Will move back to {ReplState.P}: when you exit the bash session.')
 
     def exit_bash(self):
-        if self.bash_session and self.bash_session.device:
-            self.device = self.bash_session.device
+        self.pop()
+        self.bash_session = None
+
+    def push(self, pod_targetted=False):
+        if not self.original_state:
+            self.original_state = copy(self)
+            self.pod_targetted = pod_targetted
+
+    def pop(self):
+        if o := self.original_state:
+            self.device = o.device
+            self.sts = o.sts
+            self.pod = o.pod
+            self.pg_path = o.pg_path
+            self.app_env = o.app_env
+            self.app_app = o.app_app
+            self.app_pod = o.app_pod
+            # self.export_session = o.export_session
+            self.namespace = o.namespace
+
+            self.original_state = None
+            self.pod_targetted = False
+
+    def pg_host_n_database(self):
+        host = None
+        database = None
+
+        if self.pg_path:
+            host_n_db = self.pg_path.split('/')
+            host = host_n_db[0]
+            if len(host_n_db) > 1:
+                database = host_n_db[1]
+
+        return host, database
+
+    def with_no_pod(self):
+        state1 = copy(self)
+        state1.pod = None
+
+        return state1
+
+    def with_pod(self, pod: str):
+        state1 = copy(self)
+        state1.pod = pod
+
+        return state1
+
+class DevicePodService:
+    def __init__(self, handler: 'DeviceExecHandler'):
+        self.handler = handler
+
+    def exec(self, command: str, action='bash', show_out = True):
+        state = self.handler.state
+
+        rs = None
+        if state.device == ReplState.A and state.app_app:
+            if state.app_pod:
+                rs = [AppPods.exec(state.app_pod, state.namespace, command, show_out=show_out)]
+            else:
+                pods = AppPods.pod_names(state.namespace, state.app_env, state.app_app)
+                rs = AppClusters.exec(pods, state.namespace, command, show_out=show_out)
+        elif state.pod:
+            rs = [CassandraNodes.exec(state.pod, state.namespace, command, show_out=show_out)]
+        elif state.sts:
+            rs = CassandraClusters.exec(state.sts, state.namespace, command, action=action, show_out=show_out)
+        # assume that pg-agent or ops pod is single pod
+
+        dir = None
+        if rs:
+            for r in rs:
+                if r.exit_code(): # if fails to read the session file, ignore
+                    continue
+
+                dir0 = r.stdout.strip(' \r\n')
+                if dir:
+                    if dir != dir0:
+                        log2('Inconsitent working dir found across multiple pods.')
+                        return None
+                else:
+                    dir = dir0
+
+        return dir
+
+class DeviceExecHandler:
+    def __init__(self, state: ReplState):
+        self.state = state
+
+    def __enter__(self):
+        return DevicePodService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
 
-        self.bash_session = None
+def device(state: ReplState):
+    return DeviceExecHandler(state)

adam/sql/async_executor.py

@@ -0,0 +1,62 @@
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from copy import copy
+import inspect
+import threading
+import time
+import traceback
+
+from adam.utils import log2, log_timing
+
+class AsyncExecutor:
+    # some lib does not handle asyncio loop properly, as sync exec submit does not work, use another async loop
+
+    lock = threading.Lock()
+    in_queue = set()
+    processed: dict[str, float] = {}
+    first_processed_at: float = None
+    last_processed_at: float = None
+
+    loop: asyncio.AbstractEventLoop = None
+    async_exec: ThreadPoolExecutor = None
+
+    def preload(action: callable, log_key: str = None):
+        with AsyncExecutor.lock:
+            if not AsyncExecutor.loop:
+                AsyncExecutor.loop = asyncio.new_event_loop()
+                AsyncExecutor.async_exec = ThreadPoolExecutor(max_workers=6, thread_name_prefix='async')
+                AsyncExecutor.loop.set_default_executor(AsyncExecutor.async_exec)
+
+            async def a():
+                try:
+                    t0 = time.time()
+
+                    arg_needed = len(action.__code__.co_varnames)
+
+                    if log_key:
+                        with log_timing(log_key):
+                            r = action(None) if arg_needed else action()
+                    else:
+                        r = action(None) if arg_needed else action()
+                    if inspect.isawaitable(r):
+                        await r
+
+                    AsyncExecutor.in_queue.remove(log_key)
+                    if log_key not in AsyncExecutor.processed:
+                        AsyncExecutor.processed[log_key] = time.time() - t0
+                    AsyncExecutor.last_processed_at = time.time()
+                except Exception as e:
+                    log2('preloading error', e, inspect.getsourcelines(action)[0][0])
+                    traceback.print_exc()
+
+            if log_key not in AsyncExecutor.in_queue:
+                AsyncExecutor.in_queue.add(log_key)
+                AsyncExecutor.async_exec.submit(lambda: AsyncExecutor.loop.run_until_complete(a()))
+
+    def entries_in_queue():
+        # no locking
+        return copy(AsyncExecutor.in_queue), copy(AsyncExecutor.processed), AsyncExecutor.first_processed_at, AsyncExecutor.last_processed_at
+
+    def reset():
+        AsyncExecutor.first_processed_at = time.time()
+        AsyncExecutor.processed.clear()

adam/sql/lark_completer.py

@@ -0,0 +1,286 @@
+import functools
+import os
+from typing import Iterable
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion, NestedCompleter, WordCompleter
+from prompt_toolkit.document import Document
+
+from adam.config import Config
+from adam.sql.async_executor import AsyncExecutor
+from adam.sql.lark_parser import LarkParser
+from adam.utils import debug, log_timing, offload
+from adam.utils_repl.appendable_completer import AppendableCompleter
+from adam.utils_repl.repl_completer import merge_completions
+
+__all__ = [
+    "LarkCompleter",
+]
+
+def default_columns(x: list[str]):
+    return 'id,x.,y.,z.'.split(',')
+
+class LarkCompleter(Completer, AppendableCompleter):
+    SYSTEM = 'system'
+
+    def __init__(self,
+                 dml: str = None,
+                 expandables: dict = {},
+                 variant: str = 'c',
+
+                 name: str = None,
+                 options_lambda: callable = None,
+                 auto: str = 'lazy',
+                 debug = False
+    ) -> None:
+        self.nested: NestedCompleter = None
+        self.options_lambda = options_lambda
+        if options_lambda and auto == 'lazy':
+            AsyncExecutor.preload(options_lambda, log_key=name)
+
+        self.variant = variant
+        self.parser = None
+        self.dml = dml
+        self.expandables = expandables
+
+        self.display_dict = {}
+        self.meta_dict = {}
+        self.WORD = None
+        self.sentence = False
+        self.match_middle = False
+        self.pattern = None
+
+        self.debug = debug
+
+        if variant:
+            self.parser = LarkCompleter.lark_parser(variant)
+            self.preload_lazy_auto_completes()
+
+    def __repr__(self):
+        return f"LarkCompleter.{self.variant}"
+
+    def preload_lazy_auto_completes(self):
+        for key, value in self.expandables.items():
+            if callable(value):
+                if self.auto_complete(key) == 'lazy':
+                    AsyncExecutor.preload(value, log_key=key)
+
+    def from_lambda(name: str, options_lambda: callable, auto: str = 'lazy'):
+        return LarkCompleter(name=name, options_lambda=options_lambda, auto=auto, variant=None)
+
+    @functools.lru_cache()
+    def lark_parser(variant: str):
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        with open(dir_path + f"/qingl.lark") as f:
+            grammar: str = None
+            with log_timing(f'lark.{variant}.file-read'):
+                grammar = f.read()
+
+            common_contexts = {
+                'cd_command.file_name': 'direct-dirs',
+            }
+
+            if variant in ['a', 'c0', 'p0', 'x0']:
+                grammar = grammar.replace('start: statement_sequence', f'start: qing_{variant}_statement')
+                contexts_by_path = {
+                } | common_contexts
+
+                debug(f'* GRAMMAR replaced to start: qing_{variant}_statement')
+
+                return LarkParser(grammar, contexts_by_path)
+            elif variant == 'system':
+                grammar = grammar.replace('start: statement_sequence', f'start: qing_{variant}_statement')
+                contexts_by_path = {
+                } | common_contexts
+
+                return LarkParser(grammar, contexts_by_path)
+
+            grammar = grammar.replace('qing_statement: qing_p_statement', f'qing_statement: qing_{variant}_statement')
+            debug(f'* GRAMMAR replaced to qing_statement: qing_{variant}_statement')
+
+            bash_contexts = {
+                'bash_statement.host_name': 'hosts',
+                'bash_statement.bash_command': 'bash-commands',
+            }
+
+            contexts_by_path = {
+                'describe_keyspace.keyspace_name': 'keyspaces',
+                'keyspace_ref.keyspace_path.namespace_ref.identifier_ref': 'tables',
+                'preview_table_statement.path.identifier_ref': 'tables',
+
+                'insert_statement.insert_select': 'column-names',
+                'update_statement.set_clause.path.identifier_ref': 'column-names',
+                'update_statement.where_clause.cond.expr.path.identifier_ref': 'column-names',
+                'delete_statement.where_clause.cond.expr.path.identifier_ref': 'column-names',
+
+                'select_clause.projection.result_expr.expr.path.identifier_ref': 'columns',
+                'select_from.where_clause.cond.expr.path.identifier_ref': 'columns',
+                'select_from.where_clause.cond.expr.logical_term.and_expr.cond.expr.path.identifier_ref': 'columns',
+                'select_from.group_by_clause.group_term.expr.path.identifier_ref': 'columns',
+                'select_statement.order_by_clause.ordering_term.expr.path.identifier_ref': 'columns',
+                'select_from.from_clause.from_terms.join_clause.ansi_join_clause.ansi_join_predicate.expr.path.identifier_ref': 'columns',
+                'select_from.from_clause.from_terms.join_clause.ansi_join_clause.ansi_join_predicate.expr.comparison_term.relational_expr.expr.path.identifier_ref': 'columns',
+
+                'select_from.from_clause.from_terms.from_generic.alias.identifier_ref': 'column-aliases',
+
+                'select_statement.limit_clause.expr.literal.nbr.digit': 'limits',
+            } | common_contexts
+
+            if variant == 'p':
+                contexts_by_path = bash_contexts | contexts_by_path
+            elif variant == 'c':
+                contexts_by_path = {
+                    'export_table.path.identifier_ref': 'tables',
+                    'show_column_counts_command.path.identifier_ref': 'tables',
+                    'export_statement.export_tables.keyspace_name': 'keyspaces',
+
+                    'alter_tables_statement.properties.property.property_name': 'table-props',
+                    'alter_cql_table_statement.properties.property.property_name': 'table-props',
+                    'alter_tables_statement.properties.property.property_value.literal': 'table-props-value',
+                    'alter_cql_table_statement.properties.property.property_value.literal': 'table-props-value',
+
+                    'select_clause.projection.result_expr.expr.path.identifier_ref': 'columns',
+                    'export_statement.export_tables.export_table.column_name_list.column_name': 'columns',
+
+                    'consistency_statement.consistency': 'consistencies',
+                    'export_statement.export_to.export_database_type': 'export-database-types',
+                    'drop_export_database.export_database_name': 'export-databases',
+                    'use_export_db_statement.export_database_name': 'export-databases',
+                    'clean_up_export_session_statement.clean_up_export_sessions.export_session_name': 'export-sessions',
+                    'show_export_command.export_session_name': 'export-sessions',
+                    'import_statement.import_session.export_session_name': 'export-sessions-incomplete',
+                    'download_session_statement.export_session_name': 'export-sessions-incomplete',
+                } | bash_contexts | contexts_by_path
+            elif variant == 'l':
+                contexts_by_path = {
+                    'add_partition_action.partition_ref.partition_name': 'partition-columns',
+                    'show_topn_statement.topn_count': 'topn-counts',
+                    'show_topn_statement.topn_type': 'topn-types',
+                    'show_topn_statement.topn_window': 'topn-windows'
+                } | contexts_by_path
+            elif variant == 'x':
+                contexts_by_path = {
+                    'show_column_counts_command.path.identifier_ref': 'tables',
+                    'drop_export_database.export_database_name': 'export-databases',
+                    'use_export_db_statement.export_database_name': 'export-databases',
+                } | contexts_by_path
+
+                grammar = grammar.replace('select_clause: "SELECT"i hint_comment? projection', 'select_clause: ("SELECT"i | "XELECT"i) hint_comment? projection')
+
+            with offload():
+                with open('/tmp/grammar.lark', 'wt') as f:
+                    f.write(grammar)
+
+            return LarkParser(grammar, contexts_by_path)
+
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        if not self.nested and self.options_lambda:
+            # for lazy completions
+            self.nested = NestedCompleter.from_nested_dict(self.options_lambda())
+
+        nested_words = set()
+
+        if self.nested:
+            # from NestedCompleter
+
+            # Split document.
+            text = document.text_before_cursor.lstrip()
+            stripped_len = len(document.text_before_cursor) - len(text)
+
+            # If there is a space, check for the first term, and use a
+            # subcompleter.
+            if " " in text:
+                first_term = text.split()[0]
+                completer = self.nested.options.get(first_term)
+
+                # If we have a sub completer, use this for the completions.
+                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):
+                        nested_words.add(c.text)
+                        yield c
+
+            # No space in the input: behave exactly like `WordCompleter`.
+            else:
+                completer = WordCompleter(
+                    list(self.nested.options.keys()), ignore_case=self.nested.ignore_case
+                )
+                for c in completer.get_completions(document, complete_event):
+                    nested_words.add(c.text)
+                    yield c
+
+        if self.parser:
+            full = document.text_before_cursor
+            if self.dml:
+                full = self.dml + ' ' + full
+
+            words0 = []
+            words1 = []
+            context = {}
+            for word in self.parser.next_terminals(full, context=context):
+                if ex := self.expandable(word):
+                    if ex in self.expandables:
+                        e = self.expandables[ex]
+                        if callable(e):
+                            if self.auto_complete(ex) != 'off':
+                                ctx = None
+                                if 'last-id' in context:
+                                    ctx = context['last-id']
+                                e = e(ctx)
+                                words0.extend(e)
+                        else:
+                            words0.extend(e)
+                else:
+                    words1.append(word)
+            words = words0 + words1
+
+            word_before_cursor = document.get_word_before_cursor(
+                WORD=self.WORD, pattern=self.pattern
+            )
+
+            word_before_cursor = word_before_cursor.lower()
+
+            def word_matches(word: str) -> bool:
+                return word.lower().startswith(word_before_cursor)
+
+            for word in words:
+                if word_matches(word) and word not in nested_words:
+                    display = self.display_dict.get(word, word)
+                    display_meta = self.meta_dict.get(word, "")
+                    yield Completion(
+                        word,
+                        -len(word_before_cursor),
+                        display=display,
+                        display_meta=display_meta,
+                    )
+
+    def completions_for_nesting(self, dml: str = None):
+        if dml:
+            return {dml: LarkCompleter(dml, expandables=self.expandables, variant=self.variant)}
+
+        return {
+            word.text.lower(): LarkCompleter(word.text, expandables=self.expandables, variant=self.variant)
+            for word in self.get_completions(Document(''), None)
+        }
+
+    def expandable(self, word: str):
+        return word.strip('`') if word.startswith('`') else None
+
+    def append_completions(self, key: str, value: dict[str, any]):
+        if isinstance(value, LarkCompleter) and self.variant == value.variant:
+            return
+
+        if self.nested:
+            self.nested = NestedCompleter.from_nested_dict(merge_completions(self.nested.options, value))
+        else:
+            self.nested = NestedCompleter.from_nested_dict(value)
+
+    def auto_complete(self, key: str, default = 'lazy'):
+        return Config().get(f'auto-complete.{self.variant}.{key}', default=default)