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

adam/commands/export/export.py

@@ -1,8 +1,11 @@
-from adam.commands.export.utils_export import export_tables
+from adam.commands import extract_options
 from adam.commands.command import Command
-from adam.commands.cql.cql_utils import cassandra_table_names
+from adam.commands.cql.utils_cql import cassandra_keyspaces, cassandra_table_names
+from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.export.export_handlers import export
 from adam.repl_state import ReplState, RequiredState
-from adam.sql.sql_completer import SqlCompleter
+from adam.sql.sql_completer import SqlCompleter, SqlVariant
+from adam.utils import log
 from adam.utils_k8s.statefulsets import StatefulSets
 
 class ExportTables(Command):
@@ -27,34 +30,24 @@ 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
-
-        if not state.pod:
-            state.push()
-            state.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
-        try:
-            export_tables(args, state)
-        finally:
-            state.pop()
-
-        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),
-                dml='export',
-                columns = lambda table: ['id', '*'],
-                variant='cql'
-            )
+            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 {ExportTables.COMMAND: sc()} | \
-                {f'@{p}': {ExportTables.COMMAND: sc()} for p in StatefulSets.pod_names(state.sts, state.namespace)}
+        if super().completion(state):
+            return {f'@{p}': {ExportTables.COMMAND: sc()} for p in StatefulSets.pod_names(state.sts, state.namespace)}
 
         return {}
 
     def help(self, _: ReplState):
-        return f'{ExportTables.COMMAND} [TABLE] [as target-name] [with consistency <level>]\t export table'
+        return f'{ExportTables.COMMAND} [* [in KEYSPACE]] | [TABLE] [as target-name] [with consistency <level>]\t export tables to Sqlite, Athena or CSV file'

adam/commands/export/export_databases.py

@@ -0,0 +1,174 @@
+from collections.abc import Callable
+import os
+import boto3
+
+from adam.commands.export.utils_export import ExportTableStatus
+from adam.config import Config
+from adam.repl_session import ReplSession
+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%_%'
+
+class ExportDatabases:
+    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'):
+            cnt0, log_file = SQLite.run_query(query, database=database, output=output)
+            cnt += cnt0
+        else:
+            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
+
+    def sessions_from_dbs(dbs: list[str]):
+
+        sessions = set()
+
+        for db in dbs:
+            sessions.add(db.split('_')[0])
+
+        return list(sessions)
+
+    def drop_export_dbs(db: str = None):
+        dbs: list[str] = []
+
+        if not db or db.startswith('s'):
+            dbs.extend(ExportDatabases.drop_sqlite_dbs(db))
+        if not db or db.startswith('e'):
+            dbs.extend(ExportDatabases.drop_athena_dbs(db))
+
+        return dbs
+
+    def drop_sqlite_dbs(db: str = None):
+        dbs = SQLite.database_names(db)
+        if dbs:
+            with ing(f'Droping {len(dbs)} SQLite databases'):
+                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
+
+        return dbs
+
+    def drop_athena_dbs(db: str = None):
+        dbs = Athena.database_names(f'{db}_%' if db else LIKE)
+        if dbs:
+            with ing(f'Droping {len(dbs)} Athena databases'):
+                for db in dbs:
+                    query = f'DROP DATABASE {db} CASCADE'
+                    debug(query)
+                    Athena.query(query)
+
+        with ing(f'Deleting s3 folder: export'):
+            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()
+
+        return dbs
+
+    def display_export_db(export_session: str):
+        if not export_session:
+            return
+
+        Athena.clear_cache()
+
+        keyspaces = {}
+        for table in ExportDatabases.table_names(export_session):
+            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'))
+
+    def database_names():
+        return ExportDatabases.copy_database_names() + ExportDatabases.export_database_names()
+
+    def copy_database_names():
+        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)})
+
+    def database_names_with_keyspace_cnt(importer: str = None):
+        r = {}
+
+        names = []
+        if not importer:
+            names = SQLite.database_names() + Athena.database_names(LIKE)
+        elif importer == 'athena':
+            names = Athena.database_names(LIKE)
+        else:
+            names = SQLite.database_names()
+
+        for n in names:
+            tokens = n.split('_')
+            name = tokens[0]
+            keyspace = None
+            if len(tokens) > 1:
+                keyspace = tokens[1].replace('.db', '')
+
+            if keyspace == 'root':
+                continue
+
+            if name in r:
+                r[name] += 1
+            else:
+                r[name] = 1
+
+        return r
+
+    def table_names(session: str):
+        tables = []
+
+        for session in ExportDatabases._session_database_names(session):
+            if session.startswith('s'):
+                for table in SQLite.table_names(database=session):
+                    tables.append(f'{SQLite.keyspace(session)}.{table}')
+            else:
+                for table in Athena.table_names(database=session, function='export'):
+                    tables.append(f'{session}.{table}')
+
+        return tables
+
+    def _session_database_names(db: str):
+        eprefix = db
+        if '_' in db:
+            eprefix = db.split('_')[0]
+
+        if db.startswith('s'):
+            return SQLite.database_names(prefix=f'{eprefix}_')
+        else:
+            return Athena.database_names(like=f'{eprefix}_%')

adam/commands/export/export_handlers.py

@@ -0,0 +1,71 @@
+from adam.commands.command import Command
+from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.export.exporter import Exporter
+from adam.commands.export.utils_export import ExportSpec, state_with_pod
+from adam.repl_state import ReplState
+from adam.utils import log, log2
+
+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
+
+                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:
+            # 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
+
+        if not args:
+            if state.in_repl:
+                log2('Specify export session name.')
+            else:
+                log2('* Export session name is missing.')
+
+                Command.display_help()
+
+            return 'command-missing'
+
+        with state_with_pod(state) as state:
+            tables, _ = Exporter.import_session(args, state)
+            if tables:
+                Exporter.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/export_select.py

@@ -1,11 +1,15 @@
+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.repl_state import ReplState, RequiredState
-from adam.sql.sql_completer import SqlCompleter
+from adam.sql.sql_completer import SqlCompleter, SqlVariant
 from adam.utils import log2
 from adam.utils_athena import Athena
 
 class ExportSelect(Command):
-    COMMAND = '&select'
+    COMMAND = '.select'
 
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
@@ -20,49 +24,71 @@ class ExportSelect(Command):
         return ExportSelect.COMMAND
 
     def required(self):
-        return [RequiredState.CLUSTER_OR_POD, ReplState.X]
+        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
+        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.')
 
-        if not args:
-            if state.in_repl:
-                log2('Use a SQL statement.')
-            else:
-                log2('* SQL statement is missing.')
+                        Command.display_help()
 
-                Command.display_help()
+                    return 'command-missing'
 
-            return 'command-missing'
+                if not args:
+                    if state.in_repl:
+                        log2('Use a SQL statement.')
+                    else:
+                        log2('* SQL statement is missing.')
 
-        query = ' '.join(args)
+                        Command.display_help()
 
-        Athena.run_query(f'select {query}', database=f'export_{state.export_session}')
+                    return 'command-missing'
 
-        return state
+                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:
             return {}
 
-        db = f'export_{state.export_session}'
+        db = state.export_session
 
         # warm up the caches first time when x: drive is accessed
-        Athena.table_names(database=db, function='export')
+        ExportDatabases.table_names(db)
         Athena.column_names(database=db, function='export')
         Athena.column_names(partition_cols_only=True, database=db, function='export')
 
         return {ExportSelect.COMMAND: SqlCompleter(
-            lambda: Athena.table_names(database=db, function='export'),
+            lambda: ExportDatabases.table_names(db),
             dml='select',
-            columns=lambda table: Athena.column_names(database=db, function='export'),
-            variant='athena'
+            expandables={
+                'columns':lambda table: Athena.column_names(database=db, function='export'),
+            },
+            variant=SqlVariant.ATHENA
         )}
 
     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_select_x.py

@@ -0,0 +1,54 @@
+from adam.commands.command import Command
+from adam.commands.export.export_databases import ExportDatabases
+from adam.repl_state import ReplState, RequiredState
+from adam.sql.sql_completer import SqlCompleter, SqlVariant
+from adam.utils_athena import Athena
+
+# No action body, only for a help entry and auto-completion
+class ExportSelectX(Command):
+    COMMAND = 'select_on_x'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportSelectX, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportSelectX.COMMAND
+
+    def required(self):
+        return RequiredState.EXPORT_DB
+
+    def completion(self, state: ReplState):
+        completions = {}
+
+        if state.device == ReplState.X:
+            completions = {'drop': SqlCompleter(
+                lambda: ExportDatabases.table_names(state.export_session),
+                dml='drop',
+                expandables={
+                    'export-dbs': lambda: ExportDatabases.database_names(),
+                    'columns':lambda _: Athena.column_names(database=state.export_session, function='export'),
+                },
+                variant=SqlVariant.ATHENA
+            )}
+
+            if state.export_session:
+                completions |= {'select': SqlCompleter(
+                    lambda: ExportDatabases.table_names(state.export_session),
+                    dml='select',
+                    expandables={
+                        'export-dbs': lambda: ExportDatabases.database_names(),
+                        'columns':lambda _: Athena.column_names(database=state.export_session, function='export'),
+                    },
+                    variant=SqlVariant.ATHENA
+                )}
+
+        return completions
+
+    def help(self, _: ReplState):
+        return f'<sql-select-statements>\t run queries on export database'

adam/commands/export/export_use.py

@@ -1,12 +1,12 @@
 from adam.commands.command import Command
-from adam.commands.export.export_select import ExportSelect
-from adam.config import Config
-from adam.repl_state import ReplState, RequiredState
+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'
+    COMMAND = 'use'
 
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
@@ -21,36 +21,32 @@ class ExportUse(Command):
         return ExportUse.COMMAND
 
     def required(self):
-        return RequiredState.CLUSTER_OR_POD
+        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
+        with self.validate(args, state) as (args, state):
+            if not args:
+                state.export_session = None
 
-        if not args:
-            if state.in_repl:
-                log2('Specify database to use.')
-            else:
-                log2('* database is missing.')
+                log2('Export database is unset.')
 
-                Command.display_help()
+                return state
 
-            return 'command-missing'
+            state.export_session = args[0]
+            if state.export_session.startswith('e'):
+                Athena.clear_cache()
+            else:
+                SQLite.clear_cache()
 
-        state.export_session = args[0].replace('export_', '') if args[0].startswith('export_') else args[0]
-        Athena.clear_cache()
+            ExportDatabases.display_export_db(state.export_session)
 
-        return state
+            return state
 
     def completion(self, state: ReplState):
-        # warm up the caches first time when l: drive is accessed
-        Athena.database_names()
-
-        return super().completion(state, {n: None for n in Athena.database_names()})
+        return super().completion(state, {n: None for n in ExportDatabases.database_names()})
 
     def help(self, _: ReplState):
-        return f'{ExportUse.COMMAND} <export db name>\t use Export Database'
+        return f'{ExportUse.COMMAND} <export-database-name>\t use export database'