kaqing 2.0.184__py3-none-any.whl → 2.0.227__py3-none-any.whl

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

@@ -5,7 +5,7 @@ from adam.checks.check_utils import all_checks, checks_from_csv, run_checks
 from adam.commands import extract_options, validate_args
 from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterOrPodCommandHelper
-from adam.commands.issues import Issues
+from adam.commands.diag.issues import Issues
 from adam.repl_state import ReplState
 from adam.utils import tabulize, log
 from adam.utils_issues import IssuesUtils
@@ -49,8 +49,8 @@ class Check(Issues):
     def completion(self, _: ReplState):
         return {Check.COMMAND: {check.name(): {'-s': None} for check in all_checks()}}
 
-    def help(self, _: ReplState):
-        return f'{Check.COMMAND} <check-name> [-s]\t run a single check'
+    def help(self, state: ReplState):
+        return super().help(state, 'run a single check -s show processing details', args='<check-name> [-s]')
 
 class CheckCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):

adam/commands/diag/generate_report.py

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

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

@@ -39,5 +39,6 @@ class Issues(Command):
     def completion(self, state: ReplState):
         return super().completion(state, {'-s': None})
 
-    def help(self, _: ReplState):
-        return f'{Issues.COMMAND} [-s]\t find all issues'
+    def help(self, state: ReplState, desc: str = None, args: str = None):
+        args1 = args if args else '[-s]'
+        return super().help(state, 'find all issues  -s show processing details', args=args1)

adam/commands/exit.py

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

adam/commands/export/clean_up_all_export_sessions.py

@@ -33,5 +33,5 @@ class CleanUpAllExportSessions(Command):
     def completion(self, _: ReplState):
         return {}
 
-    def help(self, _: ReplState):
-        return f'{CleanUpAllExportSessions.COMMAND}\t clean up all export sessions'
+    def help(self, state: ReplState):
+        return super().help(state, 'clean up all export sessions')

adam/commands/export/clean_up_export_sessions.py

@@ -35,5 +35,5 @@ class CleanUpExportSessions(Command):
     def completion(self, _: ReplState):
         return {}
 
-    def help(self, _: ReplState):
-        return f'{CleanUpExportSessions.COMMAND} <export-session-name,...>\t clean up export sessions'
+    def help(self, state: ReplState):
+        return super().help(state, 'clean up export sessions', args='<export-session-name>,...')

adam/commands/export/completions_x.py

@@ -0,0 +1,11 @@
+from adam.commands.export.export_databases import ExportDatabases
+from adam.repl_state import ReplState
+from adam.sql.lark_completer import LarkCompleter
+from adam.utils_athena import Athena
+
+def completions_x(state: ReplState):
+    return LarkCompleter(expandables={
+        'tables': lambda x: ExportDatabases.table_names(state.export_session),
+        'export-databases': lambda x: ExportDatabases.database_names(),
+        'columns': lambda x: Athena.column_names(database=state.export_session, function='export'),
+    }, variant=ReplState.X).completions_for_nesting()

adam/commands/export/download_export_session.py

@@ -1,6 +1,6 @@
 from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.export_sessions import ExportSessions, export_session
+from adam.commands.export.export_sessions import export_session
 from adam.repl_state import ReplState, RequiredState
 
 class DownloadExportSession(Command):
@@ -32,8 +32,8 @@ class DownloadExportSession(Command):
 
             return state
 
-    def completion(self, state: ReplState):
-        return super().completion(state, {session: None for session in ExportSessions.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import')})
+    def completion(self, _: ReplState):
+        return {}
 
-    def help(self, _: ReplState):
-        return f'{DownloadExportSession.COMMAND} <export-session-name>\t download csv files in export session'
+    def help(self, state: ReplState):
+        return super().help(state, 'download csv files in export session', args='<export-session-name>')

adam/commands/export/drop_export_database.py

@@ -35,5 +35,5 @@ class DropExportDatabase(Command):
     def completion(self, _: ReplState):
         return {}
 
-    def help(self, _: ReplState):
-        return f'{DropExportDatabase.COMMAND} <export-database-name>\t drop export database'
+    def help(self, state: ReplState):
+        return super().help(state, 'drop export database', args='<export-database-name>')

adam/commands/export/drop_export_databases.py

@@ -33,5 +33,5 @@ class DropExportDatabases(Command):
     def completion(self, _: ReplState):
         return {}
 
-    def help(self, _: ReplState):
-        return f'{DropExportDatabases.COMMAND}\t drop all export databases'
+    def help(self, state: ReplState):
+        return super().help(state, 'drop all export databases')

adam/commands/export/export.py

@@ -1,12 +1,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.exporter import export
 from adam.repl_state import ReplState, RequiredState
-from adam.sql.sql_completer import SqlCompleter, SqlVariant
-from adam.utils import log
-from adam.utils_k8s.statefulsets import StatefulSets
 
 class ExportTables(Command):
     COMMAND = 'export'
@@ -35,19 +30,8 @@ class ExportTables(Command):
                 with export(state) as exporter:
                     return exporter.export(args, export_only=export_only)
 
-    def completion(self, state: ReplState):
-        def sc():
-            return SqlCompleter(lambda: cassandra_table_names(state), expandables={
-                'dml':'export',
-                'columns': lambda table: ['id', '*'],
-                'keyspaces': lambda: cassandra_keyspaces(state),
-                'export-dbs': lambda: ExportDatabases.database_names(),
-            }, variant=SqlVariant.CQL)
-
-        if super().completion(state):
-            return {f'@{p}': {ExportTables.COMMAND: sc()} for p in StatefulSets.pod_names(state.sts, state.namespace)}
-
+    def completion(self, _: ReplState):
         return {}
 
-    def help(self, _: ReplState):
-        return f'{ExportTables.COMMAND} [* [in KEYSPACE]] | [TABLE] [as target-name] [with consistency <level>]\t export tables to Sqlite, Athena or CSV file'
+    def help(self, state: ReplState):
+        return super().help(state, 'export tables to Sqlite, Athena or CSV file', args='TABLE...')

adam/commands/export/export_databases.py

@@ -1,14 +1,16 @@
 from collections.abc import Callable
 from datetime import datetime
+from functools import partial
 import os
 import boto3
 
 from adam.commands.export.export_sessions import ExportSessions
 from adam.commands.export.importer import Importer
+from adam.commands.export.utils_export import export_log_dir
 from adam.config import Config
 from adam.repl_session import ReplSession
 from adam.repl_state import ReplState
-from adam.utils import debug, tabulize, log2, ing, log_exc
+from adam.utils import debug, log_timing, tabulize, log2, ing, log_exc
 from adam.utils_athena import Athena
 from adam.utils_sqlite import SQLite
 
@@ -31,7 +33,7 @@ class ExportDatabases:
 
         if Config().get('repl.history.push-cat-log-file', True):
             if log_file and ReplSession().prompt_session:
-                ReplSession().prompt_session.history.append_string(f':sh cat {log_file}')
+                ReplSession().prompt_session.history.append_string(f':cat {log_file}')
 
         return cnt
 
@@ -111,7 +113,8 @@ class ExportDatabases:
         return list({n.split('_')[0] for n in SQLite.database_names()})
 
     def export_database_names():
-        return list({n.split('_')[0] for n in Athena.database_names(LIKE)})
+        with log_timing('ExportDatabases.Athena.database_names'):
+            return list({n.split('_')[0] for n in Athena.database_names(LIKE)})
 
     def database_names_with_keyspace_cnt(importer: str = None):
         r = {}
@@ -191,20 +194,26 @@ class ExportDatabaseService:
     def __init__(self, handler: 'ExportDatabaseHandler'):
         self.handler = handler
 
-    def sql(self, query: str, database: str = None, backgrounded = False):
+    def sql(self, query: str, database: str = None, backgrounded = False, export_log: str = None):
         if not database:
             database = self.handler.state.export_session
 
-        def output(out: str):
-            log_prefix = Config().get('export.log-prefix', '/tmp/qing')
-            log_file = f'{log_prefix}-{datetime.now().strftime("%d%H%M%S")}-export.log'
-
-            with open(log_file, 'w') as f:
+        def output(export_log: str, out: str):
+            # serving for 1. export log or 2. job log for the sampling select 10 rows afterward
+            flag = 'at'
+            if not export_log:
+                export_log = f'{export_log_dir()}/{datetime.now().strftime("%d%H%M%S")}-export.log'
+                flag = 'w'
+
+            with open(export_log, flag) as f:
+                if flag == 'at':
+                    f.write(query)
+                    f.write('\n')
                 f.write(out)
 
-            return log_file
+            return export_log
 
-        ExportDatabases.run_query(query, database, output = output if backgrounded else None, show_query = not backgrounded)
+        ExportDatabases.run_query(query, database, output = partial(output, export_log) if backgrounded else None, show_query = not backgrounded)
 
     def drop(self, database: str):
         state = self.handler.state

adam/commands/export/export_select.py

@@ -1,12 +1,10 @@
-from adam.commands import extract_trailing_options, validate_args
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases, export_db
+from adam.commands.export.completions_x import completions_x
 from adam.repl_state import ReplState, RequiredState
-from adam.sql.sql_completer import SqlCompleter, SqlVariant
-from adam.utils_athena import Athena
 
+# No action body, only for a help entry and auto-completion
 class ExportSelect(Command):
-    COMMAND = '.select'
+    COMMAND = 'select_on_x'
 
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
@@ -23,37 +21,14 @@ class ExportSelect(Command):
     def required(self):
         return RequiredState.EXPORT_DB
 
-    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_trailing_options(args, '&') as (args, backgrounded):
-                with validate_args(args, state, name='SQL statement') as query:
-                    with export_db(state) as dbs:
-                        dbs.sql(f'select {query}', backgrounded=backgrounded)
-
-                    return state
-
     def completion(self, state: ReplState):
-        if not state.export_session:
+        if state.device != ReplState.X:
             return {}
 
-        db = state.export_session
-
-        # warm up the caches first time when x: drive is accessed
-        ExportDatabases.table_names(db)
-        Athena.column_names(database=db, function='export')
-        Athena.column_names(partition_cols_only=True, database=db, function='export')
+        if state.export_session:
+            return completions_x(state)
 
-        return {ExportSelect.COMMAND: SqlCompleter(
-            lambda: ExportDatabases.table_names(db),
-            dml='select',
-            expandables={
-                'columns':lambda table: Athena.column_names(database=db, function='export'),
-            },
-            variant=SqlVariant.ATHENA
-        )}
+        return {}
 
-    def help(self, _: ReplState):
-        return f'.<sql-select-statements>\t run queries on export database'
+    def help(self, state: ReplState):
+        return super().help(state, 'run queries on export database', command='<sql-select-statements>')

adam/commands/export/export_sessions.py

@@ -3,14 +3,14 @@ import os
 import re
 
 from adam.commands.export.importer import Importer
-from adam.commands.export.utils_export import ExportTableStatus, csv_dir, find_files
+from adam.commands.export.utils_export import ExportTableStatus, csv_dir, fs_exec, table_log_dir
 from adam.config import Config
 from adam.repl_state import ReplState
-from adam.utils import log2, tabulize, log, parallelize
+from adam.utils import log2, log_to_pods, tabulize, log, parallelize
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
-from adam.utils_k8s.pods import Pods, log_prefix
+from adam.utils_k8s.pods import Pods
 from adam.utils_k8s.statefulsets import StatefulSets
-from adam.utils_local import local_tmp_dir
+from adam.utils_local import local_downloads_dir
 
 class ExportSessions:
     def clear_export_session_cache():
@@ -36,13 +36,13 @@ class ExportSessions:
 
         prefix = Importer.prefix_from_importer(importer)
 
-        log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{prefix}*_*.log*')
+        log_files: list[str] = Pods.find_files(pod, 'cassandra', namespace, f'{table_log_dir(pod, namespace)}/{prefix}*_*.log*', remote=log_to_pods())
 
         if not log_files:
             return {}
 
         for log_file in log_files[:limit]:
-            m = re.match(f'{log_prefix()}-(.*?)_.*\.log?(.*)', log_file)
+            m = re.match(f'{table_log_dir(pod, namespace)}/([ces].*?)_.*\.log?(.*)', log_file)
             if m:
                 s = m.group(1)
                 state = m.group(2) # '', '.pending_import', '.done'
@@ -66,7 +66,8 @@ class ExportSessions:
             pod = StatefulSets.pod_names(sts, namespace)[0]
 
         CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/*', show_out=Config().is_debug(), shell='bash')
-        CassandraNodes.exec(pod, namespace, f'rm -rf {log_prefix()}-*.log*', show_out=Config().is_debug(), shell='bash')
+        cmd = f'rm -rf {table_log_dir(pod, namespace)}/*.log*'
+        fs_exec(pod, namespace, cmd, show_out=Config().is_debug())
 
         return True
 
@@ -102,17 +103,18 @@ class ExportSessions:
         csv_cnt = 0
         log_cnt = 0
 
-        log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{session}_*.log*')
+        log_files: list[str] = Pods.find_files(pod, 'cassandra', namespace, f'{table_log_dir(pod, namespace)}/{session}_*.log*', remote=log_to_pods())
 
         for log_file in log_files:
-            m = re.match(f'{log_prefix()}-{session}_(.*?)\.(.*?)\.log.*', log_file)
+            m = re.match(f'{table_log_dir(pod, namespace)}/{session}_(.*?)\.(.*?)\.log.*', log_file)
             if m:
                 table = m.group(2)
 
                 CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/{session}_{table}', show_out=not multi_tables, shell='bash')
                 csv_cnt += 1
 
-                CassandraNodes.exec(pod, namespace, f'rm -rf {log_file}', show_out=not multi_tables, shell='bash')
+                cmd = f'rm -rf {log_file}'
+                fs_exec(pod, namespace, cmd, show_out=not multi_tables)
                 log_cnt += 1
 
         return csv_cnt, log_cnt
@@ -142,7 +144,7 @@ class ExportSessions:
         def download_csv(table):
             from_path: str = table.csv_file
 
-            to_path = from_path.replace(csv_dir(), local_tmp_dir())
+            to_path = from_path.replace(csv_dir(), local_downloads_dir())
             os.makedirs(os.path.dirname(to_path), exist_ok=True)
             Pods.download_file(pod, 'cassandra', namespace, from_path, to_path)
 

adam/commands/export/export_use.py

@@ -1,6 +1,6 @@
 from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases, export_db
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
 from adam.utils import log2
 
@@ -20,7 +20,7 @@ class ExportUse(Command):
         return ExportUse.COMMAND
 
     def required(self):
-        return [ReplState.C, ReplState.X]
+        return [ReplState.C]
 
     def run(self, cmd: str, state: ReplState):
         if not(args := self.args(cmd)):
@@ -42,8 +42,8 @@ class ExportUse(Command):
 
                 return state
 
-    def completion(self, state: ReplState):
-        return super().completion(state, {n: None for n in ExportDatabases.database_names()})
+    def completion(self, _: ReplState):
+        return {}
 
-    def help(self, _: ReplState):
-        return f'{ExportUse.COMMAND} <export-database-name>\t use export database'
+    def help(self, state: ReplState):
+        return super().help(state, 'use export database', args='<export-database-name>')

adam/commands/export/export_x_select.py

@@ -0,0 +1,48 @@
+from adam.commands import extract_trailing_options, validate_args
+from adam.commands.command import Command
+from adam.commands.export.completions_x import completions_x
+from adam.commands.export.export_databases import export_db
+from adam.repl_state import ReplState, RequiredState
+
+class ExportXSelect(Command):
+    COMMAND = 'xelect'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportXSelect, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportXSelect.COMMAND
+
+    def required(self):
+        return RequiredState.EXPORT_DB
+
+    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_trailing_options(args, '&') as (args, backgrounded):
+                with validate_args(args, state, name='SQL statement') as query:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select {query}', backgrounded=backgrounded)
+
+                    return state
+
+    def completion(self, state: ReplState):
+        if state.device != ReplState.C:
+            return {}
+
+        if not state.export_session:
+            return {}
+
+        # add only xelect completions to c: drive from lark
+        return {ExportXSelect.COMMAND: completions_x(state)[ExportXSelect.COMMAND]}
+
+    def help(self, state: ReplState):
+        return super().help(state, 'run queries on export database', command='xelect...')