pgq.CoopConsumer for Python

Simply wrapper around Consumer which redirects few calls to pgq_coop schema.

1 files changed, 7 insertions, 2 deletions

diff --git a/python/pgq/consumer.py b/python/pgq/consumer.py
index a23b882e..32fa9f76 100644
--- a/python/pgq/consumer.py
+++ b/python/pgq/consumer.py
@@ -245,8 +245,8 @@ class Consumer(skytools.DBScript):
         curs.execute(q, [self.queue_name, self.consumer_name])
         return curs.fetchone()[0]
 
-    def _finish_batch(self, curs, batch_id, list):
-        """Tag events and notify that the batch is done."""
+    def _flush_retry(self, curs, list):
+        """Tag retry events."""
 
         retry = 0
         if self.pgq_lazy_fetch:
@@ -269,6 +269,11 @@ class Consumer(skytools.DBScript):
         if retry:
             self.stat_increase('retry-events', retry)
 
+    def _finish_batch(self, curs, batch_id, list):
+        """Tag events and notify that the batch is done."""
+
+        self._flush_retry(curs, batch_id, list)
+
         curs.execute("select pgq.finish_batch(%s)", [batch_id])
 
     def _tag_retry(self, cx, batch_id, ev_id, retry_time):