Add basic IPv6 support PYTHON-237

behackett · behackett · commit 5aa62a3c0e00 · 2011-04-12T10:52:25.000-07:00
commit a6ee56d8bbfdf9e0ca651f5efa6fb71556c650ed
Author: behackett &lt;bernie@10gen.com&gt;
Date:   Tue Apr 12 10:16:08 2011 -0700

    Make IPv4 explicit and add a few more tests.

commit 793584a1c7032f2d2b7f5ae9ca0740ceec84ba98
Author: behackett &lt;bernie@10gen.com&gt;
Date:   Mon Apr 11 16:51:18 2011 -0700

    Preliminary support for IPv6.
diff --git a/pymongo/connection.py b/pymongo/connection.py
@@ -58,6 +58,16 @@
 _CONNECT_TIMEOUT = 20.0
 
 
+def _partition_ipv6(source):
+    if source.find(']') == -1:
+        raise InvalidURI("an IPv6 address literal must be "
+                         "enclosed in '[' and ']' characters.")
+    i = source.find(']:')
+    if i == -1:
+        return (source[1:-1], None)
+    return (source[1: i], source[i + 2:])
+
+
 def _partition(source, sub):
     """Our own string partitioning method.
 
@@ -74,7 +84,14 @@ def _str_to_node(string, default_port=27017):
 
     "localhost:27017" -> ("localhost", 27017)
     """
-    (host, port) = _partition(string, ":")
+    # IPv6 literal
+    if string[0] == '[':
+        host, port = _partition_ipv6(string)
+    elif string.count(':') > 1 or string.find(']') != -1:
+        raise InvalidURI("an IPv6 address literal must be "
+                         "enclosed in '[' and ']' characters.")
+    else:
+        host, port = _partition(string, ":")
     if port:
         port = int(port)
     else:
@@ -240,9 +257,11 @@ def __init__(self, host=None, port=None, pool_size=None,
         username, and password present will be used.
 
         :Parameters:
-          - `host` (optional): hostname or IPv4 address of the
+          - `host` (optional): hostname or IP address of the
             instance to connect to, or a mongodb URI, or a list of
-            hostnames / mongodb URIs
+            hostnames / mongodb URIs. If `host` is an IPv6 literal
+            it must be enclosed in '[' and ']' characters following
+            the RFC2732 URL syntax (e.g. '[::1]' for localhost)
           - `port` (optional): port number on which to connect
           - `pool_size` (optional): DEPRECATED
           - `auto_start_request` (optional): DEPRECATED
@@ -585,12 +604,23 @@ def __connect(self):
             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((host, port))
-            sock.settimeout(self.__network_timeout)
-            return sock
+            try:
+                # Prefer IPv4. If there is demand for an option
+                # to specify one or the other we can add it later.
+                sock = socket.socket(socket.AF_INET)
+                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)
+                return sock
+            except socket.gaierror:
+                # If that fails try IPv6
+                sock = socket.socket(socket.AF_INET6)
+                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)
+                return sock
         except socket.error:
             self.disconnect()
             raise AutoReconnect("could not connect to %r" % list(self.__nodes))
diff --git a/test/test_connection.py b/test/test_connection.py
@@ -463,6 +463,37 @@ def test_tz_aware(self):
                 aware.pymongo_test.test.find_one()["x"].replace(tzinfo=None),
                 naive.pymongo_test.test.find_one()["x"])
 
+    def test_ipv6(self):
+        self.assertRaises(InvalidURI, _parse_uri, "::1", 27017)
+        self.assertRaises(InvalidURI, _parse_uri, "[::1", 27017)
+        self.assertRaises(InvalidURI, _parse_uri, "::1]:27017")
+        self.assertRaises(InvalidURI, _parse_uri, "mongodb://::1", 27017)
+        self.assert_(_parse_uri, "mongodb://[::1]:27017/?slaveOk=true")
+        self.assert_(_parse_uri,
+                     "[::1]:27017,[2001:0db8:85a3:0000:0000:8a2e:0370:7334]"
+                     ":27018,192.168.0.212:27019,localhost:27020")
+        self.assert_(_parse_uri,
+                     "mongodb://[2001:0db8:85a3:0000:0000:8a2e:0370:7334]"
+                     ":27017/?slaveOk=true")
+        try:
+            connection = Connection("[::1]")
+        except:
+            # Either mongod was started without --ipv6
+            # or the OS doesn't support it (or both).
+            raise SkipTest()
+
+        # Try a few simple things
+        connection = Connection("mongodb://[::1]:27017")
+        connection = Connection("mongodb://[::1]:27017/?slaveOk=true")
+        connection = Connection("[::1]:27017,localhost:27017")
+        connection = Connection("localhost:27017,[::1]:27017")
+        connection.pymongo_test.test.save({"dummy": u"object"})
+        connection.pymongo_test_bernie.test.save({"dummy": u"object"})
+
+        dbs = connection.database_names()
+        self.assert_("pymongo_test" in dbs)
+        self.assert_("pymongo_test_bernie" in dbs)
+
 
 if __name__ == "__main__":
     unittest.main()
