pgloader

#! /usr/bin/env python
# Author: Dimitri Fontaine <dim@tapoueh.org>

"""
PostgreSQL data import tool, see included man page.
"""

import os, sys, os.path, time, codecs, logging, threading, signal
from cStringIO import StringIO
from tempfile import gettempdir

import pgloader.options
import pgloader.tools
import pgloader.logger
from pgloader.tools import PGLoader_Error

def parse_options():
    """ Parse given options """
    import ConfigParser
    from optparse import OptionParser

    usage  = "%prog [-c <config_filename>] Section [Section ...]"
    parser = OptionParser(usage = usage)
    
    parser.add_option("--version", action = "store_true",
                      dest    = "version",
                      default = False,
                      help    = "show pgloader version")

    parser.add_option("-c", "--config", dest = "config",
                      default = "pgloader.conf",
                      help    = "configuration file, defauts to pgloader.conf")

    parser.add_option("-p", "--pedantic", action = "store_true",
                      dest    = "pedantic",
                      default = False,
                      help    = "pedantic mode, stop processing on warning")

    parser.add_option("-d", "--debug", action = "store_true",
                      dest    = "debug",
                      default = False,
                      help    = "add some debug information (a lot of)")

    parser.add_option("-v", "--verbose", action = "store_true",
                      dest    = "verbose",
                      default = False,
                      help    = "be verbose and about processing progress")

    parser.add_option("-q", "--quiet", action = "store_true",
                      dest    = "quiet",
                      default = False,
                      help    = "be quiet, only print out errors")

    parser.add_option("-l", "--level", dest = "loglevel",
                      default = None,
                      help    = "loglevel to use: ERROR, WARNING, INFO, DEBUG")

    parser.add_option("-L", "--logfile", dest = "logfile",
                      default = os.path.join (gettempdir(), "pgloader.log"),
                      help    = "log file, defauts to $TMPDIR/pgloader.log")

    parser.add_option("-r", "--reject-log", dest = "reject_log",
                      default = None,
                      help    = "log file for rejected data error messages")

    parser.add_option("-j", "--reject-data", dest = "reject_data",
                      default = None,
                      help    = "log file for rejected data, bad input data")

    parser.add_option("-s", "--summary", action = "store_true",
                      dest    = "summary",
                      default = False,
                      help    = "print a summary")

    parser.add_option("-n", "--dry-run", action = "store_true",
                      dest    = "dryrun",
                      default = False,
                      help    = "simulate operations, don't connect to the db")

    parser.add_option("-T", "--truncate", action = "store_true",
                      dest    = "truncate",
                      default = False,
                      help    = "truncate tables before importing data")

    parser.add_option("-D", "--disable-triggers", action = "store_true",
                      dest    = "triggers",
                      default = False,
                      help    = "Disable triggers before loading, Enable them again after")

    parser.add_option("-V", "--vacuum", action = "store_true",
                      dest = "vacuum",
                      default = False,
                      help    = "vacuum tables after data loading")

    parser.add_option("-C", "--count", dest = "count",
                      default = None, type = "int",
                      help    = "number of input lines to process")
    
    parser.add_option("-F", "--from", dest = "fromcount",
                      default = 0, type = "int",
                      help    = "number of input lines to skip")

    parser.add_option("-I", "--from-id", dest = "fromid",
                      default = None,
                      help    = "wait for given id on input to begin")

    parser.add_option("-f", "--field-separator", dest = "fsep",
                      default = pgloader.options.FIELD_SEP,
                      help    = "default field separator to use")

    parser.add_option("-E", "--encoding", dest = "encoding",
                      default = None,
                      help    = "input files encoding")

    parser.add_option("-o", "--pg-options", dest = "pg_options", action = "append",
                      help    = "list of PostgreSQL options you want to SET")

    parser.add_option("-t", "--section-threads", dest = "section_threads",
                      default = pgloader.options.SECTION_THREADS, 
                      type    = "int",
                      help    = "number of threads used per sections, default is 1")

    parser.add_option("-m", "--max-parallel-sections", dest = "parallel",
                      default = pgloader.options.MAX_PARALLEL_SECTIONS,
                      type    = "int",
                      help    = "number of sections to load in parralel, " +\
                          "default is 1")

    parser.add_option("-R", "--reformat_path", dest = "reformat_path",
                      default = None,
                      help    = "PATH where to find reformat python modules")

    parser.add_option("-1", "--psycopg1", action = "store_true",
                      dest    = "psycopg1",
                      default = False,
                      help    = "Force usage of psycopg1")

    parser.add_option("-2", "--psycopg2", action = "store_true",
                      dest    = "psycopg2",
                      default = False,
                      help    = "Force usage of psycopg2")

    parser.add_option("--psycopg-version", dest = "psycopg_version",
                      default = None,
                      help    = "Force usage of given version of psycopg")

    parser.add_option("--load-from-stdin", action = "store_true",
                      dest = "stdin",
                      default = False,
                      help    = "Load standard input data into given table name")

    parser.add_option("--load-to-table", dest = "table",
                      default = None,
                      help    = "Load to given table when --load-from-stdin")

    parser.add_option("--boundaries", dest = "boundaries",
                      default = None,
                      help    = "Load only in given boundaries, Start..End")

    (opts, args) = parser.parse_args()

    if opts.version:
        print "PGLoader version %s" % pgloader.options.PGLOADER_VERSION
        sys.exit(0)

    # check existence and read ability of config file
    if not os.path.exists(opts.config):
        print >>sys.stderr, \
              "Error: Configuration file %s does not exists" % opts.config
        print >>sys.stderr, parser.format_help()
        sys.exit(1)

    if not os.access(opts.config, os.R_OK):
        print >>sys.stderr, \
              "Error: Can't read configuration file %s" % opts.config
        print >>sys.stderr, parser.format_help()
        sys.exit(1)

    if opts.fromcount != 0 and opts.fromid is not None:
        print >>sys.stderr, \
              "Error: Can't set both options fromcount (-F) AND fromid (-I)"
        sys.exit(1)

    if opts.quiet and (opts.verbose or opts.debug):
        print >>sys.stderr, \
              "Error: Can't be verbose and quiet at the same time!"
        sys.exit(1)

    psyco_opts = 0
    for opt in [opts.psycopg1, opts.psycopg2, opts.psycopg_version]:
        if opt:
            psyco_opts += 1

    if psyco_opts > 1:
        print >>sys.stderr, \
              "Error: please use only one of the psycopg options"
        sys.exit(1)

    if opts.psycopg_version is not None:
        if opts.psycopg_version in ("1", "2"):
            opts.psycopg_version = int(opts.psycopg_version)

        else:
            print >>sys.stderr, \
                  "Error: psycopg_version can only be set to either 1 or 2"
            sys.exit(1)
        
    # if debug, then verbose
    if opts.debug:
        opts.verbose = True

    pgloader.options.DRY_RUN    = opts.dryrun
    pgloader.options.DEBUG      = opts.debug
    pgloader.options.VERBOSE    = opts.verbose
    pgloader.options.QUIET      = opts.quiet
    pgloader.options.SUMMARY    = opts.summary    
    pgloader.options.PEDANTIC   = opts.pedantic

    pgloader.options.TRUNCATE   = opts.truncate
    pgloader.options.VACUUM     = opts.vacuum
    pgloader.options.TRIGGERS   = opts.triggers
    
    pgloader.options.COUNT      = opts.count
    pgloader.options.FROM_COUNT = opts.fromcount
    pgloader.options.FROM_ID    = opts.fromid
    pgloader.options.FIELD_SEP  = opts.fsep

    pgloader.options.SECTION_THREADS       = opts.section_threads
    pgloader.options.MAX_PARALLEL_SECTIONS = opts.parallel

    pgloader.options.LOAD_FROM_STDIN = opts.stdin
    pgloader.options.LOAD_TO_TABLE   = opts.table

    if pgloader.options.MAX_PARALLEL_SECTIONS is None:
        from pgloader.options import DEFAULT_MAX_PARALLEL_SECTIONS
        pgloader.options.MAX_PARALLEL_SECTIONS = DEFAULT_MAX_PARALLEL_SECTIONS

    pgloader.options.INPUT_ENCODING = opts.encoding

    if opts.reformat_path:
        pgloader.options.REFORMAT_PATH = opts.reformat_path

    pgloader.options.LOG_FILE = opts.logfile

    # reject file names must contain one %s which will get replaced by the
    # section name
    if opts.reject_log:
        try:
            unused = opts.reject_log % "that would be a section name"
        except TypeError, e:
            # TypeError: not all arguments converted during string formatting
            # TypeError: not enough arguments for format string
            print >>sys.stderr, \
                "Error: reject log must contain a '%s' place holder for section name"
            sys.exit(1)
        pgloader.options.REJECT_LOG_FILE = opts.reject_log

    if opts.reject_data:
        try:
            unused = opts.reject_data % "that would be a section name"
        except TypeError, e:
            # TypeError: not all arguments converted during string formatting
            # TypeError: not enough arguments for format string
            print >>sys.stderr, \
                "Error: reject data must contain a '%s' place holder for section name"
            sys.exit(1)
        pgloader.options.REJECT_DATA_FILE = opts.reject_data

    if opts.loglevel:
        loglevel = pgloader.logger.level(opts.loglevel)
        pgloader.options.CLIENT_MIN_MESSAGES = loglevel
    elif opts.debug:
        pgloader.options.CLIENT_MIN_MESSAGES = logging.DEBUG
    elif opts.verbose:
        pgloader.options.CLIENT_MIN_MESSAGES = logging.INFO
    elif opts.quiet:
        pgloader.options.CLIENT_MIN_MESSAGES = logging.ERROR

    if opts.pg_options:
        pgloader.options.PG_OPTIONS = {}
        for o in opts.pg_options:
            try:
                n, v = [x.strip() for x in o.split('=')]
                if v == "":
                    raise ValueError
                pgloader.options.PG_OPTIONS[n] = v
            except ValueError, e:
                print >>sys.stderr, \
                    "Error: PostgreSQL options must have the form 'name=value'"
                sys.exit(1)

    if opts.psycopg1:
        pgloader.options.PSYCOPG_VERSION = 1
    elif opts.psycopg2:
        pgloader.options.PSYCOPG_VERSION = 2
    else:
        pgloader.options.PSYCOPG_VERSION = opts.psycopg_version

    if opts.boundaries:
        if opts.stdin:
            print >>sys.stderr, \
                  "Error: You can't set the boundaries of stdin."
            sys.exit(1)

        try:
            start, end = [int(x) for x in opts.boundaries.split("..")]
            pgloader.options.FILE_BOUNDARIES = (start, end)
        except ValueError, e:
            print >>sys.stderr, \
                  "Error: boundaries should be an integer range written X..Y"
            sys.exit(1)

    return opts.config, args

def parse_config(conffile):
    """ Parse the configuration file """
    section = 'pgsql'

    # Now read pgsql configuration section
    import ConfigParser
    config = ConfigParser.ConfigParser()

    try:
        config.read(conffile)
    except:
        print >>sys.stderr, "Error: Given file is not a configuration file"
        sys.exit(4)

    if not config.has_section(section):
        print >>sys.stderr, "Error: Please provide a [%s] section" % section
        sys.exit(5)

    # load some options
    # this has to be done after command line parsing
    from pgloader.options  import DRY_RUN, VERBOSE, DEBUG, PEDANTIC
    from pgloader.options  import NULL, EMPTY_STRING
    from pgloader.options  import CLIENT_MIN_MESSAGES, LOG_FILE
    from pgloader.options  import PG_OPTIONS
    from pgloader.tools    import check_dirname

    # first read the logging configuration
    if not CLIENT_MIN_MESSAGES:
        if config.has_option(section, 'client_min_messages'):
            cmm = config.get(section, 'client_min_messages')
            pgloader.options.CLIENT_MIN_MESSAGES = pgloader.logger.level(cmm)
        else:
            # CLIENT_MIN_MESSAGES has not been set at all
            pgloader.options.CLIENT_MIN_MESSAGES = logging.INFO

    if config.has_option(section, 'log_min_messages'):
        lmm = config.get(section, 'log_min_messages')
        pgloader.options.LOG_MIN_MESSAGES = pgloader.logger.level(lmm)
    else:
        pgloader.options.LOG_MIN_MESSAGES = logging.INFO

    if config.has_option(section, 'log_file'):
        # don't overload the command line -L option if given
        if not pgloader.options.LOG_FILE:
            pgloader.options.LOG_FILE = config.get(section, 'log_file')

    if pgloader.options.LOG_FILE:
        ok, logdir_mesg = check_dirname(pgloader.options.LOG_FILE)
        if not ok:
            # force default setting
            pgloader.options.LOG_FILE = pgloader.options.DEFAULT_LOG_FILE

    try:
        log = pgloader.logger.init(pgloader.options.CLIENT_MIN_MESSAGES,
                                   pgloader.options.LOG_MIN_MESSAGES,
                                   pgloader.options.LOG_FILE)
    except PGLoader_Error, e:
        try:
            log = pgloader.logger.init(pgloader.options.CLIENT_MIN_MESSAGES,
                                       pgloader.options.LOG_MIN_MESSAGES,
                                       pgloader.options.DEFAULT_LOG_FILE)

            log.warning(e)
            log.warning("Using default logfile %s",
                        pgloader.options.DEFAULT_LOG_FILE)
        except PGLoader_Error, e:
            print e
            sys.exit(8)
        
    pgloader.logger.log = log

    log.info("Logger initialized")
    if logdir_mesg:
        log.error(logdir_mesg)
        log.error("Default logfile %s has been used instead",
                  pgloader.options.LOG_FILE)

    if config.has_option(section, 'input_encoding'):
        input_encoding = pgloader.tools.parse_config_string(
            config.get(section, 'input_encoding'))
        pgloader.options.INPUT_ENCODING = input_encoding

    # optionnal global newline_escapes
    if config.has_option(section, 'newline_escapes'):
        setting = pgloader.tools.parse_config_string(
            config.get(section, 'newline_escapes'))
        pgloader.options.NEWLINE_ESCAPES = setting

    # Then there are null and empty_string optionnal parameters
    # They canbe overriden in specific table configuration
    if config.has_option(section, 'null'):
        pgloader.options.NULL = pgloader.tools.parse_config_string(
            config.get(section, 'null'))

    if config.has_option(section, 'empty_string'):
        pgloader.options.EMPTY_STRING = pgloader.tools.parse_config_string(
            config.get(section, 'empty_string'))

    if config.has_option(section, 'reformat_path'):
        # command line value is prefered to config format one
        if not pgloader.options.REFORMAT_PATH:
            rpath = config.get(section, 'reformat_path')
            pgloader.options.REFORMAT_PATH = rpath

    if config.has_option(section, 'max_parallel_sections'):
        if not pgloader.options.MAX_PARALLEL_SECTIONS:
            mps = config.getint(section, 'max_parallel_sections')
            pgloader.options.MAX_PARALLEL_SECTIONS = mps

    return config

def myprint(l, line_prefix = "  ", cols = 78):
    """ pretty print list l elements """
    # some code for pretty print
    lines = []
    
    tmp = line_prefix
    for e in l:
        if len(tmp) + len(e) > cols:
            lines.append(tmp)
            tmp = line_prefix
            
        if tmp != line_prefix: tmp += " "
        tmp += e
        
    lines.append(tmp)

    return lines

def duration_pprint(duration):
    """ pretty print duration (human readable information) """
    if duration > 3600:
        h  = int(duration / 3600)
        m  = int((duration - 3600 * h) / 60)
        s  = duration - 3600 * h - 60 * m + 0.5
        return '%2dh%02dm%03.1f' % (h, m, s)
    
    elif duration > 60:
        m  = int(duration / 60)
        s  = duration - 60 * m
        return ' %02dm%06.3f' % (m, s)
        
    else:
        return '%10.3f' % duration

def print_summary(dbconn, sections, summary, td):
    """ print a pretty summary """
    from pgloader.options  import VERBOSE, DEBUG, QUIET, SUMMARY
    from pgloader.options  import DRY_RUN, PEDANTIC, VACUUM
    from pgloader.pgloader import PGLoader

    retcode = 0

    t= 'Table name        |    duration |    size |  copy rows |     errors '
    _= '===================================================================='

    tu = te = ts = 0 # total updates, errors, size

    if False and not DRY_RUN:
        dbconn.reset()
        cursor = dbconn.dbconn.cursor()

    s_ok = 0
    for s in sections:
        if s not in summary:
            continue

        s_ok += 1
        if s_ok == 1:
            # print pretty sumary header now
            print
            print t
            print _

        if summary[s]:
            t, d, u, e = summary[s]
            d = duration_pprint(d)
        else:
            t = s
            d = '%9s ' % '-'
            u = e = 0

        if False and not DRY_RUN:
            sql = "select pg_total_relation_size(%s), " + \
                  "pg_size_pretty(pg_total_relation_size(%s));"
            cursor.execute(sql, [t, t])
            octets, sp = cursor.fetchone()
            ts += octets

            if sp[5:] == 'bytes': sp = sp[:-5] + ' B'
        else:
            sp = '-'

        tn = s
        if len(tn) > 18:
            tn = s[0:15] + "..."

        print '%-18s| %ss | %7s | %10d | %10d' % (tn, d, sp, u, e)

        tu += u
        te += e

        if e > 0:
            retcode += 1

    if s_ok > 1:
        td = duration_pprint(td)

        # pretty size
        if False and not DRY_RUN:
            cursor.execute("select pg_size_pretty(%s);", [ts])
            [ts] = cursor.fetchone()
            if ts[5:] == 'bytes': ts = ts[:-5] + ' B'
        else:
            ts = '-'

        print _
        print 'Total             | %ss | %7s | %10d | %10d' \
              % (td, ts, tu, te)

        if False and not DRY_RUN:
            cursor.close()

    return retcode

def load_data():
    """ read option line and configuration file, then process data
    import of given section, or all sections if no section is given on
    command line """

    # first parse command line options, and set pgloader.options values
    # accordingly
    conffile, args = parse_options()

    # now init db connection
    config = parse_config(conffile)

    from pgloader.logger  import log
    from pgloader.tools   import read_path, check_path
    from pgloader.options import VERBOSE
    
    import pgloader.options
    if pgloader.options.REFORMAT_PATH:
        rpath  = read_path(pgloader.options.REFORMAT_PATH, log, check = False)
        crpath = check_path(rpath, log)
    else:
        rpath  = crpath  = None

    if not crpath:
        if rpath:
            # don't check same path entries twice
        
            default_rpath = set(crpath) \
                            - set(pgloader.options.DEFAULT_REFORMAT_PATH)
        else:
            default_rpath = pgloader.options.DEFAULT_REFORMAT_PATH
        
        pgloader.options.REFORMAT_PATH = check_path(default_rpath, log)
    else:
        pgloader.options.REFORMAT_PATH = rpath

    log.info('Reformat path is %s', pgloader.options.REFORMAT_PATH)

    # load some pgloader package modules
    from pgloader.options  import VERBOSE, DEBUG, QUIET, SUMMARY
    from pgloader.options  import DRY_RUN, PEDANTIC, VACUUM
    from pgloader.options  import MAX_PARALLEL_SECTIONS
    from pgloader.options  import LOAD_FROM_STDIN, LOAD_TO_TABLE
    from pgloader.options  import FILE_BOUNDARIES
    from pgloader.pgloader import PGLoader
    from pgloader.tools    import PGLoader_Error

    sections = []
    summary  = {}

    # args are meant to be configuration sections, or filenames, or stdin
    if LOAD_FROM_STDIN:
        if FILE_BOUNDARIES is not None:
            log.warning("Can't use --boundaries on stdin")

        if len(args) == 0:
            s = '<stdin>'
            config.add_section(s)
            config.set(s, 'table', LOAD_TO_TABLE)
            config.set(s, 'filename', 'sys.stdin')
            config.set(s, 'columns', '*')
            config.set(s, 'format', 'csv')
            sections.append(s)
            
        elif len(args) == 1:
            if config.has_section(args[0]):
                # apply given section parameters, then load from stdin
                config.set(args[0], 'filename', 'sys.stdin')
                sections.append(args[0])
            else:
                print >>sys.stderr, \
                      "Error: Please provide a [%s] section" % args[0]
                sys.exit(5)
        else:
            print >>sys.stderr, \
                  "Error: can't read several sections all from stdin"
            sys.exit(5)

    elif len(args) > 0:
        for s in args:
            if config.has_section(s):
                sections.append(s)
            else:
                log.info("Creating a section for file '%s'" % s)
                # a filename was given, apply [pgsql] defaults
                # set the tablename as the filename sans extension
                # consider columns = *
                if not os.path.exists(s):
                    print >>sys.stderr, \
                        "Error: '%s' does not exists as a section nor as a file" % s
                    sys.exit(2)

                config.add_section(s)
                config.set(s, 'table', os.path.splitext(os.path.basename(s))[0])
                config.set(s, 'filename', s)
                config.set(s, 'columns', '*')
                config.set(s, 'format', 'csv')
                sections.append(s)

    else:
        if not LOAD_FROM_STDIN:
            # don't load all sections first when asked to load stdin
            log.debug("No argument on CLI, will consider all sections")
            for s in config.sections():
                if s != 'pgsql':
                    sections.append(s)

            # we run through sorted section list, unless we got the section list
            # from command line
            sections.sort()

    if FILE_BOUNDARIES is not None and len(sections) > 1:
        print >>sys.stderr, \
              "Error: will not apply boundaries on more than one file"
        sys.exit(5)

    log.info('Will consider following sections:')
    for line in myprint(sections):
        log.info(line)

    # we count time passed from now on
    begin = time.time()

    threads  = {}
    started  = {}
    finished = {}
    current  = 0

    max_running = MAX_PARALLEL_SECTIONS
    if max_running == -1:
        max_running = len(sections)

    log.info('Will load %d section at a time' % max_running)

    sem = threading.BoundedSemaphore(max_running)
    
    while current < len(sections):
        s = sections[current]

        try:
            loader      = None
            summary [s] = []
            started [s] = threading.Event()
            finished[s] = threading.Event()

            try:
                loader = PGLoader(s, config, sem,
                                  (started[s], finished[s]), summary[s])
            except PGLoader_Error, e:
                # could not initialize properly this loader, don't
                # ever wait for it
                started[s] .set()
                finished[s].set()
                log.error(e)
                if DEBUG:
                    raise

            except IOError, e:
                # No space left on device?  can't log it
                break

            if loader:
                if not loader.template:
                    if FILE_BOUNDARIES is not None and len(sections) == 1:
                        loader.reader.set_boundaries(FILE_BOUNDARIES)
                    filename       = loader.filename
                    input_encoding = loader.input_encoding
                    threads[s]     = loader

                    # .start() will sem.aquire(), so we won't have more
                    # than max_running threads running at any time.
                    log.debug("Starting a thread for %s" % s)
                    threads[s].start()
                else:
                    log.info("Skipping section %s, which is a template" % s)

                    for d in (summary, started, finished):
                        d.pop(s)

        except PGLoader_Error, e:
            if e == '':
                log.error('[%s] Please correct previous errors' % s)
            else:
                log.error('%s' % e)

            if DEBUG:
                raise

            if PEDANTIC:
                # was: threads[s].print_stats()
                # but now thread[s] is no more alive
                pass

        except UnicodeDecodeError, e:
            log.error("can't open '%s' with given input encoding '%s'" \
                               % (filename, input_encoding))
                                    
        except IOError, e:
            # typically, No Space Left On Device, can't report nor continue
            break

        current += 1

    # get sure each thread is started, then each one is done
    from pgloader.tools import check_events

    check_events(started, log, "is running")
    log.info("All threads are started, wait for them to terminate")
    check_events(finished, log, "processing is over")

    # check whether any thread failed
    for section, loader in threads.iteritems():
        if not loader.success:
            return 1

    # total duration
    td = time.time() - begin

    if SUMMARY:
        try:
            print_summary(None, sections, summary, td)
            print
        except PGLoader_Error, e:
            log.error("Can't print summary: %s" % e)
            return 1

        except KeyboardInterrupt:
            return 1

    return 0

if __name__ == "__main__":
    # default exit value
    ret = 1

    # Disable KeyboardInterrupt, just exit on CTRL+C
    signal.signal(signal.SIGINT, signal.SIG_DFL)

    try:
        ret = load_data()

    # SystemExit inherits from Exception in python < 2.5, so we have to
    # handle it separately.
    except SystemExit, e:
        sys.exit(e.code)

    except Exception, e:
        from pgloader.options import DEBUG
        if DEBUG:
            raise
        sys.stderr.write(str(e) + '\n')
        sys.exit(1)

    except IOError, e:
        sys.stderr.write(str(e) + '\n')
        sys.exit(1)

    sys.exit(ret)