path: root/python/skytools/scripting.py

"""Useful functions and classes for database scripts."""

import sys, os, signal, psycopg, optparse, traceback, time
import logging, logging.handlers, logging.config

from skytools.config import *
import skytools.skylog

__all__ = ['daemonize', 'run_single_process', 'DBScript',
    'I_AUTOCOMMIT', 'I_READ_COMMITTED', 'I_SERIALIZABLE']

#
# daemon mode
#

def daemonize():
    """Turn the process into daemon.
    
    Goes background and disables all i/o.
    """

    # launch new process, kill parent
    pid = os.fork()
    if pid != 0:
        os._exit(0)

    # start new session
    os.setsid()

    # stop i/o
    fd = os.open("/dev/null", os.O_RDWR)
    os.dup2(fd, 0)
    os.dup2(fd, 1)
    os.dup2(fd, 2)
    if fd > 2:
        os.close(fd)

#
# Pidfile locking+cleanup & daemonization combined
#

def _write_pidfile(pidfile):
    pid = os.getpid()
    f = open(pidfile, 'w')
    f.write(str(pid))
    f.close()

def run_single_process(runnable, daemon, pidfile):
    """Run runnable class, possibly daemonized, locked on pidfile."""

    # check if another process is running
    if pidfile and os.path.isfile(pidfile):
        print "Pidfile exists, another process running?"
        sys.exit(1)

    # daemonize if needed and write pidfile
    if daemon:
        daemonize()
    if pidfile:
        _write_pidfile(pidfile)
    
    # Catch SIGTERM to cleanup pidfile
    def sigterm_hook(signum, frame):
        try:
            os.remove(pidfile)
        except: pass
        sys.exit(0)
    # attach it to signal
    if pidfile:
        signal.signal(signal.SIGTERM, sigterm_hook)

    # run
    try:
        runnable.run()
    finally:
        # another try of cleaning up
        if pidfile:
            try:
                os.remove(pidfile)
            except: pass

#
# logging setup
#

_log_config_done = 0
_log_init_done = {}

def _init_log(job_name, cf, log_level):
    """Logging setup happens here."""
    global _log_init_done, _log_config_done

    got_skylog = 0
    use_skylog = cf.getint("use_skylog", 0)

    # load logging config if needed
    if use_skylog and not _log_config_done:
        # python logging.config braindamage:
        # cannot specify external classess without such hack
        logging.skylog = skytools.skylog

        # load general config
        list = ['skylog.ini', '~/.skylog.ini', '/etc/skylog.ini']
        for fn in list:
            fn = os.path.expanduser(fn)
            if os.path.isfile(fn):
                defs = {'job_name': job_name}
                logging.config.fileConfig(fn, defs)
                got_skylog = 1
                break
        _log_config_done = 1
        if not got_skylog:
            sys.stderr.write("skylog.ini not found!\n")
            sys.exit(1)

    # avoid duplicate logging init for job_name
    log = logging.getLogger(job_name)
    if job_name in _log_init_done:
        return log
    _log_init_done[job_name] = 1

    # compatibility: specify ini file in script config
    logfile = cf.getfile("logfile", "")
    if logfile:
        fmt = logging.Formatter('%(asctime)s %(process)s %(levelname)s %(message)s')
        size = cf.getint('log_size', 10*1024*1024)
        num = cf.getint('log_count', 3)
        hdlr = logging.handlers.RotatingFileHandler(
                    logfile, 'a', size, num)
        hdlr.setFormatter(fmt)
        log.addHandler(hdlr)

    # if skylog.ini is disabled or not available, log at least to stderr
    if not got_skylog:
        hdlr = logging.StreamHandler()
        fmt = logging.Formatter('%(asctime)s %(process)s %(levelname)s %(message)s')
        hdlr.setFormatter(fmt)
        log.addHandler(hdlr)

    log.setLevel(log_level)

    return log

#: how old connections need to be closed
DEF_CONN_AGE = 20*60  # 20 min

#: isolation level not set
I_DEFAULT = -1

#: isolation level constant for AUTOCOMMIT
I_AUTOCOMMIT = 0
#: isolation level constant for READ COMMITTED
I_READ_COMMITTED = 1
#: isolation level constant for SERIALIZABLE
I_SERIALIZABLE = 2

class DBCachedConn(object):
    """Cache a db connection."""
    def __init__(self, name, loc, max_age = DEF_CONN_AGE):
        self.name = name
        self.loc = loc
        self.conn = None
        self.conn_time = 0
        self.max_age = max_age
        self.autocommit = -1
        self.isolation_level = -1

    def get_connection(self, autocommit = 0, isolation_level = -1):
        # autocommit overrider isolation_level
        if autocommit:
            isolation_level = I_AUTOCOMMIT

        # default isolation_level is READ COMMITTED
        if isolation_level < 0:
            isolation_level = I_READ_COMMITTED

        # new conn?
        if not self.conn:
            self.isolation_level = isolation_level
            self.conn = psycopg.connect(self.loc)

            self.conn.set_isolation_level(isolation_level)
            self.conn_time = time.time()
        else:
            if self.isolation_level != isolation_level:
                raise Exception("Conflict in isolation_level")

        # done
        return self.conn

    def refresh(self):
        if not self.conn:
            return
        #for row in self.conn.notifies():
        #    if row[0].lower() == "reload":
        #        self.reset()
        #        return
        if not self.max_age:
            return
        if time.time() - self.conn_time >= self.max_age:
            self.reset()

    def reset(self):
        if not self.conn:
            return

        # drop reference
        conn = self.conn
        self.conn = None

        if self.isolation_level == I_AUTOCOMMIT:
            return

        # rollback & close
        try:
            conn.rollback()
        except: pass
        try:
            conn.close()
        except: pass

class DBScript(object):
    """Base class for database scripts.

    Handles logging, daemonizing, config, errors.
    """
    service_name = None
    job_name = None
    cf = None
    log = None

    def __init__(self, service_name, args):
        """Script setup.

        User class should override work() and optionally __init__(), startup(),
        reload(), reset() and init_optparse().

        NB: in case of daemon, the __init__() and startup()/work() will be
        run in different processes.  So nothing fancy should be done in __init__().
        
        @param service_name: unique name for script.
            It will be also default job_name, if not specified in config.
        @param args: cmdline args (sys.argv[1:]), but can be overrided
        """
        self.service_name = service_name
        self.db_cache = {}
        self.go_daemon = 0
        self.do_single_loop = 0
        self.looping = 1
        self.need_reload = 1
        self.stat_dict = {}
        self.log_level = logging.INFO
        self.work_state = 1

        # parse command line
        parser = self.init_optparse()
        self.options, self.args = parser.parse_args(args)

        # check args
        if self.options.daemon:
            self.go_daemon = 1
        if self.options.quiet:
            self.log_level = logging.WARNING
        if self.options.verbose:
            self.log_level = logging.DEBUG
        if len(self.args) < 1:
            print "need config file"
            sys.exit(1)
        conf_file = self.args[0]

        # load config
        self.cf = Config(self.service_name, conf_file)
        self.job_name = self.cf.get("job_name", self.service_name)
        self.pidfile = self.cf.getfile("pidfile", '')
        
        self.reload()

        # init logging
        self.log = _init_log(self.job_name, self.cf, self.log_level)

        # send signal, if needed
        if self.options.cmd == "kill":
            self.send_signal(signal.SIGTERM)
        elif self.options.cmd == "stop":
            self.send_signal(signal.SIGINT)
        elif self.options.cmd == "reload":
            self.send_signal(signal.SIGHUP)

    def init_optparse(self, parser = None):
        """Initialize a OptionParser() instance that will be used to 
        parse command line arguments.

        Note that it can be overrided both directions - either DBScript
        will initialize a instance and passes to user code or user can
        initialize and then pass to DBScript.init_optparse().

        @param parser: optional OptionParser() instance,
               where DBScript should attachs its own arguments.
        @return: initialized OptionParser() instance.
        """
        if parser:
            p = parser
        else:
            p = optparse.OptionParser()
            p.set_usage("%prog [options] INI")
        # generic options
        p.add_option("-q", "--quiet", action="store_true",
                     help = "make program silent")
        p.add_option("-v", "--verbose", action="store_true",
                     help = "make program verbose")
        p.add_option("-d", "--daemon", action="store_true",
                     help = "go background")

        # control options
        g = optparse.OptionGroup(p, 'control running process')
        g.add_option("-r", "--reload",
                     action="store_const", const="reload", dest="cmd",
                     help = "reload config (send SIGHUP)")
        g.add_option("-s", "--stop",
                     action="store_const", const="stop", dest="cmd",
                     help = "stop program safely (send SIGINT)")
        g.add_option("-k", "--kill",
                     action="store_const", const="kill", dest="cmd",
                     help = "kill program immidiately (send SIGTERM)")
        p.add_option_group(g)

        return p

    def send_signal(self, sig):
        if not self.pidfile:
            self.log.warning("No pidfile in config, nothing todo")
            sys.exit(0)
        if not os.path.isfile(self.pidfile):
            self.log.warning("No pidfile, process not running")
            sys.exit(0)
        pid = int(open(self.pidfile, "r").read())
        os.kill(pid, sig)
        sys.exit(0)

    def set_single_loop(self, do_single_loop):
        """Changes whether the script will loop or not."""
        self.do_single_loop = do_single_loop

    def start(self):
        """This will launch main processing thread."""
        if self.go_daemon:
            if not self.pidfile:
                self.log.error("Daemon needs pidfile")
                sys.exit(1)
        run_single_process(self, self.go_daemon, self.pidfile)

    def stop(self):
        """Safely stops processing loop."""
        self.looping = 0

    def reload(self):
        "Reload config."
        self.cf.reload()
        self.loop_delay = self.cf.getfloat("loop_delay", 1.0)

    def hook_sighup(self, sig, frame):
        "Internal SIGHUP handler.  Minimal code here."
        self.need_reload = 1

    def hook_sigint(self, sig, frame):
        "Internal SIGINT handler.  Minimal code here."
        self.stop()

    def stat_add(self, key, value):
        self.stat_put(key, value)

    def stat_put(self, key, value):
        """Sets a stat value."""
        self.stat_dict[key] = value

    def stat_increase(self, key, increase = 1):
        """Increases a stat value."""
        if key in self.stat_dict:
            self.stat_dict[key] += increase
        else:
            self.stat_dict[key] = increase

    def send_stats(self):
        "Send statistics to log."

        res = []
        for k, v in self.stat_dict.items():
            res.append("%s: %s" % (k, str(v)))

        if len(res) == 0:
            return

        logmsg = "{%s}" % ", ".join(res)
        self.log.info(logmsg)
        self.stat_dict = {}

    def get_database(self, dbname, autocommit = 0, isolation_level = -1,
                     cache = None, max_age = DEF_CONN_AGE):
        """Load cached database connection.
        
        User must not store it permanently somewhere,
        as all connections will be invalidated on reset.
        """

        if not cache:
            cache = dbname
        if cache in self.db_cache:
            dbc = self.db_cache[cache]
        else:
            loc = self.cf.get(dbname)
            dbc = DBCachedConn(cache, loc, max_age)
            self.db_cache[cache] = dbc

        return dbc.get_connection(autocommit, isolation_level)

    def close_database(self, dbname):
        """Explicitly close a cached connection.
        
        Next call to get_database() will reconnect.
        """
        if dbname in self.db_cache:
            dbc = self.db_cache[dbname]
            dbc.reset()

    def reset(self):
        "Something bad happened, reset all connections."
        for dbc in self.db_cache.values():
            dbc.reset()
        self.db_cache = {}

    def run(self):
        "Thread main loop."

        # run startup, safely
        try:
            self.startup()
        except KeyboardInterrupt, det:
            raise
        except SystemExit, det:
            raise
        except Exception, det:
            exc, msg, tb = sys.exc_info()
            self.log.fatal("Job %s crashed: %s: '%s' (%s: %s)" % (
                       self.job_name, str(exc), str(msg).rstrip(),
                       str(tb), repr(traceback.format_tb(tb))))
            del tb
            self.reset()
            sys.exit(1)

        while self.looping:
            # reload config, if needed
            if self.need_reload:
                self.reload()
                self.need_reload = 0

            # do some work
            work = self.run_once()

            # send stats that was added
            self.send_stats()

            # reconnect if needed
            for dbc in self.db_cache.values():
                dbc.refresh()

            # exit if needed
            if self.do_single_loop:
                self.log.debug("Only single loop requested, exiting")
                break

            # remember work state
            self.work_state = work
            # should sleep?
            if not work:
                try:
                    time.sleep(self.loop_delay)
                except Exception, d:
                    self.log.debug("sleep failed: "+str(d))
                    sys.exit(0)

    def run_once(self):
        "Run users work function, safely."
        try:
            return self.work()
        except SystemExit, d:
            self.send_stats()
            self.log.info("got SystemExit(%s), exiting" % str(d))
            self.reset()
            raise d
        except KeyboardInterrupt, d:
            self.send_stats()
            self.log.info("got KeyboardInterrupt, exiting")
            self.reset()
            sys.exit(1)
        except Exception, d:
            self.send_stats()
            exc, msg, tb = sys.exc_info()
            self.log.fatal("Job %s crashed: %s: '%s' (%s: %s)" % (
                       self.job_name, str(exc), str(msg).rstrip(),
                       str(tb), repr(traceback.format_tb(tb))))
            del tb
            self.reset()
            if self.looping:
                time.sleep(20)
                return 1

    def work(self):
        "Here should user's processing happen."
        raise Exception("Nothing implemented?")

    def startup(self):
        """Will be called just before entering main loop.

        In case of daemon, if will be called in same process as work(),
        unlike __init__().
        """

        # set signals
        signal.signal(signal.SIGHUP, self.hook_sighup)
        signal.signal(signal.SIGINT, self.hook_sigint)