dm-run-profiler

#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-

import getopt
import json
import os
import sys
import traceback

from hecatoncheir import DbProfilerBase
from hecatoncheir import DbProfilerFormatter
from hecatoncheir import DbProfilerRepository
from hecatoncheir.bigquery import BigQueryProfiler
from hecatoncheir.oracle import OraProfiler
from hecatoncheir.pgsql import PgProfiler
from hecatoncheir.mysql import MyProfiler
from hecatoncheir.mssql import MSSQLProfiler
from hecatoncheir import logger as log
from hecatoncheir.exception import (DbProfilerException, DriverError,
                                    InternalError, QueryError, QueryTimeout)
from hecatoncheir.msgutil import gettext as _
from hecatoncheir.table import Table2
from hecatoncheir.validation import get_validation_rules


def usage():
    print '''
Usage: %s [option...] [schema[.table]] ...

Options:
    --dbtype=TYPE              Database type
    --host=STRING              Host name
    --port=INTEGER             Port number
    --dbname=STRING            Database name
    --tnsname=STRING           TNS name (Oracle only)
    --user=STRING              User name
    --pass=STRING              User password
    --credential=STRING        Credential file name (BigQuery only)
    -P=INTEGER                 Parallel degree of table scan
    -o=FILENAME                Output file
    --batch=FILENAME           Batch execution

    --enable-validation        Enable record/column/SQL validations

    --enable-sample-rows       Enable collecting sample rows. (default)
    --disable-sample-rows      Disable collecting sample rows.

    --skip-table-profiling     Skip table (and column) profiling
    --skip-column-profiling    Skip column profiling
    --column-profiling-threshold=INTEGER
                               Threshold number of rows to skip profiling
                               columns

    --timeout=NUMBER           Query timeout in seconds (default:no timeout)

    --help                     Print this help.

''' % os.path.basename(sys.argv[0])


class Config():
    def __init__(self, dbtype):
        assert dbtype in ['pgsql', 'mysql', 'oracle', 'mssql', 'bigquery']

        self.host = None
        self.port = None
        self.dbname = None
        self.dbuser = None
        self.dbpass = None
        self.tnsname = None
        self.credential_file = None

        if dbtype == 'pgsql':
            self.init_pgsql()
        elif dbtype == 'mysql':
            self.init_mysql()
        elif dbtype == 'oracle':
            self.init_oracle()
        elif dbtype == 'mssql':
            self.init_mssql()
        elif dbtype == 'bigquery':
            self.init_bigquery()

    def init_pgsql(self):
        self.host = os.environ.get('PGHOST', u'localhost')
        self.port = os.environ.get('PGPORT', 5432)
        self.dbname = os.environ.get('PGDATABASE', u'postgres')
        self.dbuser = os.environ.get('PGUSER', u'postgres')
        self.dbpass = os.environ.get('PGPASSWORD', None)

    def init_mysql(self):
        self.host = u'localhost'
        self.port = 3306
        self.dbname = u'mysql'
        self.dbuser = u'mysql'
        self.dbpass = None

    def init_oracle(self):
        self.host = u'localhost'
        self.port = 1521
        self.dbname = u'ORCL'
        self.dbuser = u'scott'
        self.dbpass = None
        self.tnsname = None

    def init_mssql(self):
        self.host = u'localhost'
        self.port = 1433
        self.dbname = u'db'
        self.dbuser = u'sa'
        self.dbpass = None

    def init_bigquery(self):
        self.credential_file = None


def get_dbtype(argv):
    dbtype = None
    for i, o in enumerate(sys.argv[1:]):
        if o == '--dbtype':
            try:
                dbtype = sys.argv[i+2]
            except IndexError as ex:
                if str(ex) == 'list index out of range':
                    pass
                else:
                    raise ex
    return dbtype


def print_config(config):
    log.info("------ Connection settings ------")
    for k in sorted(vars(config)):
        if k == 'dbpass':
            log.info(u'%s: ****' % k)
        else:
            log.info(u'%s: %s' % (k,
                                  (vars(config)[k] if vars(config)[k]
                                   else '(none)')))
    log.info("---------------------------------")


def get_profiler(config):
    profiler = None

    if dbtype == 'pgsql':
        if config.dbpass is None:
            log.error(_("Specify database password with --pass."))
            sys.exit(1)
        profiler = PgProfiler.PgProfiler(config.host, config.port,
                                         config.dbname,
                                         config.dbuser, config.dbpass)
    elif dbtype == 'oracle':
        if config.dbpass is None:
            log.error(_("Specify database password with --pass."))
            sys.exit(1)
        if config.tnsname:
            config.dbname = config.tnsname
            config.host = None
            config.port = None
            log.info(_("TNS info: %s@%s") % (config.dbuser, config.tnsname))
        else:
            info = "%s@%s:%s/%s" % (config.dbuser, config.host,
                                    config.port, config.dbname)
            log.info(_("Conn info: %s") % info)
        profiler = OraProfiler.OraProfiler(config.host, config.port,
                                           config.dbname,
                                           config.dbuser, config.dbpass)
    elif dbtype == 'mysql':
        if config.dbpass is None:
            log.error(_("Specify database password with --pass."))
            sys.exit(1)
        profiler = MyProfiler.MyProfiler(config.host, config.port,
                                         config.dbname,
                                         config.dbuser, config.dbpass)
    elif dbtype == 'mssql':
        if config.dbpass is None:
            log.error(_("Specify database password with --pass."))
            sys.exit(1)
        profiler = MSSQLProfiler.MSSQLProfiler(config.host, config.port,
                                               config.dbname,
                                               config.dbuser, config.dbpass)
    elif dbtype == 'bigquery':
        if not config.credential_file:
            log.error(_("Specify credential file with --credential."))
            sys.exit(1)
        with open(config.credential_file) as f:
            c = json.loads(f.read())
        config.dbname = c['project_id']
        profiler = BigQueryProfiler.BigQueryProfiler(config.credential_file)

    return profiler


def table_list_from_file(profiler, filename):
    tables = []
    try:
        for line in open(filename):
            a = line.strip()
            if a.startswith("#"):
                continue
            t = a.split(".")
            if len(t) == 2 and t[0] and t[1]:
                tables.append([t[0], t[1]])
            elif len(t) == 1 and t[0]:
                tables.extend(table_list_from_schema(profiler, t[0]))
            else:
                log.error(_("Specify schema/table name in "
                            "[SCHEMA[.TABLE]] format: <%s>") % a)
    except Exception as ex:
        log.error("Could not read a table list file `%s'." % filename)
    return tables


def table_list_from_args(profiler, args):
    tables = []
    for a in args:
        t = a.split(".")
        if len(t) == 2 and t[0] and t[1]:
            tables.append([t[0], t[1]])
        elif len(t) == 1 and t[0]:
            tables.extend(table_list_from_schema(profiler, t[0]))
        else:
            log.error(_("Specify schema/table name in "
                        "[SCHEMA[.TABLE]] format: <%s>") % a)
    return tables


def table_list_from_schema(profiler, schema):
    tables = []
    for t in profiler.get_table_names(schema):
        tables.append([schema, t])
    return tables


def table_list_dedup(tables):
    tables2 = []
    for t in tables:
        if t not in tables2:
            tables2.append(t)
    return tables2


input_encoding = 'utf-8'

if __name__ == "__main__":
    dbtype = get_dbtype(sys.argv[1:])

    if not dbtype:
        log.error(_("Specify your database type with --dbtype."))
        usage()
        sys.exit(1)

    if dbtype not in ['pgsql', 'mysql', 'oracle', 'mssql', 'bigquery']:
        log.error("Database type `%s' is not supported." % dbtype)
        sys.exit(1)

    config = Config(dbtype)

    try:
        opts, args = getopt.getopt(sys.argv[1:], "P:o:",
                                   ["help", "dbtype=", "host=", "port=",
                                    "dbname=", "tnsname=", "user=", "pass=",
                                    "credential=",
                                    "batch=", "enable-validation",
                                    "enable-sample-rows",
                                    "disable-sample-rows",
                                    "skip-table-profiling",
                                    "skip-column-profiling",
                                    "skip-record-validation",
                                    "column-profiling-threshold=",
                                    "timeout="])
    except getopt.GetoptError as err:
        log.error(unicode(err))
        usage()
        sys.exit(1)

    parallel_degree = None
    output_file = None
    table_list = None
    enable_validation = False
    enable_sample_rows = True
    skip_table_profiling = False
    skip_column_profiling = False
    column_profiling_threshold = None
    skip_record_validation = False
    debug = None
    timeout = None

    for o, a in opts:
        if o in ("-P"):
            parallel_degree = int(a)
        elif o in ("-o"):
            output_file = a
        elif o in ("--dbtype"):
            assert dbtype == a
        elif o in ("--host"):
            config.host = a.decode(input_encoding)
        elif o in ("--port"):
            config.port = int(a)
        elif o in ("--dbname"):
            config.dbname = a.decode(input_encoding)
        elif o in ("--tnsname"):
            config.tnsname = a.decode(input_encoding)
        elif o in ("--user"):
            config.dbuser = a.decode(input_encoding)
        elif o in ("--pass"):
            config.dbpass = a.decode(input_encoding)
        elif o in ("--credential"):
            config.credential_file = a.decode(input_encoding)
        elif o in ("--batch"):
            table_list = a
        elif o in ("--enable-validation"):
            enable_validation = True
        elif o in ("--enable-sample-rows"):
            enable_sample_rows = True
        elif o in ("--disable-sample-rows"):
            enable_sample_rows = False
        elif o in ("--skip-table-profiling"):
            skip_table_profiling = True
        elif o in ("--skip-column-profiling"):
            skip_column_profiling = True
        elif o in ("--column-profiling-threshold"):
            column_profiling_threshold = a
        elif o in ("--skip-record-validation"):
            skip_record_validation = True
        elif o in ("--timeout"):
            timeout = int(a)
        elif o in ("--help"):
            usage()
            sys.exit(0)
        else:
            log.error("unexpected option. internal error.")
            sys.exit(1)

    if enable_validation:
        timeout = None
        log.info(_("Query timeout is disabled because data validation is enabled."))

    print_config(config)

    profiler = get_profiler(config)
    assert profiler

    profiler.profile_sample_rows = enable_sample_rows

    if parallel_degree > 1:
        profiler.parallel_degree = parallel_degree
        log.info(_("Setting paralell degree to %d for table scan.") %
                 profiler.parallel_degree)

    if column_profiling_threshold:
        profiler.column_profiling_threshold = int(column_profiling_threshold)

    tables = []
    try:
        profiler.connect()

        # Getting table names from arguments.
        if args:
            tables.extend(table_list_from_args(profiler, args))

        # Using a table list file.
        if table_list:
            tables.extend(table_list_from_file(profiler, table_list))

        # If nothing specified, show the list of schemas, then exit.
        if len(tables) == 0:
            for s in profiler.get_schema_names():
                log.info(u'Please specify one or more schemas:')
                log.info(u'  %s' % s)
                sys.exit(1)
    except DbProfilerException as e:
        error, = e.source.args
        log.error(u"%s" % unicode(e), detail=error.message)
        log.error(_("Abort."))
        sys.exit(1)
    except Exception as e:
        log.error(_("Internal error (%d)") % 4, detail=repr(e))
        log.error(_("Abort."))
        traceback.print_exc()
        sys.exit(3)

    tables = table_list_dedup(tables)

    try:
        repo = DbProfilerRepository.DbProfilerRepository(output_file)
        repo.init()
        repo.open()
    except DbProfilerException as e:
        error, = e.source.args
        log.error(u"%s" % unicode(e), detail=error.message)
        log.error(_("Abort."))
        sys.exit(1)
    except KeyboardInterrupt, e:
        log.error(_("Interrupted by the user."))
        log.error(_("Abort."))
        sys.exit(2)
    except Exception, e:
        log.error(_("Internal error (%d)") % 2, detail=repr(e))
        log.error(_("Abort."))
        traceback.print_exc()
        sys.exit(3)

    log.info(u"----------------------------------------------")
    log.info(_("Parallel degree for table scan: %d") %
             profiler.parallel_degree)
    log.info(_("Skipping table profiling: %s") % profiler.skip_table_profiling)
    log.info(_("Row count profiling: %s") % profiler.profile_row_count_enabled)
    log.info(_("Skippig column profiling: %s") %
             skip_column_profiling)
    tmp = "{:,d}".format(profiler.column_profiling_threshold)
    log.info(_("Maximum row count to enable column profiling: %s rows") % tmp)
    log.info(_("Min/Max values: %s") % profiler.profile_min_max_enabled)
    log.info(_("Number of null values: %s") % profiler.profile_nulls_enabled)
    log.info(_("Top-N most/least freq values: %d values") %
             profiler.profile_most_freq_values_enabled)
    tmp = profiler.profile_column_cardinality_enabled
    log.info(_("Column cardinality: %s") % tmp)
    log.info(_("Data validation: %s") % enable_validation)
    log.info(_("Obtaining sample records: %s") % enable_sample_rows)
    log.info(_("Query Timeout: %s") %
             (_("%d second(s)") % timeout if timeout else _('Disabled')))
    log.info(u"----------------------------------------------")

    log.info(_("Profiling on %d tables.") % len(tables))
    count = 0
    failed_count = 0
    for t in tables:
        assert t[0] and t[1]

        validation_rules = None
        if enable_validation:
            validation_rules = get_validation_rules(config.dbname, t[0], t[1])

        profiler.skip_table_profiling = skip_table_profiling
        profiler.skip_column_profiling = skip_column_profiling

        try:
            count += 1
            tmp = skip_record_validation
            newdata = profiler.run(unicode(t[0]), unicode(t[1]),
                                   validation_rules=validation_rules,
                                   skip_record_validation=tmp,
                                   timeout=timeout)
        except DriverError as e:
            log.error(_("Abort by driver error."))
            sys.exit(1)
        except QueryError as e:
            log.error(_("Profiling failed on %s.%s") %
                      (t[0], t[1]), detail=u'QueryError: '+e.value)
            failed_count += 1
            continue
        except InternalError as e:
            log.error(_("Profiling failed on %s.%s (internal error)") %
                      (t[0], t[1]), detail=u'InternalError: '+e.value)
            failed_count += 1
            continue
        except KeyboardInterrupt as e:
            log.error(_("Interrupted by the user."))
            log.error(_("Abort."))
            sys.exit(2)
        except NotImplementedError as e:
            log.error(_("This feature is not implemented."), detail=e)
            failed_count += 1
            continue
        except Exception as e:
            log.error(_("Profiling failed on %s.%s (internal error)") %
                      (t[0], t[1]), detail=repr(e))
            failed_count += 1
            traceback.print_exc()
            log.info(_("Continuing profiling."))
            continue

        # copy nls_name and comment from previous record.
        tab = Table2.find(config.dbname, t[0], t[1])
        if len(tab) == 1:
            DbProfilerBase.migrate_table_meta(tab[0].data, newdata)
        try:
            Table2.create(config.dbname, t[0], t[1], newdata)
        except InternalError as ex:
            log.error(_("Could not append table data (%s.%s) to "
                        "the repository.") % (t[0], t[1]),
                      detail=str(ex))

    try:
        repo.close()
    except Exception as e:
        log.error(_("Internal error (%d)") % 4, detail=repr(e))
        log.error(_("Abort."))
        traceback.print_exc()
        sys.exit(3)

    log.info(_("Profiling errors have occured on %d/%d tables.") %
             (count, failed_count))
    log.info(_("Completed profiling %d tables.") % (count - failed_count))

    sys.exit(0)