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

adam/commands/export/completions_x.py

@@ -0,0 +1,11 @@
+from adam.commands.export.export_databases import ExportDatabases
+from adam.repl_state import ReplState
+from adam.sql.lark_completer import LarkCompleter
+from adam.utils_athena import Athena
+
+def completions_x(state: ReplState):
+    return LarkCompleter(expandables={
+        'tables': lambda x: ExportDatabases.table_names(state.export_session),
+        'export-databases': lambda x: ExportDatabases.database_names(),
+        'columns': lambda x: Athena.column_names(database=state.export_session, function='export'),
+    }, variant=ReplState.X).completions_for_nesting()

adam/commands/export/download_export_session.py

@@ -0,0 +1,40 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.export.export_sessions import ExportSessions, export_session
+from adam.repl_state import ReplState, RequiredState
+
+class DownloadExportSession(Command):
+    COMMAND = 'download export session'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DownloadExportSession, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DownloadExportSession.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='export session') as session:
+                with export_session(state) as sessions:
+                    sessions.download_session(session)
+
+            return state
+
+    def completion(self, state: ReplState):
+        return {}
+        # return super().completion(state, {session: None for session in ExportSessions.export_session_names(state.sts, state.pod, state.namespace, export_state='pending_import')})
+
+    def help(self, _: ReplState):
+        return f'{DownloadExportSession.COMMAND} <export-session-name>\t download csv files in export session'

adam/commands/export/drop_export_database.py

@@ -1,9 +1,7 @@
+from adam.commands import validate_args
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+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 DropExportDatabase(Command):
     COMMAND = 'drop export database'
@@ -27,29 +25,12 @@ class DropExportDatabase(Command):
         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
+        with self.validate(args, state) as (args, state):
+            with validate_args(args, state, name='database name') as db:
+                with export_db(state) as dbs:
+                    dbs.drop(args[0])
 
-        if not len(args):
-            if state.in_repl:
-                log2('Database name is required.')
-                log2()
-            else:
-                log2('* Database name is missing.')
-                Command.display_help()
-
-            return 'command-missing'
-
-        Exporter.drop_databases(state.sts, state.pod, state.namespace, args[0])
-
-        SQLite.clear_cache()
-        Athena.clear_cache()
-
-        if state.export_session == args[0]:
-            state.export_session = None
-
-        return state
+                return state
 
     def completion(self, _: ReplState):
         return {}

adam/commands/export/drop_export_databases.py

@@ -1,8 +1,6 @@
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
-from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class DropExportDatabases(Command):
     COMMAND = 'drop all export databases'
@@ -26,18 +24,11 @@ class DropExportDatabases(Command):
         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
-
-        Exporter.drop_databases(state.sts, state.pod, state.namespace)
-
-        SQLite.clear_cache()
-        Athena.clear_cache()
+        with self.validate(args, state) as (args, state):
+            with export_db(state) as dbs:
+                dbs.drop_all()
 
-        state.export_session = None
-
-        return state
+            return state
 
     def completion(self, _: ReplState):
         return {}

adam/commands/export/export.py

@@ -1,12 +1,7 @@
+from adam.commands import extract_options
 from adam.commands.command import Command
-from adam.commands.cql.cql_utils import cassandra_keyspaces, cassandra_table_names
-from adam.commands.export.export_databases import ExportDatabases
-from adam.commands.export.exporter import Exporter
-from adam.commands.export.utils_export import ExportSpec
+from adam.commands.export.exporter import export
 from adam.repl_state import ReplState, RequiredState
-from adam.sql.sql_completer import SqlCompleter, SqlVariant
-from adam.utils import log
-from adam.utils_k8s.statefulsets import StatefulSets
 
 class ExportTables(Command):
     COMMAND = 'export'
@@ -30,53 +25,12 @@ class ExportTables(Command):
         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
-
-        # --export-only for testing only
-        args, export_only = Command.extract_options(args, ['--export-only'])
-
-        if not state.pod:
-            state.push()
-            state.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
-
-        export_session = state.export_session
-        spec: ExportSpec = None
-        try:
-            statuses, spec = Exporter.export_tables(args, state, export_only=export_only)
-            if not statuses:
-                return state
-
-            Exporter.clear_export_session_cache()
-
-            if spec.importer == 'csv' or export_only:
-                ExportDatabases.disply_export_session(state.sts, state.pod, state.namespace, spec.session)
-            else:
-                log()
-                ExportDatabases.display_export_db(state.export_session)
-
-        finally:
-            state.pop()
-
-            # 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
+        with self.validate(args, state) as (args, state):
+            with extract_options(args, '--export-only') as (args, export_only):
+                with export(state) as exporter:
+                    return exporter.export(args, export_only=export_only)
 
     def completion(self, state: ReplState):
-        def sc():
-            return SqlCompleter(lambda: cassandra_table_names(state), expandables={
-                'dml':'export',
-                'columns': lambda table: ['id', '*'],
-                'keyspaces': lambda: cassandra_keyspaces(state),
-                'export-dbs': lambda: ExportDatabases.database_names(),
-            }, variant=SqlVariant.CQL)
-
-        if state.sts:
-            return {f'@{p}': {ExportTables.COMMAND: sc()} for p in StatefulSets.pod_names(state.sts, state.namespace)}
-
         return {}
 
     def help(self, _: ReplState):

adam/commands/export/export_databases.py

@@ -1,26 +1,37 @@
+from collections.abc import Callable
+from datetime import datetime
 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.utils import lines_to_tabular, log, log2, ing
+from adam.repl_session import ReplSession
+from adam.repl_state import ReplState
+from adam.utils import debug, log_timing, tabulize, 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%_%'
 
 class ExportDatabases:
-    def run_query(query: str, database: str, show_query=False) -> int:
+    def run_query(query: str, database: str, show_query=False, output: Callable[[str], str] = None) -> int:
         cnt: int = 0
 
         if show_query:
             log2(query)
 
+        log_file = None
         if database.startswith('s'):
-            cnt += SQLite.run_query(query, database=database)
+            cnt0, log_file = SQLite.run_query(query, database=database, output=output)
+            cnt += cnt0
         else:
-            cnt += Athena.run_query(query, database=database)
+            cnt0, log_file = Athena.run_query(query, database=database, output=output)
+            cnt += cnt0
+
+        if Config().get('repl.history.push-cat-log-file', True):
+            if log_file and ReplSession().prompt_session:
+                ReplSession().prompt_session.history.append_string(f':sh cat {log_file}')
 
         return cnt
 
@@ -47,15 +58,13 @@ class ExportDatabases:
         dbs = SQLite.database_names(db)
         if dbs:
             with ing(f'Droping {len(dbs)} SQLite databases'):
-                try:
+                with log_exc():
                     for db in dbs:
                         file_path = f'{SQLite.local_db_dir()}/{db}'
                         try:
                             os.remove(file_path)
                         except OSError as e:
                             pass
-                except:
-                    pass
 
         return dbs
 
@@ -65,49 +74,35 @@ class ExportDatabases:
             with ing(f'Droping {len(dbs)} Athena databases'):
                 for db in dbs:
                     query = f'DROP DATABASE {db} CASCADE'
-                    if Config().is_debug():
-                        log2(query)
+                    debug(query)
                     Athena.query(query)
 
         with ing(f'Deleting s3 folder: export'):
-            try:
+            with log_exc():
                 if not db:
                     db = ''
 
                 s3 = boto3.resource('s3')
                 bucket = s3.Bucket(Config().get('export.bucket', 'c3.ops--qing'))
                 bucket.objects.filter(Prefix=f'export/{db}').delete()
-            except:
-                pass
 
         return dbs
 
-    def display_export_db(export_session: str):
-        if not export_session:
+    def show_database(database: str):
+        if not database:
             return
 
-        Athena.clear_cache()
+        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=','))
-
-    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'))
+        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()
@@ -116,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 = {}
@@ -167,4 +163,84 @@ 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()]
+        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, backgrounded = False):
+        if not database:
+            database = self.handler.state.export_session
+
+        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
+
+        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)
+
+    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
+
+    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,13 +1,10 @@
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases
+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
-from adam.utils_sqlite import SQLite
 
+# 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):
@@ -24,62 +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)
-
-        state, args = self.apply_state(args, state)
-        if not self.validate_state(state):
-            return state
-
-        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)
-
-        ExportDatabases.run_query(f'select {query}', database=state.export_session)
-
-        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'