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

adam/commands/postgres/utils_postgres.py

@@ -1,36 +1,45 @@
 import functools
 
-from adam.commands.postgres.postgres_context import PostgresContext
+from adam.commands.postgres.postgres_databases import PostgresDatabases, pg_path
 from adam.repl_state import ReplState
 from adam.utils import log2, wait_log
 from adam.utils_k8s.pods import Pods
 
 TestPG = [False]
 
+def direct_dirs(state: ReplState) -> list[str]:
+    with pg_path(state) as (host, database):
+        if database:
+            return ['..']
+        elif host:
+            return ['..'] + pg_database_names(state)
+        else:
+            return PostgresDatabases.host_names(state.namespace)
+
+def pg_database_names(state: ReplState):
+    # cache on pg_path
+    return _pg_database_names(state, state.pg_path)
+
 @functools.lru_cache()
-def pg_database_names(ns: str, pg_path: str):
+def _pg_database_names(state: ReplState, pg_path: str):
     if TestPG[0]:
         return ['azops88_c3ai_c3']
 
     wait_log('Inspecting Postgres Databases...')
 
-    pg = PostgresContext.apply(ns, pg_path)
-    return [db['name'] for db in pg.databases() if db['owner'] == PostgresContext.default_owner()]
+    return [db['name'] for db in PostgresDatabases.databases(state, default_owner=True)]
+
+def pg_table_names(state: ReplState):
+    # cache on pg_path
+    return _pg_table_names(state, state.pg_path)
 
 @functools.lru_cache()
-def pg_table_names(ns: str, pg_path: str):
+def _pg_table_names(state: ReplState, pg_path: str):
     if TestPG[0]:
         return ['C3_2_XYZ1']
 
     wait_log('Inspecting Postgres Database...')
-    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 = PostgresContext.apply(ns, pg_path)
-    if pg.db:
-        return pg.tables()
-
-    return []
+    return [table['name'] for table in PostgresDatabases.tables(state, default_schema=True)]
 
 class PostgresPodService:
     def __init__(self, handler: 'PostgresExecHandler'):
@@ -39,22 +48,26 @@ class PostgresPodService:
     def exec(self, command: str, show_out=True):
         state = self.handler.state
 
-        pod, container = PostgresContext.pod_and_container(state.namespace)
+        pod, container = PostgresDatabases.pod_and_container(state.namespace)
         if not pod:
             log2('Cannot locate postgres agent or ops pod.')
             return state
 
         return Pods.exec(pod, container, state.namespace, command, show_out=show_out)
 
-    def sql(self, args: list[str], background=False):
+    def sql(self, args: list[str], backgrounded=False):
         state = self.handler.state
 
-        PostgresContext.apply(state.namespace, state.pg_path).run_sql(' '.join(args), background=background)
+        query = args
+        if isinstance(args, list):
+            query = ' '.join(args)
+
+        PostgresDatabases.run_sql(state, query, backgrounded=backgrounded)
 
 class PostgresExecHandler:
-    def __init__(self, state: ReplState, background=False):
+    def __init__(self, state: ReplState, backgrounded=False):
         self.state = state
-        self.background = background
+        self.backgrounded = backgrounded
 
     def __enter__(self):
         return PostgresPodService(self)
@@ -62,5 +75,5 @@ class PostgresExecHandler:
     def __exit__(self, exc_type, exc_val, exc_tb):
         return False
 
-def postgres(state: ReplState, background=False):
-    return PostgresExecHandler(state, background=background)
+def postgres(state: ReplState, backgrounded=False):
+    return PostgresExecHandler(state, backgrounded=backgrounded)

adam/commands/preview_table.py

@@ -1,3 +1,4 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.devices.devices import Devices
 from adam.repl_state import ReplState, RequiredState
@@ -18,14 +19,15 @@ class PreviewTable(Command):
         return PreviewTable.COMMAND
 
     def required(self):
-        return [RequiredState.CLUSTER_OR_POD, RequiredState.PG_DATABASE, ReplState.L]
+        return [RequiredState.CLUSTER_OR_POD, RequiredState.PG_DATABASE, ReplState.L, RequiredState.EXPORT_DB]
 
     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):
-            Devices.device(state).preview(args[0] if args else None, state)
+            with validate_args(args, state, at_least=1) as table:
+                Devices.of(state).preview(table, state)
 
             return state
 

adam/commands/pwd.py

@@ -2,7 +2,7 @@ from adam.app_session import AppSession
 from adam.commands.command import Command
 from adam.commands.devices.devices import Devices
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log
+from adam.utils import tabulize, log, log_exc
 
 class Pwd(Command):
     COMMAND = 'pwd'
@@ -25,17 +25,15 @@ class Pwd(Command):
 
         with self.validate(args, state) as (_, state):
             host = "unknown"
-            try:
+            with log_exc():
                 app_session: AppSession = AppSession.create('c3', 'c3')
                 host = app_session.host
-            except:
-                pass
 
-            log(lines_to_tabular([device.pwd(state) for device in Devices.all()] + [
+            tabulize([device.pwd(state) for device in Devices.all()] + [
                 f'',
                 f'HOST\t{host}',
                 f'NAMESPACE\t{state.namespace if state.namespace else "/"}',
-            ], 'DEVICE\tLOCATION', separator='\t'))
+            ], header='DEVICE\tLOCATION', separator='\t')
             log()
 
             return state

adam/commands/reaper/reaper_forward.py

@@ -8,7 +8,7 @@ from adam.commands.reaper.utils_reaper import Reapers, port_forwarding
 from adam.config import Config
 from adam.repl_session import ReplSession
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import lines_to_tabular, log2
+from adam.utils import tabulize, log2
 
 class ReaperForward(Command):
     COMMAND = 'reaper forward'
@@ -56,7 +56,7 @@ class ReaperForward(Command):
                     'reaper-password': spec["password"]
                 }
                 log2()
-                log2(lines_to_tabular([f'{k},{v}' for k, v in d.items()], separator=','))
+                tabulize(d.items(), lambda a: f'{a[0]},{a[1]}', separator=',')
 
                 for k, v in d.items():
                     ReplSession().prompt_session.history.append_string(f'cp {k}')

adam/commands/reaper/reaper_run_abort.py

@@ -1,3 +1,4 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.reaper.utils_reaper import reaper
 from adam.repl_state import ReplState, RequiredState
@@ -26,19 +27,12 @@ class ReaperRunAbort(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not args:
-                if state.in_repl:
-                    log2('Specify run id to abort.')
-                else:
-                    Command.display_help()
+            with validate_args(args, state, name='run id') as run_id:
+                with reaper(state) as http:
+                    http.put(f'repair_run/{run_id}/state/ABORTED')
 
                 return state
 
-            with reaper(state) as http:
-                http.put(f'repair_run/{args[0]}/state/ABORTED')
-
-            return state
-
     def completion(self, state: ReplState):
         return super().completion(state)
 

adam/commands/reaper/reaper_runs.py

@@ -2,7 +2,7 @@ from adam.commands.command import Command
 from adam.commands.reaper.utils_reaper import reaper
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import convert_seconds, epoch, lines_to_tabular, log, log2
+from adam.utils import convert_seconds, epoch, tabulize, log, log2
 
 class ReaperRuns(Command):
     COMMAND = 'reaper show runs'
@@ -54,7 +54,7 @@ class ReaperRuns(Command):
 
                 runs = response.json()
                 if runs:
-                    log(lines_to_tabular(sorted([line(run) for run in runs], reverse=True), header, separator=","))
+                    tabulize(sorted([line(run) for run in runs], reverse=True), header=header, separator=",")
                 else:
                     log2('No running runs found.')
                     log2()
@@ -66,7 +66,7 @@ class ReaperRuns(Command):
 
                 runs = response.json()
                 if runs:
-                    log(lines_to_tabular(sorted([line(run) for run in runs], reverse=True), header, separator=","))
+                    tabulize(sorted([line(run) for run in runs], reverse=True), header=header, separator=",")
                 else:
                     log2('No runs found.')
 

adam/commands/reaper/reaper_schedule_activate.py

@@ -1,7 +1,12 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.reaper.utils_reaper import Reapers, reaper
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
+
+import nest_asyncio
+nest_asyncio.apply()
+
+import asyncio
 
 class ReaperScheduleActivate(Command):
     COMMAND = 'reaper activate schedule'
@@ -26,20 +31,15 @@ class ReaperScheduleActivate(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not args:
-                log2('Specify schedule to activate.')
-
-                return state
-
-            schedule_id = args[0]
-            with reaper(state) as http:
-                http.put(f'repair_schedule/{schedule_id}?state=ACTIVE')
-                Reapers.show_schedule(state, schedule_id)
+            with validate_args(args, state, name='schedule') as schedule_id:
+                with reaper(state) as http:
+                    http.put(f'repair_schedule/{schedule_id}?state=ACTIVE')
+                    Reapers.show_schedule(state, schedule_id)
 
-            return schedule_id
+                return schedule_id
 
     def completion(self, state: ReplState):
-        return super().completion(state, lambda: {id: None for id in Reapers.cached_schedule_ids(state)})
+        return super().completion(state, lambda: {id: None for id in Reapers.cached_schedule_ids(state)}, auto_key='reaper.schedules')
 
     def help(self, _: ReplState):
         return f'{ReaperScheduleActivate.COMMAND} <schedule-id>\t resume reaper schedule'

adam/commands/reaper/reaper_schedule_start.py

@@ -1,7 +1,7 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.reaper.utils_reaper import Reapers, reaper
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
 
 class ReaperScheduleStart(Command):
     COMMAND = 'reaper start schedule'
@@ -26,20 +26,15 @@ class ReaperScheduleStart(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not args:
-                log2('Specify schedule to activate.')
+            with validate_args(args, state, name='schedule') as schedule_id:
+                with reaper(state) as http:
+                    http.post(f'repair_schedule/start/{schedule_id}')
+                    Reapers.show_schedule(state, schedule_id)
 
-                return state
-
-            schedule_id = args[0]
-            with reaper(state) as http:
-                http.post(f'repair_schedule/start/{schedule_id}')
-                Reapers.show_schedule(state, schedule_id)
-
-            return schedule_id
+                return schedule_id
 
     def completion(self, state: ReplState):
-        return super().completion(state, lambda: {id: None for id in Reapers.cached_schedule_ids(state)})
+        return super().completion(state, lambda: {id: None for id in Reapers.cached_schedule_ids(state)}, auto_key='reaper.schedules')
 
     def help(self, _: ReplState):
         return f'{ReaperScheduleStart.COMMAND} <schedule-id>\t start reaper runs for schedule'

adam/commands/reaper/reaper_schedule_stop.py

@@ -1,7 +1,7 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.reaper.utils_reaper import Reapers, reaper
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
 
 class ReaperScheduleStop(Command):
     COMMAND = 'reaper stop schedule'
@@ -26,20 +26,15 @@ class ReaperScheduleStop(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not args:
-                log2('Specify run schedule to stop.')
+            with validate_args(args, state, name='schedule') as schedule_id:
+                with reaper(state) as http:
+                    http.put(f'repair_schedule/{schedule_id}?state=PAUSED')
+                    Reapers.show_schedule(state, schedule_id)
 
-                return state
-
-            with reaper(state) as http:
-                schedule_id = args[0]
-                http.put(f'repair_schedule/{schedule_id}?state=PAUSED')
-                Reapers.show_schedule(state, schedule_id)
-
-            return schedule_id
+                return schedule_id
 
     def completion(self, state: ReplState):
-        return super().completion(state, lambda: {id: None for id in Reapers.cached_schedule_ids(state)})
+        return super().completion(state, lambda: {id: None for id in Reapers.cached_schedule_ids(state)}, auto_key='reaper.schedules')
 
     def help(self, _: ReplState):
         return f'{ReaperScheduleStop.COMMAND} <schedule-id>\t pause reaper schedule'

adam/commands/reaper/utils_reaper.py

@@ -7,7 +7,7 @@ import re
 import requests
 from adam.config import Config
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log2, wait_log
+from adam.utils import tabulize, log2, wait_log
 from adam.utils_k8s.k8s import port_forwarding
 
 class ReaperService:
@@ -100,7 +100,8 @@ class Reapers:
         # forced refresh of schedule list
         if not filter:
             Reapers.schedules_ids_by_cluster[state.sts] = [schedule['id'] for schedule in schedules]
-        Reapers.show_schedules_tabular(schedules)
+
+        tabulize(schedules, lambda s: f"{s['id']} {s['state']} {s['cluster_name']} {s['keyspace_name']}", header='ID STATE CLUSTER KEYSPACE', to=2)
 
     def schedule_ids(state: ReplState, show_output = True, filter: Callable[[list[dict]], dict] = None):
         schedules = Reapers.list_schedules(state, show_output=show_output, filter=filter)
@@ -117,9 +118,6 @@ class Reapers:
 
             return res
 
-    def show_schedules_tabular(schedules: list[dict]):
-        log2(lines_to_tabular([f"{schedule['id']} {schedule['state']} {schedule['cluster_name']} {schedule['keyspace_name']}" for schedule in schedules], 'ID STATE CLUSTER KEYSPACE'))
-
     def list_reaper_pods(sts_name: str, namespace: str) -> List[client.V1Pod]:
         v1 = client.CoreV1Api()
 
@@ -193,4 +191,13 @@ class Reapers:
         if is_service:
             return Reapers.svc_name()
         else:
-            return Reapers.pod_name(state)
+            return Reapers.pod_name(state)
+
+    def schedules_auto_completion(ids: callable):
+        auto = Config().get('reaper.schedules-auto-complete', 'off')
+
+        leaf = None
+        if auto == 'on':
+            leaf = {id: None for id in ids()}
+
+        return (leaf, auto == 'lazy')

adam/commands/repair/repair_log.py

@@ -31,10 +31,7 @@ class RepairLog(Command):
             return state
 
     def completion(self, state: ReplState):
-        if state.sts:
-            return super().completion(state)
-
-        return {}
+        return super().completion(state)
 
     def help(self, _: ReplState):
         return f'{RepairLog.COMMAND}\t get repair job logs'

adam/commands/repair/repair_run.py

@@ -6,7 +6,7 @@ from adam.repl_state import ReplState, RequiredState
 from adam.config import Config
 from adam.commands.reaper.reaper_runs_abort import ReaperRunsAbort
 from adam.commands.reaper.reaper_schedule_stop import ReaperScheduleStop
-from adam.utils import log2
+from adam.utils import log2, log_exc
 
 class RepairRun(Command):
     COMMAND = 'repair run'
@@ -35,14 +35,12 @@ class RepairRun(Command):
             if len(args) == 1:
                 replace = args[0] == 'replace'
 
-            try:
+            with log_exc():
                 log2("Stopping all reaper schedules...")
                 for schedule_id in Reapers.cached_schedule_ids(state):
                     ReaperScheduleStop().run(f'reaper stop schedule {schedule_id}', state)
                 log2("Aborting all reaper runs...")
                 state = ReaperRunsAbort().run('reaper abort runs', state)
-            except:
-                pass
 
             image = Config().get('repair.image', 'ci-registry.c3iot.io/cloudops/cassrepair:2.0.11')
             secret = Config().get('repair.secret', 'ciregistryc3iotio')
@@ -61,10 +59,7 @@ class RepairRun(Command):
             return state
 
     def completion(self, state: ReplState):
-        if state.sts:
-            return super().completion(state)
-
-        return {}
+        return super().completion(state)
 
     def help(self, _: ReplState):
         return f'{RepairRun.COMMAND} [replace]\t start a repair job, default not replacing'

adam/commands/repair/repair_scan.py

@@ -1,5 +1,3 @@
-import time
-
 from adam.commands.command import Command
 from adam.utils_k8s.pods import Pods
 from adam.repl_state import ReplState, RequiredState
@@ -62,10 +60,7 @@ class RepairScan(Command):
             return state
 
     def completion(self, state: ReplState):
-        if state.sts:
-            return super().completion(state)
-
-        return {}
+        return super().completion(state)
 
     def help(self, _: ReplState):
         return f'{RepairScan.COMMAND} [n]\t scan last n days repair log, default 7 days'

adam/commands/repair/repair_stop.py

@@ -1,7 +1,6 @@
 from adam.commands.command import Command
 from adam.utils_k8s.jobs import Jobs
 from adam.repl_state import ReplState, RequiredState
-from adam.config import Config
 
 class RepairStop(Command):
     COMMAND = 'repair stop'
@@ -32,10 +31,7 @@ class RepairStop(Command):
             return state
 
     def completion(self, state: ReplState):
-        if state.sts:
-            return super().completion(state)
-
-        return {}
+        return super().completion(state)
 
     def help(self, _: ReplState):
         return f'{RepairStop.COMMAND}\t delete a repair job'

adam/commands/restart_cluster.py

@@ -0,0 +1,47 @@
+from adam.commands import extract_options
+from adam.commands.command import Command
+from adam.utils_k8s.pods import Pods
+from adam.utils_k8s.statefulsets import StatefulSets
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+
+class RestartCluster(Command):
+    COMMAND = 'restart cluster'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(RestartCluster, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return RestartCluster.COMMAND
+
+    def required(self):
+        return RequiredState.CLUSTER
+
+    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 extract_options(args, '--force') as (args, forced):
+                if not forced:
+                    log2('Please add --force for restarting all nodes in a cluster.')
+
+                    return 'force-needed'
+
+                log2(f'Restarting all pods from {state.sts}...')
+                for pod_name in StatefulSets.pod_names(state.sts, state.namespace):
+                    Pods.delete(pod_name, state.namespace)
+
+                return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state, {'--force': None})
+
+    def help(self, _: ReplState):
+        return f"{RestartCluster.COMMAND} --force\t restart all the nodes in the cluster"

adam/commands/restart_node.py

@@ -0,0 +1,51 @@
+from adam.commands import extract_options
+from adam.commands.command import Command
+from adam.commands.devices.devices import Devices
+from adam.utils_k8s.pods import Pods
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+
+class RestartNode(Command):
+    COMMAND = 'restart node'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(RestartNode, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return RestartNode.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):
+            if not state.pod:
+                log2("'pod' is required")
+
+                return 'pod-needed'
+
+            with extract_options(args, '--force') as (args, forced):
+                if not forced:
+                    log2('Please add --force for restarting pod.')
+
+                    return 'force-needed'
+
+                log2(f'Restarting {state.pod}...')
+                Pods.delete(state.pod, state.namespace)
+
+                return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state, {'--force': None}, pods=Devices.of(state).pods(state, '-'))
+
+    def help(self, _: ReplState):
+        return f"{RestartNode.COMMAND} --force\t restart the node"

adam/commands/restart_nodes.py

@@ -0,0 +1,47 @@
+from adam.commands import extract_options, validate_args
+from adam.commands.command import Command
+from adam.utils_k8s.pods import Pods
+from adam.utils_k8s.statefulsets import StatefulSets
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+
+class RestartNodes(Command):
+    COMMAND = 'restart nodes'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(RestartNodes, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return RestartNodes.COMMAND
+
+    def required(self):
+        return RequiredState.CLUSTER
+
+    def run(self, cmd: str, state: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, state)
+
+        with self.validate(args, state, apply=False) as (args, state):
+            with extract_options(args, '--force') as (args, forced):
+                with validate_args(args, state, name='pod name'):
+                    if not forced:
+                        log2('Please add --force for restarting nodes.')
+
+                        return 'force-needed'
+
+                    for arg in args:
+                        Pods.delete(arg, state.namespace)
+
+                    return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state, lambda: {p: {'--force': None} for p in StatefulSets.pod_names(state.sts, state.namespace)})
+
+    def help(self, _: ReplState):
+        return f"{RestartNodes.COMMAND} <pod-name>... --force\t restart Cassandra nodes"

adam/commands/shell.py

@@ -1,7 +1,9 @@
 import os
 
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.repl_state import ReplState
+from adam.utils import log2
 
 class Shell(Command):
     COMMAND = ':sh'
@@ -22,8 +24,13 @@ class Shell(Command):
         if not(args := self.args(cmd)):
             return super().run(cmd, state)
 
-        with self.validate(args, state):
-            os.system('QING_DROPPED=true bash')
+        with self.validate(args, state) as (args, _):
+            with validate_args(args, state, at_least=0) as args_str:
+                if args_str:
+                    os.system(args_str)
+                    log2()
+                else:
+                    os.system('QING_DROPPED=true bash')
 
             return state