PYTHON-1577 Allow applications to register a custom server selector (mongodb#371)

PYTHON-1577 Allow applications to register a custom server selector

Author: prashantmital

doc/examples/index.rst

@@ -27,5 +27,6 @@ MongoDB, you can start it like so:
  gridfs
  high_availability
  mod_wsgi
+ server_selection
  tailable
  tls

doc/examples/server_selection.rst

@@ -0,0 +1,108 @@
+Server Selector Example
+=======================
+
+Users can exert fine-grained control over the `server selection algorithm`_
+by setting the `server_selector` option on the :class:`~pymongo.MongoClient`
+to an appropriate callable. This example shows how to use this functionality
+to prefer servers running on ``localhost``.
+
+
+.. warning::
+
+ Use of custom server selector functions is a power user feature. Misusing
+ custom server selectors can have unintended consequences such as degraded
+ read/write performance.
+
+
+.. testsetup::
+
+ from pymongo import MongoClient
+
+
+.. _server selection algorithm: https://docs.mongodb.com/manual/core/read-preference-mechanics/
+
+
+Example: Selecting Servers Running on ``localhost``
+---------------------------------------------------
+
+To start, we need to write the server selector function that will be used.
+The server selector function should accept a list of
+:class:`~pymongo.server_description.ServerDescription` objects and return a
+list of server descriptions that are suitable for the read or write operation.
+A server selector must not create or modify
+:class:`~pymongo.server_description.ServerDescription` objects, and must return
+the selected instances unchanged.
+
+In this example, we write a server selector that prioritizes servers running on
+``localhost``. This can be desirable when using a sharded cluster with multiple
+``mongos``, as locally run queries are likely to see lower latency and higher
+throughput. Please note, however, that it is highly dependent on the
+application if preferring ``localhost`` is beneficial or not.
+
+In addition to comparing the hostname with ``localhost``, our server selector
+function accounts for the edge case when no servers are running on
+``localhost``. In this case, we allow the default server selection logic to
+prevail by passing through the received server description list unchanged.
+Failure to do this would render the client unable to communicate with MongoDB
+in the event that no servers were running on ``localhost``.
+
+
+The described server selection logic is implemented in the following server
+selector function:
+
+
+.. doctest::
+
+ >>> def server_selector(server_descriptions):
+ ... servers = [
+ ... server for server in server_descriptions
+ ... if server.address[0] == 'localhost'
+ ... ]
+ ... if not servers:
+ ... return server_descriptions
+ ... return servers
+
+
+
+Finally, we can create a :class:`~pymongo.MongoClient` instance with this
+server selector.
+
+
+.. doctest::
+
+ >>> client = MongoClient(server_selector=server_selector)
+
+
+
+Server Selection Process
+------------------------
+
+This section dives deeper into the server selection process for reads and
+writes. In the case of a write, the driver performs the following operations
+(in order) during the selection process:
+
+
+#. Select all writeable servers from the list of known hosts. For a replica set
+ this is the primary, while for a sharded cluster this is all the known mongoses.
+
+#. Apply the user-defined server selector function. Note that the custom server
+ selector is **not** called if there are no servers left from the previous
+ filtering stage.
+
+#. Apply the ``localThresholdMS`` setting to the list of remaining hosts. This
+ whittles the host list down to only contain servers whose latency is at most
+ ``localThresholdMS`` milliseconds higher than the lowest observed latency.
+
+#. Select a server at random from the remaining host list. The desired
+ operation is then performed against the selected server.
+
+
+In the case of **reads** the process is identical except for the first step.
+Here, instead of selecting all writeable servers, we select all servers
+matching the user's :class:`~pymongo.read_preferences.ReadPreference` from the
+list of known hosts. As an example, for a 3-member replica set with a
+:class:`~pymongo.read_preferences.Secondary` read preference, we would select
+all available secondaries.
+
+
+.. _server selection algorithm: https://docs.mongodb.com/manual/core/read-preference-mechanics/

pymongo/client_options.py

@@ -25,6 +25,7 @@
 from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import (make_read_preference,
  read_pref_mode_from_name)
+from pymongo.server_selectors import any_server_selector
 from pymongo.ssl_support import get_ssl_context
 from pymongo.write_concern import WriteConcern
 
@@ -163,6 +164,8 @@ def __init__(self, username, password, database, options):
  self.__heartbeat_frequency = options.get(
  'heartbeatfrequencyms', common.HEARTBEAT_FREQUENCY)
  self.__retry_writes = options.get('retrywrites', common.RETRY_WRITES)
+ self.__server_selector = options.get(
+ 'server_selector', any_server_selector)
 
  @property
  def _options(self):
@@ -194,6 +197,10 @@ def server_selection_timeout(self):
  """The server selection timeout for this instance in seconds."""
  return self.__server_selection_timeout
 
+ @property
+ def server_selector(self):
+ return self.__server_selector
+
  @property
  def heartbeat_frequency(self):
  """The monitoring frequency in seconds."""

pymongo/common.py

@@ -473,6 +473,15 @@ def validate_driver_or_none(option, value):
  return value
 
 
+def validate_is_callable_or_none(option, value):
+ """Validates that 'value' is a callable."""
+ if value is None:
+ return value
+ if not callable(value):
+ raise ValueError("%s must be a callable" % (option,))
+ return value
+
+
 def validate_ok_for_replace(replacement):
  """Validate a replacement document."""
  validate_is_mapping("replacement", replacement)
@@ -552,7 +561,7 @@ def validate_tzinfo(dummy, value):
  'unicode_decode_error_handler': validate_unicode_decode_error_handler,
  'retrywrites': validate_boolean_or_string,
  'compressors': validate_compressors,
- 'zlibcompressionlevel': validate_zlib_compression_level
+ 'zlibcompressionlevel': validate_zlib_compression_level,
 }
 
 TIMEOUT_VALIDATORS = {
@@ -572,6 +581,7 @@ def validate_tzinfo(dummy, value):
  'tzinfo': validate_tzinfo,
  'username': validate_string_or_none,
  'password': validate_string_or_none,
+ 'server_selector': validate_is_callable_or_none,
 }
 
 URI_VALIDATORS.update(TIMEOUT_VALIDATORS)

pymongo/mongo_client.py

@@ -215,6 +215,12 @@ def __init__(
  milliseconds) the driver will wait during server monitoring when
  connecting a new socket to a server before concluding the server
  is unavailable. Defaults to ``20000`` (20 seconds).
+ - `server_selector`: (callable or None) Optional, user-provided
+ function that augments server selection rules. The function should
+ accept as an argument a list of
+ :class:`~pymongo.server_description.ServerDescription` objects and
+ return a list of server descriptions that should be considered
+ suitable for the desired operation.
  - `serverSelectionTimeoutMS`: (integer) Controls how long (in
  milliseconds) the driver will wait to find an available,
  appropriate server to carry out a database operation; while it is
@@ -331,6 +337,8 @@ def __init__(
  is set, it must be a positive integer greater than or equal to
  90 seconds.
 
+ .. seealso:: :doc:`/examples/server_selection`
+
  | **Authentication:**
 
  - `username`: A string.
@@ -411,13 +419,16 @@ def __init__(
 
  .. mongodoc:: connections
 
- .. versionchanged:: 3.6
- Added support for mongodb+srv:// URIs.
- Added the ``retryWrites`` keyword argument and URI option.
+ .. versionchanged:: 3.8
+ Added the ``server_selector`` keyword argument.
 
  .. versionchanged:: 3.7
  Added the ``driver`` keyword argument.
 
+ .. versionchanged:: 3.6
+ Added support for mongodb+srv:// URIs.
+ Added the ``retryWrites`` keyword argument and URI option.
+
  .. versionchanged:: 3.5
  Add ``username`` and ``password`` options. Document the
  ``authSource``, ``authMechanism``, and ``authMechanismProperties ``
@@ -572,6 +583,7 @@ def __init__(
  condition_class=condition_class,
  local_threshold_ms=options.local_threshold_ms,
  server_selection_timeout=options.server_selection_timeout,
+ server_selector=options.server_selector,
  heartbeat_frequency=options.heartbeat_frequency)
 
  self._topology = Topology(self._topology_settings)

pymongo/settings.py

@@ -35,7 +35,8 @@ def __init__(self,
  condition_class=None,
  local_threshold_ms=LOCAL_THRESHOLD_MS,
  server_selection_timeout=SERVER_SELECTION_TIMEOUT,
- heartbeat_frequency=common.HEARTBEAT_FREQUENCY):
+ heartbeat_frequency=common.HEARTBEAT_FREQUENCY,
+ server_selector=None):
  """Represent MongoClient's configuration.
 
  Take a list of (host, port) pairs and optional replica set name.
@@ -53,6 +54,7 @@ def __init__(self,
  self._condition_class = condition_class or threading.Condition
  self._local_threshold_ms = local_threshold_ms
  self._server_selection_timeout = server_selection_timeout
+ self._server_selector = server_selector
  self._heartbeat_frequency = heartbeat_frequency
  self._direct = (len(self._seeds) == 1 and not replica_set_name)
  self._topology_id = ObjectId()
@@ -90,6 +92,10 @@ def local_threshold_ms(self):
  def server_selection_timeout(self):
  return self._server_selection_timeout
 
+ @property
+ def server_selector(self):
+ return self._server_selector
+
  @property
  def heartbeat_frequency(self):
  return self._heartbeat_frequency

pymongo/topology.py

@@ -190,7 +190,7 @@ def _select_servers_loop(self, selector, timeout, address):
  now = _time()
  end_time = now + timeout
  server_descriptions = self._description.apply_selector(
- selector, address)
+ selector, address, custom_selector=self._settings.server_selector)
 
  while not server_descriptions:
  # No suitable servers.
@@ -209,7 +209,8 @@ def _select_servers_loop(self, selector, timeout, address):
  self._description.check_compatible()
  now = _time()
  server_descriptions = self._description.apply_selector(
- selector, address)
+ selector, address,
+ custom_selector=self._settings.server_selector)
 
  self._description.check_compatible()
  return server_descriptions

pymongo/topology_description.py

@@ -214,7 +214,7 @@ def common_wire_version(self):
  def heartbeat_frequency(self):
  return self._topology_settings.heartbeat_frequency
 
- def apply_selector(self, selector, address):
+ def apply_selector(self, selector, address, custom_selector=None):
 
  def apply_local_threshold(selection):
  if not selection:
@@ -239,18 +239,23 @@ def apply_local_threshold(selection):
  common_wv))
 
  if self.topology_type == TOPOLOGY_TYPE.Single:
- # Ignore the selector.
+ # Ignore selectors for standalone.
  return self.known_servers
  elif address:
+ # Ignore selectors when explicit address is requested.
  description = self.server_descriptions().get(address)
  return [description] if description else []
  elif self.topology_type == TOPOLOGY_TYPE.Sharded:
- # Ignore the read preference, but apply localThresholdMS.
- return apply_local_threshold(
- Selection.from_topology_description(self))
+ # Ignore read preference.
+ selection = Selection.from_topology_description(self)
  else:
- return apply_local_threshold(
- selector(Selection.from_topology_description(self)))
+ selection = selector(Selection.from_topology_description(self))
+
+ # Apply custom selector followed by localThresholdMS.
+ if custom_selector is not None and selection:
+ selection = selection.with_server_descriptions(
+ custom_selector(selection.server_descriptions))
+ return apply_local_threshold(selection)
 
  def has_readable_server(self, read_preference=ReadPreference.PRIMARY):
  """Does this topology have any readable servers available matching the