Initialize Event properly in BaseConsumer

2 files changed, 7 insertions, 5 deletions

diff --git a/python/pgq/baseconsumer.py b/python/pgq/baseconsumer.py
index 7648a160..9159106e 100644
--- a/python/pgq/baseconsumer.py
+++ b/python/pgq/baseconsumer.py
@@ -24,8 +24,6 @@ class BaseBatchWalker(object):
      - len() after that
     """
 
-    _event_class = Event
-
     def __init__(self, curs, batch_id, queue_name, fetch_size = 300, consumer_filter = None):
         self.queue_name = queue_name
         self.fetch_size = fetch_size
@@ -36,6 +34,9 @@ class BaseBatchWalker(object):
         self.fetch_status = 0 # 0-not started, 1-in-progress, 2-done
         self.consumer_filter = consumer_filter
 
+    def _make_event(self, queue_name, row):
+        return Event(queue_name, row)
+
     def __iter__(self):
         if self.fetch_status:
             raise Exception("BatchWalker: double fetch? (%d)" % self.fetch_status)
@@ -53,7 +54,7 @@ class BaseBatchWalker(object):
 
             self.length += len(rows)
             for row in rows:
-                ev = self._event_class(self, self.queue_name, row)
+                ev = self._make_event(self.queue_name, row)
                 yield ev
 
             # if less rows than requested, it was final block
diff --git a/python/pgq/consumer.py b/python/pgq/consumer.py
index 19fbf5ef..3d78105f 100644
--- a/python/pgq/consumer.py
+++ b/python/pgq/consumer.py
@@ -31,12 +31,13 @@ class RetriableBatchWalker(BaseBatchWalker):
     """BatchWalker that returns RetriableEvents
     """
 
-    _event_class = RetriableWalkerEvent
-
     def __init__(self, curs, batch_id, queue_name, fetch_size = 300, consumer_filter = None):
         super(RetriableBatchWalker, self).__init__(self, curs, batch_id, queue_name, fetch_size, consumer_filter)
         self.status_map = {}
 
+    def _make_event(self, queue_name, row):
+        return RetriableWalkerEvent(self, queue_name, row)
+
     def tag_event_done(self, event):
         if event.id in self.status_map:
             del self.status_map[event.id]