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

adam/commands/export/export_databases.py

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+from datetime import datetime
 import os
 import boto3
 
@@ -7,7 +8,7 @@ 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 import debug, log_timing, tabulize, log2, ing, log_exc
 from adam.utils_athena import Athena
 from adam.utils_sqlite import SQLite
 
@@ -87,21 +88,21 @@ class ExportDatabases:
 
         return dbs
 
-    def display_export_db(export_session: str):
-        if not export_session:
+    def show_database(database: str):
+        if not database:
             return
 
         ExportDatabases.clear_cache()
 
         keyspaces = {}
-        for table in ExportDatabases.table_names(export_session):
+        for table in ExportDatabases.table_names(database):
             keyspace = table.split('.')[0]
             if keyspace in keyspaces:
                 keyspaces[keyspace] += 1
             else:
                 keyspaces[keyspace] = 1
 
-        log(lines_to_tabular([f'{k},{v}' for k, v in keyspaces.items()], header='SCHEMA,# of TABLES', separator=','))
+        tabulize(keyspaces.items(), lambda a: f'{a[0]},{a[1]}', header='SCHEMA,# of TABLES', separator=',')
 
     def database_names():
         return ExportDatabases.copy_database_names() + ExportDatabases.export_database_names()
@@ -110,7 +111,8 @@ class ExportDatabases:
         return list({n.split('_')[0] for n in SQLite.database_names()})
 
     def export_database_names():
-        return list({n.split('_')[0] for n in Athena.database_names(LIKE)})
+        with log_timing('ExportDatabases.Athena.database_names'):
+            return list({n.split('_')[0] for n in Athena.database_names(LIKE)})
 
     def database_names_with_keyspace_cnt(importer: str = None):
         r = {}
@@ -184,17 +186,26 @@ class ExportDatabases:
 
     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'))
+        tabulize(lines, header='NAME\tKEYSPACES', separator='\t')
 
 class ExportDatabaseService:
     def __init__(self, handler: 'ExportDatabaseHandler'):
         self.handler = handler
 
-    def sql(self, query: str, database: str = None):
+    def sql(self, query: str, database: str = None, backgrounded = False):
         if not database:
             database = self.handler.state.export_session
 
-        ExportDatabases.run_query(query, database, show_query = True)
+        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")}-export.log'
+
+            with open(log_file, 'w') as f:
+                f.write(out)
+
+            return log_file
+
+        ExportDatabases.run_query(query, database, output = output if backgrounded else None, show_query = not backgrounded)
 
     def drop(self, database: str):
         state = self.handler.state
@@ -215,6 +226,12 @@ class ExportDatabaseService:
     def show_databases(self, importer: str = None):
         ExportDatabases.show_databases(importer)
 
+    def show_database(self, database: str = None):
+        if not database:
+            database = self.handler.state.export_session
+
+        ExportDatabases.show_database(database)
+
 class ExportDatabaseHandler:
     def __init__(self, state: ReplState = None):
         self.state = state

adam/commands/export/export_select.py

@@ -1,15 +1,10 @@
-from datetime import datetime
-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.commands.command import Command
+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,58 +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()
-
-                    raise InvalidArgumentsException()
-
-                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)
-
-                        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

@@ -1,13 +1,16 @@
 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.utils import lines_to_tabular, log, parallelize
+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 log_prefix
+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():
@@ -74,7 +77,9 @@ class ExportSessions:
         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:
+        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
@@ -112,7 +117,7 @@ class ExportSessions:
 
         return csv_cnt, log_cnt
 
-    def disply_export_session(sts: str, pod: str, namespace: str, session: str):
+    def show_session(sts: str, pod: str, namespace: str, session: str):
         if not pod:
             pod = StatefulSets.pod_names(sts, namespace)[0]
 
@@ -121,4 +126,84 @@ class ExportSessions:
 
         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'))
+        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,5 +1,6 @@
+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
 
@@ -19,28 +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]
+                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'

adam/commands/export/exporter.py

@@ -1,10 +1,9 @@
 from datetime import datetime
 import time
 
-from adam.commands import validate_args
-from adam.commands.command import Command
+from adam.commands.command import InvalidArgumentsException
 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_databases import export_db
 from adam.commands.export.export_sessions import ExportSessions
 from adam.commands.export.importer import Importer
 from adam.commands.export.importer_athena import AthenaImporter
@@ -69,10 +68,10 @@ class Exporter:
 
         return spec
 
-    def import_session(args: list[str], state: ReplState, max_workers = 0) -> tuple[list[str], ExportSpec]:
+    def import_session(spec_str: str, state: ReplState, max_workers = 0) -> tuple[list[str], ExportSpec]:
         import_spec: ImportSpec = None
         with log_exc(True):
-            import_spec = Exporter.import_spec(' '.join(args), state)
+            import_spec = Exporter.import_spec(spec_str, state)
             tables, status_in_whole = ExportTableStatus.from_session(state.sts, state.pod, state.namespace, import_spec.session)
             if status_in_whole == 'done':
                 log2(f'The session has been completely done - no more csv files are found.')
@@ -84,8 +83,30 @@ class Exporter:
 
         return [], None
 
-    def import_spec(spec_str: str, state: ReplState):
-        spec: ImportSpec = ImportSpec.parse_specs(spec_str)
+    def import_local_csv_files(spec_str: str, state: ReplState, max_workers = 0) -> tuple[list[str], ExportSpec]:
+        spec: ImportSpec = None
+        with log_exc(True):
+            spec = Exporter.import_spec(spec_str, state, files=True)
+            if not spec.table_name:
+                log2(f"Use 'as <database-name>.<table-name>'.")
+                raise InvalidArgumentsException()
+
+            d_t = spec.table_name.split('.')
+            if len(d_t) != 2:
+                log2(f'Need <database-name>.<table-name> format for target table.')
+                raise InvalidArgumentsException()
+
+            database = d_t[0]
+            table = d_t[1]
+            im = AthenaImporter() if spec.importer == 'athena' else SqliteImporter()
+
+            with parallelize(spec.files, max_workers, msg='Importing|Imported {size} csv files') as exec:
+                return exec.map(lambda f: im.import_from_local_csv(state, database, table, f, len(spec.files) > 1, True)), spec
+
+        return [], None
+
+    def import_spec(spec_str: str, state: ReplState, files = False):
+        spec: ImportSpec = ImportSpec.parse_specs(spec_str, files=files)
 
         session = state.export_session
         if session:
@@ -104,15 +125,19 @@ class Exporter:
             if not spec.importer:
                 spec.importer = Importer.importer_from_session(spec.session)
 
-            if spec.importer == 'csv':
+            if not spec.importer or spec.importer == 'csv':
                 spec.importer = Config().get('export.default-importer', 'sqlite')
 
             if spec.importer == 'athena' and not AthenaImporter.ping():
                 raise Exception('Credentials for Athena is not present.')
 
             prefix = Importer.prefix_from_importer(spec.importer)
-            session = f'{prefix}{spec.session[1:]}'
-            state.export_session = session
+            if spec.session:
+                spec.session = f'{prefix}{spec.session[1:]}'
+            else:
+                spec.session = f'{prefix}{datetime.now().strftime("%Y%m%d%H%M%S")[3:]}'
+
+            state.export_session = spec.session
 
         return spec
 
@@ -196,7 +221,7 @@ class Exporter:
         queries.append(f"COPY {spec.keyspace}.{table}({columns}) TO '{csv_file}' WITH HEADER = TRUE")
         r: PodExecResult = ing(
             f'[{session}] Dumping table {spec.keyspace}.{table}{f" with consistency {consistency}" if consistency else ""}',
-            lambda: run_cql(state, ';'.join(queries), show_out=Config().is_debug(), background=True, log_file=log_file),
+            lambda: run_cql(state, ';'.join(queries), show_out=Config().is_debug(), backgrounded=True, log_file=log_file),
             suppress_log=suppress_ing_log)
 
         return log_file
@@ -255,10 +280,11 @@ class ExportService:
                 ExportSessions.clear_export_session_cache()
 
                 if spec.importer == 'csv' or export_only:
-                    ExportSessions.disply_export_session(state.sts, state.pod, state.namespace, spec.session)
+                    ExportSessions.show_session(state.sts, state.pod, state.namespace, spec.session)
                 else:
                     log()
-                    ExportDatabases.display_export_db(state.export_session)
+                    with export_db(state) as dbs:
+                        dbs.show_database()
         finally:
             # if exporting to csv, do not bind the new session id to repl state
             if spec and spec.importer == 'csv':
@@ -266,19 +292,31 @@ class ExportService:
 
         return state
 
-    def import_sesion(self, args: list[str]):
+    def import_session(self, spec_str: 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()
+        tables, _ = Exporter.import_session(spec_str, state)
+        if tables:
+            ExportSessions.clear_export_session_cache()
 
-                    log()
-                    ExportDatabases.display_export_db(state.export_session)
+            log()
+            with export_db(state) as dbs:
+                dbs.show_database()
 
-            return state
+        return state
+
+    def import_files(self, spec_str: str):
+        state = self.handler.state
+
+        tables, _ = Exporter.import_local_csv_files(spec_str, state)
+        if tables:
+            ExportSessions.clear_export_session_cache()
+
+            log()
+            with export_db(state) as dbs:
+                dbs.show_database()
+
+        return state
 
 class ExporterHandler:
     def __init__(self, state: ReplState):

adam/commands/export/import_files.py

@@ -0,0 +1,44 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.export.export_sessions import ExportSessions
+from adam.commands.export.exporter import export
+from adam.commands.export.utils_export import state_with_pod
+from adam.repl_state import ReplState, RequiredState
+
+class ImportCSVFiles(Command):
+    COMMAND = 'import files'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ImportCSVFiles, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ImportCSVFiles.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)
+
+        with self.validate(args, state) as (args, state):
+            with validate_args(args, state, name='file') as spec:
+                with state_with_pod(state) as state:
+                    with export(state) as exporter:
+                        return exporter.import_files(spec)
+
+    def completion(self, state: ReplState):
+        # warm up cache
+        # 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 {}
+
+    def help(self, _: ReplState):
+        return f'{ImportCSVFiles.COMMAND} <file-names,...>\t import files in session to Athena or SQLite'

adam/commands/export/import_session.py

@@ -1,6 +1,8 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
 from adam.commands.export.export_sessions import ExportSessions
 from adam.commands.export.exporter import export
+from adam.commands.export.utils_export import state_with_pod
 from adam.repl_state import ReplState, RequiredState
 
 class ImportSession(Command):
@@ -26,13 +28,15 @@ class ImportSession(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            with export(state) as exporter:
-                return exporter.import_sesion(args)
+            with validate_args(args, state, name='export session') as spec:
+                with state_with_pod(state) as state:
+                    with export(state) as exporter:
+                        return exporter.import_session(spec)
 
     def completion(self, state: ReplState):
         # warm up cache
-        ExportSessions.export_session_names(state.sts, state.pod, state.namespace)
-        ExportSessions.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

@@ -16,6 +16,13 @@ class Importer:
     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
 
+    @abstractmethod
+    def import_from_local_csv(self, state: ReplState,
+                        keyspace: str, table: str, target_table: str, columns: str,
+                        csv_file: str,
+                        multi_tables = True, create_db = False):
+        pass
+
     def move_to_done(self, state: ReplState, from_session: str, keyspace: str, target_table: str):
         pod = state.pod
         namespace = state.namespace