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

adam/commands/devices/device_cass.py

@@ -1,10 +1,11 @@
-from adam.commands.bash.bash import Bash
-from adam.commands.command import Command
+from adam.commands.bash.bash_completer import BashCompleter
+from adam.commands.command import Command, InvalidState
 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
 from adam.config import Config
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log, log2
+from adam.utils import lines_to_tabular, log, log2, wait_log
 from adam.utils_k8s.cassandra_clusters import CassandraClusters
 from adam.utils_k8s.custom_resources import CustomResources
 from adam.utils_k8s.kube_context import KubeContext
@@ -41,7 +42,7 @@ class DeviceCass(Command, Device):
 
     def ls(self, cmd: str, state: ReplState):
         if state.pod:
-            return Bash().run('bash ' + cmd, state)
+            return self.bash(state, state, cmd.split(' '))
         elif state.sts and state.namespace:
             show_pods(StatefulSets.pods(state.sts, state.namespace), state.namespace, show_namespace=not KubeContext.in_cluster_namespace())
             show_rollout(state.sts, state.namespace)
@@ -57,6 +58,18 @@ class DeviceCass(Command, Device):
         else:
             return {cmd: {n: None for n in StatefulSets.list_sts_names()}}
 
+    def cat(self, cmd: str, state: ReplState):
+        if state.pod:
+            return self.bash(state, state, cmd.split(' '))
+
+        return InvalidState()
+
+    def cat_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.sts:
+            return {f'@{p}': {cmd: None} for p in StatefulSets.pod_names(state.sts, state.namespace)}
+
+        return default
+
     def show_statefulsets(self):
         ss = StatefulSets.list_sts_names()
         if len(ss) == 0:
@@ -101,7 +114,7 @@ class DeviceCass(Command, Device):
                     if dir in names:
                         state.pod = names[dir]
                     else:
-                        log2('Not a valid pod name or host id.')
+                        state.pod = dir
 
     def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
         if state.pod:
@@ -140,6 +153,33 @@ class DeviceCass(Command, Device):
                 if auto_enter == 'first-pod':
                     state.pod = f'{state.sts}-0'
                 if KubeContext().in_cluster_namespace:
-                    Config().wait_log(f'Moving to the only Cassandra cluster: {state.sts}...')
+                    wait_log(f'Moving to the only Cassandra cluster: {state.sts}...')
                 else:
-                    Config().wait_log(f'Moving to the only Cassandra cluster: {state.sts}@{state.namespace}...')
+                    wait_log(f'Moving to the only Cassandra cluster: {state.sts}@{state.namespace}...')
+
+    def show_tables(self, state: ReplState):
+        log(lines_to_tabular(cassandra_table_names(state), separator=','))
+
+    def show_table_preview(self, state: ReplState, table: str, rows: int):
+        with cassandra(state) as pods:
+            pods.cql(f'select * from {table} limit {rows}', show_out=True, use_single_quotes=True, on_any=True)
+
+    def bash_target_changed(self, s0: ReplState, s1: ReplState):
+        return s0.sts != s1.sts or s0.pod != s1.pod
+
+    def exec_no_dir(self, command: str, state: ReplState):
+        with cassandra(state) as pods:
+            return pods.exec(command, action='bash', show_out=True, shell='bash')
+
+    def exec_with_dir(self, command: str, session_just_created: bool, state: ReplState):
+        with cassandra(state) as pods:
+            return pods.exec(command, action='bash', show_out=not session_just_created, shell='bash')
+
+    def bash_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        completions = {cmd: BashCompleter(lambda: [])}
+
+        if state.sts and state.namespace:
+            completions |= {f'@{p}': {cmd: BashCompleter(lambda: [])} for p in StatefulSets.pod_names(state.sts, state.namespace)}
+
+        return completions
+

adam/commands/devices/device_export.py

@@ -3,7 +3,7 @@ from adam.commands.devices.device import Device
 from adam.commands.export.export_databases import ExportDatabases
 from adam.config import Config
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log, log2
+from adam.utils import lines_to_tabular, log, log2, wait_log
 
 class DeviceExport(Command, Device):
     COMMAND = f'{ReplState.X}:'
@@ -78,7 +78,7 @@ class DeviceExport(Command, Device):
     def enter(self, state: ReplState):
         if auto_enter := Config().get('repl.x.auto-enter', 'no'):
             if auto_enter == 'latest':
-                Config().wait_log(f'Moving to the latest export database...')
+                wait_log(f'Moving to the latest export database...')
                 if dbs := ExportDatabases.database_names():
                     state.export_session = sorted(dbs)[-1]
                 else:

adam/commands/devices/device_postgres.py

@@ -1,10 +1,11 @@
-from adam.commands.command import Command
+from adam.commands.bash.bash_completer import BashCompleter
+from adam.commands.command import Command, InvalidState
 from adam.commands.devices.device import Device
 from adam.commands.postgres.postgres_context import PostgresContext
-from adam.commands.postgres.postgres_utils import pg_database_names, pg_table_names
-from adam.config import Config
+from adam.commands.postgres.utils_postgres import pg_database_names, pg_table_names, postgres
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log
+from adam.utils import lines_to_tabular, log, log2, wait_log
+from adam.utils_k8s.pods import Pods
 
 class DevicePostgres(Command, Device):
     COMMAND = f'{ReplState.P}:'
@@ -45,6 +46,12 @@ class DevicePostgres(Command, Device):
         else:
             self.show_pg_hosts(state)
 
+    def cat(self, cmd: str, state: ReplState):
+        if state.pod:
+            return self.bash(state, state, cmd.split(' '))
+
+        return InvalidState()
+
     def show_pg_hosts(self, state: ReplState):
         if state.namespace:
             def line(pg: PostgresContext):
@@ -105,5 +112,33 @@ class DevicePostgres(Command, Device):
 
         return False, None
 
-    def enter(self, state: ReplState):
-        Config().wait_log('Inspecting postgres database instances...')
+    def enter(self, _: ReplState):
+        wait_log('Inspecting postgres database instances...')
+
+    def show_tables(self, state: ReplState):
+        pg = PostgresContext.apply(state.namespace, state.pg_path)
+        lines = [db["name"] for db in pg.tables() if db["schema"] == PostgresContext.default_schema()]
+        log(lines_to_tabular(lines, separator=','))
+
+    def show_table_preview(self, state: ReplState, table: str, rows: int):
+        PostgresContext.apply(state.namespace, state.pg_path).run_sql(f'select * from {table} limit {rows}')
+
+    def bash(self, s0: ReplState, s1: ReplState, args: list[str]):
+        pod, container = PostgresContext.pod_and_container(s1.namespace)
+        log2(f'Running on {pod}(container:{container})...')
+
+        return super().bash(s0, s1, args)
+
+    def bash_target_changed(self, s0: ReplState, s1: ReplState):
+        return s0.pg_path != s1.pg_path
+
+    def exec_no_dir(self, command: str, state: ReplState):
+        with postgres(state) as pod:
+            return pod.exec(command, show_out=True)
+
+    def exec_with_dir(self, command: str, session_just_created: bool, state: ReplState):
+        with postgres(state) as pod:
+            return pod.exec(command, show_out=not session_just_created)
+
+    def bash_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        return {cmd: BashCompleter(lambda: [])}

adam/commands/exit.py

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

adam/commands/export/clean_up_all_export_sessions.py

@@ -0,0 +1,37 @@
+from adam.commands.command import Command
+from adam.commands.export.exporter import Exporter
+from adam.repl_state import ReplState, RequiredState
+
+class CleanUpAllExportSessions(Command):
+    COMMAND = 'clean up all export sessions'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(CleanUpAllExportSessions, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return CleanUpAllExportSessions.COMMAND
+
+    def required(self):
+        return RequiredState.CLUSTER_OR_POD
+
+    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):
+            if Exporter.clean_up_all_sessions(state.sts, state.pod, state.namespace):
+                Exporter.clear_export_session_cache()
+
+            return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{CleanUpAllExportSessions.COMMAND}\t clean up all export sessions'

adam/commands/export/clean_up_export_sessions.py

@@ -1,9 +1,10 @@
 from adam.commands.command import Command
 from adam.commands.export.exporter import Exporter
 from adam.repl_state import ReplState, RequiredState
+from adam.utils import log, log2
 
 class CleanUpExportSessions(Command):
-    COMMAND = 'clean up all export sessions'
+    COMMAND = 'clean up export sessions'
 
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
@@ -24,17 +25,27 @@ class CleanUpExportSessions(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return 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.')
+
+                    Command.display_help()
+
+                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.')
 
-        if Exporter.clean_up_all_sessions(state.sts, state.pod, state.namespace):
             Exporter.clear_export_session_cache()
 
-        return state
+            return state
 
     def completion(self, _: ReplState):
         return {}
 
     def help(self, _: ReplState):
-        return f'{CleanUpExportSessions.COMMAND}\t clean up all export sessions'
+        return f'{CleanUpExportSessions.COMMAND} <export-session-name>\t clean up export session'

adam/commands/export/drop_export_database.py

@@ -27,29 +27,26 @@ class DropExportDatabase(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return 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 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()
 
-            return 'command-missing'
+                return 'command-missing'
 
-        Exporter.drop_databases(state.sts, state.pod, state.namespace, args[0])
+            Exporter.drop_databases(state.sts, state.pod, state.namespace, args[0])
 
-        SQLite.clear_cache()
-        Athena.clear_cache()
+            SQLite.clear_cache()
+            Athena.clear_cache()
 
-        if state.export_session == args[0]:
-            state.export_session = None
+            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

@@ -26,18 +26,15 @@ class DropExportDatabases(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return state
-
-        Exporter.drop_databases(state.sts, state.pod, state.namespace)
+        with self.validate(args, state) as (args, state):
+            Exporter.drop_databases(state.sts, state.pod, state.namespace)
 
-        SQLite.clear_cache()
-        Athena.clear_cache()
+            SQLite.clear_cache()
+            Athena.clear_cache()
 
-        state.export_session = None
+            state.export_session = None
 
-        return state
+            return state
 
     def completion(self, _: ReplState):
         return {}

adam/commands/export/export.py

@@ -1,8 +1,8 @@
+from adam.commands import extract_options
 from adam.commands.command import Command
-from adam.commands.cql.cql_utils import cassandra_keyspaces, cassandra_table_names
+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 Exporter
-from adam.commands.export.utils_export import ExportSpec
+from adam.commands.export.export_handlers import export
 from adam.repl_state import ReplState, RequiredState
 from adam.sql.sql_completer import SqlCompleter, SqlVariant
 from adam.utils import log
@@ -30,40 +30,10 @@ class ExportTables(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return state
-
-        # --export-only for testing only
-        args, export_only = Command.extract_options(args, ['--export-only'])
-
-        if not state.pod:
-            state.push()
-            state.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
-
-        export_session = state.export_session
-        spec: ExportSpec = None
-        try:
-            statuses, spec = Exporter.export_tables(args, state, export_only=export_only)
-            if not statuses:
-                return state
-
-            Exporter.clear_export_session_cache()
-
-            if spec.importer == 'csv' or export_only:
-                ExportDatabases.disply_export_session(state.sts, state.pod, state.namespace, spec.session)
-            else:
-                log()
-                ExportDatabases.display_export_db(state.export_session)
-
-        finally:
-            state.pop()
-
-            # if exporting to csv, do not bind the new session id to repl state
-            if spec and spec.importer == 'csv':
-                state.export_session = export_session
-
-        return state
+        with self.validate(args, state) as (args, state):
+            with extract_options(args, '--export-only') as (args, export_only):
+                with export(state) as exporter:
+                    return exporter.export(args, export_only=export_only)
 
     def completion(self, state: ReplState):
         def sc():
@@ -74,7 +44,7 @@ class ExportTables(Command):
                 'export-dbs': lambda: ExportDatabases.database_names(),
             }, variant=SqlVariant.CQL)
 
-        if state.sts:
+        if super().completion(state):
             return {f'@{p}': {ExportTables.COMMAND: sc()} for p in StatefulSets.pod_names(state.sts, state.namespace)}
 
         return {}

adam/commands/export/export_databases.py

@@ -1,9 +1,11 @@
+from collections.abc import Callable
 import os
 import boto3
 
 from adam.commands.export.utils_export import ExportTableStatus
 from adam.config import Config
-from adam.utils import lines_to_tabular, log, log2, ing
+from adam.repl_session import ReplSession
+from adam.utils import debug, lines_to_tabular, log, log2, ing, log_exc
 from adam.utils_athena import Athena
 from adam.utils_k8s.statefulsets import StatefulSets
 from adam.utils_sqlite import SQLite
@@ -11,16 +13,23 @@ from adam.utils_sqlite import SQLite
 LIKE = 'e%_%'
 
 class ExportDatabases:
-    def run_query(query: str, database: str, show_query=False) -> int:
+    def run_query(query: str, database: str, show_query=False, output: Callable[[str], str] = None) -> int:
         cnt: int = 0
 
         if show_query:
             log2(query)
 
+        log_file = None
         if database.startswith('s'):
-            cnt += SQLite.run_query(query, database=database)
+            cnt0, log_file = SQLite.run_query(query, database=database, output=output)
+            cnt += cnt0
         else:
-            cnt += Athena.run_query(query, database=database)
+            cnt0, log_file = Athena.run_query(query, database=database, output=output)
+            cnt += cnt0
+
+        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}')
 
         return cnt
 
@@ -47,15 +56,13 @@ class ExportDatabases:
         dbs = SQLite.database_names(db)
         if dbs:
             with ing(f'Droping {len(dbs)} SQLite databases'):
-                try:
+                with log_exc():
                     for db in dbs:
                         file_path = f'{SQLite.local_db_dir()}/{db}'
                         try:
                             os.remove(file_path)
                         except OSError as e:
                             pass
-                except:
-                    pass
 
         return dbs
 
@@ -65,20 +72,17 @@ class ExportDatabases:
             with ing(f'Droping {len(dbs)} Athena databases'):
                 for db in dbs:
                     query = f'DROP DATABASE {db} CASCADE'
-                    if Config().is_debug():
-                        log2(query)
+                    debug(query)
                     Athena.query(query)
 
         with ing(f'Deleting s3 folder: export'):
-            try:
+            with log_exc():
                 if not db:
                     db = ''
 
                 s3 = boto3.resource('s3')
                 bucket = s3.Bucket(Config().get('export.bucket', 'c3.ops--qing'))
                 bucket.objects.filter(Prefix=f'export/{db}').delete()
-            except:
-                pass
 
         return dbs
 

adam/commands/export/export_handlers.py

@@ -0,0 +1,71 @@
+from adam.commands.command import Command
+from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.export.exporter import Exporter
+from adam.commands.export.utils_export import ExportSpec, state_with_pod
+from adam.repl_state import ReplState
+from adam.utils import log, log2
+
+class ExportService:
+    def __init__(self, handler: 'ExporterHandler'):
+        self.handler = handler
+
+    def export(self, args: list[str], export_only=False):
+        state = self.handler.state
+        export_session = state.export_session
+        spec: ExportSpec = None
+        try:
+            with state_with_pod(state) as state:
+                # --export-only for testing only
+                statuses, spec = Exporter.export_tables(args, state, export_only=export_only)
+                if not statuses:
+                    return state
+
+                Exporter.clear_export_session_cache()
+
+                if spec.importer == 'csv' or export_only:
+                    ExportDatabases.disply_export_session(state.sts, state.pod, state.namespace, spec.session)
+                else:
+                    log()
+                    ExportDatabases.display_export_db(state.export_session)
+        finally:
+            # if exporting to csv, do not bind the new session id to repl state
+            if spec and spec.importer == 'csv':
+                state.export_session = export_session
+
+        return state
+
+    def import_sesion(self, args: list[str]):
+        state = self.handler.state
+
+        if not args:
+            if state.in_repl:
+                log2('Specify export session name.')
+            else:
+                log2('* Export session name is missing.')
+
+                Command.display_help()
+
+            return 'command-missing'
+
+        with state_with_pod(state) as state:
+            tables, _ = Exporter.import_session(args, state)
+            if tables:
+                Exporter.clear_export_session_cache()
+
+                log()
+                ExportDatabases.display_export_db(state.export_session)
+
+        return state
+
+class ExporterHandler:
+    def __init__(self, state: ReplState):
+        self.state = state
+
+    def __enter__(self):
+        return ExportService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+def export(state: ReplState):
+    return ExporterHandler(state)

adam/commands/export/export_select.py

@@ -1,10 +1,12 @@
+from datetime import datetime
+from adam.commands import extract_trailing_options
 from adam.commands.command import Command
 from adam.commands.export.export_databases import ExportDatabases
+from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
 from adam.sql.sql_completer import SqlCompleter, SqlVariant
 from adam.utils import log2
 from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class ExportSelect(Command):
     COMMAND = '.select'
@@ -28,38 +30,45 @@ class ExportSelect(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return state
+        with self.validate(args, state) as (args, state):
+            with extract_trailing_options(args, '&') as (args, backgrounded):
+                if not state.export_session:
+                    if state.in_repl:
+                        if state.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.')
 
-        if not state.export_session:
-            if state.in_repl:
-                if state.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.')
+                        Command.display_help()
+
+                    return 'command-missing'
+
+                if not args:
+                    if state.in_repl:
+                        log2('Use a SQL statement.')
+                    else:
+                        log2('* SQL statement is missing.')
 
-                Command.display_help()
+                        Command.display_help()
 
-            return 'command-missing'
+                    return 'command-missing'
 
-        if not args:
-            if state.in_repl:
-                log2('Use a SQL statement.')
-            else:
-                log2('* SQL statement is missing.')
+                query = ' '.join(args)
 
-                Command.display_help()
+                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")}-sqlite.log'
 
-            return 'command-missing'
+                    with open(log_file, 'w') as f:
+                        f.write(out)
 
-        query = ' '.join(args)
+                    return log_file
 
-        ExportDatabases.run_query(f'select {query}', database=state.export_session)
+                ExportDatabases.run_query(f'select {query}', database=state.export_session, output=output if backgrounded else None)
 
-        return state
+                return state
 
     def completion(self, state: ReplState):
         if not state.export_session: