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

adam/commands/export/export_databases.py

@@ -2,12 +2,13 @@ from collections.abc import Callable
 import os
 import boto3
 
-from adam.commands.export.utils_export import ExportTableStatus
+from adam.commands.export.export_sessions import ExportSessions
+from adam.commands.export.importer import Importer
 from adam.config import Config
 from adam.repl_session import ReplSession
+from adam.repl_state import ReplState
 from adam.utils import debug, lines_to_tabular, log, log2, ing, log_exc
 from adam.utils_athena import Athena
-from adam.utils_k8s.statefulsets import StatefulSets
 from adam.utils_sqlite import SQLite
 
 LIKE = 'e%_%'
@@ -90,7 +91,7 @@ class ExportDatabases:
         if not export_session:
             return
 
-        Athena.clear_cache()
+        ExportDatabases.clear_cache()
 
         keyspaces = {}
         for table in ExportDatabases.table_names(export_session):
@@ -102,17 +103,6 @@ class ExportDatabases:
 
         log(lines_to_tabular([f'{k},{v}' for k, v in keyspaces.items()], header='SCHEMA,# of TABLES', separator=','))
 
-    def disply_export_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()
-        log(lines_to_tabular([f'{table.keyspace}\t{table.table}\t{table.target_table}\t{"export_completed_pending_import" if table.status == "pending_import" else table.status}' for table in tables], header='KEYSPACE\tTABLE\tTARGET_TABLE\tSTATUS', separator='\t'))
-
     def database_names():
         return ExportDatabases.copy_database_names() + ExportDatabases.export_database_names()
 
@@ -171,4 +161,69 @@ class ExportDatabases:
         if db.startswith('s'):
             return SQLite.database_names(prefix=f'{eprefix}_')
         else:
-            return Athena.database_names(like=f'{eprefix}_%')
+            return Athena.database_names(like=f'{eprefix}_%')
+
+    def drop_databases(sts: str, pod: str, namespace: str, database: str = None):
+        importer = None
+        if database:
+            importer = Importer.importer_from_session(database)
+
+        sessions_done = ExportSessions.export_session_names(sts, pod, namespace, importer=importer, export_state='done')
+        sessions = ExportDatabases.sessions_from_dbs(ExportDatabases.drop_export_dbs(database))
+        if sessions_done and sessions:
+            intersects = list(set(sessions_done) & set(sessions))
+            with ing(f'Cleaning up {len(intersects)} completed sessions'):
+                ExportSessions.clean_up_sessions(sts, pod, namespace, list(intersects))
+                ExportSessions.clear_export_session_cache()
+
+    def clear_cache(database: str = None):
+        if not database or database.startswith('s'):
+            SQLite.clear_cache()
+        if not database or database.startswith('e'):
+            Athena.clear_cache()
+
+    def show_databases(importer: str = None):
+        lines = [f'{k}\t{v}' for k, v in ExportDatabases.database_names_with_keyspace_cnt(importer).items()]
+        log(lines_to_tabular(lines, 'NAME\tKEYSPACES', separator='\t'))
+
+class ExportDatabaseService:
+    def __init__(self, handler: 'ExportDatabaseHandler'):
+        self.handler = handler
+
+    def sql(self, query: str, database: str = None):
+        if not database:
+            database = self.handler.state.export_session
+
+        ExportDatabases.run_query(query, database, show_query = True)
+
+    def drop(self, database: str):
+        state = self.handler.state
+
+        ExportDatabases.drop_databases(state.sts, state.pod, state.namespace, database)
+        ExportDatabases.clear_cache(database)
+        if state.export_session == database:
+            state.export_session = None
+
+    def drop_all(self):
+        state = self.handler.state
+
+        ExportDatabases.drop_databases(state.sts, state.pod, state.namespace)
+        ExportDatabases.clear_cache()
+
+        state.export_session = None
+
+    def show_databases(self, importer: str = None):
+        ExportDatabases.show_databases(importer)
+
+class ExportDatabaseHandler:
+    def __init__(self, state: ReplState = None):
+        self.state = state
+
+    def __enter__(self):
+        return ExportDatabaseService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+def export_db(state: ReplState = None):
+    return ExportDatabaseHandler(state)

adam/commands/export/export_select.py

@@ -1,6 +1,6 @@
 from datetime import datetime
-from adam.commands import extract_trailing_options
-from adam.commands.command import Command
+from adam.commands import extract_trailing_options, validate_args
+from adam.commands.command import Command, InvalidArgumentsException
 from adam.commands.export.export_databases import ExportDatabases
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
@@ -43,32 +43,21 @@ class ExportSelect(Command):
 
                         Command.display_help()
 
-                    return 'command-missing'
+                    raise InvalidArgumentsException()
 
-                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 validate_args(args, state, name='SQL statement') as query:
+                    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)
+                        with open(log_file, 'w') as f:
+                            f.write(out)
 
-                    return log_file
+                        return log_file
 
-                ExportDatabases.run_query(f'select {query}', database=state.export_session, output=output if backgrounded else None)
+                    ExportDatabases.run_query(f'select {query}', database=state.export_session, output=output if backgrounded else None)
 
-                return state
+                    return state
 
     def completion(self, state: ReplState):
         if not state.export_session:

adam/commands/export/export_select_x.py

@@ -38,15 +38,15 @@ class ExportSelectX(Command):
             )}
 
             if state.export_session:
-                completions |= {'select': SqlCompleter(
+                completions |= {dml: SqlCompleter(
                     lambda: ExportDatabases.table_names(state.export_session),
-                    dml='select',
+                    dml=dml,
                     expandables={
                         'export-dbs': lambda: ExportDatabases.database_names(),
                         'columns':lambda _: Athena.column_names(database=state.export_session, function='export'),
                     },
                     variant=SqlVariant.ATHENA
-                )}
+                ) for dml in ['select', 'preview']}
 
         return completions
 

adam/commands/export/export_sessions.py

@@ -0,0 +1,124 @@
+import functools
+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.utils import lines_to_tabular, log, parallelize
+from adam.utils_k8s.cassandra_nodes import CassandraNodes
+from adam.utils_k8s.pods import log_prefix
+from adam.utils_k8s.statefulsets import StatefulSets
+
+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 disply_export_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()
+        log(lines_to_tabular([f'{table.keyspace}\t{table.table}\t{table.target_table}\t{"export_completed_pending_import" if table.status == "pending_import" else table.status}' for table in tables], header='KEYSPACE\tTABLE\tTARGET_TABLE\tSTATUS', separator='\t'))

adam/commands/export/export_use.py

@@ -2,8 +2,6 @@ from adam.commands.command import Command
 from adam.commands.export.export_databases import ExportDatabases
 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'
@@ -36,10 +34,6 @@ class ExportUse(Command):
                 return state
 
             state.export_session = args[0]
-            if state.export_session.startswith('e'):
-                Athena.clear_cache()
-            else:
-                SQLite.clear_cache()
 
             ExportDatabases.display_export_db(state.export_session)
 

adam/commands/export/exporter.py

@@ -1,21 +1,21 @@
 from datetime import datetime
-import functools
-import re
 import time
 
+from adam.commands import validate_args
+from adam.commands.command import Command
 from adam.commands.cql.utils_cql import cassandra_table_names, run_cql, table_spec
 from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.export.export_sessions import ExportSessions
 from adam.commands.export.importer import Importer
 from adam.commands.export.importer_athena import AthenaImporter
 from adam.commands.export.importer_sqlite import SqliteImporter
-from adam.commands.export.utils_export import ExportSpec, ExportTableStatus, ExportTableSpec, ImportSpec, csv_dir, find_files
+from adam.commands.export.utils_export import ExportSpec, ExportTableStatus, ExportTableSpec, ImportSpec, csv_dir, find_files, state_with_pod
 from adam.config import Config
 from adam.pod_exec_result import PodExecResult
 from adam.repl_state import ReplState
-from adam.utils import debug, parallelize, log2, ing, log_exc
+from adam.utils import debug, log, parallelize, log2, ing, log_exc
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.pods import log_prefix
-from adam.utils_k8s.statefulsets import StatefulSets
 
 class Exporter:
     def export_tables(args: list[str], state: ReplState, export_only: bool = False, max_workers = 0) -> tuple[list[str], ExportSpec]:
@@ -211,109 +211,7 @@ class Exporter:
 
     def import_from_csv(spec: ExportTableSpec, state: ReplState, session: str, importer: str, table: str, target_table: str, columns: str, multi_tables = True, create_db = False):
         im = AthenaImporter() if importer == 'athena' else SqliteImporter()
-        return im.import_from_csv(state.pod, state.namespace, state.export_session, session if session else state.export_session, spec.keyspace, table, target_table, columns, multi_tables, create_db)
-
-    def clear_export_session_cache():
-        Exporter.find_export_sessions.cache_clear()
-        Exporter.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 Exporter.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: Exporter.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
+        return im.import_from_csv(state, session if session else state.export_session, spec.keyspace, table, target_table, columns, multi_tables, create_db)
 
     def resove_table_n_columns(spec: ExportTableSpec, state: ReplState, include_ks_in_target = False, importer = 'sqlite'):
         table = spec.table
@@ -339,15 +237,58 @@ class Exporter:
 
         return table, target_table, columns
 
-    def drop_databases(sts: str, pod: str, namespace: str, db: str = None):
-        importer = None
-        if db:
-            importer = Importer.importer_from_session(db)
-
-        sessions_done = Exporter.export_session_names(sts, pod, namespace, importer=importer, export_state='done')
-        sessions = ExportDatabases.sessions_from_dbs(ExportDatabases.drop_export_dbs(db))
-        if sessions_done and sessions:
-            intersects = list(set(sessions_done) & set(sessions))
-            with ing(f'Cleaning up {len(intersects)} completed sessions'):
-                Exporter.clean_up_sessions(sts, pod, namespace, list(intersects))
-                Exporter.clear_export_session_cache()
+class ExportService:
+    def __init__(self, handler: 'ExporterHandler'):
+        self.handler = handler
+
+    def export(self, args: list[str], export_only=False):
+        state = self.handler.state
+        export_session = state.export_session
+        spec: ExportSpec = None
+        try:
+            with state_with_pod(state) as state:
+                # --export-only for testing only
+                statuses, spec = Exporter.export_tables(args, state, export_only=export_only)
+                if not statuses:
+                    return state
+
+                ExportSessions.clear_export_session_cache()
+
+                if spec.importer == 'csv' or export_only:
+                    ExportSessions.disply_export_session(state.sts, state.pod, state.namespace, spec.session)
+                else:
+                    log()
+                    ExportDatabases.display_export_db(state.export_session)
+        finally:
+            # if exporting to csv, do not bind the new session id to repl state
+            if spec and spec.importer == 'csv':
+                state.export_session = export_session
+
+        return state
+
+    def import_sesion(self, args: list[str]):
+        state = self.handler.state
+
+        with validate_args(args, state, name='export session') as args_str:
+            with state_with_pod(state) as state:
+                tables, _ = Exporter.import_session(args_str, state)
+                if tables:
+                    ExportSessions.clear_export_session_cache()
+
+                    log()
+                    ExportDatabases.display_export_db(state.export_session)
+
+            return state
+
+class ExporterHandler:
+    def __init__(self, state: ReplState):
+        self.state = state
+
+    def __enter__(self):
+        return ExportService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+def export(state: ReplState):
+    return ExporterHandler(state)

adam/commands/export/import_session.py

@@ -1,6 +1,6 @@
 from adam.commands.command import Command
-from adam.commands.export.export_handlers import export
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_sessions import ExportSessions
+from adam.commands.export.exporter import export
 from adam.repl_state import ReplState, RequiredState
 
 class ImportSession(Command):
@@ -31,8 +31,8 @@ class ImportSession(Command):
 
     def completion(self, state: ReplState):
         # warm up cache
-        Exporter.export_session_names(state.sts, state.pod, state.namespace)
-        Exporter.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import')
+        ExportSessions.export_session_names(state.sts, state.pod, state.namespace)
+        ExportSessions.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import')
 
         return {}
 

adam/commands/export/importer.py

@@ -2,6 +2,7 @@ from abc import abstractmethod
 
 from adam.commands.export.utils_export import csv_dir
 from adam.config import Config
+from adam.repl_state import ReplState
 from adam.utils import ing
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.pods import log_prefix
@@ -12,10 +13,13 @@ class Importer:
         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):
+    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):
         pass
 
-    def move_to_done(self, pod: str, namespace: str, to_session: str, from_session: str, keyspace: str, target_table: str):
+    def move_to_done(self, state: ReplState, from_session: str, keyspace: str, target_table: str):
+        pod = state.pod
+        namespace = state.namespace
+        to_session = state.export_session
         log_file = f'{log_prefix()}-{from_session}_{keyspace}.{target_table}.log.pending_import'
 
         to = f'{log_prefix()}-{to_session}_{keyspace}.{target_table}.log.done'
@@ -30,9 +34,12 @@ class Importer:
 
         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 remove_csv(self, state: ReplState, from_session: str, table: str, target_table: str, multi_tables = True):
+        pod = state.pod
+        namespace = state.namespace
+
+        with ing(f'[{from_session}] Cleaning up temporary files', suppress_log=multi_tables):
+            CassandraNodes.exec(pod, namespace, f'rm -rf {self.csv_file(from_session, table, target_table)}', show_out=Config().is_debug(), shell='bash')
 
     def db(self, session: str, keyspace: str):
         return f'{session}_{keyspace}'