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

adam/commands/export/utils_export.py

@@ -0,0 +1,366 @@
+import io
+import os
+import re
+
+from adam.config import Config
+from adam.utils import ExecResult, log2
+from adam.repl_state import ReplState
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.statefulsets import StatefulSets
+from adam.utils_local import local_exec
+
+class ImportSpec:
+    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, files = False):
+        session_or_files: str = None
+        importer: str = None
+        table_name: str = None
+
+        if specs_str:
+            importer, rest = ImportSpec._extract_importer(specs_str.strip(' '))
+
+            if rest:
+                table_name, session_or_files = ImportSpec._extract_table_name(rest)
+
+        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
+        rest = spec_str
+
+        p = re.compile(r"(.*?)to\s+(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            rest = match.group(1).strip(' ')
+            importer = match.group(2).strip(' ')
+
+        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
+
+    def __str__(self):
+        return f'keyspace: {self.keyspace}, consistency: {self.consistency}, importer: {self.importer}, tables: {",".join([t.table for t in self.tables])}, session: {self.session}'
+
+    def __eq__(self, other):
+        if not isinstance(other, ExportSpec):
+            return NotImplemented
+
+        return self.keyspace == other.keyspace and self.tables == other.tables and self.consistency == other.consistency and self.importer == other.importer and self.session == other.session
+
+    def parse_specs(specs_str: str):
+        keyspace: str = None
+        consistency: str = None
+        importer: str = None
+        specs: list[ExportTableSpec] = None
+
+        if specs_str:
+            importer, specs_str = ExportSpec._extract_importer(specs_str.strip(' '))
+            keyspace, specs_str = ExportSpec._extract_keyspace(specs_str)
+            consistency, specs = ExportSpec._extract_consisteny(specs_str)
+
+        return ExportSpec(keyspace, consistency, importer, specs)
+
+    def _extract_keyspace(spec_str: str) -> tuple[str, str]:
+        keyspace = None
+        rest = spec_str
+
+        p = re.compile(r"\s*\*\s+in\s+(\S+)(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            keyspace = match.group(1).strip(' ')
+            rest = match.group(2).strip(' ')
+        elif spec_str.startswith('*'):
+            keyspace = '*'
+            rest = spec_str[1:].strip(' ')
+
+        return keyspace, rest
+
+    def _extract_consisteny(spec_str: str) -> tuple[str, list['ExportTableSpec']]:
+        consistency = None
+
+        p = re.compile(r"(.*?)with\s+consistency\s+(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            spec_str = match.group(1).strip(' ')
+            consistency = match.group(2)
+
+        if spec_str:
+            p = r",\s*(?![^()]*\))"
+            specs = re.split(p, spec_str)
+
+            return consistency, [ExportTableSpec.parse(spec) for spec in specs]
+
+        return consistency, None
+
+class ExportTableSpec:
+    def __init__(self, keyspace: str, table: str, columns: str = None, target_table: str = None):
+        self.keyspace = keyspace
+        self.table = table
+        self.columns = columns
+        self.target_table = target_table
+
+    def __str__(self):
+        return f'{self.keyspace}.{self.table}({self.columns}) AS {self.target_table}'
+
+    def __eq__(self, other):
+        if not isinstance(other, ExportTableSpec):
+            return NotImplemented
+
+        return self.keyspace == other.keyspace and self.table == other.table and self.columns == other.columns and self.target_table == other.target_table
+
+    def from_status(status: 'ExportTableStatus'):
+        return ExportTableSpec(status.keyspace, status.table, target_table=status.target_table)
+
+    def parse(spec_str: str) -> 'ExportTableSpec':
+        target = None
+
+        p = re.compile(r"(.*?)\s+as\s+(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            spec_str = match.group(1)
+            target = match.group(2)
+
+        keyspace = None
+        table = spec_str
+        columns = None
+
+        p = re.compile('(.*?)\.(.*?)\((.*)\)')
+        match = p.match(spec_str)
+        if match:
+            keyspace = match.group(1)
+            table = match.group(2)
+            columns = match.group(3)
+        else:
+            p = re.compile('(.*?)\.(.*)')
+            match = p.match(spec_str)
+            if match:
+                keyspace = match.group(1)
+                table = match.group(2)
+
+        return ExportTableSpec(keyspace, table, columns, target)
+
+    def __eq__(self, other):
+        if isinstance(other, ExportTableSpec):
+            return self.keyspace == other.keyspace and self.table == other.table and self.columns == other.columns and self.target_table == other.target_table
+
+        return False
+
+    def __str__(self):
+        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, 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}'
+
+    def __eq__(self, other):
+        if isinstance(other, ExportTableStatus):
+            return self.keyspace == other.keyspace and self.table == other.table and self.status == other.status and self.target_table == other.target_table
+
+        return False
+
+    def from_session(sts: str, pod: str, namespace: str, export_session: str) -> tuple['ExportTableStatus', str]:
+        statuses: list[ExportTableStatus] = []
+
+        status_in_whole = 'done'
+        log_files: list[str] = find_files(pod, namespace, f'{export_log_prefix()}-{export_session}_*.log*')
+
+        for log_file in log_files:
+            status: ExportTableStatus = ExportTableStatus.from_log_file(pod, namespace, export_session, log_file)
+            statuses.append(status)
+
+            if status.status != 'done':
+                status_in_whole = status.status
+
+        return statuses, status_in_whole
+
+    def from_log_file(pod: str, namespace: str, copy_session: str, log_file: str):
+        def get_csv_files_n_table(target_table: str):
+            db = f'{copy_session}_{target_table}'
+            csv_file = f'{csv_dir()}/{db}/*.csv'
+            csv_files: list[str] = find_files(pod, namespace, csv_file, remote=True)
+            if csv_files:
+                table = target_table
+                m = re.match(f'{csv_dir()}/{db}/(.*).csv', csv_files[0])
+                if m:
+                    table = m.group(1)
+                return csv_files, table
+
+            return csv_files, target_table
+
+        m = re.match(f'{export_log_prefix()}-{copy_session}_(.*?)\.(.*?)\.log(.*)', log_file)
+        if m:
+            keyspace = m.group(1)
+            target_table = m.group(2)
+            state = m.group(3)
+            if state == '.pending_import':
+                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)
+
+            # 4 rows exported to 1 files in 0 day, 0 hour, 0 minute, and 1.335 seconds.
+            pattern = 'rows exported to'
+            r: ExecResult = local_exec(['grep', pattern, log_file], show_out=Config().is_debug())
+            # r = CassandraNodes.exec(pod, namespace, f"grep '{pattern}' {log_file}", show_out=Config().is_debug(), shell='bash')
+
+            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, csv_files[0])
+                else:
+                    return ExportTableStatus(keyspace, target_table, 'imported', target_table)
+            else:
+                return ExportTableStatus(keyspace, target_table, 'export_in_pregress')
+
+        return ExportTableStatus(None, None, 'unknown')
+
+def csv_dir():
+    return Config().get('export.csv_dir', '/c3/cassandra/tmp')
+
+def find_files(pod: str, namespace: str, pattern: str, mmin: int = 0, remote = False):
+    stdout = ''
+    if not remote:
+        dir = os.path.dirname(pattern)
+        base = os.path.basename(pattern)
+        cmd = ['find', dir, '-name', base]
+        if mmin:
+                cmd += ['-mmin',  f'-{mmin}']
+
+        stdout = local_exec(cmd, show_out=Config().is_debug()).stdout
+    else:
+        cmd = f'find {pattern}'
+        if mmin:
+                cmd = f'{cmd} -mmin -{mmin}'
+
+        stdout = CassandraNodes.exec(pod, namespace, cmd, show_out=Config().is_debug(), shell='bash').stdout
+
+    log_files = []
+    for line in stdout.split('\n'):
+        line = line.strip(' \r')
+        if line:
+            log_files.append(line)
+
+    return log_files
+
+def export_log_prefix():
+    return Config().get('export.log-prefix', '/tmp/qing')
+
+def os_system_exec(cmd: str, show_out = False):
+    if show_out: log2(cmd)
+
+    os.system(cmd)
+
+class GeneratorStream(io.RawIOBase):
+    def __init__(self, generator):
+        self._generator = generator
+        self._buffer = b''  # Buffer to store leftover bytes from generator yields
+
+    def readable(self):
+        return True
+
+    def _read_from_generator(self):
+        try:
+            chunk = next(self._generator)
+            if isinstance(chunk, str):
+                chunk = chunk.encode('utf-8')  # Encode if generator yields strings
+            self._buffer += chunk
+        except StopIteration:
+            pass  # Generator exhausted
+
+    def readinto(self, b):
+        # Fill the buffer if necessary
+        while len(self._buffer) < len(b):
+            old_buffer_len = len(self._buffer)
+            self._read_from_generator()
+            if len(self._buffer) == old_buffer_len:  # Generator exhausted and buffer empty
+                break
+
+        bytes_to_read = min(len(b), len(self._buffer))
+        b[:bytes_to_read] = self._buffer[:bytes_to_read]
+        self._buffer = self._buffer[bytes_to_read:]
+        return bytes_to_read
+
+    def read(self, size=-1):
+        if size == -1:  # Read all remaining data
+            while True:
+                old_buffer_len = len(self._buffer)
+                self._read_from_generator()
+                if len(self._buffer) == old_buffer_len:
+                    break
+            data = self._buffer
+            self._buffer = b''
+            return data
+        else:
+            # Ensure enough data in buffer
+            while len(self._buffer) < size:
+                old_buffer_len = len(self._buffer)
+                self._read_from_generator()
+                if len(self._buffer) == old_buffer_len:
+                    break
+
+            data = self._buffer[:size]
+            self._buffer = self._buffer[size:]
+            return data
+
+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()
+
+            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):
+        if self.pushed:
+            self.state.pop()
+
+        return False
+
+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.of(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.of(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/generate_report.py

@@ -0,0 +1,52 @@
+import click
+import json
+
+from adam.checks.check_result import CheckResult
+from adam.checks.check_utils import run_checks
+from adam.commands import extract_options
+from adam.commands.command import Command
+from adam.commands.commands_utils import kaqing_log_file
+from adam.repl_state import ReplState
+from adam.utils import log2
+
+class GenerateReport(Command):
+    COMMAND = 'generate report'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(GenerateReport, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return GenerateReport.COMMAND
+
+    def required(self):
+        return ReplState.NON_L
+
+    def run(self, cmd: str, state: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, state)
+
+        with self.validate(args, state) as (args, state):
+            with extract_options(args, ['-s', '--show']) as (args, show_out):
+                results = run_checks(state.sts, state.namespace, state.pod, show_out=show_out)
+                output = CheckResult.report(results)
+
+                if state.in_repl:
+                    with kaqing_log_file() as json_file:
+                        json.dump(output, json_file, indent=2)
+                        log2(f'Report stored in {json_file.name}.')
+                else:
+                    click.echo(json.dumps(output, indent=2))
+
+                return output
+
+    def completion(self, state: ReplState):
+        return super().completion(state, {'-s': None})
+
+    def help(self, _: ReplState):
+        return f"{GenerateReport.COMMAND} [-s]\t generate report"

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.of(state).bash(state, state, cmd.split(' '))
+
+    def completion(self, state: ReplState):
+        return super().completion(state, lambda: {f: None for f in Devices.of(state).files(state)}, pods=Devices.of(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'
@@ -28,20 +28,24 @@ class Help(Command):
 
         lines = []
         lines.append('NAVIGATION')
-        lines.append('  a: | c: | l: | p:\t switch to another operational device: App, Cassandra, Audit or Postgres')
+        lines.append('  a: | c: | l: | p: | x:\t switch to another operational device: App, Cassandra, Audit, Postgres or Export')
         lines.extend(section(ReplCommands.navigation()))
-        lines.append('CHECK CASSANDRA')
-        lines.extend(section(ReplCommands.cassandra_check()))
-        lines.append('CASSANDRA OPERATIONS')
+        lines.append('CASSANDRA')
         lines.extend(section(ReplCommands.cassandra_ops()))
+        lines.append('POSTGRES')
+        lines.extend(section(ReplCommands.postgres_ops()))
+        lines.append('APP')
+        lines.extend(section(ReplCommands.app_ops()))
+        lines.append('EXPORT DB')
+        lines.extend(section(ReplCommands.export_ops()))
+        lines.append('AUDIT')
+        lines.extend(section(ReplCommands.audit_ops()))
         lines.append('TOOLS')
         lines.extend(section(ReplCommands.tools()))
-        lines.append('APP')
-        lines.extend(section(ReplCommands.app()))
         lines.append('')
         lines.extend(section(ReplCommands.exit()))
 
-        log(lines_to_tabular(lines, separator='\t'))
+        tabulize(lines, separator='\t')
 
         return lines
 

adam/commands/intermediate_command.py

@@ -0,0 +1,52 @@
+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 tabulize, log, log2
+
+class IntermediateCommand(Command):
+    def run(self, cmd: str, state: ReplState):
+        if not(args := self.args(cmd)):
+            return super().run(cmd, state)
+
+        return self.intermediate_run(cmd, state, args, self.cmd_list())
+
+    def completion(self, state: ReplState):
+        return {}
+
+    @abstractmethod
+    def cmd_list(self):
+        pass
+
+    def intermediate_run(self, cmd: str, state: ReplState, args: list[str], cmds: list['Command'], separator='\t', display_help=True):
+        state, _ = self.apply_state(args, state)
+
+        if state.in_repl:
+            if display_help:
+                tabulize(cmds, lambda c: c.help(state), separator=separator)
+
+            return 'command-missing'
+        else:
+            # head with the Chain of Responsibility pattern
+            if not self.run_subcommand(cmd, state):
+                if display_help:
+                    log2('* Command is missing.')
+                    Command.display_help()
+                return 'command-missing'
+
+        return state
+
+    def run_subcommand(self, cmd: str, state: ReplState):
+        cmds = Command.chain(self.cmd_list())
+        return cmds.run(cmd, state)
+
+    def intermediate_help(super_help: str, cmd: str, cmd_list: list['Command'], separator='\t', show_cluster_help=False):
+        log(super_help)
+        log()
+        log('Sub-Commands:')
+
+        tabulize(cmd_list, lambda c: c.help(ReplState()).replace(f'{cmd} ', '  ', 1), separator=separator)
+        if show_cluster_help:
+            log()
+            ClusterCommandHelper.cluster_help()