kaqing 2.0.171__py3-none-any.whl → 2.0.204__py3-none-any.whl

adam/commands/devices/device_postgres.py

@@ -1,11 +1,10 @@
 from adam.commands.bash.bash_completer import BashCompleter
-from adam.commands.command import Command, InvalidState
+from adam.commands.command import Command
 from adam.commands.devices.device import Device
-from adam.commands.postgres.postgres_context import PostgresContext
+from adam.commands.postgres.postgres_databases import PostgresDatabases, pg_path
 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, log2, wait_log
-from adam.utils_k8s.pods import Pods
+from adam.utils import tabulize, log2, wait_log
 
 class DevicePostgres(Command, Device):
     COMMAND = f'{ReplState.P}:'
@@ -36,79 +35,98 @@ class DevicePostgres(Command, Device):
     def help(self, _: ReplState):
         return f'{DevicePostgres.COMMAND}\t move to Postgres Operations device'
 
-    def ls(self, cmd: str, state: ReplState):
-        if state.pg_path:
-            pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
-            if pg.db:
-                self.show_pg_tables(pg)
-            else:
-                self.show_pg_databases(pg)
-        else:
-            self.show_pg_hosts(state)
+    def pod(self, state: ReplState) -> str:
+        pod, _ = PostgresDatabases.pod_and_container(state.namespace)
+        return pod
 
-    def cat(self, cmd: str, state: ReplState):
-        if state.pod:
-            return self.bash(state, state, cmd.split(' '))
+    def pod_names(self, state: ReplState) -> list[str]:
+        return [self.pod(state)]
 
-        return InvalidState()
+    def default_container(self, state: ReplState) -> str:
+        _, container = PostgresDatabases.pod_and_container(state.namespace)
+        return container
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.pod_targetted:
+            self.bash(state, state, ['ls'])
+            return
+
+        with pg_path(state) as (host, database):
+            if database:
+                tabulize(pg_table_names(state), header='NAME', separator=',')
+            elif host:
+                tabulize(pg_database_names(state), header='DATABASE', separator=',')
+            else:
+                self.show_pg_hosts(state)
 
     def show_pg_hosts(self, state: ReplState):
         if state.namespace:
-            def line(pg: PostgresContext):
-                return f'{pg.path()},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
-
-            lines = [line(PostgresContext.apply(state.namespace, pg)) for pg in PostgresContext.hosts(state.namespace)]
-
-            log(lines_to_tabular(lines, 'NAME,ENDPOINT,USERNAME,PASSWORD', separator=','))
+            tabulize(PostgresDatabases.hosts(state),
+                     lambda p: f'{p.host},{p.endpoint()}:{p.port()},{p.username()},{p.password()}',
+                     header='NAME,ENDPOINT,USERNAME,PASSWORD',
+                     separator=',')
         else:
-            def line(pg: PostgresContext):
-                return f'{pg.path()},{pg.namespace},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
-
-            lines = [line(PostgresContext.apply(state.namespace, pg)) for pg in PostgresContext.hosts(state.namespace)]
-
-            log(lines_to_tabular(lines, 'NAME,NAMESPACE,ENDPOINT,USERNAME,PASSWORD', separator=','))
-
-    def show_pg_databases(self, pg: PostgresContext):
-        log(lines_to_tabular(pg_database_names(pg.namespace, pg.path()), 'DATABASE', separator=','))
-
-    def show_pg_tables(self, pg: PostgresContext):
-        log(lines_to_tabular(pg_table_names(pg.namespace, pg.path()), 'NAME', separator=','))
+            tabulize(PostgresDatabases.hosts(state),
+                     lambda p: f'{p.host},{p.namespace},{p.endpoint()}:{p.port()},{p.username()},{p.password()}',
+                     header='NAME,NAMESPACE,ENDPOINT,USERNAME,PASSWORD',
+                     separator=',')
 
     def cd(self, dir: str, state: ReplState):
         if dir == '':
             state.pg_path = None
-        else:
-            context: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path, arg=dir)
-            # patch up state.namespace from pg cd
-            if not state.namespace and context.namespace:
-                state.namespace = context.namespace
-            state.pg_path = context.path()
+            return
+
+        with pg_path(state) as (host, database):
+            if dir == '..':
+                if database:
+                    database = None
+                else:
+                    host = None
+            else:
+                tks = dir.split('@')
+                if not host:
+                    host = tks[0]
+                else:
+                    database = tks[0]
+
+                if len(tks) > 1:
+                    state.namespace = tks[1]
+
+            if database:
+                state.pg_path = f'{host}/{database}'
+            else:
+                state.pg_path = host
 
     def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
-        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path) if state.pg_path else None
-        if pg and pg.db:
-            return {cmd: {'..': None}}
-        elif pg and pg.host:
-            return {cmd: {'..': None} | {p: None for p in pg_database_names(state.namespace, pg.path())}}
-        else:
-            return {cmd: {p: None for p in PostgresContext.hosts(state.namespace)}}
+        return {cmd: {d: None for d in self.direct_dirs(cmd, state, default=default)}}
+
+    def direct_dirs(self, cmd: str, state: ReplState, default: dict = {}) -> list[str]:
+        with pg_path(state) as (host, database):
+            if database:
+                return ['..']
+            elif host:
+                return ['..'] + [p for p in pg_database_names(state)]
+            else:
+                return [p for p in PostgresDatabases.host_names(state.namespace)]
 
     def pwd(self, state: ReplState):
         words = []
 
-        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
+        with pg_path(state) as (host, database):
+            if host:
+                words.append(f'host/{host}')
+            if database:
+                words.append(f'database/{database}')
 
-        if pg.host:
-            words.append(f'host/{pg.host}')
-        if pg.db:
-            words.append(f'database/{pg.db}')
-
-        return '\t'.join([f'{ReplState.P}:>'] + (words if words else ['/']))
+            return '\t'.join([f'{ReplState.P}:>'] + (words if words else ['/']))
 
     def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
-        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
-        if pg.db:
-            return True, chain.run(f'pg {cmd}', state)
+        with pg_path(state) as (_, database):
+            if not database:
+                database = PostgresDatabases.default_db()
+
+            if database:
+                return True, chain.run(f'pg {cmd}', state)
 
         return False, None
 
@@ -116,15 +134,13 @@ class DevicePostgres(Command, Device):
         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=','))
+        tabulize(PostgresDatabases.tables(state, default_schema=True), lambda d: d['name'], 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}')
+        PostgresDatabases.run_sql(state, f'select * from {table} limit {rows}')
 
     def bash(self, s0: ReplState, s1: ReplState, args: list[str]):
-        pod, container = PostgresContext.pod_and_container(s1.namespace)
+        pod, container = PostgresDatabases.pod_and_container(s1.namespace)
         log2(f'Running on {pod}(container:{container})...')
 
         return super().bash(s0, s1, args)

adam/commands/devices/devices.py

@@ -7,7 +7,7 @@ from adam.commands.devices.device_postgres import DevicePostgres
 from adam.repl_state import ReplState
 
 class Devices:
-    def device(state: ReplState) -> Device:
+    def of(state: ReplState) -> Device:
         if state.device == ReplState.A:
             return DeviceApp()
         elif state.device == ReplState.C:

adam/commands/download_cassandra_log.py

@@ -0,0 +1,45 @@
+from adam.commands.command import Command
+from adam.commands.devices.devices import Devices
+from adam.config import Config
+from adam.utils_k8s.pod_exec_result import PodExecResult
+from adam.utils import log2
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.repl_state import ReplState, RequiredState
+from adam.utils_k8s.pods import Pods
+
+class DownloadCassandraLog(Command):
+    COMMAND = 'download cassandra log'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DownloadCassandraLog, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DownloadCassandraLog.COMMAND
+
+    def required(self):
+        return RequiredState.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):
+            path = Config().get('logs.path', '/c3/cassandra/logs/system.log')
+            r: PodExecResult = CassandraNodes.exec(state.pod, state.namespace, f'cat {path}', backgrounded=True, no_history=True)
+
+            to_file = Pods.download_file(state.pod, 'cassandra', state.namespace, path)
+            log2(f'Downloaded to {to_file}.')
+
+            return r
+
+    def completion(self, state: ReplState):
+        return super().completion(state, pods=Devices.of(state).pods(state, '-'), auto='jit')
+
+    def help(self, _: ReplState):
+        return f'{DownloadCassandraLog.COMMAND}\t download cassandra system log'

adam/commands/download_file.py

@@ -0,0 +1,47 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.devices.devices import Devices
+from adam.config import Config
+from adam.utils_k8s.pod_exec_result import PodExecResult
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+from adam.utils_k8s.pods import Pods
+
+class DownloadFile(Command):
+    COMMAND = 'download file'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DownloadFile, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DownloadFile.COMMAND
+
+    def required(self):
+        return [RequiredState.CLUSTER_OR_POD, RequiredState.APP_APP, ReplState.P]
+
+    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 validate_args(args, state, name='file'):
+                to_file = Pods.download_file(Devices.of(state).pod(state),
+                                             Devices.of(state).default_container(state),
+                                             state.namespace,
+                                             args[0],
+                                             args[1] if len(args) > 1 else None)
+                log2(f'Downloaded to {to_file}.')
+
+                return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state, lambda: {f: None for f in Devices.of(state).files(state)}, pods=Devices.of(state).pods(state, '-'), auto='jit')
+
+    def help(self, _: ReplState):
+        return f'{DownloadFile.COMMAND} from-file [to-file]\t download file from pod'

adam/commands/export/clean_up_all_export_sessions.py

@@ -1,5 +1,5 @@
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_sessions import export_session
 from adam.repl_state import ReplState, RequiredState
 
 class CleanUpAllExportSessions(Command):
@@ -25,8 +25,8 @@ class CleanUpAllExportSessions(Command):
             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()
+            with export_session(state) as sessions:
+                sessions.clean_up_all()
 
             return state
 

adam/commands/export/clean_up_export_sessions.py

@@ -1,7 +1,7 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_sessions import export_session
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log, log2
 
 class CleanUpExportSessions(Command):
     COMMAND = 'clean up export sessions'
@@ -26,26 +26,14 @@ class CleanUpExportSessions(Command):
             return super().run(cmd, 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.')
+            with validate_args(args, state, name='export session', separator=',') as session_names:
+                with export_session(state) as sessions:
+                    sessions.clean_up(session_names)
 
-                    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.')
-
-            Exporter.clear_export_session_cache()
-
-            return state
+                return state
 
     def completion(self, _: ReplState):
         return {}
 
     def help(self, _: ReplState):
-        return f'{CleanUpExportSessions.COMMAND} <export-session-name>\t clean up export session'
+        return f'{CleanUpExportSessions.COMMAND} <export-session-name,...>\t clean up export sessions'

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

@@ -0,0 +1,40 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.export.export_sessions import ExportSessions, export_session
+from adam.repl_state import ReplState, RequiredState
+
+class DownloadExportSession(Command):
+    COMMAND = 'download export session'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DownloadExportSession, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DownloadExportSession.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):
+            with validate_args(args, state, name='export session') as session:
+                with export_session(state) as sessions:
+                    sessions.download_session(session)
+
+            return state
+
+    def completion(self, state: ReplState):
+        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/drop_export_database.py

@@ -1,9 +1,7 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
-from adam.utils import log2
-from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class DropExportDatabase(Command):
     COMMAND = 'drop export database'
@@ -28,25 +26,11 @@ class DropExportDatabase(Command):
             return super().run(cmd, state)
 
         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()
+            with validate_args(args, state, name='database name') as db:
+                with export_db(state) as dbs:
+                    dbs.drop(args[0])
 
-                return 'command-missing'
-
-            Exporter.drop_databases(state.sts, state.pod, state.namespace, args[0])
-
-            SQLite.clear_cache()
-            Athena.clear_cache()
-
-            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

@@ -1,8 +1,6 @@
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
-from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class DropExportDatabases(Command):
     COMMAND = 'drop all export databases'
@@ -27,12 +25,8 @@ class DropExportDatabases(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            Exporter.drop_databases(state.sts, state.pod, state.namespace)
-
-            SQLite.clear_cache()
-            Athena.clear_cache()
-
-            state.export_session = None
+            with export_db(state) as dbs:
+                dbs.drop_all()
 
             return state
 

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.export_handlers import export
+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):