PYTHON-2671 Support loadBalanced URI option (#614)

Add workaround in test_dns until PYTHON-2679 is completed.

(cherry picked from commit 2c41c6f)

Author: ShaneHarvey

pymongo/client_options.py

@@ -173,6 +173,7 @@ def __init__(self, username, password, database, options):
         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 @@ def retry_reads(self):
     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

@@ -627,6 +627,7 @@ def validate_tzinfo(dummy, value):
     '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

@@ -73,7 +73,8 @@
 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
 
 
@@ -700,11 +701,7 @@ def __init__(
         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)
@@ -752,7 +749,9 @@ def __init__(
             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 @@ def __init__(self,
                  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 @@ def __init__(self,
             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 @@ def direct(self):
         """
         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

@@ -371,7 +371,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,
@@ -509,17 +528,17 @@ 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
         if "ssl" not in options:
             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 @@ def __init__(self):
         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:
@@ -605,6 +606,12 @@ def require_no_standalone(self, func):
                              "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
+  }
+}