kaqing 2.0.102__py3-none-any.whl → 2.0.104__py3-none-any.whl

adam/batch.py

@@ -9,7 +9,6 @@ from adam.commands.command_helpers import ClusterCommandHelper, ClusterOrPodComm
 from adam.commands.cql.cqlsh import CqlCommandHelper, Cqlsh
 from adam.commands.deploy.deploy import Deploy, DeployCommandHelper
 from adam.commands.deploy.undeploy import Undeploy, UndeployCommandHelper
-from adam.commands.describe.describe import Describe, DescribeCommandHelper
 from adam.commands.issues import Issues
 from adam.commands.login import Login
 from adam.commands.logs import Logs
@@ -97,19 +96,6 @@ def deploy(kubeconfig: str, config: str, param: list[str], namespace: str, extra
     run_command(Deploy(), kubeconfig, config, param, None, namespace, None, extra_args)
 
 
-@cli.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True), cls=DescribeCommandHelper, help='Describe keyspace(s) or table(s).')
-@click.option('--kubeconfig', '-k', required=False, metavar='path', help='path to kubeconfig file')
-@click.option('--config', default='params.yaml', metavar='path', help='path to kaqing parameters file')
-@click.option('--param', '-v', multiple=True, metavar='<key>=<value>', help='parameter override')
-@click.option('--cluster', '-c', required=False, metavar='statefulset', help='Kubernetes statefulset name')
-@click.option('--namespace', '-n', required=False, metavar='namespace', help='Kubernetes namespace')
-@click.option('--pod', '-p', required=False, metavar='pod', help='Kubernetes pod name')
-@click.option('--all-nodes', '-a', is_flag=True, help='execute on all Cassandra nodes')
-@click.argument('extra_args', nargs=-1, metavar='<cluster|pod>', type=click.UNPROCESSED)
-def describe(kubeconfig: str, config: str, param: list[str], cluster: str, namespace: str, pod: str, all_nodes: bool, extra_args):
-    run_command(Describe(), kubeconfig, config, param, cluster, namespace, pod, extra_args + ('&',) if all_nodes else extra_args)
-
-
 @cli.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True), cls=ClusterOrPodCommandHelper, help="Print Qing's issues.")
 @click.option('--kubeconfig', '-k', required=False, metavar='path', help='path to kubeconfig file')
 @click.option('--config', default='params.yaml', metavar='path', help='path to kaqing parameters file')

adam/commands/audit/audit.py

@@ -1,6 +1,7 @@
 import click
 
 from adam.commands.audit.audit_repair_tables import AuditRepairTables
+from adam.commands.audit.audit_run import AuditRun
 from adam.commands.command import Command
 from adam.config import Config
 from adam.repl_state import ReplState
@@ -54,22 +55,17 @@ class Audit(Command):
                 audit_column_names()
                 audit_column_names(partition_cols_only=True)
 
-            # def columns(_):
-            #     return audit_column_names()
-
-            return super().completion(state) | SqlCompleter.completions(
+            return super().completion(state) | SqlCompleter(
                 lambda: audit_table_names(),
                 columns=lambda table: audit_column_names(),
                 partition_columns=lambda table: audit_column_names(partition_cols_only=True),
                 variant='athena'
-            )
-        # | {
-        #         'desc': {table: None for table in audit_table_names()}}
+            ).completions_for_nesting()
 
         return {}
 
     def cmd_list():
-        return [AuditRepairTables()]
+        return [AuditRepairTables(), AuditRun()]
 
     def help(self, _: ReplState):
         return f'[{Audit.COMMAND}] <sql-statements>\t run SQL queries on Authena audit database'

adam/commands/audit/audit_repair_tables.py

@@ -6,7 +6,7 @@ from adam.commands.command import Command
 from adam.config import Config
 from adam.repl_state import ReplState
 from adam.utils import log, log2
-from adam.utils_athena import audit_query, run_audit_query
+from adam.utils_athena import AuditMeta, audit_query, get_meta, put_meta, run_audit_query
 
 class AuditRepairTables(Command):
     COMMAND = 'audit repair'
@@ -30,11 +30,12 @@ class AuditRepairTables(Command):
 
         state, args = self.apply_state(args, state)
 
-        tables = Config().get('audit.tables', 'audit').split(',')
+        tables = Config().get('audit.athena.repair-partition-tables', 'audit').split(',')
         if args:
             tables = args
 
-        self.repair(tables)
+        meta = get_meta()
+        self.repair(tables, meta)
 
         return state
 
@@ -52,43 +53,20 @@ class AuditRepairTables(Command):
     def auto_repair(self, hours: int):
         self.auto_repaired = True
 
-        state, _, rs = audit_query(f'select * from meta')
-        if state == 'SUCCEEDED':
-            do_repair = True
-            if len(rs) > 1:
-                try:
-                    checked_in = float(rs[1]['Data'][0]['VarCharValue'])
-                    do_repair = checked_in + hours * 60 * 60 < time.time()
-                except:
-                    pass
-
-            if do_repair:
-                tables = Config().get('audit.athena.tables', 'audit').split(',')
-                self.repair(tables, show_sql=True)
-                log2(f'Audit tables have been auto-repaired.')
-
-    def repair(self, tables: list[str], show_sql = False):
+        meta = get_meta()
+        if meta.checked_in + hours * 60 * 60 < time.time():
+            tables = Config().get('audit.athena.repair-partition-tables', 'audit').split(',')
+            self.repair(tables, meta, show_sql=True)
+            log2(f'Audit tables have been auto-repaired.')
+
+    def repair(self, tables: list[str], meta: AuditMeta, show_sql = False):
         with concurrent.futures.ThreadPoolExecutor(max_workers=Config().get('audit.workers', 3)) as executor:
             for table in tables:
                 if show_sql:
                     log(f'MSCK REPAIR TABLE {table}')
 
                 executor.submit(run_audit_query, f'MSCK REPAIR TABLE {table}', None,)
-            executor.submit(self.check_in,)
+            executor.submit(put_meta, 'check-in', meta,)
 
     def help(self, _: ReplState):
-        return f"{AuditRepairTables.COMMAND} \t run MSCK REPAIR command for new partition discovery"
-
-    def check_in(self):
-        payload = {
-            'action': 'check-in'
-        }
-        audit_endpoint = Config().get("audit.endpoint", "https://4psvtaxlcb.execute-api.us-west-2.amazonaws.com/prod/")
-        try:
-            response = requests.post(audit_endpoint, json=payload, timeout=Config().get("audit.timeout", 10))
-            if response.status_code in [200, 201]:
-                Config().debug(response.text)
-            else:
-                log2(f"Error: {response.status_code} {response.text}")
-        except requests.exceptions.Timeout as e:
-            log2(f"Timeout occurred: {e}")
+        return f"{AuditRepairTables.COMMAND} \t run MSCK REPAIR command for new partition discovery"

adam/commands/audit/audit_run.py

@@ -0,0 +1,49 @@
+from adam.commands.command import Command
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_athena import AuditMeta, find_new_clusters, get_meta, put_meta, run_audit_query
+
+class AuditRun(Command):
+    COMMAND = 'audit run'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(AuditRun, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+        self.auto_repaired = False
+
+    def command(self):
+        return AuditRun.COMMAND
+
+    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)
+
+        meta: AuditMeta = get_meta()
+        clusters = find_new_clusters(meta.cluster_last_checked)
+        if clusters:
+            put_meta('add-clusters', meta, clusters=clusters)
+            log2(f'Added {len(clusters)} new clusters.')
+            tables = Config().get('audit.athena.repair-cluster-tables', 'cluster').split(',')
+            for table in tables:
+                run_audit_query(f'MSCK REPAIR TABLE {table}')
+        else:
+            log2(f'No new clusters were found.')
+
+        return state
+
+    def completion(self, state: ReplState):
+        if state.device == ReplState.L:
+            return super().completion(state)
+
+        return {}
+
+    def help(self, _: ReplState):
+        return f"{AuditRun.COMMAND} \t run"

adam/commands/bash.py

@@ -3,6 +3,8 @@ from adam.utils_k8s.cassandra_clusters import CassandraClusters
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.pod_exec_result import PodExecResult
 from adam.repl_state import BashSession, ReplState, RequiredState
+from adam.utils_repl.automata_completer import AutomataCompleter
+from adam.utils_repl.state_machine import StateMachine
 
 class Bash(Command):
     COMMAND = 'bash'
@@ -84,9 +86,65 @@ class Bash(Command):
 
     def completion(self, state: ReplState):
         if state.pod or state.sts:
-            return {Bash.COMMAND: None}
+            return {Bash.COMMAND: AutomataCompleter(BashStateMachine())}
 
         return {}
 
     def help(self, _: ReplState):
-        return f'{Bash.COMMAND} [bash-commands]\t run bash on the Cassandra nodes'
+        return f'{Bash.COMMAND} [bash-commands]\t run bash on the Cassandra nodes'
+
+BASH_SPEC = [
+    # <command> ::= <simple_command> | <pipeline> | <conditional_command>
+    # <simple_command> ::= <word> <argument>* <redirection>*
+    # <pipeline> ::= <command> '|' <command>
+    # <conditional_command> ::= <command> '&&' <command> | <command> '||' <command>
+    # <word> ::= <letter> <letter_or_digit>*
+    # <argument> ::= <word>
+    # <redirection> ::= '>' <filename> | '<' <filename>
+    # <filename> ::= <word>
+    # <letter> ::= 'a' | 'b' | ... | 'z' | 'A' | 'B' | ... | 'Z'
+    # <digit> ::= '0' | '1' | ... | '9'
+    # <letter_or_digit> ::= <letter> | <digit>
+
+    '                                > word           > word',
+    'word                            > word           > word                                ^ |,>,2>,<,&,&&,||',
+    '-                               > pipe           > word_pipe',
+    '-                               > _rdr0_         > word_rdr0',
+    '-                               > _rdr1_         > word_rdr1',
+    '-                               > _rdr2_         > word_rdr2',
+    '-                               > &              > word_bg                             ^ |,>,2>,<,&,&&,||',
+    '-                               > &&|_or_        > word',
+    'word_a                          > word           > word',
+    'word_pipe                       > word           > word',
+    'word_rdr0                       > word           > word_rdr0_f',
+    'word_rdr1                       > word           > word_rdr1_f',
+    'word_rdr2                       > word           > word_rdr2_f',
+    'word_rdr1_f                     > pipe           > word_pipe                           ^ |,2>,<,&,&&,||',
+    '-                               > _rdr2_         > word_rdr2',
+    '-                               > _rdr0_         > word_rdr0',
+    'word_rdr2_f                     > pipe           > word_pipe                           ^ |,<,&,&&,||',
+    '-                               > _rdr0_         > word_rdr0',
+    '-                               > &              > word_bg                             ^ |,>,2>,<,&,&&,||',
+    '-                               > &&|_or_        > word',
+    'word_rdr0_f                     > pipe           > word_pipe                           ^ |,&,&&,||',
+    '-                               > &              > word_bg                             ^ |,>,2>,<,&,&&,||',
+    '-                               > &&|_or_        > word',
+    'word_bg                         > &&|_or_        >                                     ^ &&,||',
+]
+
+BASH_KEYWORDS = [
+    '&',
+    '&&',
+    '|',
+    '||',
+    '>',
+    '2>',
+    '>>',
+    '<'
+]
+class BashStateMachine(StateMachine[str]):
+    def spec(self) -> str:
+        return BASH_SPEC
+
+    def keywords(self) -> list[str]:
+        return BASH_KEYWORDS

adam/commands/cd.py

@@ -1,6 +1,6 @@
 from adam.commands.command import Command
 from adam.commands.postgres.postgres_utils import pg_database_names
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.utils_k8s.cassandra_clusters import CassandraClusters
 from adam.utils_k8s.kube_context import KubeContext
 from adam.utils_k8s.statefulsets import StatefulSets
@@ -36,11 +36,11 @@ class Cd(Command):
                 if dir == '':
                     state.pg_path = None
                 else:
-                    session = PostgresSession(state.namespace, state.pg_path)
+                    context: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path, arg=dir)
                     # patch up state.namespace from pg cd
-                    if not state.namespace and (ns := session.find_namespace(arg)):
-                        state.namespace = ns
-                    state.pg_path = session.directory(arg)
+                    if not state.namespace and context.namespace:
+                        state.namespace = context.namespace
+                    state.pg_path = context.path()
             elif state.device == ReplState.A:
                 if dir == '':
                     state.app_env = None
@@ -91,13 +91,13 @@ class Cd(Command):
 
     def completion(self, state: ReplState):
         if state.device == ReplState.P:
-            pg = PostgresSession(state.namespace, state.pg_path) if state.pg_path else None
+            pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path) if state.pg_path else None
             if pg and pg.db:
                 return {Cd.COMMAND: {'..': None}}
             elif pg and pg.host:
-                return {Cd.COMMAND: {'..': None} | {p: None for p in pg_database_names(state.namespace, pg.directory())}}
+                return {Cd.COMMAND: {'..': None} | {p: None for p in pg_database_names(state.namespace, pg.path())}}
             else:
-                return {Cd.COMMAND: {p: None for p in PostgresSession.hosts(state.namespace)}}
+                return {Cd.COMMAND: {p: None for p in PostgresContext.hosts(state.namespace)}}
         elif state.device == ReplState.A:
             if state.app_app:
                 return {Cd.COMMAND: {'..': None}}

adam/commands/commands_utils.py

@@ -1,5 +1,4 @@
 from concurrent.futures import ThreadPoolExecutor
-import time
 from kubernetes import client
 from typing import List
 
@@ -10,7 +9,7 @@ from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.pods import Pods
 from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState
-from adam.utils import convert_seconds, duration, lines_to_tabular, log, log2
+from adam.utils import duration, lines_to_tabular, log, log2
 
 def show_pods(pods: List[client.V1Pod], ns: str, show_namespace = True, show_host_id = True):
     if len(pods) == 0:

adam/commands/cql/cql_completions.py

@@ -1,4 +1,4 @@
-from adam.commands.cql.cql_utils import table_names
+from adam.commands.cql.cql_utils import cassandra_table_names
 from adam.config import Config
 from adam.repl_state import ReplState
 from adam.sql.sql_completer import SqlCompleter
@@ -8,8 +8,8 @@ def cql_completions(state: ReplState) -> dict[str, any]:
     return {
         'describe': {
             'keyspaces': None,
-            'table': {t: None for t in table_names(state)},
+            'table': {t: None for t in cassandra_table_names(state)},
             'tables': None},
-    } | SqlCompleter.completions(lambda: table_names(state), table_props=lambda: {
+    } | SqlCompleter(lambda: cassandra_table_names(state), table_props=lambda: {
         'GC_GRACE_SECONDS': ps
-    }, variant='cql')
+    }, variant='cql').completions_for_nesting()

adam/commands/cql/cql_utils.py

@@ -23,7 +23,7 @@ def keyspaces(state: ReplState, on_any=False):
 
     return parse_cql_desc_keyspaces(r.stdout if state.pod else r[0].stdout)
 
-def table_names(state: ReplState):
+def cassandra_table_names(state: ReplState):
     return [f'{k}.{t}' for k, ts in tables(state, on_any=True).items() for t in ts]
 
 @functools.lru_cache()

adam/commands/cql/cqlsh.py

@@ -3,7 +3,7 @@ import click
 from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterOrPodCommandHelper
 from adam.commands.cql.cql_completions import cql_completions
-from .cql_utils import run_cql, table_names
+from .cql_utils import run_cql, cassandra_table_names
 from adam.repl_state import ReplState, RequiredState
 from adam.utils import log, log2
 

adam/commands/deploy/deploy_pg_agent.py

@@ -1,5 +1,5 @@
 from adam.commands.command import Command
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
 
@@ -29,7 +29,7 @@ class DeployPgAgent(Command):
         if not self.validate_state(state):
             return state
 
-        PostgresSession.deploy_pg_agent(Config().get('pg.agent.name', 'ops-pg-agent'), state.namespace)
+        PostgresContext.deploy_pg_agent(Config().get('pg.agent.name', 'ops-pg-agent'), state.namespace)
 
     def completion(self, state: ReplState):
         return super().completion(state)

adam/commands/deploy/undeploy_pg_agent.py

@@ -1,6 +1,6 @@
 from adam.commands.command import Command
 from adam.commands.deploy.deploy_utils import deleting
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
 
@@ -30,7 +30,7 @@ class UndeployPgAgent(Command):
         if not self.validate_state(state):
             return state
 
-        deleting('pod', lambda: PostgresSession.undeploy_pg_agent(Config().get('pg.agent.name', 'ops-pg-agent'), state.namespace))
+        deleting('pod', lambda: PostgresContext.undeploy_pg_agent(Config().get('pg.agent.name', 'ops-pg-agent'), state.namespace))
 
         return state
 

adam/commands/ls.py

@@ -4,7 +4,7 @@ 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.utils_k8s.custom_resources import CustomResources
 from adam.utils_k8s.ingresses import Ingresses
@@ -45,7 +45,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:
@@ -110,25 +110,25 @@ 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=','))

adam/commands/nodetool.py

@@ -47,7 +47,7 @@ class NodeTool(Command):
 
     def completion(self, state: ReplState):
         if state.pod or state.sts:
-            return {NodeTool.COMMAND: {'help': None} | {c: None for c in NODETOOL_COMMANDS}}
+            return {NodeTool.COMMAND: {'help': None} | {c: {'&': None} for c in NODETOOL_COMMANDS}}
 
         return {}
 

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)

adam/commands/postgres/{postgres_session.py → postgres_context.py}

@@ -8,7 +8,28 @@ from adam.utils_k8s.pods import Pods
 from adam.utils_k8s.secrets import Secrets
 from adam.utils import log2
 
-class PostgresSession:
+class PostgresContext:
+    def apply(namespace: str, path: str, arg: str = None) -> 'PostgresContext':
+        context = PostgresContext(namespace, path)
+
+        if arg:
+            if arg == '..':
+                if context.db:
+                    context.db = None
+                else:
+                    context.host = None
+            else:
+                tks = arg.split('@')
+                if not context.host:
+                    context.host = tks[0]
+                else:
+                    context.db = tks[0]
+
+                if not namespace and tks[1]:
+                    context.namespace = tks[1]
+
+        return context
+
     def __init__(self, ns: str, path: str):
         self.namespace = ns
         self.conn_details = None
@@ -25,29 +46,7 @@ class PostgresSession:
             if len(tks) > 1:
                 self.db = tks[1]
 
-    def find_namespace(self, arg: str):
-        if arg:
-            tks = arg.split('@')
-            if len(tks) > 1:
-                return tks[1]
-
-        return None
-
-    def directory(self, arg: str = None):
-        if arg:
-            if arg == '..':
-                if self.db:
-                    self.db = None
-                else:
-                    self.host = None
-            else:
-                tks = arg.split('@')
-                arg = tks[0]
-                if not self.host:
-                    self.host = arg
-                else:
-                    self.db = arg
-
+    def path(self):
         if not self.host:
             return None
 
@@ -58,7 +57,7 @@ class PostgresSession:
         return f'{self.host}/{self.db}'
 
     def hosts(ns: str):
-        return PostgresSession.hosts_for_namespace(ns)
+        return PostgresContext.hosts_for_namespace(ns)
 
     @functools.lru_cache()
     def hosts_for_namespace(ns: str):
@@ -133,7 +132,7 @@ class PostgresSession:
         return dbs
 
     def run_sql(self, sql: str, show_out = True):
-        db = self.db if self.db else PostgresSession.default_db()
+        db = self.db if self.db else PostgresContext.default_db()
 
         if KubeContext.in_cluster():
             cmd1 = f'env PGPASSWORD={self.password()} psql -h {self.endpoint()} -p {self.port()} -U {self.username()} {db} --pset pager=off -c'
@@ -151,7 +150,7 @@ class PostgresSession:
             pod_name = Config().get('pg.agent.name', 'ops-pg-agent')
 
             if Config().get('pg.agent.just-in-time', False):
-                if not PostgresSession.deploy_pg_agent(pod_name, ns):
+                if not PostgresContext.deploy_pg_agent(pod_name, ns):
                     return
 
             real_pod_name = pod_name

adam/commands/postgres/postgres_utils.py

@@ -1,6 +1,6 @@
 import functools
 
-from adam.commands.postgres.postgres_session import PostgresSession
+from adam.commands.postgres.postgres_context import PostgresContext
 from adam.config import Config
 
 TestPG = [False]
@@ -12,8 +12,8 @@ def pg_database_names(ns: str, pg_path: str):
 
     Config().wait_log('Inspecting Postgres Databases...')
 
-    pg = PostgresSession(ns, pg_path)
-    return [db['name'] for db in pg.databases() if db['owner'] == PostgresSession.default_owner()]
+    pg = PostgresContext.apply(ns, pg_path)
+    return [db['name'] for db in pg.databases() if db['owner'] == PostgresContext.default_owner()]
 
 @functools.lru_cache()
 def pg_table_names(ns: str, pg_path: str):
@@ -21,10 +21,10 @@ def pg_table_names(ns: str, pg_path: str):
         return ['C3_2_XYZ1']
 
     Config().wait_log('Inspecting Postgres Database...')
-    return [table['name'] for table in pg_tables(ns, pg_path) if table['schema'] == PostgresSession.default_schema()]
+    return [table['name'] for table in pg_tables(ns, pg_path) if table['schema'] == PostgresContext.default_schema()]
 
 def pg_tables(ns: str, pg_path: str):
-    pg = PostgresSession(ns, pg_path)
+    pg = PostgresContext.apply(ns, pg_path)
     if pg.db:
         return pg.tables()
 

adam/commands/postgres/psql_completions.py

@@ -7,4 +7,4 @@ def psql_completions(ns: str, pg_path: str):
         '\d': None,
         '\dt': None,
         '\du': None
-    } | SqlCompleter.completions(lambda: pg_table_names(ns, pg_path))
+    } | SqlCompleter(lambda: pg_table_names(ns, pg_path)).completions_for_nesting()