kaqing 2.0.14__py3-none-any.whl → 2.0.145__py3-none-any.whl

adam/commands/export/show_column_counts.py

@@ -0,0 +1,63 @@
+from adam.commands.command import Command
+from adam.commands.export.export_databases import ExportDatabases
+from adam.config import Config
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+from adam.utils_athena import Athena
+from adam.utils_sqlite import SQLite
+
+class ShowColumnCounts(Command):
+    COMMAND = 'show column counts on'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowColumnCounts, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowColumnCounts.COMMAND
+
+    def required(self):
+        return RequiredState.EXPORT_DB
+
+    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'
+
+        copy_or_export = 'copy'
+        if state.export_session.startswith('e'):
+            copy_or_export = 'export'
+
+        table = args[0]
+        query = Config().get(f'{copy_or_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)
+
+        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)})
+
+    def help(self, _: ReplState):
+        return f'{ShowColumnCounts.COMMAND} <export-table-name>\t show column count per id'

adam/commands/export/show_export_databases.py

@@ -0,0 +1,39 @@
+from adam.commands.command import Command
+from adam.commands.devices.device_export import DeviceExport
+from adam.repl_state import ReplState
+
+class ShowExportDatabases(Command):
+    COMMAND = 'show export databases'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowExportDatabases, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowExportDatabases.COMMAND
+
+    def required(self):
+        return [ReplState.C, 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
+
+        DeviceExport().show_export_databases()
+
+        return state
+
+    def completion(self, state: ReplState):
+        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

@@ -0,0 +1,51 @@
+from adam.commands.command import Command
+from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.export.exporter import Exporter
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+
+class ShowExportSession(Command):
+    COMMAND = 'show export session'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowExportSession, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowExportSession.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 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])
+
+        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)})
+
+    def help(self, _: ReplState):
+        return f'{ShowExportSession.COMMAND} <export-session-name>\t show export session'

adam/commands/export/show_export_sessions.py

@@ -0,0 +1,47 @@
+from adam.commands.command import Command
+from adam.commands.export.exporter import Exporter
+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'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowExportSessions, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowExportSessions.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
+
+        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'))
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{ShowExportSessions.COMMAND}\t list export sessions'

adam/commands/export/utils_export.py

@@ -0,0 +1,291 @@
+import io
+import re
+
+from adam.config import Config
+from adam.pod_exec_result import PodExecResult
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.pods import log_prefix
+
+class ImportSpec:
+    def __init__(self, session: str, importer: str):
+        self.session = session
+        self.importer = importer
+
+    def parse_specs(specs_str: str):
+        session: str = None
+        importer: str = None
+
+        if specs_str:
+            importer, session = ImportSpec._extract_importer(specs_str.strip(' '))
+
+        return ImportSpec(session, 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
+
+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.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):
+        self.keyspace = keyspace
+        self.target_table = target_table
+        self.status = status
+        self.table = table
+
+    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):
+        statuses: list[ExportTableStatus] = []
+
+        status_in_whole = 'done'
+        log_files: list[str] = find_files(pod, namespace, f'{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)
+            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'{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':
+                _, table = get_csv_files_n_table(target_table)
+                return ExportTableStatus(keyspace, target_table, 'pending_import', table)
+            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: PodExecResult = 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)
+                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):
+    if mmin:
+        r = CassandraNodes.exec(pod, namespace, f'find {pattern} -mmin -{mmin}', show_out=Config().is_debug(), shell='bash')
+    else:
+        r = CassandraNodes.exec(pod, namespace, f'find {pattern}', show_out=Config().is_debug(), shell='bash')
+
+    log_files = []
+    for line in r.stdout.split('\n'):
+        line = line.strip(' \r')
+        if line:
+            log_files.append(line)
+
+    return log_files
+
+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

adam/commands/help.py

@@ -23,20 +23,25 @@ class Help(Command):
             return super().run(cmd, state)
 
         def section(cmds : list[ReplCommands]):
-            return [f'  {c.help(state)}' for c in cmds if c.help(state)]
+            sorted_cmds = sorted(cmds, key=lambda cmd: cmd.command())
+            return [f'  {c.help(state)}' for c in sorted_cmds if c.help(state)]
 
         lines = []
         lines.append('NAVIGATION')
-        lines.append('  a: | c: | p:\t switch to another operational device: App, Cassandra 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()))
 

adam/commands/issues.py

@@ -21,11 +21,17 @@ class Issues(Command):
     def command(self):
         return Issues.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)
 
         state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
         args, show = Command.extract_options(args, ['-s', '--show'])
 
         results = run_checks(state.sts, state.namespace, state.pod, show_output=show)

adam/commands/kubectl.py

@@ -0,0 +1,41 @@
+import subprocess
+
+from adam.commands.command import Command
+from adam.repl_state import ReplState, RequiredState
+
+class Kubectl(Command):
+    COMMAND = 'k'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(Kubectl, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return Kubectl.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)
+
+        state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
+        subprocess.run(["kubectl"] + args)
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+
+    def help(self, _: ReplState):
+        return f'{Kubectl.COMMAND} \t run a kubectl command'

adam/commands/login.py

@@ -8,7 +8,7 @@ from adam.config import Config
 from adam.sso.idp import Idp
 from adam.sso.idp_login import IdpLogin
 from adam.commands.command import Command
-from adam.repl_state import ReplState
+from adam.repl_state import ReplState, RequiredState
 from adam.utils import log, log2
 
 class Login(Command):
@@ -26,6 +26,9 @@ class Login(Command):
     def command(self):
         return Login.COMMAND
 
+    def required(self):
+        return ReplState.NON_L
+
     def run(self, cmd: str, state: ReplState):
         def custom_handler(signum, frame):
             AppSession.ctrl_c_entered = True
@@ -38,7 +41,7 @@ class Login(Command):
         state, args = self.apply_state(args, state)
         args, debug = Command.extract_options(args, ['d'])
         if debug:
-            Config().set('debug.show-out', True)
+            Config().set('debug', True)
 
         username: str = os.getenv('USERNAME')
         if len(args) > 0:
@@ -59,8 +62,8 @@ class Login(Command):
 
         return state
 
-    def completion(self, _: ReplState):
-        return {}
+    def completion(self, state: ReplState):
+        return super().completion(state)
 
     def help(self, _: ReplState):
         return f'{Login.COMMAND}\t SSO login'

adam/commands/logs.py

@@ -1,6 +1,6 @@
 from adam.commands.command import Command
 from adam.config import Config
-from adam.k8s_utils.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.repl_state import ReplState, RequiredState
 
 class Logs(Command):
@@ -33,6 +33,7 @@ class Logs(Command):
         return CassandraNodes.exec(state.pod, state.namespace, f'cat {path}')
 
     def completion(self, _: ReplState):
+        # available only on cli
         return {}
 
     def help(self, _: ReplState):