--- /dev/null
+
+"""
+Consumer that stores last applied position in local file.
+
+For cases where the consumer cannot use single database for remote tracking.
+
+"""
+
+import sys
+import os
+import errno
+import skytools
+import pgq
+
+__all__ = ['LocalConsumer']
+
+class LocalConsumer(pgq.Consumer):
+ """Consumer that applies batches sequentially in second database.
+
+ Requirements:
+ - Whole batch in one TX.
+ - Must not use retry queue.
+
+ Features:
+ - Can detect if several batches are already applied to dest db.
+ - If some ticks are lost. allows to seek back on queue.
+ Whether it succeeds, depends on pgq configuration.
+
+ Config options::
+
+ ## Parameters for LocalConsumer ##
+
+ # file location where last applied tick is tracked
+ local_tracking_file = ~/state/%(job_name)s.tick
+ """
+
+ def reload(self):
+ super(LocalConsumer, self).reload()
+
+ self.local_tracking_file = self.cf.getfile('local_tracking_file')
+
+ def init_optparse(self, parser = None):
+ p = super(LocalConsumer, self).init_optparse(parser)
+ p.add_option("--rewind", action = "store_true",
+ help = "change queue position according to local tick")
+ p.add_option("--reset", action = "store_true",
+ help = "reset local tick based on queue position")
+ return p
+
+ def startup(self):
+ if self.options.rewind:
+ self.rewind()
+ sys.exit(0)
+ if self.options.reset:
+ self.dst_reset()
+ sys.exit(0)
+ super(LocalConsumer, self).startup()
+
+ self.check_queue()
+
+ def check_queue(self):
+ queue_tick = -1
+ local_tick = self.load_local_tick()
+
+ db = self.get_database(self.db_name)
+ curs = db.cursor()
+ q = "select last_tick from pgq.get_consumer_info(%s, %s)"
+ curs.execute(q, [self.queue_name, self.consumer_name])
+ rows = curs.fetchall()
+ if len(rows) == 1:
+ queue_tick = rows[0]['last_tick']
+ db.commit()
+
+ if queue_tick < 0:
+ if local_tick >= 0:
+ self.log.info("Registering consumer at tick %d", local_tick)
+ q = "select * from pgq.register_consumer_at(%s, %s, %s)"
+ curs.execute(q, [self.queue_name, self.consumer_name, local_tick])
+ else:
+ self.log.info("Registering consumer at queue top")
+ q = "select * from pgq.register_consumer(%s, %s)"
+ curs.execute(q, [self.queue_name, self.consumer_name])
+ elif local_tick < 0:
+ self.log.info("Local tick missing, storing queueu tick %d", queue_tick)
+ self.save_local_tick(queue_tick)
+ elif local_tick > queue_tick:
+ self.log.warning("Tracking out of sync: queue=%d local=%d. Repositioning on queue. [Database failure?]",
+ queue_tick, local_tick)
+ q = "select * from pgq.register_consumer_at(%s, %s, %s)"
+ curs.execute(q, [self.queue_name, self.consumer_name, local_tick])
+ elif local_tick < queue_tick:
+ self.log.warning("Tracking out of sync: queue=%d local=%d. Rewinding queue. [Lost file data?]",
+ queue_tick, local_tick)
+ q = "select * from pgq.register_consumer_at(%s, %s, %s)"
+ curs.execute(q, [self.queue_name, self.consumer_name, local_tick])
+ else:
+ self.log.info("Ticks match: Queue=%d Local=%d", queue_tick, local_tick)
+
+ def work(self):
+ if self.work_state < 0:
+ self.check_queue()
+ return super(LocalConsumer, self).work()
+
+ def process_batch(self, db, batch_id, event_list):
+ """Process all events in batch.
+ """
+
+ # check if done
+ if self.is_batch_done():
+ return
+
+ # actual work
+ self.process_local_batch(db, batch_id, event_list)
+
+ # finish work
+ self.set_batch_done()
+
+ def process_local_batch(self, db, batch_id, event_list):
+ for ev in event_list:
+ self.process_local_event(db, batch_id, ev)
+
+ def process_local_event(self, db, batch_id, ev):
+ raise Exception('process_remote_batch not implemented')
+
+ def is_batch_done(self):
+ """Helper function to keep track of last successful batch
+ in external database.
+ """
+
+ local_tick = self.load_local_tick()
+
+ cur_tick = self.batch_info['tick_id']
+ prev_tick = self.batch_info['prev_tick_id']
+
+ if local_tick < 0:
+ # seems this consumer has not run yet?
+ return False
+
+ if prev_tick == local_tick:
+ # on track
+ return False
+
+ if cur_tick == local_tick:
+ # current batch is already applied, skip it
+ return True
+
+ # anything else means problems
+ raise Exception('Lost position: batch %d..%d, dst has %d' % (
+ prev_tick, cur_tick, local_tick))
+
+ def set_batch_done(self):
+ """Helper function to set last successful batch
+ in external database.
+ """
+ tick_id = self.batch_info['tick_id']
+ self.save_local_tick(tick_id)
+
+ def register_consumer(self):
+ new = super(LocalConsumer, self).register_consumer()
+ if new: # fixme
+ self.dst_reset()
+
+ def unregister_consumer(self):
+ """If unregistering, also clean completed tick table on dest."""
+
+ super(LocalConsumer, self).unregister_consumer()
+ self.dst_reset()
+
+ def rewind(self):
+ dst_tick = self.load_local_tick()
+ if dst_tick >= 0:
+ src_db = self.get_database(self.db_name)
+ src_curs = src_db.cursor()
+
+ self.log.info("Rewinding queue to tick local tick %d", dst_tick)
+ q = "select pgq.register_consumer_at(%s, %s, %s)"
+ src_curs.execute(q, [self.queue_name, self.consumer_name, dst_tick])
+
+ src_db.commit()
+ else:
+ self.log.error('Cannot rewind, no tick found in local file')
+
+ def dst_reset(self):
+ self.log.info("Removing local tracking file")
+ try:
+ os.remove(self.local_tracking_file)
+ except:
+ pass
+
+ def load_local_tick(self):
+ """Reads stored tick or -1."""
+ try:
+ f = open(self.local_tracking_file, 'r')
+ buf = f.read()
+ f.close()
+ data = buf.strip()
+ if data:
+ tick_id = int(data)
+ else:
+ tick_id = -1
+ return tick_id
+ except IOError, ex:
+ if ex.errno == errno.ENOENT:
+ return -1
+ raise
+
+ def save_local_tick(self, tick_id):
+ """Store tick in local file."""
+ data = str(tick_id)
+ skytools.write_atomic(self.local_tracking_file, data)
+
projects / skytools.git / commitdiff