Add connect_timeout to cluster and Connection.factory

aholmberg · aholmberg · commit b94e31be846c · 2015-04-21T14:55:37.000-05:00
PYTHON-206
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -427,6 +427,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
@@ -465,7 +473,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.
@@ -518,6 +527,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()
@@ -707,11 +717,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
@@ -240,20 +240,19 @@ def handle_fork(self):
         pass
 
     @classmethod
-    def factory(cls, *args, **kwargs):
+    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).
         """
-        timeout = kwargs.pop('timeout', 5.0)
-        conn = cls(*args, **kwargs)
+        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")
+            raise OperationTimedOut("Timed out creating connection (%s seconds)" % timeout)
         else:
             return conn
 
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
