PYTHON-2671 Support loadBalanced URI option (#614)

Add workaround in test_dns until PYTHON-2679 is completed.

pymongo/client_options.py

@@ -173,6 +173,7 @@ class ClientOptions(object):
         self.__server_selector = options.get(
             'server_selector', any_server_selector)
         self.__auto_encryption_opts = options.get('auto_encryption_opts')
+        self.__load_balanced = options.get('loadbalanced')
 
     @property
     def _options(self):
@@ -257,3 +258,8 @@ class ClientOptions(object):
     def auto_encryption_opts(self):
         """A :class:`~pymongo.encryption.AutoEncryptionOpts` or None."""
         return self.__auto_encryption_opts
+
+    @property
+    def load_balanced(self):
+        """True if the client was configured to connect to a load balancer."""
+        return self.__load_balanced

pymongo/common.py

@@ -623,6 +623,7 @@ URI_OPTIONS_VALIDATOR_MAP = {
     'replicaset': validate_string_or_none,
     'retryreads': validate_boolean_or_string,
     'retrywrites': validate_boolean_or_string,
+    'loadbalanced': validate_boolean_or_string,
     'serverselectiontimeoutms': validate_timeout_or_zero,
     'sockettimeoutms': validate_timeout_or_none_or_zero,
     'ssl_keyfile': validate_readable,

pymongo/mongo_client.py

@@ -70,7 +70,8 @@ from pymongo.topology_description import TOPOLOGY_TYPE
 from pymongo.settings import TopologySettings
 from pymongo.uri_parser import (_handle_option_deprecations,
                                 _handle_security_options,
-                                _normalize_options)
+                                _normalize_options,
+                                _check_options)
 from pymongo.write_concern import DEFAULT_WRITE_CONCERN
 
 
@@ -692,11 +693,7 @@ class MongoClient(common.BaseObject):
         opts = _handle_security_options(opts)
         # Normalize combined options.
         opts = _normalize_options(opts)
-
-        # Ensure directConnection was not True if there are multiple seeds.
-        if len(seeds) > 1 and opts.get('directconnection'):
-            raise ConfigurationError(
-                "Cannot specify multiple hosts with directConnection=true")
+        _check_options(seeds, opts)
 
         # Username and password passed as kwargs override user info in URI.
         username = opts.get("username", username)
@@ -739,7 +736,9 @@ class MongoClient(common.BaseObject):
             server_selector=options.server_selector,
             heartbeat_frequency=options.heartbeat_frequency,
             fqdn=fqdn,
-            direct_connection=options.direct_connection)
+            direct_connection=options.direct_connection,
+            load_balanced=options.load_balanced,
+        )
 
         self._topology = Topology(self._topology_settings)
 

pymongo/settings.py

@@ -39,7 +39,8 @@ class TopologySettings(object):
                  heartbeat_frequency=common.HEARTBEAT_FREQUENCY,
                  server_selector=None,
                  fqdn=None,
-                 direct_connection=None):
+                 direct_connection=None,
+                 load_balanced=None):
         """Represent MongoClient's configuration.
 
         Take a list of (host, port) pairs and optional replica set name.
@@ -65,6 +66,7 @@ class TopologySettings(object):
             self._direct = (len(self._seeds) == 1 and not self.replica_set_name)
         else:
             self._direct = direct_connection
+        self._load_balanced = load_balanced
 
         self._topology_id = ObjectId()
         # Store the allocation traceback to catch unclosed clients in the
@@ -124,6 +126,11 @@ class TopologySettings(object):
         """
         return self._direct
 
+    @property
+    def load_balanced(self):
+        """True if the client was configured to connect to a load balancer."""
+        return self._load_balanced
+
     def get_topology_type(self):
         if self.direct:
             return TOPOLOGY_TYPE.Single

pymongo/uri_parser.py

@@ -362,7 +362,26 @@ def split_hosts(hosts, default_port=DEFAULT_PORT):
 _BAD_DB_CHARS = re.compile('[' + re.escape(r'/ "$') + ']')
 
 _ALLOWED_TXT_OPTS = frozenset(
-    ['authsource', 'authSource', 'replicaset', 'replicaSet'])
+    ['authsource', 'authSource', 'replicaset', 'replicaSet', 'loadbalanced',
+     'loadBalanced'])
+
+
+def _check_options(nodes, options):
+    # Ensure directConnection was not True if there are multiple seeds.
+    if len(nodes) > 1 and options.get('directconnection'):
+        raise ConfigurationError(
+            'Cannot specify multiple hosts with directConnection=true')
+
+    if options.get('loadbalanced'):
+        if len(nodes) > 1:
+            raise ConfigurationError(
+                'Cannot specify multiple hosts with loadBalanced=true')
+        if options.get('directconnection'):
+            raise ConfigurationError(
+                'Cannot specify directConnection=true with loadBalanced=true')
+        if options.get('replicaset'):
+            raise ConfigurationError(
+                'Cannot specify replicaSet with loadBalanced=true')
 
 
 def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
@@ -500,7 +519,8 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
                 dns_options, validate, warn, normalize)
             if set(parsed_dns_options) - _ALLOWED_TXT_OPTS:
                 raise ConfigurationError(
-                    "Only authSource and replicaSet are supported from DNS")
+                    "Only authSource, replicaSet, and loadBalanced are "
+                    "supported from DNS")
             for opt, val in parsed_dns_options.items():
                 if opt not in options:
                     options[opt] = val
@@ -508,9 +528,8 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
             options["ssl"] = True if validate else 'true'
     else:
         nodes = split_hosts(hosts, default_port=default_port)
-        if len(nodes) > 1 and options.get('directConnection'):
-            raise ConfigurationError(
-                "Cannot specify multiple hosts with directConnection=true")
+
+    _check_options(nodes, options)
 
     return {
         'nodelist': nodes,

test/__init__.py

@@ -216,6 +216,7 @@ class ClientContext(object):
         self.client = None
         self.conn_lock = threading.Lock()
         self.is_data_lake = False
+        self.load_balancer = False
         if COMPRESSORS:
             self.default_client_options["compressors"] = COMPRESSORS
         if MONGODB_API_VERSION:
@@ -616,6 +617,12 @@ class ClientContext(object):
                              "Must be connected to a replica set or mongos",
                              func=func)
 
+    def require_load_balancer(self, func):
+        """Run a test only if the client is connected to a load balancer."""
+        return self._require(lambda: self.load_balancer,
+                             "Must be connected to a load balancer",
+                             func=func)
+
     def check_auth_with_sharding(self, func):
         """Skip a test when connected to mongos < 2.0 and running with auth."""
         condition = lambda: not (self.auth_enabled and

test/srv_seedlist/load-balanced/loadBalanced-directConnection.json

@@ -0,0 +1,14 @@
+{
+  "uri": "mongodb+srv://test20.test.build.10gen.cc/?directConnection=false",
+  "seeds": [
+    "localhost.test.build.10gen.cc:27017"
+  ],
+  "hosts": [
+    "localhost.test.build.10gen.cc:27017"
+  ],
+  "options": {
+    "loadBalanced": true,
+    "ssl": true,
+    "directConnection": false
+  }
+}

test/srv_seedlist/load-balanced/loadBalanced-replicaSet-errors.json

@@ -0,0 +1,7 @@
+{
+  "uri": "mongodb+srv://test20.test.build.10gen.cc/?replicaSet=replset",
+  "seeds": [],
+  "hosts": [],
+  "error": true,
+  "comment": "Should fail because loadBalanced=true is incompatible with replicaSet"
+}

test/srv_seedlist/load-balanced/loadBalanced-true-multiple-hosts.json

@@ -0,0 +1,7 @@
+{
+  "uri": "mongodb+srv://test1.test.build.10gen.cc/?loadBalanced=true",
+  "seeds": [],
+  "hosts": [],
+  "error": true,
+  "comment": "Should fail because loadBalanced is true but the SRV record resolves to multiple hosts"
+}

test/srv_seedlist/load-balanced/loadBalanced-true-txt.json

@@ -0,0 +1,13 @@
+{
+  "uri": "mongodb+srv://test20.test.build.10gen.cc/",
+  "seeds": [
+    "localhost.test.build.10gen.cc:27017"
+  ],
+  "hosts": [
+    "localhost.test.build.10gen.cc:27017"
+  ],
+  "options": {
+    "loadBalanced": true,
+    "ssl": true
+  }
+}

test/srv_seedlist/replica-set/loadBalanced-false-txt.json

@@ -0,0 +1,15 @@
+{
+  "uri": "mongodb+srv://test21.test.build.10gen.cc/",
+  "seeds": [
+    "localhost.test.build.10gen.cc:27017"
+  ],
+  "hosts": [
+    "localhost:27017",
+    "localhost:27018",
+    "localhost:27019"
+  ],
+  "options": {
+    "loadBalanced": false,
+    "ssl": true
+  }
+}

test/test_dns.py

@@ -30,16 +30,28 @@ from test import client_context, unittest
 from test.utils import wait_until
 
 
-TEST_PATH = os.path.join(
-    os.path.dirname(os.path.realpath(__file__)), 'srv_seedlist')
+class TestDNSRepl(unittest.TestCase):
+    TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)),
+                             'srv_seedlist', 'replica-set')
+    load_balanced = False
 
-class TestDNS(unittest.TestCase):
-    pass
+    @client_context.require_replica_set
+    def setUp(self):
+        pass
+
+
+class TestDNSLoadBalanced(unittest.TestCase):
+    TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)),
+                             'srv_seedlist', 'load-balanced')
+    load_balanced = True
+
+    @client_context.require_load_balancer
+    def setUp(self):
+        pass
 
 
 def create_test(test_case):
 
-    @client_context.require_replica_set
     def run_test(self):
         if not _HAVE_DNSPYTHON:
             raise unittest.SkipTest("DNS tests require the dnspython module")
@@ -91,6 +103,12 @@ def create_test(test_case):
                     # tests.
                     copts['tlsAllowInvalidHostnames'] = True
 
+                # The SRV spec tests assume drivers auto discover replica set
+                # members. This should be removed during PYTHON-2679.
+                if not self.load_balanced and (
+                        'directconnection' not in result['options']):
+                    copts['directConnection'] = False
+
                 client = MongoClient(uri, **copts)
                 wait_until(
                     lambda: hosts == client.nodes,
@@ -106,15 +124,17 @@ def create_test(test_case):
     return run_test
 
 
-def create_tests():
-    for filename in glob.glob(os.path.join(TEST_PATH, '*.json')):
+def create_tests(cls):
+    for filename in glob.glob(os.path.join(cls.TEST_PATH, '*.json')):
         test_suffix, _ = os.path.splitext(os.path.basename(filename))
         with open(filename) as dns_test_file:
             test_method = create_test(json.load(dns_test_file))
-        setattr(TestDNS, 'test_' + test_suffix, test_method)
+        setattr(cls, 'test_' + test_suffix, test_method)
 
 
-create_tests()
+create_tests(TestDNSRepl)
+create_tests(TestDNSLoadBalanced)
+
 
 class TestParsingErrors(unittest.TestCase):