kaqing 2.0.93__py3-none-any.whl → 2.0.115__py3-none-any.whl

adam/commands/export/utils_export.py

@@ -0,0 +1,253 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime
+import io
+import re
+import time
+from typing import Callable
+import boto3
+
+from adam.commands.cql.cql_utils import cassandra_table_names, run_cql, table_spec
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import elapsed_time, log2
+from adam.utils_athena import Athena
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.pods import Pods
+
+def export_tables(args: list[str], state: ReplState, max_workers = 0):
+    consistency = None
+    specs = None
+
+    if args:
+        consistency, specs = ExportSpec.parse_multiple(' '.join(args))
+
+    if not specs:
+        specs = [ExportSpec(t) for t in cassandra_table_names(state, keyspace=f'{state.namespace}_db')]
+
+    if not max_workers:
+        max_workers = Config().action_workers('export', 8)
+
+    if max_workers > 1 and len(specs) > 1:
+        log2(f'Executing on {len(specs)} Cassandra tables in parallel...')
+        start_time = time.time()
+        try:
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                futures = [executor.submit(export_table, spec, state, True, consistency=consistency) for spec in specs]
+                if len(futures) == 0:
+                    return []
+
+            return [future.result() for future in as_completed(futures)]
+        finally:
+            log2(f"{len(specs)} parallel table export elapsed time: {elapsed_time(start_time)} with {max_workers} workers")
+    else:
+        return [export_table(spec, state, multi_tables=len(specs) > 1, consistency=consistency) for spec in specs]
+
+class ExportSpec:
+    def __init__(self, table: str, columns: str = None, target_table: str = None):
+        self.table = table
+        self.columns = columns
+        self.target_table = target_table
+
+    def parse(spec_str: str) -> 'ExportSpec':
+        target = None
+
+        p = re.compile(r"(.*?)\s+as\s+(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            spec_str = match.group(1)
+            target = match.group(2)
+
+        table = spec_str
+        columns = None
+
+        p = re.compile('(.*?)\((.*)\)')
+        match = p.match(spec_str)
+        if match:
+            table = match.group(1)
+            columns = match.group(2)
+
+        return ExportSpec(table, columns, target)
+
+    def __eq__(self, other):
+        if isinstance(other, ExportSpec):
+            return self.table == other.table and self.columns == other.columns and self.target_table == other.target_table
+
+        return False
+
+    def __str__(self):
+        return f'{self.table}, {self.columns}, {self.target_table}'
+
+    def parse_multiple(spec_str: str) -> tuple[str, list['ExportSpec']]:
+        consistency = None
+
+        p = re.compile(r"(.*?)with\s+consistency\s+(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            spec_str = match.group(1).strip(' ')
+            consistency = match.group(2)
+
+        if spec_str:
+            p = r",\s*(?![^()]*\))"
+            specs = re.split(p, spec_str)
+
+            return consistency, [ExportSpec.parse(spec) for spec in specs]
+
+        return consistency, []
+
+def export_table(spec: ExportSpec, state: ReplState, multi_tables = True, consistency: str = None):
+    table = spec.table
+    columns = spec.columns
+    if not columns:
+        columns = Config().get('export.columns', f'<keys>')
+
+    if columns == '<keys>':
+        columns = ','.join(table_spec(state, table, on_any=True).keys())
+    elif columns == '<row-key>':
+        columns = table_spec(state, table, on_any=True).row_key()
+    elif columns == '*':
+        columns = ','.join([c.name for c in table_spec(state, table, on_any=True).columns])
+
+    if not columns:
+        log2(f'ERROR: Empty columns on {table}.')
+        return table
+
+    athena_table = spec.target_table if spec.target_table else table
+    if '.' in athena_table:
+        athena_table = athena_table.split('.')[-1]
+
+    temp_dir = Config().get('export.temp_dir', '/c3/cassandra/tmp')
+    session = state.export_session
+    create_db = not session
+    if create_db:
+        session = datetime.now().strftime("%Y%m%d%H%M%S")
+        state.export_session = session
+    db = f'export_{session}'
+
+    CassandraNodes.exec(state.pod, state.namespace, f'mkdir -p {temp_dir}/{session}', show_out=not multi_tables, shell='bash')
+    csv_file = f'{temp_dir}/{session}/{table}.csv'
+    succeeded = False
+    try:
+        suppress_ing_log = Config().is_debug() or multi_tables
+        queries = []
+        if consistency:
+            queries.append(f'CONSISTENCY {consistency}')
+        queries.append(f"COPY {table}({columns}) TO '{csv_file}' WITH HEADER = TRUE")
+        ing(f'Dumping table {table}{f" with consistency {consistency}" if consistency else ""}',
+            lambda: run_cql(state, ';'.join(queries), show_out=False),
+            suppress_log=suppress_ing_log)
+
+        def upload_to_s3():
+            bytes = Pods.read_file(state.pod, 'cassandra', state.namespace, csv_file)
+
+            s3 = boto3.client('s3')
+            s3.upload_fileobj(GeneratorStream(bytes), 'c3.ops--qing', f'export/{session}/{athena_table}/{table}.csv')
+
+        ing(f'Uploading to S3', upload_to_s3, suppress_log=suppress_ing_log)
+
+        def create_schema():
+            query = f'CREATE DATABASE IF NOT EXISTS {db};'
+            if Config().is_debug():
+                log2(query)
+            Athena.query(query, 'default')
+
+            query = f'DROP TABLE IF EXISTS {athena_table};'
+            if Config().is_debug():
+                log2(query)
+            Athena.query(query, db)
+
+            # columns = ', '.join([f'{h.strip(" ")} string' for h in header[0].split(',')])
+            athena_columns = ', '.join([f'{c} string' for c in columns.split(',')])
+            query = f'CREATE EXTERNAL TABLE IF NOT EXISTS {athena_table}(\n' + \
+                    f'    {athena_columns})\n' + \
+                        "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'\n" + \
+                        'WITH SERDEPROPERTIES (\n' + \
+                        '    "separatorChar" = ",",\n' + \
+                        '    "quoteChar"     = "\\"")\n' + \
+                    f"LOCATION 's3://c3.ops--qing/export/{session}/{athena_table}'\n" + \
+                        'TBLPROPERTIES ("skip.header.line.count"="1");'
+            if Config().is_debug():
+                log2(query)
+            try:
+                Athena.query(query, db)
+            except Exception as e:
+                log2(f'*** Failed query:\n{query}')
+                raise e
+
+        ing(f"Creating database {db}" if create_db else f"Creating table {athena_table}", create_schema, suppress_log=suppress_ing_log)
+
+        succeeded = True
+    except Exception as e:
+        log2(e)
+    finally:
+        ing('Cleaning up temporary files',
+            lambda: CassandraNodes.exec(state.pod, state.namespace, f'rm -rf {csv_file}', show_out=False, shell='bash'),
+            suppress_log=suppress_ing_log)
+
+    if succeeded:
+        Athena.clear_cache()
+
+        if not suppress_ing_log:
+            query = f'select * from {athena_table} limit 10'
+            log2(query)
+            Athena.run_query(query, db)
+
+    return table
+
+def ing(msg: str, body: Callable[[], None], suppress_log=False):
+    if not suppress_log:
+        log2(f'{msg}...', nl=False)
+    body()
+    if not suppress_log:
+        log2(' OK')
+
+class GeneratorStream(io.RawIOBase):
+    def __init__(self, generator):
+        self._generator = generator
+        self._buffer = b''  # Buffer to store leftover bytes from generator yields
+
+    def readable(self):
+        return True
+
+    def _read_from_generator(self):
+        try:
+            chunk = next(self._generator)
+            if isinstance(chunk, str):
+                chunk = chunk.encode('utf-8')  # Encode if generator yields strings
+            self._buffer += chunk
+        except StopIteration:
+            pass  # Generator exhausted
+
+    def readinto(self, b):
+        # Fill the buffer if necessary
+        while len(self._buffer) < len(b):
+            old_buffer_len = len(self._buffer)
+            self._read_from_generator()
+            if len(self._buffer) == old_buffer_len:  # Generator exhausted and buffer empty
+                break
+
+        bytes_to_read = min(len(b), len(self._buffer))
+        b[:bytes_to_read] = self._buffer[:bytes_to_read]
+        self._buffer = self._buffer[bytes_to_read:]
+        return bytes_to_read
+
+    def read(self, size=-1):
+        if size == -1:  # Read all remaining data
+            while True:
+                old_buffer_len = len(self._buffer)
+                self._read_from_generator()
+                if len(self._buffer) == old_buffer_len:
+                    break
+            data = self._buffer
+            self._buffer = b''
+            return data
+        else:
+            # Ensure enough data in buffer
+            while len(self._buffer) < size:
+                old_buffer_len = len(self._buffer)
+                self._read_from_generator()
+                if len(self._buffer) == old_buffer_len:
+                    break
+
+            data = self._buffer[:size]
+            self._buffer = self._buffer[size:]
+            return data

adam/commands/help.py

@@ -30,14 +30,18 @@ class Help(Command):
         lines.append('NAVIGATION')
         lines.append('  a: | c: | l: | p:\t switch to another operational device: App, Cassandra, Audit or Postgres')
         lines.extend(section(ReplCommands.navigation()))
-        lines.append('CHECK CASSANDRA')
-        lines.extend(section(ReplCommands.cassandra_check()))
-        lines.append('CASSANDRA OPERATIONS')
+        lines.append('CASSANDRA')
         lines.extend(section(ReplCommands.cassandra_ops()))
+        lines.append('POSTGRES')
+        lines.extend(section(ReplCommands.postgres_ops()))
+        lines.append('APP')
+        lines.extend(section(ReplCommands.app_ops()))
+        lines.append('EXPORT DB')
+        lines.extend(section(ReplCommands.export_ops()))
+        lines.append('AUDIT')
+        lines.extend(section(ReplCommands.audit_ops()))
         lines.append('TOOLS')
         lines.extend(section(ReplCommands.tools()))
-        lines.append('APP')
-        lines.extend(section(ReplCommands.app()))
         lines.append('')
         lines.extend(section(ReplCommands.exit()))
 

adam/commands/issues.py

@@ -21,11 +21,17 @@ class Issues(Command):
     def command(self):
         return Issues.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)
 
         state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
         args, show = Command.extract_options(args, ['-s', '--show'])
 
         results = run_checks(state.sts, state.namespace, state.pod, show_output=show)

adam/commands/kubectl.py

@@ -0,0 +1,41 @@
+import subprocess
+
+from adam.commands.command import Command
+from adam.repl_state import ReplState, RequiredState
+
+class Kubectl(Command):
+    COMMAND = 'k'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(Kubectl, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return Kubectl.COMMAND
+
+    def required(self):
+        return RequiredState.NAMESPACE
+
+    def run(self, cmd: str, state: ReplState):
+        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
+
+        subprocess.run(["kubectl"] + args)
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+
+    def help(self, _: ReplState):
+        return f'{Kubectl.COMMAND} \t run a kubectl command'

adam/commands/login.py

@@ -8,7 +8,7 @@ from adam.config import Config
 from adam.sso.idp import Idp
 from adam.sso.idp_login import IdpLogin
 from adam.commands.command import Command
-from adam.repl_state import ReplState
+from adam.repl_state import ReplState, RequiredState
 from adam.utils import log, log2
 
 class Login(Command):
@@ -26,6 +26,9 @@ class Login(Command):
     def command(self):
         return Login.COMMAND
 
+    def required(self):
+        return ReplState.NON_L
+
     def run(self, cmd: str, state: ReplState):
         def custom_handler(signum, frame):
             AppSession.ctrl_c_entered = True
@@ -59,8 +62,8 @@ class Login(Command):
 
         return state
 
-    def completion(self, _: ReplState):
-        return {}
+    def completion(self, state: ReplState):
+        return super().completion(state)
 
     def help(self, _: ReplState):
         return f'{Login.COMMAND}\t SSO login'

adam/commands/logs.py

@@ -1,6 +1,6 @@
 from adam.commands.command import Command
 from adam.config import Config
-from adam.k8s_utils.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.repl_state import ReplState, RequiredState
 
 class Logs(Command):
@@ -33,6 +33,7 @@ class Logs(Command):
         return CassandraNodes.exec(state.pod, state.namespace, f'cat {path}')
 
     def completion(self, _: ReplState):
+        # available only on cli
         return {}
 
     def help(self, _: ReplState):

adam/commands/ls.py

@@ -1,20 +1,21 @@
 import copy
 
+from adam.commands.bash.bash import Bash
 from adam.commands.command import Command
 from adam.commands.commands_utils import show_pods, show_rollout
-from adam.commands.cql.cqlsh import Cqlsh
 from adam.commands.postgres.postgres_utils import pg_database_names, pg_table_names
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
-from adam.k8s_utils.custom_resources import CustomResources
-from adam.k8s_utils.ingresses import Ingresses
-from adam.k8s_utils.kube_context import KubeContext
-from adam.k8s_utils.statefulsets import StatefulSets
-from adam.pod_exec_result import PodExecResult
+from adam.utils_athena import Athena
+from adam.utils_k8s.app_pods import AppPods
+from adam.utils_k8s.custom_resources import CustomResources
+from adam.utils_k8s.ingresses import Ingresses
+from adam.utils_k8s.kube_context import KubeContext
+from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState
 from adam.utils import lines_to_tabular, log, log2
 from adam.apps import Apps
-from adam.utils_athena import audit_table_names
+from adam.utils_audits import Audits
 
 class Ls(Command):
     COMMAND = 'ls'
@@ -45,7 +46,7 @@ class Ls(Command):
 
         if state.device == ReplState.P:
             if state.pg_path:
-                pg = PostgresSession(state.namespace, state.pg_path)
+                pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
                 if pg.db:
                     self.show_pg_tables(pg)
                 else:
@@ -53,7 +54,13 @@ class Ls(Command):
             else:
                 self.show_pg_hosts(state)
         elif state.device == ReplState.A:
-            if state.app_env:
+            if state.app_pod:
+                return Bash().run('bash ' + cmd, state)
+            elif state.app_app:
+                pods = AppPods.pod_names(state.namespace, state.app_env, state.app_app)
+
+                log(lines_to_tabular(pods, 'POD_NAME'))
+            elif state.app_env:
                 def line(n: str, ns: str):
                     host = Ingresses.get_host(Config().get('app.login.ingress', '{app_id}-k8singr-appleader-001').replace('{app_id}', f'{ns}-{n}'), ns)
                     if not host:
@@ -76,10 +83,7 @@ class Ls(Command):
             self.show_audit_log_tables()
         else:
             if state.pod:
-                r: PodExecResult = Cqlsh().run(f'cql describe tables', state)
-                if r.stderr:
-                    log(r.stderr)
-                log(r.stdout)
+                return Bash().run('bash ' + cmd, state)
             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)
@@ -91,7 +95,7 @@ class Ls(Command):
     def show_statefulsets(self):
         ss = StatefulSets.list_sts_names()
         if len(ss) == 0:
-            log2('No cassandra statefulsets found.')
+            log2('No Cassandra clusters found.')
             return
 
         app_ids = CustomResources.get_app_ids()
@@ -110,37 +114,45 @@ class Ls(Command):
 
     def show_pg_hosts(self, state: ReplState):
         if state.namespace:
-            def line(pg: PostgresSession):
-                return f'{pg.directory()},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
+            def line(pg: PostgresContext):
+                return f'{pg.path()},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
 
-            lines = [line(PostgresSession(state.namespace, pg)) for pg in PostgresSession.hosts(state.namespace)]
+            lines = [line(PostgresContext.apply(state.namespace, pg)) for pg in PostgresContext.hosts(state.namespace)]
 
             log(lines_to_tabular(lines, 'NAME,ENDPOINT,USERNAME,PASSWORD', separator=','))
         else:
-            def line(pg: PostgresSession):
-                return f'{pg.directory()},{pg.namespace},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
+            def line(pg: PostgresContext):
+                return f'{pg.path()},{pg.namespace},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
 
-            lines = [line(PostgresSession(state.namespace, pg)) for pg in PostgresSession.hosts(state.namespace)]
+            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: PostgresSession):
-        log(lines_to_tabular(pg_database_names(pg.namespace, pg.directory()), 'DATABASE', 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: PostgresSession):
-        log(lines_to_tabular(pg_table_names(pg.namespace, pg.directory()), 'NAME', separator=','))
+    def show_pg_tables(self, pg: PostgresContext):
+        log(lines_to_tabular(pg_table_names(pg.namespace, pg.path()), 'NAME', separator=','))
 
     def show_audit_log_tables(self):
-        log(lines_to_tabular(audit_table_names(), 'NAME', separator=','))
+        log(lines_to_tabular(Athena.table_names(), 'NAME', separator=','))
 
     def completion(self, state: ReplState):
-        if state.pod:
-            return {}
+        if state.device == ReplState.C:
+            def pod_names():
+                return [p for p in StatefulSets.pod_names(state.sts, state.namespace)]
+
+            if state.sts:
+                return super().completion(state) | {f'@{p}': {'ls': None} for p in pod_names()}
+            else:
+                return {Ls.COMMAND: {n: None for n in StatefulSets.list_sts_names()}}
+        elif state.device == ReplState.A and state.app_app:
+            def pod_names():
+                return [p for p in AppPods.pod_names(state.namespace, state.app_env, state.app_app)]
 
-        if not state.sts:
-            return {Ls.COMMAND: {n: None for n in StatefulSets.list_sts_names()}}
+            return super().completion(state) | {f'@{p}': {'ls': None} for p in pod_names()}
 
-        return {Ls.COMMAND: None}
+        return super().completion(state)
 
     def help(self, _: ReplState):
         return f'{Ls.COMMAND} [device:]\t list apps, envs, clusters, nodes, pg hosts or pg databases'

adam/commands/medusa/medusa_backup.py

@@ -2,9 +2,9 @@ from datetime import datetime
 import re
 
 from adam.commands.command import Command
-from adam.k8s_utils.statefulsets import StatefulSets
+from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState, RequiredState
-from adam.k8s_utils.custom_resources import CustomResources
+from adam.utils_k8s.custom_resources import CustomResources
 from adam.utils import log2
 
 

adam/commands/medusa/medusa_restore.py

@@ -1,9 +1,9 @@
 from datetime import datetime
 
 from adam.commands.command import Command
-from adam.k8s_utils.statefulsets import StatefulSets
+from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState, RequiredState
-from adam.k8s_utils.custom_resources import CustomResources
+from adam.utils_k8s.custom_resources import CustomResources
 from adam.config import Config
 from adam.utils import lines_to_tabular, log2
 

adam/commands/medusa/medusa_show_backupjobs.py

@@ -1,7 +1,7 @@
 from adam.commands.command import Command
-from adam.k8s_utils.statefulsets import StatefulSets
+from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState, RequiredState
-from adam.k8s_utils.custom_resources import CustomResources
+from adam.utils_k8s.custom_resources import CustomResources
 from adam.utils import lines_to_tabular, log2
 
 
@@ -29,6 +29,7 @@ class MedusaShowBackupJobs(Command):
         state, args = self.apply_state(args, state)
         if not self.validate_state(state):
             return state
+
         ns = state.namespace
         dc = StatefulSets.get_datacenter(state.sts, ns)
         if not dc:

adam/commands/medusa/medusa_show_restorejobs.py

@@ -1,7 +1,7 @@
 from adam.commands.command import Command
-from adam.k8s_utils.statefulsets import StatefulSets
+from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState, RequiredState
-from adam.k8s_utils.custom_resources import CustomResources
+from adam.utils_k8s.custom_resources import CustomResources
 from adam.utils import lines_to_tabular, log2
 
 class MedusaShowRestoreJobs(Command):

adam/commands/nodetool.py

@@ -4,12 +4,11 @@ from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterOrPodCommandHelper
 from adam.commands.nodetool_commands import NODETOOL_COMMANDS
 from adam.config import Config
-from adam.k8s_utils.cassandra_clusters import CassandraClusters
-from adam.k8s_utils.cassandra_nodes import CassandraNodes
-from adam.pod_exec_result import PodExecResult
-from adam.repl_session import ReplSession
+from adam.utils_k8s.cassandra_clusters import CassandraClusters
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log, random_alphanumeric
+from adam.utils import log
+from adam.utils_k8s.statefulsets import StatefulSets
 
 class NodeTool(Command):
     COMMAND = 'nodetool'
@@ -41,26 +40,22 @@ class NodeTool(Command):
         command = f"nodetool -u {user} -pw {pw} {' '.join(args)}"
 
         if state.pod:
-            results: PodExecResult = CassandraNodes.exec(state.pod, state.namespace, command, show_out=True)
-            if results and results.log_file and Config().get('repl.history.push-cat-remote-log-file', True):
-                ReplSession().prompt_session.history.append_string(f'bash cat {results.log_file}')
+            return CassandraNodes.exec(state.pod, state.namespace, command, show_out=True)
         elif state.sts:
-            results: list[PodExecResult] = CassandraClusters.exec(state.sts, state.namespace, command, action='nodetool', show_out=True)
-            if results and Config().get('repl.history.push-cat-remote-log-file', True):
-                for result in results:
-                    if result.log_file:
-                        ReplSession().prompt_session.history.append_string(f'bash {result.pod} cat {result.log_file}')
+            return CassandraClusters.exec(state.sts, state.namespace, command, action='nodetool', show_out=True)
 
-        return results
+        return state
 
     def completion(self, state: ReplState):
         if state.pod or state.sts:
-            return {NodeTool.COMMAND: {'help': None} | {c: None for c in NODETOOL_COMMANDS}}
+            d = {c: {'&': None} for c in NODETOOL_COMMANDS}
+            return {NodeTool.COMMAND: {'help': None} | d} | \
+                   {f'@{p}': {NodeTool.COMMAND: d} for p in StatefulSets.pod_names(state.sts, state.namespace)}
 
         return {}
 
     def help(self, _: ReplState):
-        return f'{NodeTool.COMMAND} <sub-command>\t run nodetool with arguments'
+        return f'{NodeTool.COMMAND} <sub-command> [&]\t run nodetool with arguments'
 
 class NodeToolCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):

adam/commands/postgres/postgres.py

@@ -5,7 +5,7 @@ from adam.commands.postgres.psql_completions import psql_completions
 from adam.commands.postgres.postgres_utils import pg_table_names
 from .postgres_ls import PostgresLs
 from .postgres_preview import PostgresPreview
-from .postgres_session import PostgresSession
+from .postgres_context import PostgresContext
 from adam.repl_state import ReplState
 from adam.utils import log, log2
 
@@ -62,7 +62,7 @@ class Postgres(Command):
 
             return state
 
-        PostgresSession(state.namespace, state.pg_path).run_sql(' '.join(args))
+        PostgresContext.apply(state.namespace, state.pg_path).run_sql(' '.join(args))
 
     def completion(self, state: ReplState):
         if state.device != state.P:
@@ -70,7 +70,7 @@ class Postgres(Command):
             return {}
 
         leaf = {}
-        session = PostgresSession(state.namespace, state.pg_path)
+        session = PostgresContext.apply(state.namespace, state.pg_path)
         if session.db:
           if pg_table_names(state.namespace, state.pg_path):
             leaf = psql_completions(state.namespace, state.pg_path)
@@ -86,7 +86,7 @@ class Postgres(Command):
             return {}
 
     def help(self, _: ReplState):
-        return f'[{Postgres.COMMAND}] <sql-statements>\t run psql with queries'
+        return f'<sql-statements>\t run queries on Postgres databases'
 
 class PostgresCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):