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

adam/commands/export/importer_athena.py

@@ -0,0 +1,157 @@
+import os
+import boto3
+
+from adam.commands.export.export_databases import export_db
+from adam.commands.export.importer import Importer
+from adam.config import Config
+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
+
+class AthenaImporter(Importer):
+    def ping():
+        session = boto3.session.Session()
+        credentials = session.get_credentials()
+
+        return credentials is not None
+
+    def prefix(self):
+        return 'e'
+
+    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,
+                        job_log: str = None):
+        csv_file = self.csv_file(from_session, table, target_table)
+        pod = state.pod
+        namespace = state.namespace
+        to_session = state.export_session
+        database = self.db(to_session, keyspace)
+
+        succeeded = False
+        try:
+            bucket = Config().get('export.bucket', 'c3.ops--qing')
+
+            with ing(f'[{to_session}] Uploading to S3', suppress_log=multi_tables):
+                bytes = Pods.read_file(pod, 'cassandra', namespace, csv_file)
+
+                s3 = boto3.client('s3')
+                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(state, from_session, table, target_table, multi_tables, job_log=job_log)
+                Athena.clear_cache()
+
+                if multi_tables:
+                    log2(f'[{to_session}] {keyspace}.{target_table} OK', file=job_log)
+                else:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {keyspace}.{target_table} limit 10', backgrounded=True, export_log=job_log)
+
+    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

@@ -0,0 +1,78 @@
+from typing import TextIO
+import pandas
+
+from adam.commands.export.export_databases import export_db
+from adam.commands.export.importer import Importer
+from adam.repl_state import ReplState
+from adam.utils import GeneratorStream, bytes_generator_from_file, ing, log2
+from adam.utils_k8s.pods import Pods
+from adam.utils_sqlite import SQLite, sqlite
+
+class SqliteImporter(Importer):
+    def prefix(self):
+        return 's'
+
+    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,
+                        job_log: str = None):
+
+        csv_file = self.csv_file(from_session, table, target_table)
+        pod = state.pod
+        namespace = state.namespace
+        to_session = state.export_session
+
+        succeeded = False
+        try:
+            with ing(f'[{to_session}] Uploading to Sqlite', suppress_log=multi_tables, job_log=job_log):
+                # 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(state, from_session, keyspace, target_table)
+
+            succeeded = True
+
+            return to, to_session
+        finally:
+            if succeeded:
+                self.remove_csv(state, from_session, table, target_table, multi_tables, job_log=job_log)
+                SQLite.clear_cache()
+
+                if multi_tables:
+                    log2(f'[{to_session}] {keyspace}.{target_table} OK', file=job_log)
+                else:
+                    with export_db(state) as dbs:
+                        dbs.sql(f'select * from {keyspace}.{target_table} limit 10', backgrounded=True, export_log=job_log)
+
+    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')
+
+            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

@@ -0,0 +1,45 @@
+from adam.commands import extract_trailing_options, validate_args
+from adam.commands.command import Command
+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
+
+class ShowColumnCounts(Command):
+    COMMAND = 'show column counts on'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowColumnCounts, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowColumnCounts.COMMAND
+
+    def required(self):
+        return RequiredState.EXPORT_DB
+
+    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_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)
+
+            return state
+
+    def completion(self, state: ReplState):
+        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

@@ -0,0 +1,39 @@
+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):
+    COMMAND = 'show export databases'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowExportDatabases, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowExportDatabases.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.show_databases()
+
+            return state
+
+    def completion(self, state: ReplState):
+        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

@@ -0,0 +1,39 @@
+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 ShowExportSession(Command):
+    COMMAND = 'show export session'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowExportSession, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowExportSession.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.show_session(session)
+
+            return state
+
+    def completion(self, state: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{ShowExportSession.COMMAND} <export-session-name>\t show export session'

adam/commands/export/show_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 ShowExportSessions(Command):
+    COMMAND = 'show export sessions'
+
+    # the singleton pattern
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(cls, 'instance'): cls.instance = super(ShowExportSessions, cls).__new__(cls)
+
+        return cls.instance
+
+    def __init__(self, successor: Command=None):
+        super().__init__(successor)
+
+    def command(self):
+        return ShowExportSessions.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.show_all_sessions()
+
+            return state
+
+    def completion(self, state: ReplState):
+        return {}
+
+    def help(self, _: ReplState):
+        return f'{ShowExportSessions.COMMAND}\t list export sessions'