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

adam/commands/devices/device_export.py

@@ -0,0 +1,86 @@
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+from adam.commands.export.export_databases import ExportDatabases
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import lines_to_tabular, log, log2
+
+class DeviceExport(Command, Device):
+    COMMAND = f'{ReplState.X}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DeviceExport, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DeviceExport.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.X
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DeviceExport.COMMAND}\t move to Export Database Operations device'
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.export_session:
+            self.show_export_tables(state.export_session)
+        else:
+            self.show_export_databases()
+
+    def show_export_databases(self, 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'))
+
+    def show_export_tables(self, export_session: str):
+        log(lines_to_tabular(ExportDatabases.table_names(export_session), 'NAME', separator=','))
+
+    def cd(self, dir: str, state: ReplState):
+        if dir in ['', '..']:
+            state.export_session = None
+        else:
+            state.export_session = dir
+
+    def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        if state.export_session:
+            return {cmd: {'..': None} | {n: None for n in ExportDatabases.database_names()}}
+        else:
+            return {cmd: {n: None for n in ExportDatabases.database_names()}}
+
+    def pwd(self, state: ReplState):
+        words = []
+
+        if state.export_session:
+            words.append(state.export_session)
+
+        return '\t'.join([f'{ReplState.X}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        result = chain.run(f'.{cmd}', state)
+        if type(result) is ReplState:
+            if state.export_session and not result.export_session:
+                state.export_session = None
+
+        return True, result
+
+    def enter(self, state: ReplState):
+        if auto_enter := Config().get('repl.x.auto-enter', 'no'):
+            if auto_enter == 'latest':
+                Config().wait_log(f'Moving to the latest export database...')
+                if dbs := ExportDatabases.database_names():
+                    state.export_session = sorted(dbs)[-1]
+                else:
+                    log2('No export database found.')
+

adam/commands/devices/device_postgres.py

@@ -0,0 +1,109 @@
+from adam.commands.command import Command
+from adam.commands.devices.device import Device
+from adam.commands.postgres.postgres_context import PostgresContext
+from adam.commands.postgres.postgres_utils import pg_database_names, pg_table_names
+from adam.config import Config
+from adam.repl_state import ReplState
+from adam.utils import lines_to_tabular, log
+
+class DevicePostgres(Command, Device):
+    COMMAND = f'{ReplState.P}:'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DevicePostgres, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DevicePostgres.COMMAND
+
+    def run(self, cmd: str, state: ReplState):
+        if not self.args(cmd):
+            return super().run(cmd, state)
+
+        state.device = ReplState.P
+
+        return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state)
+
+    def help(self, _: ReplState):
+        return f'{DevicePostgres.COMMAND}\t move to Postgres Operations device'
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.pg_path:
+            pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
+            if pg.db:
+                self.show_pg_tables(pg)
+            else:
+                self.show_pg_databases(pg)
+        else:
+            self.show_pg_hosts(state)
+
+    def show_pg_hosts(self, state: ReplState):
+        if state.namespace:
+            def line(pg: PostgresContext):
+                return f'{pg.path()},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
+
+            lines = [line(PostgresContext.apply(state.namespace, pg)) for pg in PostgresContext.hosts(state.namespace)]
+
+            log(lines_to_tabular(lines, 'NAME,ENDPOINT,USERNAME,PASSWORD', separator=','))
+        else:
+            def line(pg: PostgresContext):
+                return f'{pg.path()},{pg.namespace},{pg.endpoint()}:{pg.port()},{pg.username()},{pg.password()}'
+
+            lines = [line(PostgresContext.apply(state.namespace, pg)) for pg in PostgresContext.hosts(state.namespace)]
+
+            log(lines_to_tabular(lines, 'NAME,NAMESPACE,ENDPOINT,USERNAME,PASSWORD', separator=','))
+
+    def show_pg_databases(self, pg: PostgresContext):
+        log(lines_to_tabular(pg_database_names(pg.namespace, pg.path()), 'DATABASE', separator=','))
+
+    def show_pg_tables(self, pg: PostgresContext):
+        log(lines_to_tabular(pg_table_names(pg.namespace, pg.path()), 'NAME', separator=','))
+
+    def cd(self, dir: str, state: ReplState):
+        if dir == '':
+            state.pg_path = None
+        else:
+            context: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path, arg=dir)
+            # patch up state.namespace from pg cd
+            if not state.namespace and context.namespace:
+                state.namespace = context.namespace
+            state.pg_path = context.path()
+
+    def cd_completion(self, cmd: str, state: ReplState, default: dict = {}):
+        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path) if state.pg_path else None
+        if pg and pg.db:
+            return {cmd: {'..': None}}
+        elif pg and pg.host:
+            return {cmd: {'..': None} | {p: None for p in pg_database_names(state.namespace, pg.path())}}
+        else:
+            return {cmd: {p: None for p in PostgresContext.hosts(state.namespace)}}
+
+    def pwd(self, state: ReplState):
+        words = []
+
+        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
+
+        if pg.host:
+            words.append(f'host/{pg.host}')
+        if pg.db:
+            words.append(f'database/{pg.db}')
+
+        return '\t'.join([f'{ReplState.P}:>'] + (words if words else ['/']))
+
+    def try_fallback_action(self, chain: Command, state: ReplState, cmd: str):
+        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
+        if pg.db:
+            return True, chain.run(f'pg {cmd}', state)
+
+        return False, None
+
+    def enter(self, state: ReplState):
+        Config().wait_log('Inspecting postgres database instances...')

adam/commands/devices/devices.py

@@ -0,0 +1,25 @@
+from adam.commands.devices.device import Device
+from adam.commands.devices.device_app import DeviceApp
+from adam.commands.devices.device_auit_log import DeviceAuditLog
+from adam.commands.devices.device_cass import DeviceCass
+from adam.commands.devices.device_export import DeviceExport
+from adam.commands.devices.device_postgres import DevicePostgres
+from adam.repl_state import ReplState
+
+class Devices:
+    def device(state: ReplState) -> Device:
+        if state.device == ReplState.A:
+            return DeviceApp()
+        elif state.device == ReplState.C:
+            return DeviceCass()
+        elif state.device == ReplState.L:
+            return DeviceAuditLog()
+        elif state.device == ReplState.P:
+            return DevicePostgres()
+        elif state.device == ReplState.X:
+            return DeviceExport()
+
+        return DeviceCass()
+
+    def all():
+        return [DeviceApp(), DeviceCass(), DeviceAuditLog(), DevicePostgres(), DeviceExport()]

adam/commands/export/clean_up_export_session.py

@@ -0,0 +1,53 @@
+from adam.commands.command import Command
+from adam.commands.export.exporter import Exporter
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log, log2
+
+class CleanUpExportSession(Command):
+    COMMAND = 'clean up export session'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(CleanUpExportSession, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return CleanUpExportSession.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)
+
+        state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
+        if not args:
+            if state.in_repl:
+                log2('Specify export session name.')
+            else:
+                log2('* Session name is missing.')
+
+                Command.display_help()
+
+            return 'command-missing'
+
+        csv_cnt, log_cnt = Exporter.clean_up_session(state.sts, state.pod, state.namespace, args[0])
+        log(f'Removed {csv_cnt} csv and {log_cnt} log files.')
+
+        Exporter.clear_export_session_cache()
+
+        return state
+
+    def completion(self, _: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{CleanUpExportSession.COMMAND} <export-session-name>\t clean up export session'

adam/commands/{frontend/teardown_frontend.py → export/clean_up_export_sessions.py}

@@ -1,14 +1,13 @@
 from adam.commands.command import Command
-from adam.k8s_utils.ingresses import Ingresses
-from adam.k8s_utils.services import Services
+from adam.commands.export.exporter import Exporter
 from adam.repl_state import ReplState, RequiredState
 
-class TearDownFrontend(Command):
-    COMMAND = 'teardown frontend'
+class CleanUpExportSessions(Command):
+    COMMAND = 'clean up all export sessions'
 
     # the singleton pattern
     def __new__(cls, *args, **kwargs):
-        if not hasattr(cls, 'instance'): cls.instance = super(TearDownFrontend, cls).__new__(cls)
+        if not hasattr(cls, 'instance'): cls.instance = super(CleanUpExportSessions, cls).__new__(cls)
 
         return cls.instance
 
@@ -16,10 +15,10 @@ class TearDownFrontend(Command):
         super().__init__(successor)
 
     def command(self):
-        return TearDownFrontend.COMMAND
+        return CleanUpExportSessions.COMMAND
 
     def required(self):
-        return RequiredState.NAMESPACE
+        return RequiredState.CLUSTER_OR_POD
 
     def run(self, cmd: str, state: ReplState):
         if not(args := self.args(cmd)):
@@ -29,9 +28,8 @@ class TearDownFrontend(Command):
         if not self.validate_state(state):
             return state
 
-        name = 'ops'
-        Ingresses.delete_ingress(name, state.namespace)
-        Services.delete_service(name, state.namespace)
+        if Exporter.clean_up_all_sessions(state.sts, state.pod, state.namespace):
+            Exporter.clear_export_session_cache()
 
         return state
 
@@ -39,4 +37,4 @@ class TearDownFrontend(Command):
         return {}
 
     def help(self, _: ReplState):
-        return f'{TearDownFrontend.COMMAND}\t tear down frontend'
+        return f'{CleanUpExportSessions.COMMAND}\t clean up all export sessions'

adam/commands/export/drop_export_database.py

@@ -0,0 +1,58 @@
+from adam.commands.command import Command
+from adam.commands.export.exporter import Exporter
+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'
+
+    # 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)
+
+        state, args = self.apply_state(args, state)
+        if not self.validate_state(state):
+            return state
+
+        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
+
+    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,46 @@
+from adam.commands.command import Command
+from adam.commands.export.exporter import Exporter
+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'
+
+    # 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)
+
+        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()
+
+        state.export_session = None
+
+        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,83 @@
+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.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'
+
+    # 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)
+
+        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
+
+    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):
+        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,170 @@
+import os
+import boto3
+
+from adam.commands.export.utils_export import ExportTableStatus
+from adam.config import Config
+from adam.utils import lines_to_tabular, log, log2, ing
+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:
+        cnt: int = 0
+
+        if show_query:
+            log2(query)
+
+        if database.startswith('s'):
+            cnt += SQLite.run_query(query, database=database)
+        else:
+            cnt += Athena.run_query(query, database=database)
+
+        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'):
+                try:
+                    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
+
+    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'
+                    if Config().is_debug():
+                        log2(query)
+                    Athena.query(query)
+
+        with ing(f'Deleting s3 folder: export'):
+            try:
+                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:
+            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}_%')