Rework connection handling, and add support for replSets

Mike Dirolf · Mike Dirolf · commit 0726f2154337 · 2010-07-15T11:58:10.000-04:00
diff --git a/pymongo/connection.py b/pymongo/connection.py
@@ -186,6 +186,8 @@ def __init__(self, host=None, port=None, pool_size=None,
             database = db or database
             username = u or username
             password = p or password
+        if not nodes:
+            raise ConfigurationError("need to specify at least one host")
         self.__nodes = nodes
         if database and username is None:
             raise InvalidURI("cannot specify database without "
@@ -206,6 +208,9 @@ def __init__(self, host=None, port=None, pool_size=None,
         self.__port = None
 
         self.__slave_okay = slave_okay
+        if slave_okay and len(self.__nodes) > 1:
+            raise ConfigurationError("cannot specify slave_okay for a paired "
+                                     "or replica set connection")
 
         self.__cursor_manager = CursorManager(self)
 
@@ -309,16 +314,13 @@ def paired(cls, left, right=None, **connection_args):
         The remaining keyword arguments are the same as those accepted
         by :meth:`~Connection`.
         """
+        if isinstance(left, str) or isinstance(right, str):
+            raise TypeError("arguments to paired must be tuples")
         if right is None:
             right = (cls.HOST, cls.PORT)
         return cls([":".join(map(str, left)), ":".join(map(str, right))],
                    **connection_args)
 
-    def __master(self, sock):
-        """Is this socket connected to a master server?
-        """
-        return self["admin"].command("ismaster", _sock=sock)["ismaster"]
-
     def _cache_index(self, database, collection, index, ttl):
         """Add an index to the index cache for ensure_index operations.
 
@@ -420,63 +422,79 @@ def tz_aware(self):
         """
         return self.__tz_aware
 
+    def __add_hosts_and_get_primary(self, response):
+        if "hosts" in response:
+            self.__nodes.update([h.split(":") for h in response["hosts"]])
+        return response.get("primary", False)
+
+    def __try_node(self, node):
+        self.disconnect()
+        self.__host, self.__port = node
+        try:
+            response = self.admin.command("isMaster")
+            self.end_request()
+
+            primary = self.__add_hosts_and_get_primary(response)
+            if response["ismaster"]:
+                return True
+            return primary
+        except:
+            self.end_request()
+            return None
+
     def __find_master(self):
         """Create a new socket and use it to figure out who the master is.
 
         Sets __host and __port so that :attr:`host` and :attr:`port`
         will return the address of the master. Also (possibly) updates
         any replSet information.
         """
-        self.__host = None
-        self.__port = None
-        sock = None
-        sock_error = False
-        close = True
-
+        # Special case the first node to try to get the primary or any
+        # additional hosts from a replSet:
+        first = iter(self.__nodes).next()
+
+        primary = self.__try_node(first)
+        if primary is True:
+            return first
+        if self.__slave_okay and primary is not None: # no network error
+            return first
+
+        # Wasn't the first node, but we got a primary - let's try it:
+        tried = [first]
+        if primary:
+            if self.__try_node(primary) is True:
+                return primary
+            tried.append(primary)
+
+        nodes = self.__nodes - set(tried)
+
+        # Just scan
         # TODO parallelize these to minimize connect time?
-        for (host, port) in self.__nodes:
-            try:
-                try:
-                    sock = socket.socket()
-                    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-                    sock.settimeout(self.__network_timeout or _CONNECT_TIMEOUT)
-                    sock.connect((host, port))
-                    sock.settimeout(self.__network_timeout)
-                    master = self.__master(sock)
-                    if master or self.__slave_okay:
-                        self.__host = host
-                        self.__port = port
-                        self.__pool.return_unowned(sock)
-                        close = False
-                        return
-                except socket.error, e:
-                    sock_error = True
-            finally:
-                if sock is not None and close:
-                    sock.close()
-        if sock_error or self.__host is None:
-            raise AutoReconnect("could not find master")
-        raise ConfigurationError("No master node in %r. You must specify "
-                                 "slave_okay to connect to "
-                                 "slaves." % self.__nodes)
+        for node in nodes:
+            if self.__try_node(node) is True:
+                return node
+
+        raise AutoReconnect("could not find master/primary")
 
     def __connect(self):
         """(Re-)connect to Mongo and return a new (connected) socket.
 
         Connect to the master if this is a paired connection.
         """
-        if self.__host is None or self.__port is None:
-            self.__find_master()
+        host, port = (self.__host, self.__port)
+        if host is None or port is None:
+            host, port = self.__find_master()
 
         try:
             sock = socket.socket()
             sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
             sock.settimeout(self.__network_timeout or _CONNECT_TIMEOUT)
-            sock.connect((self.__host, self.__port))
+            sock.connect((host, port))
             sock.settimeout(self.__network_timeout)
             return sock
         except socket.error:
-            raise AutoReconnect("could not connect to %r" % self.__nodes)
+            self.disconnect()
+            raise AutoReconnect("could not connect to %r" % list(self.__nodes))
 
     def disconnect(self):
         """Disconnect from MongoDB.
@@ -492,17 +510,8 @@ def disconnect(self):
         .. versionadded:: 1.3
         """
         self.__pool = Pool(self.__connect)
-
-    def _reset(self):
-        """Reset everything and start connecting again.
-
-        Closes all open sockets and resets them to None. Re-finds the master.
-
-        This should be done in case of a connection failure or a "not master"
-        error.
-        """
-        self.disconnect()
-        self.__find_master()
+        self.__host = None
+        self.__port = None
 
     def set_cursor_manager(self, manager_class):
         """Set this connection's cursor manager.
@@ -540,7 +549,7 @@ def __check_response_to_last_error(self, response):
         if error.get("err", 0) is None:
             return error
         if error["err"] == "not master":
-            self._reset()
+            self.disconnect()
 
         if "code" in error:
             if error["code"] in [11000, 11001]:
@@ -580,7 +589,7 @@ def _send_message(self, message, with_last_error=False):
                 return self.__check_response_to_last_error(response)
             return None
         except (ConnectionFailure, socket.error), e:
-            self._reset()
+            self.disconnect()
             raise AutoReconnect(str(e))
 
     def __receive_data_on_socket(self, length, sock):
@@ -642,7 +651,7 @@ def _send_message_with_response(self, message, _sock=None,
                 return self.__send_and_receive(message, _sock)
             except (ConnectionFailure, socket.error), e:
                 if reset:
-                    self._reset()
+                    self.disconnect()
                 raise AutoReconnect(str(e))
         finally:
             if "network_timeout" in kwargs:
@@ -687,12 +696,8 @@ def __cmp__(self, other):
     def __repr__(self):
         if len(self.__nodes) == 1:
             return "Connection(%r, %r)" % (self.__host, self.__port)
-        elif len(self.__nodes) == 2:
-            return ("Connection.paired((%r, %r), (%r, %r))" %
-                    (self.__nodes[0][0],
-                     self.__nodes[0][1],
-                     self.__nodes[1][0],
-                     self.__nodes[1][1]))
+        else:
+            return "Connection(%r)" % ["%s:%d" % n for n in self.__nodes]
 
     def __getattr__(self, name):
         """Get a database by name.
diff --git a/pymongo/cursor.py b/pymongo/cursor.py
@@ -495,7 +495,7 @@ def __send_message(self, message):
             response = helpers._unpack_response(response, self.__id,
                                                 self.__as_class, self.__tz_aware)
         except AutoReconnect:
-            db.connection._reset()
+            db.connection.disconnect()
             raise
         self.__id = response["cursor_id"]
 
diff --git a/pymongo/database.py b/pymongo/database.py
@@ -391,7 +391,7 @@ def error(self):
         if error.get("err", 0) is None:
             return None
         if error["err"] == "not master":
-            self.__connection._reset()
+            self.__connection.disconnect()
         return error
 
     def last_status(self):
diff --git a/test/test_connection.py b/test/test_connection.py
@@ -56,6 +56,8 @@ def test_types(self):
         self.assertRaises(TypeError, Connection, "localhost", 1.14)
         self.assertRaises(TypeError, Connection, "localhost", [])
 
+        self.assertRaises(ConfigurationError, Connection, [])
+
     def test_constants(self):
         Connection.HOST = self.host
         Connection.PORT = self.port
diff --git a/test/test_paired.py b/test/test_paired.py
@@ -28,7 +28,7 @@
 import warnings
 sys.path[0:0] = [""]
 
-from pymongo.errors import ConnectionFailure, ConfigurationError
+from pymongo.errors import ConnectionFailure
 from pymongo.connection import Connection
 
 skip_tests = True
@@ -83,17 +83,17 @@ def test_connect(self):
         self.assertEqual(port, connection.port)
 
         slave = self.left == (host, port) and self.right or self.left
-        self.assertRaises(ConfigurationError, Connection.paired,
+        self.assertRaises(ConnectionFailure, Connection.paired,
                           slave, self.bad)
-        self.assertRaises(ConfigurationError, Connection.paired,
+        self.assertRaises(ConnectionFailure, Connection.paired,
                           self.bad, slave)
 
     def test_repr(self):
         self.skip()
         connection = Connection.paired(self.left, self.right)
 
         self.assertEqual(repr(connection),
-                         "Connection.paired(('%s', %s), ('%s', %s))" %
+                         "Connection(['%s:%s', '%s:%s'])" %
                          (self.left[0],
                           self.left[1],
                           self.right[0],
diff --git a/tools/auto_reconnect_test.py b/tools/auto_reconnect_test.py
@@ -17,7 +17,7 @@
 import threading
 import time
 
-from pymongo.errors import ConnectionFailure
+from pymongo.errors import AutoReconnect
 from pymongo.connection import Connection
 
 db = Connection.paired(("localhost", 27018)).test
@@ -28,12 +28,12 @@ def run(self):
         while True:
             time.sleep(1)
             try:
-                id = db.test.save({"x": 1})
+                id = db.test.save({"x": 1}, safe=True)
                 assert db.test.find_one(id)["x"] == 1
                 db.test.remove(id)
                 db.connection.end_request()
                 print "Y"
-            except ConnectionFailure, e:
+            except Exception, e:
                 print e
                 print "N"
 
