summaryrefslogtreecommitdiff
path: root/python/skytools/scripting.py
diff options
context:
space:
mode:
Diffstat (limited to 'python/skytools/scripting.py')
-rw-r--r--python/skytools/scripting.py523
1 files changed, 523 insertions, 0 deletions
diff --git a/python/skytools/scripting.py b/python/skytools/scripting.py
new file mode 100644
index 00000000..cf976801
--- /dev/null
+++ b/python/skytools/scripting.py
@@ -0,0 +1,523 @@
+
+"""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)
+
+