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

adam/commands/devices/device_postgres.py

@@ -1,11 +1,10 @@
 from adam.commands.bash.bash_completer import BashCompleter
-from adam.commands.command import Command, InvalidState
+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_databases import PostgresDatabases, pg_path
 from adam.commands.postgres.utils_postgres import pg_database_names, pg_table_names, postgres
 from adam.repl_state import ReplState
-from adam.utils import lines_to_tabular, log, log2, wait_log
-from adam.utils_k8s.pods import Pods
+from adam.utils import tabulize, log2, wait_log
 
 class DevicePostgres(Command, Device):
     COMMAND = f'{ReplState.P}:'
@@ -36,79 +35,98 @@ class DevicePostgres(Command, Device):
     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 pod(self, state: ReplState) -> str:
+        pod, _ = PostgresDatabases.pod_and_container(state.namespace)
+        return pod
 
-    def cat(self, cmd: str, state: ReplState):
-        if state.pod:
-            return self.bash(state, state, cmd.split(' '))
+    def pod_names(self, state: ReplState) -> list[str]:
+        return [self.pod(state)]
 
-        return InvalidState()
+    def default_container(self, state: ReplState) -> str:
+        _, container = PostgresDatabases.pod_and_container(state.namespace)
+        return container
+
+    def ls(self, cmd: str, state: ReplState):
+        if state.pod_targetted:
+            self.bash(state, state, ['ls'])
+            return
+
+        with pg_path(state) as (host, database):
+            if database:
+                tabulize(pg_table_names(state), header='NAME', separator=',')
+            elif host:
+                tabulize(pg_database_names(state), header='DATABASE', separator=',')
+            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=','))
+            tabulize(PostgresDatabases.hosts(state),
+                     lambda p: f'{p.host},{p.endpoint()}:{p.port()},{p.username()},{p.password()}',
+                     header='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=','))
+            tabulize(PostgresDatabases.hosts(state),
+                     lambda p: f'{p.host},{p.namespace},{p.endpoint()}:{p.port()},{p.username()},{p.password()}',
+                     header='NAME,NAMESPACE,ENDPOINT,USERNAME,PASSWORD',
+                     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()
+            return
+
+        with pg_path(state) as (host, database):
+            if dir == '..':
+                if database:
+                    database = None
+                else:
+                    host = None
+            else:
+                tks = dir.split('@')
+                if not host:
+                    host = tks[0]
+                else:
+                    database = tks[0]
+
+                if len(tks) > 1:
+                    state.namespace = tks[1]
+
+            if database:
+                state.pg_path = f'{host}/{database}'
+            else:
+                state.pg_path = host
 
     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)}}
+        return {cmd: {d: None for d in self.direct_dirs(cmd, state, default=default)}}
+
+    def direct_dirs(self, cmd: str, state: ReplState, default: dict = {}) -> list[str]:
+        with pg_path(state) as (host, database):
+            if database:
+                return ['..']
+            elif host:
+                return ['..'] + [p for p in pg_database_names(state)]
+            else:
+                return [p for p in PostgresDatabases.host_names(state.namespace)]
 
     def pwd(self, state: ReplState):
         words = []
 
-        pg: PostgresContext = PostgresContext.apply(state.namespace, state.pg_path)
+        with pg_path(state) as (host, database):
+            if host:
+                words.append(f'host/{host}')
+            if database:
+                words.append(f'database/{database}')
 
-        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 ['/']))
+            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)
+        with pg_path(state) as (_, database):
+            if not database:
+                database = PostgresDatabases.default_db()
+
+            if database:
+                return True, chain.run(f'pg {cmd}', state)
 
         return False, None
 
@@ -116,15 +134,13 @@ class DevicePostgres(Command, Device):
         wait_log('Inspecting postgres database instances...')
 
     def show_tables(self, state: ReplState):
-        pg = PostgresContext.apply(state.namespace, state.pg_path)
-        lines = [db["name"] for db in pg.tables() if db["schema"] == PostgresContext.default_schema()]
-        log(lines_to_tabular(lines, separator=','))
+        tabulize(PostgresDatabases.tables(state, default_schema=True), lambda d: d['name'], separator=',')
 
     def show_table_preview(self, state: ReplState, table: str, rows: int):
-        PostgresContext.apply(state.namespace, state.pg_path).run_sql(f'select * from {table} limit {rows}')
+        PostgresDatabases.run_sql(state, f'select * from {table} limit {rows}')
 
     def bash(self, s0: ReplState, s1: ReplState, args: list[str]):
-        pod, container = PostgresContext.pod_and_container(s1.namespace)
+        pod, container = PostgresDatabases.pod_and_container(s1.namespace)
         log2(f'Running on {pod}(container:{container})...')
 
         return super().bash(s0, s1, args)

adam/commands/download_file.py

@@ -0,0 +1,47 @@
+from adam.commands import validate_args
+from adam.commands.command import Command
+from adam.commands.devices.devices import Devices
+from adam.config import Config
+from adam.pod_exec_result import PodExecResult
+from adam.repl_state import ReplState, RequiredState
+from adam.utils import log2
+from adam.utils_k8s.pods import Pods
+
+class DownloadFile(Command):
+    COMMAND = 'download file'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(DownloadFile, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return DownloadFile.COMMAND
+
+    def required(self):
+        return [RequiredState.CLUSTER_OR_POD, RequiredState.APP_APP, ReplState.P]
+
+    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'):
+                to_file = Pods.download_file(Devices.device(state).pod(state),
+                                             Devices.device(state).default_container(state),
+                                             state.namespace,
+                                             args[0],
+                                             args[1] if len(args) > 1 else None)
+                log2(f'Downloaded to {to_file}.')
+
+                return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state, lambda: {f: None for f in Devices.device(state).files(state)}, pods=Devices.device(state).pods(state, '-'), auto='jit')
+
+    def help(self, _: ReplState):
+        return f'{DownloadFile.COMMAND} from-file [to-file]\t download file from pod'

adam/commands/export/clean_up_all_export_sessions.py

@@ -1,5 +1,5 @@
 from adam.commands.command import Command
-from adam.commands.export.exporter import Exporter
+from adam.commands.export.export_sessions import export_session
 from adam.repl_state import ReplState, RequiredState
 
 class CleanUpAllExportSessions(Command):
@@ -25,8 +25,8 @@ class CleanUpAllExportSessions(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if Exporter.clean_up_all_sessions(state.sts, state.pod, state.namespace):
-                Exporter.clear_export_session_cache()
+            with export_session(state) as sessions:
+                sessions.clean_up_all()
 
             return state
 

adam/commands/export/clean_up_export_sessions.py

@@ -1,7 +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_sessions import export_session
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log, log2
 
 class CleanUpExportSessions(Command):
     COMMAND = 'clean up export sessions'
@@ -26,26 +26,14 @@ class CleanUpExportSessions(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not args:
-                if state.in_repl:
-                    log2('Specify export session name.')
-                else:
-                    log2('* Session name is missing.')
+            with validate_args(args, state, name='export session', separator=',') as session_names:
+                with export_session(state) as sessions:
+                    sessions.clean_up(session_names)
 
-                    Command.display_help()
-
-                return 'command-missing'
-
-            sessions = [arg.strip(' ') for arg in ' '.join(args).split(',')]
-            csv_cnt, log_cnt = Exporter.clean_up_sessions(state.sts, state.pod, state.namespace, sessions)
-            log(f'Removed {csv_cnt} csv and {log_cnt} log files.')
-
-            Exporter.clear_export_session_cache()
-
-            return state
+                return state
 
     def completion(self, _: ReplState):
         return {}
 
     def help(self, _: ReplState):
-        return f'{CleanUpExportSessions.COMMAND} <export-session-name>\t clean up export session'
+        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

@@ -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'
@@ -28,25 +26,11 @@ class DropExportDatabase(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            if not len(args):
-                if state.in_repl:
-                    log2('Database name is required.')
-                    log2()
-                else:
-                    log2('* Database name is missing.')
-                    Command.display_help()
+            with validate_args(args, state, name='database name') as db:
+                with export_db(state) as dbs:
+                    dbs.drop(args[0])
 
-                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'
@@ -27,12 +25,8 @@ class DropExportDatabases(Command):
             return super().run(cmd, state)
 
         with self.validate(args, state) as (args, state):
-            Exporter.drop_databases(state.sts, state.pod, state.namespace)
-
-            SQLite.clear_cache()
-            Athena.clear_cache()
-
-            state.export_session = None
+            with export_db(state) as dbs:
+                dbs.drop_all()
 
             return state
 

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.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.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'
@@ -36,17 +31,6 @@ class ExportTables(Command):
                     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 super().completion(state):
-            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,13 +1,15 @@
 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.repl_session import ReplSession
-from adam.utils import debug, lines_to_tabular, log, log2, ing, log_exc
+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%_%'
@@ -86,32 +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
 
-        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()
@@ -120,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 = {}
@@ -171,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)