From 540df5829ebad62f6d7b54df302d1abc8f475d6a Mon Sep 17 00:00:00 2001 From: Tyler Hobbs Date: Thu, 17 Apr 2014 11:00:15 -0500 Subject: [PATCH] Add missing timeout to register_watcher for gevent --- cassandra/io/geventreactor.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/cassandra/io/geventreactor.py b/cassandra/io/geventreactor.py index 322b9c00..27b92979 100644 --- a/cassandra/io/geventreactor.py +++ b/cassandra/io/geventreactor.py @@ -194,11 +194,15 @@ class GeventConnection(Connection): for i in xrange(0, len(data), chunk_size): self._write_queue.put(data[i:i + chunk_size]) - def register_watcher(self, event_type, callback): + def register_watcher(self, event_type, callback, register_timeout=None): self._push_watchers[event_type].add(callback) - self.wait_for_response(RegisterMessage(event_list=[event_type])) + self.wait_for_response( + RegisterMessage(event_list=[event_type]), + timeout=register_timeout) - def register_watchers(self, type_callback_dict): + def register_watchers(self, type_callback_dict, register_timeout=None): for event_type, callback in type_callback_dict.items(): self._push_watchers[event_type].add(callback) - self.wait_for_response(RegisterMessage(event_list=type_callback_dict.keys())) + self.wait_for_response( + RegisterMessage(event_list=type_callback_dict.keys()), + timeout=register_timeout)