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

adam/commands/export/importer_athena.py

@@ -1,9 +1,11 @@
+import os
 import boto3
 
+from adam.commands.export.export_databases import export_db
 from adam.commands.export.importer import Importer
-from adam.commands.export.utils_export import GeneratorStream
 from adam.config import Config
-from adam.utils import log2, ing
+from adam.repl_state import ReplState
+from adam.utils import GeneratorStream, bytes_generator_from_file, debug, log2, ing
 from adam.utils_athena import Athena
 from adam.utils_k8s.pods import Pods
 
@@ -17,9 +19,14 @@ class AthenaImporter(Importer):
     def prefix(self):
         return 'e'
 
-    def import_from_csv(self, pod: str, namespace: str, to_session: str, from_session: str, keyspace: str, table: str, target_table: str, columns: str, multi_tables = True, create_db = False):
+    def import_from_csv(self, state: ReplState, from_session: str,
+                        keyspace: str, table: str, target_table: str, columns: str,
+                        multi_tables = True, create_db = False):
         csv_file = self.csv_file(from_session, table, target_table)
-        db = self.db(to_session, keyspace)
+        pod = state.pod
+        namespace = state.namespace
+        to_session = state.export_session
+        database = self.db(to_session, keyspace)
 
         succeeded = False
         try:
@@ -29,52 +36,113 @@ class AthenaImporter(Importer):
                 bytes = Pods.read_file(pod, 'cassandra', namespace, csv_file)
 
                 s3 = boto3.client('s3')
-                s3.upload_fileobj(GeneratorStream(bytes), bucket, f'export/{db}/{keyspace}/{target_table}/{table}.csv')
-
-            msg: str = None
-            if create_db:
-                msg = f"[{to_session}] Creating database {db}"
-            else:
-                msg = f"[{to_session}] Creating table {target_table}"
-            with ing(msg, suppress_log=multi_tables):
-                query = f'CREATE DATABASE IF NOT EXISTS {db};'
-                if Config().is_debug():
-                    log2(query)
-                Athena.query(query, 'default')
-
-                query = f'DROP TABLE IF EXISTS {target_table};'
-                if Config().is_debug():
-                    log2(query)
-                Athena.query(query, db)
-
-                athena_columns = ', '.join([f'{c} string' for c in columns.split(',')])
-                query = f'CREATE EXTERNAL TABLE IF NOT EXISTS {target_table}(\n' + \
-                        f'    {athena_columns})\n' + \
-                            "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'\n" + \
-                            'WITH SERDEPROPERTIES (\n' + \
-                            '    "separatorChar" = ",",\n' + \
-                            '    "quoteChar"     = "\\"")\n' + \
-                        f"LOCATION 's3://{bucket}/export/{db}/{keyspace}/{target_table}'\n" + \
-                            'TBLPROPERTIES ("skip.header.line.count"="1");'
-                if Config().is_debug():
-                    log2(query)
-                try:
-                    Athena.query(query, db)
-                except Exception as e:
-                    log2(f'*** Failed query:\n{query}')
-                    raise e
-
-            to, _ = self.move_to_done(pod, namespace, to_session, from_session, keyspace, target_table)
+                s3.upload_fileobj(GeneratorStream(bytes), bucket, f'export/{database}/{keyspace}/{target_table}/{table}.csv')
+
+            self.create_schema(to_session, bucket, database, keyspace, table, columns, multi_tables, create_db)
+
+            to, _ = self.move_to_done(state, from_session, keyspace, target_table)
 
             succeeded = True
 
             return to, to_session
         finally:
             if succeeded:
-                self.remove_csv(pod, namespace, from_session, table, target_table, multi_tables)
+                self.remove_csv(state, from_session, table, target_table, multi_tables)
                 Athena.clear_cache()
 
                 if not multi_tables:
-                    query = f'select * from {target_table} limit 10'
-                    log2(query)
-                    Athena.run_query(query, db)
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {database}.{target_table} limit 10')
+
+    def import_from_local_csv(self, state: ReplState,
+                        keyspace: str, table: str, csv_file: str, multi_tables = True, create_db = False):
+        to_session = state.export_session
+        database = self.db(to_session, keyspace)
+
+        succeeded = False
+        try:
+            columns = None
+            with open(csv_file, 'r') as f:
+                columns = f.readline()
+
+            bucket = Config().get('export.bucket', 'c3.ops--qing')
+
+            with ing(f'[{to_session}] Uploading to S3', suppress_log=multi_tables):
+                bytes = bytes_generator_from_file(csv_file)
+
+                s3 = boto3.client('s3')
+                s3.upload_fileobj(GeneratorStream(bytes), bucket, f'export/{database}/{keyspace}/{table}/{os.path.basename(csv_file)}')
+
+            self.create_schema(to_session, bucket, database, keyspace, table, columns, multi_tables, create_db)
+            # msg: str = None
+            # if create_db:
+            #     msg = f"[{to_session}] Creating database {database}"
+            # else:
+            #     msg = f"[{to_session}] Creating table {target_table}"
+            # with ing(msg, suppress_log=multi_tables):
+            #     query = f'CREATE DATABASE IF NOT EXISTS {database};'
+            #     debug(query)
+            #     Athena.query(query, 'default')
+
+            #     query = f'DROP TABLE IF EXISTS {target_table};'
+            #     debug(query)
+            #     Athena.query(query, database)
+
+            #     athena_columns = ', '.join([f'{c} string' for c in columns.split(',')])
+            #     query = f'CREATE EXTERNAL TABLE IF NOT EXISTS {target_table}(\n' + \
+            #             f'    {athena_columns})\n' + \
+            #                 "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'\n" + \
+            #                 'WITH SERDEPROPERTIES (\n' + \
+            #                 '    "separatorChar" = ",",\n' + \
+            #                 '    "quoteChar"     = "\\"")\n' + \
+            #             f"LOCATION 's3://{bucket}/export/{database}/{keyspace}/{target_table}'\n" + \
+            #                 'TBLPROPERTIES ("skip.header.line.count"="1");'
+            #     debug(query)
+            #     try:
+            #         Athena.query(query, database)
+            #     except Exception as e:
+            #         log2(f'*** Failed query:\n{query}')
+            #         raise e
+
+            succeeded = True
+
+            return csv_file, to_session
+        finally:
+            if succeeded:
+                Athena.clear_cache()
+
+                if not multi_tables:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {database}.{table} limit 10')
+
+    def create_schema(self, to_session: str, bucket: str, database: str, keyspace: str, table: str, columns: list[str], multi_tables: bool, create_db = False):
+        msg: str = None
+        if create_db:
+            msg = f"[{to_session}] Creating database {database}"
+        else:
+            msg = f"[{to_session}] Creating table {table}"
+
+        with ing(msg, suppress_log=multi_tables):
+            query = f'CREATE DATABASE IF NOT EXISTS {database};'
+            debug(query)
+            Athena.query(query, 'default')
+
+            query = f'DROP TABLE IF EXISTS {table};'
+            debug(query)
+            Athena.query(query, database)
+
+            athena_columns = ', '.join([f'{c} string' for c in columns.split(',')])
+            query = f'CREATE EXTERNAL TABLE IF NOT EXISTS {table}(\n' + \
+                    f'    {athena_columns})\n' + \
+                        "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'\n" + \
+                        'WITH SERDEPROPERTIES (\n' + \
+                        '    "separatorChar" = ",",\n' + \
+                        '    "quoteChar"     = "\\"")\n' + \
+                    f"LOCATION 's3://{bucket}/export/{database}/{keyspace}/{table}'\n" + \
+                        'TBLPROPERTIES ("skip.header.line.count"="1");'
+            debug(query)
+            try:
+                Athena.query(query, database)
+            except Exception as e:
+                log2(f'*** Failed query:\n{query}')
+                raise e

adam/commands/export/importer_sqlite.py

@@ -1,47 +1,67 @@
-import os
-import sqlite3
 import pandas
 
+from adam.commands.export.export_databases import export_db
 from adam.commands.export.importer import Importer
-from adam.commands.export.utils_export import GeneratorStream
-from adam.utils import log2, ing
+from adam.repl_state import ReplState
+from adam.utils import GeneratorStream, bytes_generator_from_file, ing
 from adam.utils_k8s.pods import Pods
-from adam.utils_sqlite import SQLite
+from adam.utils_sqlite import SQLite, sqlite
 
 class SqliteImporter(Importer):
     def prefix(self):
         return 's'
 
-    def import_from_csv(self, pod: str, namespace: str, to_session: str, from_session: str, keyspace: str, table: str, target_table: str, columns: str, multi_tables = True, create_db = False):
+    def import_from_csv(self, state: ReplState, from_session: str,
+                        keyspace: str, table: str, target_table: str, columns: str,
+                        multi_tables = True, create_db = False):
         csv_file = self.csv_file(from_session, table, target_table)
-        db = self.db(to_session, keyspace)
+        pod = state.pod
+        namespace = state.namespace
+        to_session = state.export_session
 
         succeeded = False
-        conn = None
         try:
-            os.makedirs(SQLite.local_db_dir(), exist_ok=True)
-            conn = sqlite3.connect(f'{SQLite.local_db_dir()}/{db}.db')
-
             with ing(f'[{to_session}] Uploading to Sqlite', suppress_log=multi_tables):
-                bytes = Pods.read_file(pod, 'cassandra', namespace, csv_file)
-                df = pandas.read_csv(GeneratorStream(bytes))
-
-                df.to_sql(target_table, conn, index=False, if_exists='replace')
+                # create a connection to single keyspace
+                with sqlite(to_session, keyspace) as conn:
+                    bytes = Pods.read_file(pod, 'cassandra', namespace, csv_file)
+                    df = pandas.read_csv(GeneratorStream(bytes))
+                    df.to_sql(target_table, conn, index=False, if_exists='replace')
 
-            to, _ = self.move_to_done(pod, namespace, to_session, from_session, keyspace, target_table)
+            to, _ = self.move_to_done(state, from_session, keyspace, target_table)
 
             succeeded = True
 
             return to, to_session
         finally:
             if succeeded:
-                self.remove_csv(pod, namespace, from_session, table, target_table, multi_tables)
+                self.remove_csv(state, from_session, table, target_table, multi_tables)
                 SQLite.clear_cache()
 
                 if not multi_tables:
-                    query = f'select * from {target_table} limit 10'
-                    log2(query)
-                    SQLite.run_query(query, conn_passed=conn)
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {keyspace}.{target_table} limit 10')
+
+    def import_from_local_csv(self, state: ReplState,
+                        keyspace: str, table: str, csv_file: str, multi_tables = True, create_db = False):
+        to_session = state.export_session
+
+        succeeded = False
+        try:
+            with ing(f'[{to_session}] Uploading to Sqlite', suppress_log=multi_tables):
+                # create a connection to single keyspace
+                with sqlite(to_session, keyspace) as conn:
+                    bytes = bytes_generator_from_file(csv_file)
+                    df = pandas.read_csv(GeneratorStream(bytes))
+                    df.to_sql(table, conn, index=False, if_exists='replace')
 
-            if conn:
-                conn.close()
+            succeeded = True
+
+            return csv_file, to_session
+        finally:
+            if succeeded:
+                SQLite.clear_cache()
+
+                if not multi_tables:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {keyspace}.{table} limit 10')

adam/commands/export/show_column_counts.py

@@ -1,10 +1,9 @@
+from adam.commands import extract_trailing_options, validate_args
 from adam.commands.command import Command
-from adam.commands.export.export_databases import ExportDatabases
+from adam.commands.cql.utils_cql import cassandra_table_names
+from adam.commands.export.export_databases import ExportDatabases, export_db
 from adam.config import Config
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
-from adam.utils_athena import Athena
-from adam.utils_sqlite import SQLite
 
 class ShowColumnCounts(Command):
     COMMAND = 'show column counts on'
@@ -28,36 +27,19 @@ class ShowColumnCounts(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 args:
-            if state.in_repl:
-                log2('Use a SQL statement.')
-            else:
-                log2('* SQL statement is missing.')
-
-                Command.display_help()
-
-            return 'command-missing'
+        with self.validate(args, state) as (args, state):
+            with extract_trailing_options(args, '&') as (args, backgrounded):
+                with validate_args(args, state, name='SQL statement') as table:
+                    with export_db(state) as dbs:
+                        query = Config().get(f'export.column_counts_query', 'select id, count(id) as columns from {table} group by id')
+                        query = query.replace('{table}', table)
+                        dbs.sql(query, state.export_session, backgrounded=backgrounded)
 
-        copy_or_export = 'copy'
-        if state.export_session.startswith('e'):
-            copy_or_export = 'export'
-
-        table = args[0]
-        query = Config().get(f'{copy_or_export}.column_counts_query', 'select id, count(id) as columns from {table} group by id')
-        query = query.replace('{table}', table)
-        ExportDatabases.run_query(query, state.export_session)
-
-        return state
+            return state
 
     def completion(self, state: ReplState):
-        if not state.export_session:
-            return {}
-
-        return super().completion(state, lambda: {t: None for t in ExportDatabases.table_names(state.export_session)})
+        return super().completion(state, lambda: {t: None for t in ExportDatabases.table_names(state.export_session)}, auto_key='x.tables')
+        # return {}
 
     def help(self, _: ReplState):
         return f'{ShowColumnCounts.COMMAND} <export-table-name>\t show column count per id'

adam/commands/export/show_export_databases.py

@@ -1,5 +1,6 @@
 from adam.commands.command import Command
 from adam.commands.devices.device_export import DeviceExport
+from adam.commands.export.export_databases import export_db
 from adam.repl_state import ReplState
 
 class ShowExportDatabases(Command):
@@ -24,16 +25,15 @@ class ShowExportDatabases(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
-
-        DeviceExport().show_export_databases()
+        with self.validate(args, state) as (args, state):
+            with export_db(state) as dbs:
+                dbs.show_databases()
 
-        return state
+            return state
 
     def completion(self, state: ReplState):
-        return DeviceExport().ls_completion(ShowExportDatabases.COMMAND, state, default = super().completion(state))
+        return {}
+        # return DeviceExport().ls_completion(ShowExportDatabases.COMMAND, state, default = super().completion(state))
 
     def help(self, _: ReplState):
         return f'{ShowExportDatabases.COMMAND}\t list export databases'

adam/commands/export/show_export_session.py

@@ -1,8 +1,7 @@
+from adam.commands import validate_args
 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.export_sessions import ExportSessions, export_session
 from adam.repl_state import ReplState, RequiredState
-from adam.utils import log2
 
 class ShowExportSession(Command):
     COMMAND = 'show export session'
@@ -26,26 +25,15 @@ class ShowExportSession(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 args:
-            if state.in_repl:
-                log2('Specify export database name.')
-            else:
-                log2('* Database name is missing.')
-
-                Command.display_help()
+        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.show_session(session)
 
-            return 'command-missing'
-
-        ExportDatabases.disply_export_session(state.sts, state.pod, state.namespace, args[0])
-
-        return state
+            return state
 
     def completion(self, state: ReplState):
-        return super().completion(state, {session: None for session in Exporter.export_session_names(state.sts, state.pod, state.namespace)})
+        return {}
 
     def help(self, _: ReplState):
         return f'{ShowExportSession.COMMAND} <export-session-name>\t show export session'

adam/commands/export/show_export_sessions.py

@@ -1,8 +1,6 @@
 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 lines_to_tabular, log
-from adam.utils_k8s.statefulsets import StatefulSets
 
 class ShowExportSessions(Command):
     COMMAND = 'show export sessions'
@@ -26,22 +24,14 @@ class ShowExportSessions(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
-
-        pod = state.pod
-        if not pod:
-            pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
+        with self.validate(args, state) as (args, state):
+            with export_session(state) as sessions:
+                sessions.show_all_sessions()
 
-        sessions: dict[str, str] = Exporter.find_export_sessions(pod, state.namespace)
-        log(lines_to_tabular([f'{session}\t{export_state}' for session, export_state in sorted(sessions.items(), reverse=True)],
-                                header='EXPORT_SESSION\tSTATUS', separator='\t'))
-
-        return state
+            return state
 
     def completion(self, state: ReplState):
-        return super().completion(state)
+        return {}
 
     def help(self, _: ReplState):
         return f'{ShowExportSessions.COMMAND}\t list export sessions'

adam/commands/export/utils_export.py

@@ -3,22 +3,33 @@ import re
 
 from adam.config import Config
 from adam.pod_exec_result import PodExecResult
+from adam.repl_state import ReplState
 from adam.utils_k8s.cassandra_nodes import CassandraNodes
 from adam.utils_k8s.pods import log_prefix
+from adam.utils_k8s.statefulsets import StatefulSets
 
 class ImportSpec:
-    def __init__(self, session: str, importer: str):
+    def __init__(self, table_name: str, session: str = None, files: list[str] = None, importer: str = None):
+        self.table_name = table_name
         self.session = session
+        self.files = files
         self.importer = importer
 
-    def parse_specs(specs_str: str):
-        session: str = None
+    def parse_specs(specs_str: str, files = False):
+        session_or_files: str = None
         importer: str = None
+        table_name: str = None
 
         if specs_str:
-            importer, session = ImportSpec._extract_importer(specs_str.strip(' '))
+            importer, rest = ImportSpec._extract_importer(specs_str.strip(' '))
+
+            if rest:
+                table_name, session_or_files = ImportSpec._extract_table_name(rest)
+
+        if not files:
+            return ImportSpec(table_name, session=session_or_files, importer=importer)
 
-        return ImportSpec(session, importer)
+        return ImportSpec(table_name, files=[f.strip(' ') for f in session_or_files.split(',')], importer=importer)
 
     def _extract_importer(spec_str: str) -> tuple[str, str]:
         importer = None
@@ -32,12 +43,25 @@ class ImportSpec:
 
         return importer, rest
 
+    def _extract_table_name(spec_str: str) -> tuple[str, str]:
+        table_name = None
+        rest = spec_str
+
+        p = re.compile(r"(.*?)as\s+(.*)", re.IGNORECASE)
+        match = p.match(spec_str)
+        if match:
+            rest = match.group(1).strip(' ')
+            table_name = match.group(2).strip(' ')
+
+        return table_name, rest
+
 class ExportSpec(ImportSpec):
     def __init__(self, keyspace: str, consistency: str, importer: str, tables: list['ExportTableSpec'], session: str = None):
         super().__init__(None, importer)
 
         self.keyspace = keyspace
         self.consistency = consistency
+        self.importer = importer
         self.tables = tables
         self.session = session
 
@@ -152,11 +176,12 @@ class ExportTableSpec:
         return f'{self.keyspace}.{self.table}({self.columns}) as {self.target_table}'
 
 class ExportTableStatus:
-    def __init__(self, keyspace: str, target_table: str, status: str, table: str = None):
+    def __init__(self, keyspace: str, target_table: str, status: str, table: str = None, csv_file: str = ''):
         self.keyspace = keyspace
         self.target_table = target_table
         self.status = status
         self.table = table
+        self.csv_file = csv_file
 
     def __str__(self):
         return f'{self.keyspace}.{self.table} as {self.target_table} = {self.status}'
@@ -167,7 +192,7 @@ class ExportTableStatus:
 
         return False
 
-    def from_session(sts: str, pod: str, namespace: str, export_session: str):
+    def from_session(sts: str, pod: str, namespace: str, export_session: str) -> tuple['ExportTableStatus', str]:
         statuses: list[ExportTableStatus] = []
 
         status_in_whole = 'done'
@@ -202,8 +227,8 @@ class ExportTableStatus:
             target_table = m.group(2)
             state = m.group(3)
             if state == '.pending_import':
-                _, table = get_csv_files_n_table(target_table)
-                return ExportTableStatus(keyspace, target_table, 'pending_import', table)
+                csv_files, table = get_csv_files_n_table(target_table)
+                return ExportTableStatus(keyspace, target_table, 'pending_import', table, csv_files[0] if csv_files else '')
             elif state == '.done':
                 return ExportTableStatus(keyspace, target_table, 'done', target_table)
 
@@ -213,7 +238,7 @@ class ExportTableStatus:
             if r.exit_code() == 0:
                 csv_files, table = get_csv_files_n_table(target_table)
                 if csv_files:
-                    return ExportTableStatus(keyspace, target_table, 'exported', table)
+                    return ExportTableStatus(keyspace, target_table, 'exported', table, csv_files[0])
                 else:
                     return ExportTableStatus(keyspace, target_table, 'imported', target_table)
             else:
@@ -288,4 +313,32 @@ class GeneratorStream(io.RawIOBase):
 
             data = self._buffer[:size]
             self._buffer = self._buffer[size:]
-            return data
+            return data
+
+class PodPushHandler:
+    def __init__(self, state: ReplState, pod: str = None):
+        self.state = state
+        self.pushed = False
+        self.pod = pod
+
+    def __enter__(self):
+        state = self.state
+
+        if not state.pod:
+            self.pushed = True
+            state.push()
+
+            if not self.pod:
+                self.pod = StatefulSets.pod_names(state.sts, state.namespace)[0]
+            state.pod = self.pod
+
+        return state
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.pushed:
+            self.state.pop()
+
+        return False
+
+def state_with_pod(state: ReplState, pod: str = None):
+    return PodPushHandler(state, pod=pod)

adam/commands/find_files.py

@@ -0,0 +1,51 @@
+import os
+
+from adam.commands.command import Command
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_local import local_tmp_dir
+
+class FindLocalFiles(Command):
+    COMMAND = 'find local'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(FindLocalFiles, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return FindLocalFiles.COMMAND
+
+    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):
+            cmd = 'find'
+
+            if not args:
+                cmd = f'find {local_tmp_dir()}'
+            elif len(args) == 1:
+                cmd = f"find {local_tmp_dir()} -name '{args[0]}'"
+            else:
+                new_args = [f"'{arg}'" if '*' in arg else arg for arg in args]
+                cmd = 'find ' + ' '.join(new_args)
+
+            log2(cmd)
+            os.system(cmd)
+
+            return state
+
+    def completion(self, state: ReplState):
+        return super().completion(state, {
+            '*.csv': None,
+            '*.db': None,
+            '*': None
+        })
+
+    def help(self, _: ReplState):
+        return f'{FindLocalFiles.COMMAND} [linux-find-arguments]\t find files from local machine'