kaqing 2.0.110__py3-none-any.whl → 2.0.214__py3-none-any.whl

adam/commands/export/clean_up_all_export_sessions.py

@@ -0,0 +1,37 @@
+from adam.commands.command import Command
+from adam.commands.export.export_sessions import export_session
+from adam.repl_state import ReplState, RequiredState
+
+class CleanUpAllExportSessions(Command):
+    COMMAND = 'clean up all export sessions'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(CleanUpAllExportSessions, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return CleanUpAllExportSessions.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 export_session(state) as sessions:
+                sessions.clean_up_all()
+
+            return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{CleanUpAllExportSessions.COMMAND}\t clean up all export sessions'

adam/commands/export/clean_up_export_sessions.py

@@ -0,0 +1,39 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.export.export_sessions import export_session
+from adam.repl_state import ReplState, RequiredState
+
+class CleanUpExportSessions(Command):
+    COMMAND = 'clean up export sessions'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(CleanUpExportSessions, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return CleanUpExportSessions.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', separator=',') as session_names:
+                with export_session(state) as sessions:
+                    sessions.clean_up(session_names)
+
+                return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{CleanUpExportSessions.COMMAND} <export-session-name,...>\t clean up export sessions'

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

@@ -0,0 +1,39 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.export.export_databases import export_db
+from adam.repl_state import ReplState
+
+class DropExportDatabase(Command):
+    COMMAND = 'drop export database'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DropExportDatabase, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DropExportDatabase.COMMAND
+
+    def required(self):
+        return [ReplState.C, ReplState.X]
+
+    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='database name') as db:
+                with export_db(state) as dbs:
+                    dbs.drop(args[0])
+
+                return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{DropExportDatabase.COMMAND} <export-database-name>\t drop export database'

adam/commands/export/drop_export_databases.py

@@ -0,0 +1,37 @@
+from adam.commands.command import Command
+from adam.commands.export.export_databases import export_db
+from adam.repl_state import ReplState
+
+class DropExportDatabases(Command):
+    COMMAND = 'drop all export databases'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DropExportDatabases, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DropExportDatabases.COMMAND
+
+    def required(self):
+        return [ReplState.C, ReplState.X]
+
+    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 export_db(state) as dbs:
+                dbs.drop_all()
+
+            return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{DropExportDatabases.COMMAND}\t drop all export databases'

adam/commands/export/export.py

@@ -0,0 +1,37 @@
+from adam.commands import extract_options
+from adam.commands.command import Command
+from adam.commands.export.exporter import export
+from adam.repl_state import ReplState, RequiredState
+
+class ExportTables(Command):
+    COMMAND = 'export'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportTables, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportTables.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 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):
+        return {}
+
+    def help(self, _: ReplState):
+        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,251 @@
+from collections.abc import Callable
+from datetime import datetime
+from functools import partial
+import os
+import boto3
+
+from adam.commands.export.export_sessions import ExportSessions
+from adam.commands.export.importer import Importer
+from adam.commands.export.utils_export import export_log_dir
+from adam.config import Config
+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_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':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 show_database(database: str):
+        if not database:
+            return
+
+        ExportDatabases.clear_cache()
+
+        keyspaces = {}
+        for table in ExportDatabases.table_names(database):
+            keyspace = table.split('.')[0]
+            if keyspace in keyspaces:
+                keyspaces[keyspace] += 1
+            else:
+                keyspaces[keyspace] = 1
+
+        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()
+
+    def copy_database_names():
+        return list({n.split('_')[0] for n in SQLite.database_names()})
+
+    def export_database_names():
+        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 = {}
+
+        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}_%')
+
+    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, export_log: str = None):
+        if not database:
+            database = self.handler.state.export_session
+
+        def output(export_log: str, out: str):
+            # serving for 1. export log or 2. job log for the sampling select 10 rows afterward
+            flag = 'at'
+            if not export_log:
+                export_log = f'{export_log_dir()}/{datetime.now().strftime("%d%H%M%S")}-export.log'
+                flag = 'w'
+
+            with open(export_log, flag) as f:
+                f.write(out)
+
+            return export_log
+
+        ExportDatabases.run_query(query, database, output = partial(output, export_log) 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

@@ -0,0 +1,34 @@
+from adam.commands.command import Command
+from adam.commands.export.completions_x import completions_x
+from adam.repl_state import ReplState, RequiredState
+
+# No action body, only for a help entry and auto-completion
+class ExportSelect(Command):
+    COMMAND = 'select_on_x'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ExportSelect, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ExportSelect.COMMAND
+
+    def required(self):
+        return RequiredState.EXPORT_DB
+
+    def completion(self, state: ReplState):
+        if state.device != ReplState.X:
+            return {}
+
+        if state.export_session:
+            return completions_x(state)
+
+        return {}
+
+    def help(self, _: ReplState):
+        return f'<sql-select-statements>\t run queries on export database'