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

adam/commands/cql/cqlsh.py

@@ -3,7 +3,8 @@ 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 adam.utils_k8s.statefulsets import StatefulSets
+from .cql_utils import run_cql
 from adam.repl_state import ReplState, RequiredState
 from adam.utils import log, log2
 
@@ -33,11 +34,14 @@ class Cqlsh(Command):
         if not self.validate_state(state):
             return state
 
+        background = False
         opts = []
         cqls = []
-        for arg in args:
+        for index, arg in enumerate(args):
             if arg.startswith('--'):
                 opts.append(arg)
+            elif index == len(args) -1 and arg == '&':
+                background = True
             elif arg != '-e':
                 cqls.append(arg)
         if not cqls:
@@ -51,7 +55,7 @@ class Cqlsh(Command):
             return 'no-cql'
 
         cql = ' '.join(cqls)
-        return run_cql(state, cql, opts, show_out=True)
+        return run_cql(state, cql, opts, show_out=True, background=background)
 
     def completion(self, state: ReplState) -> dict[str, any]:
         if state.device != state.C:
@@ -59,12 +63,13 @@ class Cqlsh(Command):
             return {}
 
         if state.sts or state.pod:
-            return {Cqlsh.COMMAND: None} | cql_completions(state)
+            return cql_completions(state) | \
+                   {f'@{p}': cql_completions(state) for p in StatefulSets.pod_names(state.sts, state.namespace)}
 
         return {}
 
     def help(self, _: ReplState) -> str:
-        return f'[{Cqlsh.COMMAND}] <cql-statements>\t run cqlsh with queries'
+        return f'<cql-statements> [&]\t run cqlsh with queries'
 
 class CqlCommandHelper(click.Command):
     def get_help(self, ctx: click.Context):

adam/commands/deploy/deploy.py

@@ -4,7 +4,7 @@ from adam.commands.command import Command
 from adam.commands.deploy.deploy_pg_agent import DeployPgAgent
 from adam.commands.deploy.deploy_pod import DeployPod
 from .deploy_frontend import DeployFrontend
-from adam.repl_state import ReplState
+from adam.repl_state import ReplState, RequiredState
 
 class Deploy(Command):
     COMMAND = 'deploy'
@@ -22,10 +22,16 @@ class Deploy(Command):
     def command(self):
         return Deploy.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)
 
+        if not self.validate_state(state):
+            return state
+
         return super().intermediate_run(cmd, state, args, Deploy.cmd_list())
 
     def cmd_list():

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.py

@@ -4,7 +4,7 @@ from adam.commands.command import Command
 from adam.commands.deploy.undeploy_frontend import UndeployFrontend
 from adam.commands.deploy.undeploy_pg_agent import UndeployPgAgent
 from adam.commands.deploy.undeploy_pod import UndeployPod
-from adam.repl_state import ReplState
+from adam.repl_state import ReplState, RequiredState
 
 class Undeploy(Command):
     COMMAND = 'undeploy'
@@ -22,10 +22,16 @@ class Undeploy(Command):
     def command(self):
         return Undeploy.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)
 
+        if not self.validate_state(state):
+            return state
+
         return super().intermediate_run(cmd, state, args, Undeploy.cmd_list())
 
     def cmd_list():

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/devices.py

@@ -1,6 +1,35 @@
 from adam.commands.command import Command
 from adam.repl_state import ReplState
 
+class DeviceExport(Command):
+    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'{DeviceAuditLog.COMMAND}\t move to Export Database Operations device'
+
 class DeviceAuditLog(Command):
     COMMAND = f'{ReplState.L}:'
 

adam/commands/export/export.py

@@ -0,0 +1,60 @@
+from adam.commands.export.utils_export import export_tables
+from adam.commands.command import Command
+from adam.commands.cql.cql_utils import cassandra_table_names
+from adam.repl_state import ReplState, RequiredState
+from adam.sql.sql_completer import SqlCompleter
+from adam.utils_k8s.statefulsets import StatefulSets
+
+class ExportTables(Command):
+    COMMAND = 'export'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportTables, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportTables.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)
+
+        state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
+        if not state.pod:
+            state.push()
+            state.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
+        try:
+            export_tables(args, state)
+        finally:
+            state.pop()
+
+        return state
+
+    def completion(self, state: ReplState):
+        def sc():
+            return SqlCompleter(
+                lambda: cassandra_table_names(state),
+                dml='export',
+                columns = lambda table: ['id', '*'],
+                variant='cql'
+            )
+
+        if state.sts:
+            return {ExportTables.COMMAND: sc()} | \
+                {f'@{p}': {ExportTables.COMMAND: sc()} for p in StatefulSets.pod_names(state.sts, state.namespace)}
+
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{ExportTables.COMMAND} [TABLE] [as target-name] [with consistency <level>]\t export table'

adam/commands/export/export_on_x.py

@@ -0,0 +1,76 @@
+from adam.commands.command import Command
+from adam.repl_state import ReplState, RequiredState
+from adam.sql.sql_completer import SqlCompleter
+from adam.utils import log2
+from adam.utils_athena import Athena
+
+class ExportUseX(Command):
+    COMMAND = 'use'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportUseX, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportUseX.COMMAND
+
+    def required(self):
+        return ReplState.X
+
+    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
+
+        if not args:
+            if state.in_repl:
+                log2('Specify database to use.')
+            else:
+                log2('* database is missing.')
+
+                Command.display_help()
+
+            return 'command-missing'
+
+        state.export_session = args[0].replace('export_', '') if args[0].startswith('export_') else args[0]
+        Athena.clear_cache()
+
+        return state
+
+    def completion(self, state: ReplState):
+        dict = super().completion(state, {n: None for n in Athena.database_names()})
+
+        if state.export_session:
+            db = f'export_{state.export_session}'
+            dict = dict | {'select': SqlCompleter(
+                lambda: Athena.table_names(database=db, function='export'),
+                dml='select',
+                columns=lambda table: Athena.column_names(database=db, function='export'),
+                variant='athena'
+            )}
+
+        return dict
+
+    def help(self, _: ReplState):
+        return f'{ExportUseX.COMMAND} <export db name>\t use Export Database'
+
+# No action body, only for a help entry and auto-completion
+class ExportSelectX(Command):
+    COMMAND = 'select'
+
+    def command(self):
+        return ExportSelectX.COMMAND
+
+    def required(self):
+        return RequiredState.EXPORT_DB
+
+    def help(self, _: ReplState):
+        return f'<sql-select-statements>\t run queries on export database'

adam/commands/export/export_rmdbs.py

@@ -0,0 +1,65 @@
+import boto3
+
+from adam.commands.command import Command
+from adam.commands.export.utils_export import ing
+from adam.config import Config
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+from adam.utils_athena import Athena
+
+class RemoveExportDatabases(Command):
+    COMMAND = '&rmdbs'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(RemoveExportDatabases, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return RemoveExportDatabases.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)
+
+        state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
+        dbs = Athena.database_names('export_')
+        def drop_all_exports():
+            for db in dbs:
+                query = f'DROP DATABASE {db} CASCADE'
+                if Config().is_debug():
+                    log2(query)
+                Athena.query(query)
+
+        if Config().is_debug():
+            drop_all_exports()
+        else:
+            ing(f'Droping {len(dbs)} databases', drop_all_exports)
+
+        def delete_s3_folder():
+            s3 = boto3.resource('s3')
+            bucket = s3.Bucket('c3.ops--qing')
+            bucket.objects.filter(Prefix='export/').delete()
+
+        if Config().is_debug():
+            delete_s3_folder()
+        else:
+            ing(f'Deleting s3 folder: export', delete_s3_folder)
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{RemoveExportDatabases.COMMAND}\t remove all export databases'

adam/commands/export/export_select.py

@@ -0,0 +1,68 @@
+from adam.commands.command import Command
+from adam.repl_state import ReplState, RequiredState
+from adam.sql.sql_completer import SqlCompleter
+from adam.utils import log2
+from adam.utils_athena import Athena
+
+class ExportSelect(Command):
+    COMMAND = '&select'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportSelect, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportSelect.COMMAND
+
+    def required(self):
+        return [RequiredState.CLUSTER_OR_POD, ReplState.X]
+
+    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
+
+        if not args:
+            if state.in_repl:
+                log2('Use a SQL statement.')
+            else:
+                log2('* SQL statement is missing.')
+
+                Command.display_help()
+
+            return 'command-missing'
+
+        query = ' '.join(args)
+
+        Athena.run_query(f'select {query}', database=f'export_{state.export_session}')
+
+        return state
+
+    def completion(self, state: ReplState):
+        if not state.export_session:
+            return {}
+
+        db = f'export_{state.export_session}'
+
+        # warm up the caches first time when x: drive is accessed
+        Athena.table_names(database=db, function='export')
+        Athena.column_names(database=db, function='export')
+        Athena.column_names(partition_cols_only=True, database=db, function='export')
+
+        return {ExportSelect.COMMAND: SqlCompleter(
+            lambda: Athena.table_names(database=db, function='export'),
+            dml='select',
+            columns=lambda table: Athena.column_names(database=db, function='export'),
+            variant='athena'
+        )}
+
+    def help(self, _: ReplState):
+        return f'<sql-select-statements>\t run queries on export database'

adam/commands/export/export_use.py

@@ -0,0 +1,56 @@
+from adam.commands.command import Command
+from adam.commands.export.export_select import ExportSelect
+from adam.config import Config
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+from adam.utils_athena import Athena
+
+class ExportUse(Command):
+    COMMAND = '&use'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportUse, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportUse.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)
+
+        state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
+        if not args:
+            if state.in_repl:
+                log2('Specify database to use.')
+            else:
+                log2('* database is missing.')
+
+                Command.display_help()
+
+            return 'command-missing'
+
+        state.export_session = args[0].replace('export_', '') if args[0].startswith('export_') else args[0]
+        Athena.clear_cache()
+
+        return state
+
+    def completion(self, state: ReplState):
+        # warm up the caches first time when l: drive is accessed
+        Athena.database_names()
+
+        return super().completion(state, {n: None for n in Athena.database_names()})
+
+    def help(self, _: ReplState):
+        return f'{ExportUse.COMMAND} <export db name>\t use Export Database'