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

adam/commands/export/importer.py

@@ -0,0 +1,67 @@
+from abc import abstractmethod
+
+from adam.commands.export.utils_export import csv_dir
+from adam.config import Config
+from adam.utils import ing
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.pods import log_prefix
+
+class Importer:
+    @abstractmethod
+    def prefix(self):
+        pass
+
+    @abstractmethod
+    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):
+        pass
+
+    def move_to_done(self, pod: str, namespace: str, to_session: str, from_session: str, keyspace: str, target_table: str):
+        log_file = f'{log_prefix()}-{from_session}_{keyspace}.{target_table}.log.pending_import'
+
+        to = f'{log_prefix()}-{to_session}_{keyspace}.{target_table}.log.done'
+
+        CassandraNodes.exec(pod, namespace, f'mv {log_file} {to}', show_out=Config().is_debug(), shell='bash')
+
+        return to, to_session
+
+    def prefix_adjusted_session(self, session: str):
+        if not session.startswith(self.prefix()):
+            return f'{self.prefix()}{session[1:]}'
+
+        return session
+
+    def remove_csv(self, pod: str, namespace: str, session: str, table: str, target_table: str, multi_tables = True):
+        with ing(f'[{session}] Cleaning up temporary files', suppress_log=multi_tables):
+            CassandraNodes.exec(pod, namespace, f'rm -rf {self.csv_file(session, table, target_table)}', show_out=Config().is_debug(), shell='bash')
+
+    def db(self, session: str, keyspace: str):
+        return f'{session}_{keyspace}'
+
+    def csv_file(self, session: str, table: str, target_table: str):
+        return f'{csv_dir()}/{session}_{target_table}/{table}.csv'
+
+    def prefix_from_importer(importer: str = ''):
+        if not importer:
+            return ''
+
+        prefix = 's'
+
+        if importer == 'athena':
+            prefix = 'e'
+        elif importer == 'csv':
+            prefix = 'c'
+
+        return prefix
+
+    def importer_from_session(session: str):
+        if not session:
+            return None
+
+        importer = 'csv'
+
+        if session.startswith('s'):
+            importer = 'sqlite'
+        elif session.startswith('e'):
+            importer = 'athena'
+
+        return importer

adam/commands/export/importer_athena.py

@@ -0,0 +1,80 @@
+import boto3
+
+from adam.commands.export.importer import Importer
+from adam.commands.export.utils_export import GeneratorStream
+from adam.config import Config
+from adam.utils import log2, ing
+from adam.utils_athena import Athena
+from adam.utils_k8s.pods import Pods
+
+class AthenaImporter(Importer):
+    def ping():
+        session = boto3.session.Session()
+        credentials = session.get_credentials()
+
+        return credentials is not None
+
+    def prefix(self):
+        return 'e'
+
+    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):
+        csv_file = self.csv_file(from_session, table, target_table)
+        db = self.db(to_session, keyspace)
+
+        succeeded = False
+        try:
+            bucket = Config().get('export.bucket', 'c3.ops--qing')
+
+            with ing(f'[{to_session}] Uploading to S3', suppress_log=multi_tables):
+                bytes = Pods.read_file(pod, 'cassandra', namespace, csv_file)
+
+                s3 = boto3.client('s3')
+                s3.upload_fileobj(GeneratorStream(bytes), bucket, f'export/{db}/{keyspace}/{target_table}/{table}.csv')
+
+            msg: str = None
+            if create_db:
+                msg = f"[{to_session}] Creating database {db}"
+            else:
+                msg = f"[{to_session}] Creating table {target_table}"
+            with ing(msg, suppress_log=multi_tables):
+                query = f'CREATE DATABASE IF NOT EXISTS {db};'
+                if Config().is_debug():
+                    log2(query)
+                Athena.query(query, 'default')
+
+                query = f'DROP TABLE IF EXISTS {target_table};'
+                if Config().is_debug():
+                    log2(query)
+                Athena.query(query, db)
+
+                athena_columns = ', '.join([f'{c} string' for c in columns.split(',')])
+                query = f'CREATE EXTERNAL TABLE IF NOT EXISTS {target_table}(\n' + \
+                        f'    {athena_columns})\n' + \
+                            "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'\n" + \
+                            'WITH SERDEPROPERTIES (\n' + \
+                            '    "separatorChar" = ",",\n' + \
+                            '    "quoteChar"     = "\\"")\n' + \
+                        f"LOCATION 's3://{bucket}/export/{db}/{keyspace}/{target_table}'\n" + \
+                            'TBLPROPERTIES ("skip.header.line.count"="1");'
+                if Config().is_debug():
+                    log2(query)
+                try:
+                    Athena.query(query, db)
+                except Exception as e:
+                    log2(f'*** Failed query:\n{query}')
+                    raise e
+
+            to, _ = self.move_to_done(pod, namespace, to_session, from_session, keyspace, target_table)
+
+            succeeded = True
+
+            return to, to_session
+        finally:
+            if succeeded:
+                self.remove_csv(pod, namespace, from_session, table, target_table, multi_tables)
+                Athena.clear_cache()
+
+                if not multi_tables:
+                    query = f'select * from {target_table} limit 10'
+                    log2(query)
+                    Athena.run_query(query, db)

adam/commands/export/importer_sqlite.py

@@ -0,0 +1,47 @@
+import os
+import sqlite3
+import pandas
+
+from adam.commands.export.importer import Importer
+from adam.commands.export.utils_export import GeneratorStream
+from adam.utils import log2, ing
+from adam.utils_k8s.pods import Pods
+from adam.utils_sqlite import SQLite
+
+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):
+        csv_file = self.csv_file(from_session, table, target_table)
+        db = self.db(to_session, keyspace)
+
+        succeeded = False
+        conn = None
+        try:
+            os.makedirs(SQLite.local_db_dir(), exist_ok=True)
+            conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{db}.db')
+
+            with ing(f'[{to_session}] Uploading to Sqlite', suppress_log=multi_tables):
+                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)
+
+            succeeded = True
+
+            return to, to_session
+        finally:
+            if succeeded:
+                self.remove_csv(pod, namespace, from_session, table, target_table, multi_tables)
+                SQLite.clear_cache()
+
+                if not multi_tables:
+                    query = f'select * from {target_table} limit 10'
+                    log2(query)
+                    SQLite.run_query(query, conn_passed=conn)
+
+            if conn:
+                conn.close()

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