kaqing 2.0.98__py3-none-any.whl → 2.0.203__py3-none-any.whl

adam/commands/devices/device_auit_log.py

@@ -0,0 +1,49 @@
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+from adam.repl_state import ReplState
+from adam.utils import tabulize, log
+from adam.utils_athena import Athena
+
+class DeviceAuditLog(Command, Device):
+    COMMAND = f'{ReplState.L}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DeviceAuditLog, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DeviceAuditLog.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.L
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DeviceAuditLog.COMMAND}\t move to Audit Log Operations device'
+
+    def ls(self, cmd: str, _: ReplState):
+        tabulize(Athena.table_names(), header='NAME', separator=',')
+
+    def pwd(self, _: ReplState):
+        return '\t'.join([f'{ReplState.L}:>', '/'])
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        return True, chain.run(f'audit {cmd}', state)
+
+    def show_tables(self, _: ReplState):
+        tabulize(Athena.table_names(), separator=',')
+
+    def show_table_preview(self, _: ReplState, table: str, rows: int):
+        Athena.run_query(f'select * from {table} limit {rows}')

adam/commands/devices/device_cass.py

@@ -0,0 +1,179 @@
+from adam.commands.bash.bash_completer import BashCompleter
+from adam.commands.command import Command, InvalidStateException
+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 tabulize, 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
+from adam.utils_k8s.statefulsets import StatefulSets
+
+class DeviceCass(Command, Device):
+    COMMAND = f'{ReplState.C}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DeviceCass, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DeviceCass.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.C
+
+        return state
+
+    def pod(self, state: ReplState) -> str:
+        return state.pod
+
+    def pod_names(self, state: ReplState) -> list[str]:
+        return StatefulSets.pod_names(state.sts, state.namespace)
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DeviceCass.COMMAND}\t move to Cassandra Operations device'
+
+    def default_container(self, _: ReplState) -> str:
+        return 'cassandra'
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.pod:
+            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)
+        else:
+            self.show_statefulsets()
+
+    def ls_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.sts:
+            return super().completion(state) | {f'@{p}': {'ls': None} for p in self.pod_names(state)}
+        else:
+            return {cmd: {n: None for n in StatefulSets.list_sts_names()}}
+
+    def show_statefulsets(self):
+        ss = StatefulSets.list_sts_names()
+        if len(ss) == 0:
+            log2('No Cassandra clusters found.')
+            return
+
+        app_ids = CustomResources.get_app_ids()
+        list = []
+        for s in ss:
+            cr_name = CustomResources.get_cr_name(s)
+            app_id = 'Unknown'
+            if cr_name in app_ids:
+                app_id = app_ids[cr_name]
+            list.append(f"{s} {app_id}")
+
+        header = 'STATEFULSET_NAME@NAMESPACE APP_ID'
+        if KubeContext.in_cluster_namespace():
+            header = 'STATEFULSET_NAME APP_ID'
+        tabulize(list, header=header)
+
+    def cd(self, dir: str, state: ReplState):
+        if dir == '':
+            state.sts = None
+            state.pod = None
+        elif dir == '..':
+            if state.pod:
+                state.pod = None
+            else:
+                state.sts = None
+        else:
+            if not state.sts:
+                ss_and_ns = dir.split('@')
+                state.sts = ss_and_ns[0]
+                if len(ss_and_ns) > 1:
+                    state.namespace = ss_and_ns[1]
+            elif not state.pod:
+                p, _ = KubeContext.is_pod_name(dir)
+                if p:
+                    state.pod = p
+                else:
+                    names = CassandraClusters.pod_names_by_host_id(state.sts, state.namespace)
+                    if dir in names:
+                        state.pod = names[dir]
+                    else:
+                        state.pod = dir
+
+    def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.pod:
+            return {cmd: {'..': None}}
+        elif state.sts:
+            return {cmd: {'..': None} | {p: None for p in self.pod_names(state)}}
+        else:
+            return {cmd: {p: None for p in StatefulSets.list_sts_names()}}
+
+    def pwd(self, state: ReplState):
+        words = []
+
+        if state.sts:
+            words.append(f'sts/{state.sts}')
+        if state.pod:
+            words.append(f'pod/{state.pod}')
+
+        return '\t'.join([f'{ReplState.C}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        if state.sts:
+            return True, chain.run(f'cql {cmd}', state)
+
+        return False, None
+
+    def enter(self, state: ReplState):
+        auto_enter = Config().get('repl.c.auto-enter', 'cluster')
+        if auto_enter and auto_enter in ['cluster', 'first-pod']:
+            sts = StatefulSets.list_sts_name_and_ns()
+            if not sts:
+                log2("No Cassandra clusters found.")
+            elif not state.sts and len(sts) == 1:
+                cluster = sts[0]
+                state.sts = cluster[0]
+                state.namespace = cluster[1]
+                if auto_enter == 'first-pod':
+                    state.pod = f'{state.sts}-0'
+                if KubeContext().in_cluster_namespace:
+                    wait_log(f'Moving to the only Cassandra cluster: {state.sts}...')
+                else:
+                    wait_log(f'Moving to the only Cassandra cluster: {state.sts}@{state.namespace}...')
+
+    def show_tables(self, state: ReplState):
+        tabulize(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 self.pod_names(state)}
+
+        return completions
+

adam/commands/devices/device_export.py

@@ -0,0 +1,87 @@
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+from adam.commands.export.export_databases import ExportDatabases, export_db
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import tabulize, log2, wait_log
+
+class DeviceExport(Command, Device):
+    COMMAND = f'{ReplState.X}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DeviceExport, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DeviceExport.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.X
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DeviceExport.COMMAND}\t move to Export Database Operations device'
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.export_session:
+            tabulize(ExportDatabases.table_names(state.export_session), header='NAME', separator=',')
+        else:
+            ExportDatabases.show_databases()
+
+    def show_table_preview(self, state: ReplState, table: str, rows: int):
+        if state.export_session:
+            with export_db(state) as dbs:
+                dbs.sql(f'select * from {table} limit {rows}')
+
+    def cd(self, dir: str, state: ReplState):
+        if dir in ['', '..']:
+            state.export_session = None
+        else:
+            state.export_session = dir
+
+    def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.export_session:
+            return {cmd: {'..': None} | {n: None for n in ExportDatabases.database_names()}}
+        else:
+            return {cmd: {n: None for n in ExportDatabases.database_names()}}
+
+    def pwd(self, state: ReplState):
+        words = []
+
+        if state.export_session:
+            words.append(state.export_session)
+
+        return '\t'.join([f'{ReplState.X}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        if cmd.startswith('select '):
+            cmd = f'xelect {cmd[7:]}'
+
+        result = chain.run(cmd, state)
+        if type(result) is ReplState:
+            if state.export_session and not result.export_session:
+                state.export_session = None
+
+        return True, result
+
+    def enter(self, state: ReplState):
+        if auto_enter := Config().get('repl.x.auto-enter', 'no'):
+            if auto_enter == 'latest':
+                wait_log(f'Moving to the latest export database...')
+                if dbs := ExportDatabases.database_names():
+                    state.export_session = sorted(dbs)[-1]
+                else:
+                    log2('No export database found.')
+

adam/commands/devices/device_postgres.py

@@ -0,0 +1,160 @@
+from adam.commands.bash.bash_completer import BashCompleter
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+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 tabulize, log2, wait_log
+
+class DevicePostgres(Command, Device):
+    COMMAND = f'{ReplState.P}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DevicePostgres, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DevicePostgres.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.P
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DevicePostgres.COMMAND}\t move to Postgres Operations device'
+
+    def pod(self, state: ReplState) -> str:
+        pod, _ = PostgresDatabases.pod_and_container(state.namespace)
+        return pod
+
+    def pod_names(self, state: ReplState) -> list[str]:
+        return [self.pod(state)]
+
+    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:
+            tabulize(PostgresDatabases.hosts(state),
+                     lambda p: f'{p.host},{p.endpoint()}:{p.port()},{p.username()},{p.password()}',
+                     header='NAME,ENDPOINT,USERNAME,PASSWORD',
+                     separator=',')
+        else:
+            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
+            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 = {}):
+        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 = []
+
+        with pg_path(state) as (host, database):
+            if host:
+                words.append(f'host/{host}')
+            if database:
+                words.append(f'database/{database}')
+
+            return '\t'.join([f'{ReplState.P}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        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
+
+    def enter(self, _: ReplState):
+        wait_log('Inspecting postgres database instances...')
+
+    def show_tables(self, state: ReplState):
+        tabulize(PostgresDatabases.tables(state, default_schema=True), lambda d: d['name'], separator=',')
+
+    def show_table_preview(self, state: ReplState, table: str, rows: int):
+        PostgresDatabases.run_sql(state, f'select * from {table} limit {rows}')
+
+    def bash(self, s0: ReplState, s1: ReplState, args: list[str]):
+        pod, container = PostgresDatabases.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/devices/devices.py

@@ -0,0 +1,25 @@
+from adam.commands.devices.device import Device
+from adam.commands.devices.device_app import DeviceApp
+from adam.commands.devices.device_auit_log import DeviceAuditLog
+from adam.commands.devices.device_cass import DeviceCass
+from adam.commands.devices.device_export import DeviceExport
+from adam.commands.devices.device_postgres import DevicePostgres
+from adam.repl_state import ReplState
+
+class Devices:
+    def of(state: ReplState) -> Device:
+        if state.device == ReplState.A:
+            return DeviceApp()
+        elif state.device == ReplState.C:
+            return DeviceCass()
+        elif state.device == ReplState.L:
+            return DeviceAuditLog()
+        elif state.device == ReplState.P:
+            return DevicePostgres()
+        elif state.device == ReplState.X:
+            return DeviceExport()
+
+        return DeviceCass()
+
+    def all():
+        return [DeviceApp(), DeviceCass(), DeviceAuditLog(), DevicePostgres(), DeviceExport()]

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/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.export_sessions import export_session
+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):
+            with export_session(state) as sessions:
+                sessions.clean_up_all()
+
+            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

@@ -0,0 +1,39 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.export.export_sessions import export_session
+from adam.repl_state import ReplState, RequiredState
+
+class CleanUpExportSessions(Command):
+    COMMAND = 'clean up export sessions'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(CleanUpExportSessions, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return CleanUpExportSessions.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', separator=',') as session_names:
+                with export_session(state) as sessions:
+                    sessions.clean_up(session_names)
+
+                return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{CleanUpExportSessions.COMMAND} <export-session-name,...>\t clean up export sessions'