Update loopquery: allow any number of queries, support reading parameters from config file.
authorRadek Brich <radek.brich@devl.cz>
Thu, 31 Jan 2013 11:02:04 +0100
changeset 62 af637235ca81
parent 61 703bba757605
child 63 8c7f0a51ba50
Update loopquery: allow any number of queries, support reading parameters from config file.
loopquery.py
pgtoolkit/toolbase.py
--- a/loopquery.py	Fri Jan 25 18:04:17 2013 +0100
+++ b/loopquery.py	Thu Jan 31 11:02:04 2013 +0100
@@ -2,29 +2,91 @@
 """
 loopquery
 
-Execute some query in loop, waiting for some time interval before next run.
+Execute some queries in loop. Execute interval can be set.
 
 """
 
-from pgtoolkit import toolbase
+from pgtoolkit import toolbase, config
 
+import logging.handlers
 import time
+from datetime import datetime, timedelta
 
 
-class LoopQueryTool(toolbase.SimpleTool):
+class LoopQueryTool(toolbase.ToolBase):
     def __init__(self):
-        toolbase.SimpleTool.__init__(self, name='loopquery', desc='Run query in loop.')
-        self.parser.add_argument('-q', dest='query', type=str, help='Query to run.')
-        self.parser.add_argument('-i', dest='interval', type=float, help='Run interval in seconds.')
+        toolbase.ToolBase.__init__(self, name='loopquery', desc='Run query in loop.')
+        self.parser.add_argument('target', nargs='?', metavar='target', type=str, help='Target database')
+        self.parser.add_argument('-c', dest='config', type=str, help='Additional config file (besides pgtoolkit.conf).')
+        self.parser.add_argument('-q', dest='queries', nargs='*', help='Queries to run.')
+        self.parser.add_argument('--mins', dest='delay_mins', type=int, help='Delay between queries in minutes.')
+        self.parser.add_argument('--secs', dest='delay_secs', type=int, help='Delay between queries in seconds.')
+
+        self.config.add_argument('target', type=str, default=None)
+        self.config.add_argument('queries', type=list, default=[])
+        self.config.add_argument('delay_mins', type=int, default=0)
+        self.config.add_argument('delay_secs', type=int, default=0)
+        self.config.add_argument('log_path', type=str)
+
+        self.target_isolation_level = 'autocommit'
+
         self.init()
 
+    def init(self):
+        toolbase.ToolBase.init(self)
+        if self.args.config:
+            self.config.load(self.args.config)
+        self.queries = self.args.queries or self.config.queries
+        self.delay_mins = self.config.delay_mins
+        self.delay_secs = self.args.delay_secs or self.config.delay_secs
+        if self.config.log_path:
+            self.init_file_logs(self.config.log_path)
+        self.prepare_conns(target = self.args.target or self.config.target)
+
+    def init_file_logs(self, path):
+        format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')
+        handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)
+        handler.setFormatter(format)
+        handler.setLevel(logging.DEBUG)
+        logging.getLogger('main').addHandler(handler)
+
+        format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S')
+        handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5)
+        handler.setFormatter(format)
+        handler.setLevel(logging.DEBUG)
+        logging.getLogger('pgmanager_notices').addHandler(handler)
+
     def main(self):
+        self.prepare()
         while True:
-            self.log.info('Executing: %s', self.args.query)
+            self.wait()
+            self.action()
+
+    def prepare(self):
+        """Check current time, set next action time."""
+        dt = datetime.today()
+        dt = dt.replace(second = 0, microsecond = 0)
+        self.next_action_time = dt + timedelta(minutes = self.delay_mins,
+                                               seconds = self.delay_secs)
+
+    def wait(self):
+        """Wait for action time, compute next action time."""
+        now = datetime.today()
+        self.log.debug('Next run %s', self.next_action_time)
+        if self.next_action_time > now:
+            td = self.next_action_time - now
+            self.log.debug('Sleep %ds', td.seconds + td.microseconds/1e6)
+            time.sleep(td.seconds + td.microseconds/1e6)
+        self.next_action_time += timedelta(minutes = self.delay_mins,
+                                           seconds = self.delay_secs)
+
+    def action(self):
+        """Execute the queries."""
+        for q in self.queries:
+            self.log.info('%s', q)
             with self.pgm.cursor('target') as curs:
-                curs.execute(self.args.query)
-            self.log.info('Done.')
-            time.sleep(self.args.interval)
+                curs.execute(q)
+        self.log.info('Done')
 
 
 tool = LoopQueryTool()
--- a/pgtoolkit/toolbase.py	Fri Jan 25 18:04:17 2013 +0100
+++ b/pgtoolkit/toolbase.py	Thu Jan 31 11:02:04 2013 +0100
@@ -36,22 +36,22 @@
 
     def init_logging(self):
         # logging
+        format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S')
         handler = logging.StreamHandler()
-        format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S')
         handler.setFormatter(format)
         handler.setLevel(logging.DEBUG)
         self.log = logging.getLogger('main')
         self.log.addHandler(handler)
         self.log.setLevel(logging.DEBUG)
 
-        logger_notices = logging.getLogger('pgmanager_notices')
-        logger_notices.addHandler(handler)
-        logger_notices.setLevel(logging.DEBUG)
+        log_notices = logging.getLogger('pgmanager_notices')
+        log_notices.addHandler(handler)
+        log_notices.setLevel(logging.DEBUG)
 
         if self.args.debug:
-            logger_sql = logging.getLogger('pgmanager_sql')
-            logger_sql.addHandler(handler)
-            logger_sql.setLevel(logging.DEBUG)
+            log_sql = logging.getLogger('pgmanager_sql')
+            log_sql.addHandler(handler)
+            log_sql.setLevel(logging.DEBUG)
 
     def prepare_conn_from_metadb(self, name, lookup_name):
         '''Create connection in pgmanager using meta DB.
@@ -80,12 +80,19 @@
                     dsn=dsn)
                 return True
 
-    def prepare_conns_from_cmdline_args(self, *pgm_names):
+    def prepare_conns(self, **kwargs):
+        """Create connections in PgManager.
+
+        Keyword arguments meaning:
+            key: connection name for use in PgManager
+            value: connection name in config or meta DB
+
+        """
         if self.config.meta_db:
             self.pgm.create_conn(name='meta', dsn=self.config.meta_db)
 
-        for name in pgm_names:
-            lookup_name = self.args.__dict__[name]
+        for name in kwargs:
+            lookup_name = kwargs[name]
             found = self.prepare_conn_from_config(name, lookup_name)
             if not found and self.config.meta_db:
                 found = self.prepare_conn_from_metadb(name, lookup_name)
@@ -103,7 +110,7 @@
 
     def init(self):
         ToolBase.init(self)
-        self.prepare_conns_from_cmdline_args('target')
+        self.prepare_conns(target=self.args.target)
 
 
 class SrcDstTool(ToolBase):
@@ -118,7 +125,7 @@
         ToolBase.init(self)
         if self.is_reversed():
             self.args.src, self.args.dst = self.args.dst, self.args.src
-        self.prepare_conns_from_cmdline_args('src', 'dst')
+        self.prepare_conns(src=self.args.src, dst=self.args.dst)
 
     def is_reversed(self):
         return 'reverse' in self.args and self.args.reverse