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

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, _: ReplState):
+        return f"{GenerateReport.COMMAND} [-s]\t generate report"

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

@@ -33,7 +33,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')})
+        return {}
+        # return super().completion(state, {session: None for session in ExportSessions.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import')})
 
     def help(self, _: ReplState):
         return f'{DownloadExportSession.COMMAND} <export-session-name>\t download csv files in export session'

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'
@@ -36,17 +31,6 @@ class ExportTables(Command):
                     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)}
-
         return {}
 
     def help(self, _: ReplState):

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,23 @@ 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'
+        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(log_file, 'w') as f:
+            with open(export_log, flag) as f:
                 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'
+        return f'<sql-select-statements>\t run queries on export database'

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, export_log_dir, find_files, os_system_exec
 from adam.config import Config
 from adam.repl_state import ReplState
 from adam.utils import log2, 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] = find_files(pod, namespace, f'{export_log_dir()}/{prefix}*_*.log*')
 
         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'{export_log_dir()}/([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 {export_log_dir()}/*.log*'
+        os_system_exec(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] = find_files(pod, namespace, f'{export_log_dir()}/{session}_*.log*')
 
         for log_file in log_files:
-            m = re.match(f'{log_prefix()}-{session}_(.*?)\.(.*?)\.log.*', log_file)
+            m = re.match(f'{export_log_dir()}/{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}'
+                os_system_exec(cmd, show_out=not multi_tables)
                 log_cnt += 1
 
         return csv_cnt, log_cnt
@@ -142,8 +144,7 @@ class ExportSessions:
         def download_csv(table):
             from_path: str = table.csv_file
 
-            to_path = from_path.replace(csv_dir(), local_tmp_dir())
-            os.makedirs(os.path.dirname(to_path), exist_ok=True)
+            to_path = from_path.replace(csv_dir(), local_downloads_dir())
             Pods.download_file(pod, 'cassandra', namespace, from_path, to_path)
 
             log2(f'[{session}] Downloaded to {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)):
@@ -43,7 +43,7 @@ class ExportUse(Command):
                 return state
 
     def completion(self, state: ReplState):
-        return super().completion(state, {n: None for n in ExportDatabases.database_names()})
+        return {}
 
     def help(self, _: ReplState):
         return f'{ExportUse.COMMAND} <export-database-name>\t use export database'

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, _: ReplState):
+        return f'xelect...\t run queries on export database'