Merge pull request apache#285 from datastax/PYTHON-206

aholmberg · aholmberg · commit 95b6c13ef4c8 · 2015-04-24T10:07:03.000-05:00
PYTHON-206 - Connect timeouts
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -426,6 +426,14 @@ def auth_provider(self, value):
     See :attr:`.schema_event_refresh_window` for discussion of rationale
     """
 
+    connect_timeout = 5
+    """
+    Timeout, in seconds, for creating new connections.
+
+    This timeout covers the entire connection negotiation, including TCP
+    establishment, options passing, and authentication.
+    """
+
     sessions = None
     control_connection = None
     scheduler = None
@@ -464,7 +472,8 @@ def __init__(self,
                  control_connection_timeout=2.0,
                  idle_heartbeat_interval=30,
                  schema_event_refresh_window=2,
-                 topology_event_refresh_window=10):
+                 topology_event_refresh_window=10,
+                 connect_timeout=5):
         """
         Any of the mutable Cluster attributes may be set as keyword arguments
         to the constructor.
@@ -517,6 +526,7 @@ def __init__(self,
         self.idle_heartbeat_interval = idle_heartbeat_interval
         self.schema_event_refresh_window = schema_event_refresh_window
         self.topology_event_refresh_window = topology_event_refresh_window
+        self.connect_timeout = connect_timeout
 
         self._listeners = set()
         self._listener_lock = Lock()
@@ -706,11 +716,11 @@ def connection_factory(self, address, *args, **kwargs):
         Intended for internal use only.
         """
         kwargs = self._make_connection_kwargs(address, kwargs)
-        return self.connection_class.factory(address, *args, **kwargs)
+        return self.connection_class.factory(address, self.connect_timeout, *args, **kwargs)
 
     def _make_connection_factory(self, host, *args, **kwargs):
         kwargs = self._make_connection_kwargs(host.address, kwargs)
-        return partial(self.connection_class.factory, host.address, *args, **kwargs)
+        return partial(self.connection_class.factory, host.address, self.connect_timeout, *args, **kwargs)
 
     def _make_connection_kwargs(self, address, kwargs_dict):
         if self._auth_provider_callable:
diff --git a/cassandra/connection.py b/cassandra/connection.py
@@ -239,6 +239,23 @@ def handle_fork(self):
         """
         pass
 
+    @classmethod
+    def factory(cls, host, timeout, *args, **kwargs):
+        """
+        A factory function which returns connections which have
+        succeeded in connecting and are ready for service (or
+        raises an exception otherwise).
+        """
+        conn = cls(host, *args, **kwargs)
+        conn.connected_event.wait(timeout)
+        if conn.last_error:
+            raise conn.last_error
+        elif not conn.connected_event.is_set():
+            conn.close()
+            raise OperationTimedOut("Timed out creating connection (%s seconds)" % timeout)
+        else:
+            return conn
+
     def close(self):
         raise NotImplementedError()
 
diff --git a/cassandra/io/asyncorereactor.py b/cassandra/io/asyncorereactor.py
@@ -156,19 +156,6 @@ def handle_fork(cls):
             cls._loop._cleanup()
             cls._loop = None
 
-    @classmethod
-    def factory(cls, *args, **kwargs):
-        timeout = kwargs.pop('timeout', 5.0)
-        conn = cls(*args, **kwargs)
-        conn.connected_event.wait(timeout)
-        if conn.last_error:
-            raise conn.last_error
-        elif not conn.connected_event.is_set():
-            conn.close()
-            raise OperationTimedOut("Timed out creating connection")
-        else:
-            return conn
-
     def __init__(self, *args, **kwargs):
         Connection.__init__(self, *args, **kwargs)
         asyncore.dispatcher.__init__(self)
diff --git a/cassandra/io/eventletreactor.py b/cassandra/io/eventletreactor.py
@@ -57,19 +57,6 @@ class EventletConnection(Connection):
     def initialize_reactor(cls):
         eventlet.monkey_patch()
 
-    @classmethod
-    def factory(cls, *args, **kwargs):
-        timeout = kwargs.pop('timeout', 5.0)
-        conn = cls(*args, **kwargs)
-        conn.connected_event.wait(timeout)
-        if conn.last_error:
-            raise conn.last_error
-        elif not conn.connected_event.is_set():
-            conn.close()
-            raise OperationTimedOut("Timed out creating connection")
-        else:
-            return conn
-
     def __init__(self, *args, **kwargs):
         Connection.__init__(self, *args, **kwargs)
 
diff --git a/cassandra/io/geventreactor.py b/cassandra/io/geventreactor.py
@@ -50,19 +50,6 @@ class GeventConnection(Connection):
     _write_watcher = None
     _socket = None
 
-    @classmethod
-    def factory(cls, *args, **kwargs):
-        timeout = kwargs.pop('timeout', 5.0)
-        conn = cls(*args, **kwargs)
-        conn.connected_event.wait(timeout)
-        if conn.last_error:
-            raise conn.last_error
-        elif not conn.connected_event.is_set():
-            conn.close()
-            raise OperationTimedOut("Timed out creating connection")
-        else:
-            return conn
-
     def __init__(self, *args, **kwargs):
         Connection.__init__(self, *args, **kwargs)
 
diff --git a/cassandra/io/libevreactor.py b/cassandra/io/libevreactor.py
@@ -236,19 +236,6 @@ def handle_fork(cls):
             cls._libevloop._cleanup()
             cls._libevloop = None
 
-    @classmethod
-    def factory(cls, *args, **kwargs):
-        timeout = kwargs.pop('timeout', 5.0)
-        conn = cls(*args, **kwargs)
-        conn.connected_event.wait(timeout)
-        if conn.last_error:
-            raise conn.last_error
-        elif not conn.connected_event.is_set():
-            conn.close()
-            raise OperationTimedOut("Timed out creating new connection")
-        else:
-            return conn
-
     def __init__(self, *args, **kwargs):
         Connection.__init__(self, *args, **kwargs)
 
diff --git a/cassandra/io/twistedreactor.py b/cassandra/io/twistedreactor.py
@@ -148,24 +148,6 @@ def initialize_reactor(cls):
         if not cls._loop:
             cls._loop = TwistedLoop()
 
-    @classmethod
-    def factory(cls, *args, **kwargs):
-        """
-        A factory function which returns connections which have
-        succeeded in connecting and are ready for service (or
-        raises an exception otherwise).
-        """
-        timeout = kwargs.pop('timeout', 5.0)
-        conn = cls(*args, **kwargs)
-        conn.connected_event.wait(timeout)
-        if conn.last_error:
-            raise conn.last_error
-        elif not conn.connected_event.is_set():
-            conn.close()
-            raise OperationTimedOut("Timed out creating connection")
-        else:
-            return conn
-
     def __init__(self, *args, **kwargs):
         """
         Initialization method.
diff --git a/tests/integration/standard/test_connection.py b/tests/integration/standard/test_connection.py
@@ -58,7 +58,7 @@ def get_connection(self):
         e = None
         for i in range(5):
             try:
-                conn = self.klass.factory(protocol_version=PROTOCOL_VERSION)
+                conn = self.klass.factory(host='127.0.0.1', timeout=5, protocol_version=PROTOCOL_VERSION)
                 break
             except (OperationTimedOut, NoHostAvailable) as e:
                 continue
