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

adam/commands/export/export_select.py

@@ -1,15 +1,10 @@
-from datetime import datetime
-from adam.commands import extract_trailing_options
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases
-from adam.config import Config
+from adam.commands.export.completions_x import completions_x
 from adam.repl_state import ReplState, RequiredState
-from adam.sql.sql_completer import SqlCompleter, SqlVariant
-from adam.utils import log2
-from adam.utils_athena import Athena
 
+# No action body, only for a help entry and auto-completion
 class ExportSelect(Command):
-    COMMAND = '.select'
+    COMMAND = 'select_on_x'
 
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
@@ -26,69 +21,14 @@ class ExportSelect(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)
-
-        with self.validate(args, state) as (args, state):
-            with extract_trailing_options(args, '&') as (args, backgrounded):
-                if not state.export_session:
-                    if state.in_repl:
-                        if state.device == ReplState.C:
-                            log2("Select an export database first with 'use' command.")
-                        else:
-                            log2('cd to an export database first.')
-                    else:
-                        log2('* export database is missing.')
-
-                        Command.display_help()
-
-                    return 'command-missing'
-
-                if not args:
-                    if state.in_repl:
-                        log2('Use a SQL statement.')
-                    else:
-                        log2('* SQL statement is missing.')
-
-                        Command.display_help()
-
-                    return 'command-missing'
-
-                query = ' '.join(args)
-
-                def output(out: str):
-                    log_prefix = Config().get('export.log-prefix', '/tmp/qing')
-                    log_file = f'{log_prefix}-{datetime.now().strftime("%d%H%M%S")}-sqlite.log'
-
-                    with open(log_file, 'w') as f:
-                        f.write(out)
-
-                    return log_file
-
-                ExportDatabases.run_query(f'select {query}', database=state.export_session, output=output if backgrounded else None)
-
-                return state
-
     def completion(self, state: ReplState):
-        if not state.export_session:
+        if state.device != ReplState.X:
             return {}
 
-        db = state.export_session
-
-        # warm up the caches first time when x: drive is accessed
-        ExportDatabases.table_names(db)
-        Athena.column_names(database=db, function='export')
-        Athena.column_names(partition_cols_only=True, database=db, function='export')
+        if state.export_session:
+            return completions_x(state)
 
-        return {ExportSelect.COMMAND: SqlCompleter(
-            lambda: ExportDatabases.table_names(db),
-            dml='select',
-            expandables={
-                'columns':lambda table: Athena.column_names(database=db, function='export'),
-            },
-            variant=SqlVariant.ATHENA
-        )}
+        return {}
 
     def help(self, _: ReplState):
-        return f'.<sql-select-statements>\t run queries on export database'
+        return f'<sql-select-statements>\t run queries on export database'

adam/commands/export/export_sessions.py

@@ -0,0 +1,209 @@
+import functools
+import os
+import re
+
+from adam.commands.export.importer import Importer
+from adam.commands.export.utils_export import ExportTableStatus, csv_dir, find_files
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import log2, tabulize, log, parallelize
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.pods import Pods, log_prefix
+from adam.utils_k8s.statefulsets import StatefulSets
+from adam.utils_local import local_tmp_dir
+
+class ExportSessions:
+    def clear_export_session_cache():
+        ExportSessions.find_export_sessions.cache_clear()
+        ExportSessions.export_session_names.cache_clear()
+
+    @functools.lru_cache()
+    def export_session_names(sts: str, pod: str, namespace: str, importer: str = None, export_state = None):
+        if not sts or not namespace:
+            return []
+
+        if not pod:
+            pod = StatefulSets.pod_names(sts, namespace)[0]
+
+        if not pod:
+            return []
+
+        return [session for session, state in ExportSessions.find_export_sessions(pod, namespace, importer).items() if not export_state or state == export_state]
+
+    @functools.lru_cache()
+    def find_export_sessions(pod: str, namespace: str, importer: str = None, limit = 100):
+        sessions: dict[str, str] = {}
+
+        prefix = Importer.prefix_from_importer(importer)
+
+        log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{prefix}*_*.log*')
+
+        if not log_files:
+            return {}
+
+        for log_file in log_files[:limit]:
+            m = re.match(f'{log_prefix()}-(.*?)_.*\.log?(.*)', log_file)
+            if m:
+                s = m.group(1)
+                state = m.group(2) # '', '.pending_import', '.done'
+                if state:
+                    state = state.strip('.')
+                else:
+                    state = 'in_export'
+
+                if s not in sessions:
+                    sessions[s] = state
+                elif sessions[s] == 'done' and state != 'done':
+                    sessions[s] = state
+
+        return sessions
+
+    def clean_up_all_sessions(sts: str, pod: str, namespace: str):
+        if not sts or not namespace:
+            return False
+
+        if not pod:
+            pod = StatefulSets.pod_names(sts, namespace)[0]
+
+        CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/*', show_out=Config().is_debug(), shell='bash')
+        CassandraNodes.exec(pod, namespace, f'rm -rf {log_prefix()}-*.log*', show_out=Config().is_debug(), shell='bash')
+
+        return True
+
+    def clean_up_sessions(sts: str, pod: str, namespace: str, sessions: list[str], max_workers = 0):
+        if not sessions:
+            return []
+
+        if not max_workers:
+            max_workers = Config().action_workers('export', 8)
+
+        with parallelize(sessions,
+                         max_workers,
+                         msg='Cleaning|Cleaned up {size} export sessions') as exec:
+            cnt_tuples = exec.map(lambda session: ExportSessions.clean_up_session(sts, pod, namespace, session, True))
+            csv_cnt = 0
+            log_cnt = 0
+            for (csv, log) in cnt_tuples:
+                csv_cnt += csv
+                log_cnt += log
+
+            return csv_cnt, log_cnt
+
+    def clean_up_session(sts: str, pod: str, namespace: str, session: str, multi_tables = True):
+        if not sts or not namespace:
+            return 0, 0
+
+        if not pod:
+            pod = StatefulSets.pod_names(sts, namespace)[0]
+
+        if not pod:
+            return 0, 0
+
+        csv_cnt = 0
+        log_cnt = 0
+
+        log_files: list[str] = find_files(pod, namespace, f'{log_prefix()}-{session}_*.log*')
+
+        for log_file in log_files:
+            m = re.match(f'{log_prefix()}-{session}_(.*?)\.(.*?)\.log.*', log_file)
+            if m:
+                table = m.group(2)
+
+                CassandraNodes.exec(pod, namespace, f'rm -rf {csv_dir()}/{session}_{table}', show_out=not multi_tables, shell='bash')
+                csv_cnt += 1
+
+                CassandraNodes.exec(pod, namespace, f'rm -rf {log_file}', show_out=not multi_tables, shell='bash')
+                log_cnt += 1
+
+        return csv_cnt, log_cnt
+
+    def show_session(sts: str, pod: str, namespace: str, session: str):
+        if not pod:
+            pod = StatefulSets.pod_names(sts, namespace)[0]
+
+        if not pod:
+            return
+
+        tables, _ = ExportTableStatus.from_session(sts, pod, namespace, session)
+        log()
+        tabulize(tables,
+                 lambda t: f'{t.keyspace}\t{t.target_table}\t{"export_completed_pending_import" if t.status == "pending_import" else t.status}\t{t.csv_file}',
+                 header='KEYSPACE\tTARGET_TABLE\tSTATUS\tCSV_FILES',
+                 separator='\t')
+
+    def download_session(sts: str, pod: str, namespace: str, session: str):
+        if not pod:
+            pod = StatefulSets.pod_names(sts, namespace)[0]
+
+        if not pod:
+            return
+
+        tables, _ = ExportTableStatus.from_session(sts, pod, namespace, session)
+        def download_csv(table):
+            from_path: str = table.csv_file
+
+            to_path = from_path.replace(csv_dir(), local_tmp_dir())
+            os.makedirs(os.path.dirname(to_path), exist_ok=True)
+            Pods.download_file(pod, 'cassandra', namespace, from_path, to_path)
+
+            log2(f'[{session}] Downloaded to {to_path}.')
+
+        with parallelize(tables,
+                         workers=Config().get('download.workers', 8),
+                         msg='Downloading|Downloaded {size} csv files') as exec:
+            exec.map(download_csv)
+
+class ExportSessionService:
+    def __init__(self, handler: 'ExportSessionHandler'):
+        self.handler = handler
+
+    def clean_up(self, sessions: list[str]):
+        state = self.handler.state
+
+        csv_cnt, log_cnt = ExportSessions.clean_up_sessions(state.sts, self.pod(), state.namespace, sessions)
+
+        log(f'Removed {csv_cnt} csv and {log_cnt} log files.')
+
+        ExportSessions.clear_export_session_cache()
+
+    def clean_up_all(self):
+        state = self.handler.state
+
+        if ExportSessions.clean_up_all_sessions(state.sts, self.pod(), state.namespace):
+            ExportSessions.clear_export_session_cache()
+
+    def show_all_sessions(self):
+        state = self.handler.state
+
+        sessions = sorted(ExportSessions.find_export_sessions(self.pod(), state.namespace).items(), reverse=True)
+        tabulize(sessions, lambda args: f'{args[0]}\t{args[1]}', header='EXPORT_SESSION\tSTATUS', separator='\t')
+
+    def show_session(self, session: str):
+        state = self.handler.state
+        ExportSessions.show_session(state.sts, self.pod(), state.namespace, session)
+
+    def download_session(self, session: str):
+        state = self.handler.state
+        ExportSessions.download_session(state.sts, self.pod(), state.namespace, session)
+
+    def pod(self):
+        state = self.handler.state
+
+        pod = state.pod
+        if not pod:
+            pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
+
+        return pod
+
+class ExportSessionHandler:
+    def __init__(self, state: ReplState = None):
+        self.state = state
+
+    def __enter__(self):
+        return ExportSessionService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+def export_session(state: ReplState = None):
+    return ExportSessionHandler(state)

adam/commands/export/export_use.py

@@ -1,9 +1,8 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
 from adam.utils import log2
-from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class ExportUse(Command):
     COMMAND = 'use'
@@ -21,32 +20,30 @@ class ExportUse(Command):
         return ExportUse.COMMAND
 
     def required(self):
-        return [ReplState.C, ReplState.X]
+        return [ReplState.C]
 
     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):
-            if not args:
-                state.export_session = None
+            with validate_args(args, state, at_least=0) as session:
+                if not session:
+                    state.export_session = None
 
-                log2('Export database is unset.')
+                    log2('Export database is unset.')
 
-                return state
+                    return state
 
-            state.export_session = args[0]
-            if state.export_session.startswith('e'):
-                Athena.clear_cache()
-            else:
-                SQLite.clear_cache()
+                state.export_session = session
 
-            ExportDatabases.display_export_db(state.export_session)
+                with export_db(state) as dbs:
+                    dbs.show_database()
 
-            return state
+                return state
 
     def completion(self, state: ReplState):
-        return super().completion(state, {n: None for n in ExportDatabases.database_names()})
+        return {}
 
     def help(self, _: ReplState):
         return f'{ExportUse.COMMAND} <export-database-name>\t use export database'

adam/commands/export/export_x_select.py

@@ -0,0 +1,48 @@
+from adam.commands import extract_trailing_options, validate_args
+from adam.commands.command import Command
+from adam.commands.export.completions_x import completions_x
+from adam.commands.export.export_databases import export_db
+from adam.repl_state import ReplState, RequiredState
+
+class ExportXSelect(Command):
+    COMMAND = 'xelect'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportXSelect, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportXSelect.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)
+
+        with self.validate(args, state) as (args, state):
+            with extract_trailing_options(args, '&') as (args, backgrounded):
+                with validate_args(args, state, name='SQL statement') as query:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select {query}', backgrounded=backgrounded)
+
+                    return state
+
+    def completion(self, state: ReplState):
+        if state.device != ReplState.C:
+            return {}
+
+        if not state.export_session:
+            return {}
+
+        # add only xelect completions to c: drive from lark
+        return {ExportXSelect.COMMAND: completions_x(state)[ExportXSelect.COMMAND]}
+
+    def help(self, _: ReplState):
+        return f'xelect...\t run queries on export database'