kaqing 2.0.172__py3-none-any.whl → 2.0.186__py3-none-any.whl

adam/commands/export/importer_sqlite.py

@@ -1,8 +1,9 @@
 import pandas
 
+from adam.commands.export.export_databases import export_db
 from adam.commands.export.importer import Importer
-from adam.commands.export.utils_export import GeneratorStream
-from adam.utils import log2, ing
+from adam.repl_state import ReplState
+from adam.utils import GeneratorStream, bytes_generator_from_file, ing
 from adam.utils_k8s.pods import Pods
 from adam.utils_sqlite import SQLite, sqlite
 
@@ -10,30 +11,57 @@ class SqliteImporter(Importer):
     def prefix(self):
         return 's'
 
-    def import_from_csv(self, pod: str, namespace: str, to_session: str, from_session: str, keyspace: str, table: str, target_table: str, columns: str, multi_tables = True, create_db = False):
+    def import_from_csv(self, state: ReplState, from_session: str,
+                        keyspace: str, table: str, target_table: str, columns: str,
+                        multi_tables = True, create_db = False):
         csv_file = self.csv_file(from_session, table, target_table)
+        pod = state.pod
+        namespace = state.namespace
+        to_session = state.export_session
 
         succeeded = False
-        with sqlite(to_session, keyspace) as conn:
-            try:
-                with ing(f'[{to_session}] Uploading to Sqlite', suppress_log=multi_tables):
+        try:
+            with ing(f'[{to_session}] Uploading to Sqlite', suppress_log=multi_tables):
+                # create a connection to single keyspace
+                with sqlite(to_session, keyspace) as conn:
                     bytes = Pods.read_file(pod, 'cassandra', namespace, csv_file)
                     df = pandas.read_csv(GeneratorStream(bytes))
-
                     df.to_sql(target_table, conn, index=False, if_exists='replace')
 
-                to, _ = self.move_to_done(pod, namespace, to_session, from_session, keyspace, target_table)
+            to, _ = self.move_to_done(state, from_session, keyspace, target_table)
+
+            succeeded = True
+
+            return to, to_session
+        finally:
+            if succeeded:
+                self.remove_csv(state, from_session, table, target_table, multi_tables)
+                SQLite.clear_cache()
+
+                if not multi_tables:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {keyspace}.{target_table} limit 10')
+
+    def import_from_local_csv(self, state: ReplState,
+                        keyspace: str, table: str, csv_file: str, multi_tables = True, create_db = False):
+        to_session = state.export_session
+
+        succeeded = False
+        try:
+            with ing(f'[{to_session}] Uploading to Sqlite', suppress_log=multi_tables):
+                # create a connection to single keyspace
+                with sqlite(to_session, keyspace) as conn:
+                    bytes = bytes_generator_from_file(csv_file)
+                    df = pandas.read_csv(GeneratorStream(bytes))
+                    df.to_sql(table, conn, index=False, if_exists='replace')
 
-                succeeded = True
+            succeeded = True
 
-                return to, to_session
-            finally:
-                if succeeded:
-                    self.remove_csv(pod, namespace, from_session, table, target_table, multi_tables)
-                    SQLite.clear_cache()
+            return csv_file, to_session
+        finally:
+            if succeeded:
+                SQLite.clear_cache()
 
-                    if not multi_tables:
-                        with sqlite(to_session) as conn:
-                            query = f'select * from {keyspace}.{target_table} limit 10'
-                            log2(query)
-                            SQLite.run_query_with_conn(conn, query)
+                if not multi_tables:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {keyspace}.{table} limit 10')

adam/commands/export/show_column_counts.py

@@ -1,8 +1,9 @@
+from adam.commands import extract_trailing_options, validate_args
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.cql.utils_cql import cassandra_table_names
+from adam.commands.export.export_databases import ExportDatabases, export_db
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
 
 class ShowColumnCounts(Command):
     COMMAND = 'show column counts on'
@@ -27,28 +28,18 @@ class ShowColumnCounts(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, 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'
-
-            table = args[0]
-            query = Config().get(f'export.column_counts_query', 'select id, count(id) as columns from {table} group by id')
-            query = query.replace('{table}', table)
-            ExportDatabases.run_query(query, state.export_session)
+            with extract_trailing_options(args, '&') as (args, backgrounded):
+                with validate_args(args, state, name='SQL statement') as table:
+                    with export_db(state) as dbs:
+                        query = Config().get(f'export.column_counts_query', 'select id, count(id) as columns from {table} group by id')
+                        query = query.replace('{table}', table)
+                        dbs.sql(query, state.export_session, backgrounded=backgrounded)
 
             return state
 
     def completion(self, state: ReplState):
-        if not state.export_session:
-            return {}
-
-        return super().completion(state, lambda: {t: None for t in ExportDatabases.table_names(state.export_session)})
+        return super().completion(state, lambda: {t: None for t in ExportDatabases.table_names(state.export_session)}, auto_key='x.tables')
+        # return {}
 
     def help(self, _: ReplState):
         return f'{ShowColumnCounts.COMMAND} <export-table-name>\t show column count per id'

adam/commands/export/show_export_databases.py

@@ -1,5 +1,6 @@
 from adam.commands.command import Command
 from adam.commands.devices.device_export import DeviceExport
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
 
 class ShowExportDatabases(Command):
@@ -25,12 +26,14 @@ class ShowExportDatabases(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            DeviceExport().show_export_databases()
+            with export_db(state) as dbs:
+                dbs.show_databases()
 
             return state
 
     def completion(self, state: ReplState):
-        return DeviceExport().ls_completion(ShowExportDatabases.COMMAND, state, default = super().completion(state))
+        return {}
+        # return DeviceExport().ls_completion(ShowExportDatabases.COMMAND, state, default = super().completion(state))
 
     def help(self, _: ReplState):
         return f'{ShowExportDatabases.COMMAND}\t list export databases'

adam/commands/export/show_export_session.py

@@ -1,8 +1,7 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_sessions import ExportSessions, export_session
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
 
 class ShowExportSession(Command):
     COMMAND = 'show export session'
@@ -27,22 +26,14 @@ class ShowExportSession(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not args:
-                if state.in_repl:
-                    log2('Specify export database name.')
-                else:
-                    log2('* Database name is missing.')
-
-                    Command.display_help()
-
-                return 'command-missing'
-
-            ExportDatabases.disply_export_session(state.sts, state.pod, state.namespace, args[0])
+            with validate_args(args, state, name='export session') as session:
+                with export_session(state) as sessions:
+                    sessions.show_session(session)
 
             return state
 
     def completion(self, state: ReplState):
-        return super().completion(state, {session: None for session in Exporter.export_session_names(state.sts, state.pod, state.namespace)})
+        return {}
 
     def help(self, _: ReplState):
         return f'{ShowExportSession.COMMAND} <export-session-name>\t show export session'

adam/commands/export/show_export_sessions.py

@@ -1,8 +1,6 @@
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_sessions import export_session
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import lines_to_tabular, log
-from adam.utils_k8s.statefulsets import StatefulSets
 
 class ShowExportSessions(Command):
     COMMAND = 'show export sessions'
@@ -27,18 +25,13 @@ class ShowExportSessions(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            pod = state.pod
-            if not pod:
-                pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
-
-            sessions: dict[str, str] = Exporter.find_export_sessions(pod, state.namespace)
-            log(lines_to_tabular([f'{session}\t{export_state}' for session, export_state in sorted(sessions.items(), reverse=True)],
-                                    header='EXPORT_SESSION\tSTATUS', separator='\t'))
+            with export_session(state) as sessions:
+                sessions.show_all_sessions()
 
             return state
 
     def completion(self, state: ReplState):
-        return super().completion(state)
+        return {}
 
     def help(self, _: ReplState):
         return f'{ShowExportSessions.COMMAND}\t list export sessions'

adam/commands/export/utils_export.py

@@ -9,18 +9,27 @@ from adam.utils_k8s.pods import log_prefix
 from adam.utils_k8s.statefulsets import StatefulSets
 
 class ImportSpec:
-    def __init__(self, session: str, importer: str):
+    def __init__(self, table_name: str, session: str = None, files: list[str] = None, importer: str = None):
+        self.table_name = table_name
         self.session = session
+        self.files = files
         self.importer = importer
 
-    def parse_specs(specs_str: str):
-        session: str = None
+    def parse_specs(specs_str: str, files = False):
+        session_or_files: str = None
         importer: str = None
+        table_name: str = None
 
         if specs_str:
-            importer, session = ImportSpec._extract_importer(specs_str.strip(' '))
+            importer, rest = ImportSpec._extract_importer(specs_str.strip(' '))
+
+            if rest:
+                table_name, session_or_files = ImportSpec._extract_table_name(rest)
 
-        return ImportSpec(session, importer)
+        if not files:
+            return ImportSpec(table_name, session=session_or_files, importer=importer)
+
+        return ImportSpec(table_name, files=[f.strip(' ') for f in session_or_files.split(',')], importer=importer)
 
     def _extract_importer(spec_str: str) -> tuple[str, str]:
         importer = None
@@ -34,12 +43,25 @@ class ImportSpec:
 
         return importer, rest
 
+    def _extract_table_name(spec_str: str) -> tuple[str, str]:
+        table_name = None
+        rest = spec_str
+
+        p = re.compile(r"(.*?)as\s+(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            rest = match.group(1).strip(' ')
+            table_name = match.group(2).strip(' ')
+
+        return table_name, rest
+
 class ExportSpec(ImportSpec):
     def __init__(self, keyspace: str, consistency: str, importer: str, tables: list['ExportTableSpec'], session: str = None):
         super().__init__(None, importer)
 
         self.keyspace = keyspace
         self.consistency = consistency
+        self.importer = importer
         self.tables = tables
         self.session = session
 
@@ -154,11 +176,12 @@ class ExportTableSpec:
         return f'{self.keyspace}.{self.table}({self.columns}) as {self.target_table}'
 
 class ExportTableStatus:
-    def __init__(self, keyspace: str, target_table: str, status: str, table: str = None):
+    def __init__(self, keyspace: str, target_table: str, status: str, table: str = None, csv_file: str = ''):
         self.keyspace = keyspace
         self.target_table = target_table
         self.status = status
         self.table = table
+        self.csv_file = csv_file
 
     def __str__(self):
         return f'{self.keyspace}.{self.table} as {self.target_table} = {self.status}'
@@ -169,7 +192,7 @@ class ExportTableStatus:
 
         return False
 
-    def from_session(sts: str, pod: str, namespace: str, export_session: str):
+    def from_session(sts: str, pod: str, namespace: str, export_session: str) -> tuple['ExportTableStatus', str]:
         statuses: list[ExportTableStatus] = []
 
         status_in_whole = 'done'
@@ -204,8 +227,8 @@ class ExportTableStatus:
             target_table = m.group(2)
             state = m.group(3)
             if state == '.pending_import':
-                _, table = get_csv_files_n_table(target_table)
-                return ExportTableStatus(keyspace, target_table, 'pending_import', table)
+                csv_files, table = get_csv_files_n_table(target_table)
+                return ExportTableStatus(keyspace, target_table, 'pending_import', table, csv_files[0] if csv_files else '')
             elif state == '.done':
                 return ExportTableStatus(keyspace, target_table, 'done', target_table)
 
@@ -215,7 +238,7 @@ class ExportTableStatus:
             if r.exit_code() == 0:
                 csv_files, table = get_csv_files_n_table(target_table)
                 if csv_files:
-                    return ExportTableStatus(keyspace, target_table, 'exported', table)
+                    return ExportTableStatus(keyspace, target_table, 'exported', table, csv_files[0])
                 else:
                     return ExportTableStatus(keyspace, target_table, 'imported', target_table)
             else:
@@ -292,23 +315,30 @@ class GeneratorStream(io.RawIOBase):
             self._buffer = self._buffer[size:]
             return data
 
-class PodHandler:
-    def __init__(self, state: ReplState):
+class PodPushHandler:
+    def __init__(self, state: ReplState, pod: str = None):
         self.state = state
+        self.pushed = False
+        self.pod = pod
 
     def __enter__(self):
         state = self.state
 
         if not state.pod:
+            self.pushed = True
             state.push()
-            state.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
+
+            if not self.pod:
+                self.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
+            state.pod = self.pod
 
         return state
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.state.pop()
+        if self.pushed:
+            self.state.pop()
 
         return False
 
-def state_with_pod(state: ReplState):
-    return PodHandler(state)
+def state_with_pod(state: ReplState, pod: str = None):
+    return PodPushHandler(state, pod=pod)

adam/commands/find_files.py

@@ -0,0 +1,51 @@
+import os
+
+from adam.commands.command import Command
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_local import local_tmp_dir
+
+class FindLocalFiles(Command):
+    COMMAND = 'find local'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(FindLocalFiles, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return FindLocalFiles.COMMAND
+
+    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):
+            cmd = 'find'
+
+            if not args:
+                cmd = f'find {local_tmp_dir()}'
+            elif len(args) == 1:
+                cmd = f"find {local_tmp_dir()} -name '{args[0]}'"
+            else:
+                new_args = [f"'{arg}'" if '*' in arg else arg for arg in args]
+                cmd = 'find ' + ' '.join(new_args)
+
+            log2(cmd)
+            os.system(cmd)
+
+            return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state, {
+            '*.csv': None,
+            '*.db': None,
+            '*': None
+        })
+
+    def help(self, _: ReplState):
+        return f'{FindLocalFiles.COMMAND} [linux-find-arguments]\t find files from local machine'

adam/commands/find_processes.py

@@ -0,0 +1,76 @@
+from adam.commands import extract_options, validate_args
+from adam.commands.command import Command
+from adam.commands.devices.devices import Devices
+from adam.commands.export.utils_export import state_with_pod
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2, tabulize
+
+class FindProcesses(Command):
+    COMMAND = 'find processes'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(FindProcesses, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return FindProcesses.COMMAND
+
+    def required(self):
+        return [RequiredState.CLUSTER_OR_POD, RequiredState.APP_APP, ReplState.P]
+
+    def run(self, cmd: str, state: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, state)
+
+        with self.validate(args, state) as (args, state):
+            with extract_options(args, '-kill') as (args, kill):
+                with validate_args(args, state, name='words to look for'):
+                    arg = ' | '.join([f'grep {a}' for a in args])
+                    awk = "awk '{ print $1, $2, $8, $NF }'"
+                    rs = Devices.device(state).bash(state, state, f"ps -ef | grep -v grep | {arg} | {awk}".split(' '))
+
+                    lines: list[list[str]] = []
+                    for r in rs:
+                        for l in r.stdout.split('\n'):
+                            l = l.strip(' \t\r\n')
+                            if not l:
+                                continue
+
+                            tokens = [r.pod] + l.split(' ')
+                            lines.append(tokens)
+
+                    pids = []
+                    for l in lines:
+                        pids.append(f'{l[2]}@{l[0]}')
+
+                    tabulize(lines, lambda l: '\t'.join(l), header = 'POD\tUSER\tPID\tCMD\tLAST_ARG', separator='\t')
+                    log2()
+                    log2(f'PIDS with {",".join(args)}: {",".join(pids)}')
+
+                    if kill:
+                        log2()
+                        for pidp in pids:
+                            pid_n_pod = pidp.split('@')
+                            pid = pid_n_pod[0]
+                            if len(pid_n_pod) < 2:
+                                continue
+
+                            pod = pid_n_pod[1]
+
+                            log2(f'@{pod} bash kill -9 {pid}')
+
+                            with state_with_pod(state, pod) as state1:
+                                Devices.device(state).bash(state, state1, ['kill', '-9', pid])
+
+                    return rs
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{FindProcesses.COMMAND} word... [-kill]\t find processes with words'

adam/commands/head.py

@@ -0,0 +1,36 @@
+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
+
+class Head(Command):
+    COMMAND = 'head'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(Head, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return Head.COMMAND
+
+    def required(self):
+        return [RequiredState.CLUSTER_OR_POD, RequiredState.APP_APP, ReplState.P]
+
+    def run(self, cmd: str, state: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, state)
+
+        with self.validate(args, state) as (args, state):
+            with validate_args(args, state, name='file'):
+                return Devices.device(state).bash(state, state, cmd.split(' '))
+
+    def completion(self, state: ReplState):
+        return super().completion(state, lambda: {f: None for f in Devices.device(state).files(state)}, pods=Devices.device(state).pods(state, '-'), auto='jit')
+
+    def help(self, _: ReplState):
+        return f'{Head.COMMAND} file [&]\t run head command on the pod'

adam/commands/help.py

@@ -1,7 +1,7 @@
 from adam.commands.command import Command
 from adam.repl_commands import ReplCommands
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log
+from adam.utils import tabulize, log
 
 class Help(Command):
     COMMAND = 'help'
@@ -45,7 +45,7 @@ class Help(Command):
         lines.append('')
         lines.extend(section(ReplCommands.exit()))
 
-        log(lines_to_tabular(lines, separator='\t'))
+        tabulize(lines, separator='\t')
 
         return lines
 

adam/commands/intermediate_command.py

@@ -3,7 +3,7 @@ from abc import abstractmethod
 from adam.commands.command import Command
 from adam.commands.command_helpers import ClusterCommandHelper
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log, log2
+from adam.utils import tabulize, log, log2
 
 class IntermediateCommand(Command):
     def run(self, cmd: str, state: ReplState):
@@ -12,6 +12,9 @@ class IntermediateCommand(Command):
 
         return self.intermediate_run(cmd, state, args, self.cmd_list())
 
+    def completion(self, state: ReplState):
+        return {}
+
     @abstractmethod
     def cmd_list(self):
         pass
@@ -21,7 +24,7 @@ class IntermediateCommand(Command):
 
         if state.in_repl:
             if display_help:
-                log(lines_to_tabular([c.help(state) for c in cmds], separator=separator))
+                tabulize(cmds, lambda c: c.help(state), separator=separator)
 
             return 'command-missing'
         else:
@@ -43,7 +46,7 @@ class IntermediateCommand(Command):
         log()
         log('Sub-Commands:')
 
-        log(lines_to_tabular([c.help(ReplState()).replace(f'{cmd} ', '  ', 1) for c in cmd_list], separator=separator))
+        tabulize(cmd_list, lambda c: c.help(ReplState()).replace(f'{cmd} ', '  ', 1), separator=separator)
         if show_cluster_help:
             log()
             ClusterCommandHelper.cluster_help()

adam/commands/ls.py

@@ -35,7 +35,7 @@ class Ls(Command):
             return state
 
     def completion(self, state: ReplState):
-        return Devices.device(state).ls_completion(Ls.COMMAND, state, default = super().completion(state))
+        return super().completion(state, {'&': None}, pods=Devices.device(state).pods(state, '-'))
 
     def help(self, _: ReplState):
         return f'{Ls.COMMAND} [device:]\t list apps, envs, clusters, nodes, pg hosts/databases or export databases'

adam/commands/medusa/medusa_backup.py

@@ -1,6 +1,6 @@
 from datetime import datetime
-import re
 
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.utils_k8s.statefulsets import StatefulSets
 from adam.repl_state import ReplState, RequiredState
@@ -33,22 +33,19 @@ class MedusaBackup(Command):
             ns = state.namespace
             sts = state.sts
             now_dtformat = datetime.now().strftime("%Y-%m-%d.%H.%M.%S")
-            bkname = 'medusa-' + now_dtformat + 'full-backup-' + sts
-            if len(args) == 1:
-                bkname = str(args[0])
-            groups = re.match(r'^(.*?-.*?-).*', sts)
-            dc = StatefulSets.get_datacenter(state.sts, ns)
-            if not dc:
-                return state
-
-            try:
-                CustomResources.create_medusa_backupjob(bkname, dc, ns)
-            except Exception as e:
-                log2("Exception: MedusaBackup failed: %s\n" % e)
-            finally:
-                CustomResources.clear_caches()
+            with validate_args(args, state, default='medusa-' + now_dtformat + 'full-backup-' + sts) as bkname:
+                dc = StatefulSets.get_datacenter(state.sts, ns)
+                if not dc:
+                    return state
+
+                try:
+                    CustomResources.create_medusa_backupjob(bkname, dc, ns)
+                except Exception as e:
+                    log2("Exception: MedusaBackup failed: %s\n" % e)
+                finally:
+                    CustomResources.clear_caches()
 
-            return state
+                return state
 
     def completion(self, state: ReplState):
         return super().completion(state)