Merge branch 'master' of git://github.com/markokr/skytools-dev

1 files changed, 87 insertions, 48 deletions

diff --git a/python/skytools/scripting.py b/python/skytools/scripting.py
index 2b151e9c..83e050eb 100644
--- a/python/skytools/scripting.py
+++ b/python/skytools/scripting.py
@@ -1,22 +1,6 @@
 
 """Useful functions and classes for database scripts.
 
-Config template::
-
-    [service_type]
-    # generic ini file example from Skytools framework
-    # if job name is omitted then ini file name without extension is used
-    job_name = %(config_name)s
-
-    # loop delay in seconds
-    loop_delay = 1.0
-
-    # location of log and pid files usually kept in ~/log and ~/pid
-    logfile = ~/log/%(job_name)s.log
-    pidfile = ~/pid/%(job_name)s.pid
-
-    # should centralized logging be used
-    use_skylog = 0
 """
 
 import sys, os, signal, optparse, time, errno
@@ -277,14 +261,42 @@ class DBScript(object):
     """Base class for database scripts.
 
     Handles logging, daemonizing, config, errors.
+
+    Config template::
+
+        ## Parameters for skytools.DBScript ##
+
+        # how many seconds to sleep between work loops
+        # if missing or 0, then instead sleeping, the script will exit
+        loop_delay = 1.0
+
+        # where to log
+        logfile = ~/log/%(job_name)s.log
+
+        # where to write pidfile
+        pidfile = ~/pid/%(job_name)s.pid
+
+        # per-process name to use in logging
+        #job_name = %(config_name)s
+
+        # whether centralized logging should be used (loaded from skylog.ini)
+        #use_skylog = 0
+
+        # default lifetime for database connections (in seconds)
+        #connection_lifetime = 1200
     """
     service_name = None
     job_name = None
     cf = None
     log = None
     pidfile = None
-    loop_delay = 1
-    doc_string = None
+    loop_delay = 0
+
+    # result from last work() call:
+    #  1 - there is probably more work, don't sleep
+    #  0 - no work, sleep before calling again
+    # -1 - exception was thrown
+    work_state = 1
 
     def __init__(self, service_name, args):
         """Script setup.
@@ -302,12 +314,9 @@ class DBScript(object):
         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()
@@ -345,14 +354,29 @@ class DBScript(object):
     def print_ini(self):
         """Prints out ini file from doc string of the script of default for dbscript
         
-        Used by --ini option on command line
+        Used by --ini option on command line.
         """
 
+        # current service name
+        print("[%s]\n" % self.service_name)
+
+        # walk class hierarchy
+        bases = [self.__class__]
+        while len(bases) > 0:
+            parents = []
+            for c in bases:
+                for p in c.__bases__:
+                    if p not in parents:
+                        parents.append(p)
+                doc = c.__doc__
+                if doc:
+                    self._print_ini_frag(doc)
+            bases = parents
+
+    def _print_ini_frag(self, doc):
         # use last '::' block as config template
-        doc = self.__doc__
         pos = doc and doc.rfind('::') or -1
         if pos < 0:
-            print 'No ini template defined.'
             return
         doc = doc[pos+2 : ].rstrip()
 
@@ -366,9 +390,10 @@ class DBScript(object):
                 pfx = ln[ : wslen]
             if pfx:
                 if ln.startswith(pfx):
-                    print ln[ len(pfx) : ]
+                    print(ln[ len(pfx) : ])
                 else:
-                    print ln
+                    print(ln)
+        print('')
 
     def load_config(self):
         """Loads and returns skytools.Config instance.
@@ -435,7 +460,8 @@ class DBScript(object):
 
     def set_single_loop(self, do_single_loop):
         """Changes whether the script will loop or not."""
-        self.do_single_loop = do_single_loop
+        if do_single_loop:
+            self.loop_delay = 0
 
     def _boot_daemon(self):
         run_single_process(self, self.go_daemon, self.pidfile)
@@ -450,7 +476,7 @@ class DBScript(object):
 
     def stop(self):
         """Safely stops processing loop."""
-        self.looping = 0
+        self.loop_delay = 0
 
     def reload(self):
         "Reload config."
@@ -506,7 +532,7 @@ class DBScript(object):
         self.log.info(logmsg)
         self.stat_dict = {}
 
-    def connection_setup(self, dbname, conn):
+    def connection_hook(self, dbname, conn):
         pass
 
     def get_database(self, dbname, autocommit = 0, isolation_level = -1,
@@ -525,7 +551,7 @@ class DBScript(object):
         else:
             if not connstr:
                 connstr = self.cf.get(dbname)
-            dbc = DBCachedConn(cache, connstr, max_age, setup_func = self.connection_setup)
+            dbc = DBCachedConn(cache, connstr, max_age, setup_func = self.connection_hook)
             self.db_cache[cache] = dbc
 
         return dbc.get_connection(autocommit, isolation_level)
@@ -552,7 +578,7 @@ class DBScript(object):
         # run startup, safely
         self.run_func_safely(self.startup)
 
-        while self.looping:
+        while 1:
             # reload config, if needed
             if self.need_reload:
                 self.reload()
@@ -568,36 +594,36 @@ class DBScript(object):
             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:
-                time.sleep(self.loop_delay)
+                if self.loop_delay:
+                    time.sleep(self.loop_delay)
+                else:
+                    break
 
     def run_once(self):
         return self.run_func_safely(self.work, True)
 
     def run_func_safely(self, func, prefer_looping = False):
         "Run users work function, safely."
+        cname = None
+        emsg = None
         try:
             return func()
-        except UsageError, ex:
-            self.log.error(str(ex))
+        except UsageError, d:
+            self.log.error(str(d))
             sys.exit(1)
         except SystemExit, d:
             self.send_stats()
-            if prefer_looping and not self.do_single_loop:
+            if prefer_looping and self.loop_delay:
                 self.log.info("got SystemExit(%s), exiting" % str(d))
             self.reset()
             raise d
         except KeyboardInterrupt, d:
             self.send_stats()
-            if prefer_looping and not self.do_single_loop:
+            if prefer_looping and self.loop_delay:
                 self.log.info("got KeyboardInterrupt, exiting")
             self.reset()
             sys.exit(1)
@@ -607,26 +633,39 @@ class DBScript(object):
                 cname = d.cursor.connection.my_name
                 dsn = d.cursor.connection.dsn
                 sql = d.cursor.query
+                emsg = str(d).strip()
                 self.log.error("Job %s got error on connection '%s': %s" % (
-                    self.job_name,
-                    d.cursor.connection.my_name,
-                    str(d).strip()))
+                    self.job_name, cname, emsg))
             else:
                 n = "psycopg2.%s" % d.__class__.__name__
+                emsg = str(d).rstrip()
                 self.log.exception("Job %s crashed: %s: %s" % (
-                       self.job_name, n, str(d).rstrip()))
+                       self.job_name, n, emsg))
         except Exception, d:
             self.send_stats()
+            emsg = str(d).rstrip()
             self.log.exception("Job %s crashed: %s" % (
-                       self.job_name, str(d).rstrip()))
+                       self.job_name, emsg))
 
         # reset and sleep
         self.reset()
-        if prefer_looping and self.looping and not self.do_single_loop:
+        self.exception_hook(d, emsg, cname)
+        if prefer_looping and self.loop_delay:
             time.sleep(20)
-            return 1
+            return -1
         sys.exit(1)
 
+    def exception_hook(self, det, emsg, cname):
+        """Called on after exception processing.
+
+        Can do additional logging.
+
+        @param det: exception details
+        @param emsg: exception msg
+        @param cname: connection name or None
+        """
+        pass
+
     def work(self):
         """Here should user's processing happen.