sensu-plugins-mongodb-mrtrotl 1.4.0

data/lib/pymongo/read_preferences.py

@@ -0,0 +1,609 @@
+# Copyright 2012-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License",
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for choosing which member of a replica set to read from."""
+
+from collections import abc
+from typing import Any, Dict, Mapping, Optional, Sequence
+
+from pymongo import max_staleness_selectors
+from pymongo.errors import ConfigurationError
+from pymongo.server_selectors import (
+    member_with_tags_server_selector,
+    secondary_with_tags_server_selector,
+)
+
+_PRIMARY = 0
+_PRIMARY_PREFERRED = 1
+_SECONDARY = 2
+_SECONDARY_PREFERRED = 3
+_NEAREST = 4
+
+
+_MONGOS_MODES = (
+    "primary",
+    "primaryPreferred",
+    "secondary",
+    "secondaryPreferred",
+    "nearest",
+)
+
+
+def _validate_tag_sets(tag_sets):
+    """Validate tag sets for a MongoClient."""
+    if tag_sets is None:
+        return tag_sets
+
+    if not isinstance(tag_sets, (list, tuple)):
+        raise TypeError(("Tag sets %r invalid, must be a sequence") % (tag_sets,))
+    if len(tag_sets) == 0:
+        raise ValueError(
+            ("Tag sets %r invalid, must be None or contain at least one set of tags") % (tag_sets,)
+        )
+
+    for tags in tag_sets:
+        if not isinstance(tags, abc.Mapping):
+            raise TypeError(
+                "Tag set %r invalid, must be an instance of dict, "
+                "bson.son.SON or other type that inherits from "
+                "collection.Mapping" % (tags,)
+            )
+
+    return list(tag_sets)
+
+
+def _invalid_max_staleness_msg(max_staleness):
+    return "maxStalenessSeconds must be a positive integer, not %s" % max_staleness
+
+
+# Some duplication with common.py to avoid import cycle.
+def _validate_max_staleness(max_staleness):
+    """Validate max_staleness."""
+    if max_staleness == -1:
+        return -1
+
+    if not isinstance(max_staleness, int):
+        raise TypeError(_invalid_max_staleness_msg(max_staleness))
+
+    if max_staleness <= 0:
+        raise ValueError(_invalid_max_staleness_msg(max_staleness))
+
+    return max_staleness
+
+
+def _validate_hedge(hedge):
+    """Validate hedge."""
+    if hedge is None:
+        return None
+
+    if not isinstance(hedge, dict):
+        raise TypeError("hedge must be a dictionary, not %r" % (hedge,))
+
+    return hedge
+
+
+_Hedge = Mapping[str, Any]
+_TagSets = Sequence[Mapping[str, Any]]
+
+
+class _ServerMode(object):
+    """Base class for all read preferences."""
+
+    __slots__ = ("__mongos_mode", "__mode", "__tag_sets", "__max_staleness", "__hedge")
+
+    def __init__(
+        self,
+        mode: int,
+        tag_sets: Optional[_TagSets] = None,
+        max_staleness: int = -1,
+        hedge: Optional[_Hedge] = None,
+    ) -> None:
+        self.__mongos_mode = _MONGOS_MODES[mode]
+        self.__mode = mode
+        self.__tag_sets = _validate_tag_sets(tag_sets)
+        self.__max_staleness = _validate_max_staleness(max_staleness)
+        self.__hedge = _validate_hedge(hedge)
+
+    @property
+    def name(self) -> str:
+        """The name of this read preference."""
+        return self.__class__.__name__
+
+    @property
+    def mongos_mode(self) -> str:
+        """The mongos mode of this read preference."""
+        return self.__mongos_mode
+
+    @property
+    def document(self) -> Dict[str, Any]:
+        """Read preference as a document."""
+        doc: Dict[str, Any] = {"mode": self.__mongos_mode}
+        if self.__tag_sets not in (None, [{}]):
+            doc["tags"] = self.__tag_sets
+        if self.__max_staleness != -1:
+            doc["maxStalenessSeconds"] = self.__max_staleness
+        if self.__hedge not in (None, {}):
+            doc["hedge"] = self.__hedge
+        return doc
+
+    @property
+    def mode(self) -> int:
+        """The mode of this read preference instance."""
+        return self.__mode
+
+    @property
+    def tag_sets(self) -> _TagSets:
+        """Set ``tag_sets`` to a list of dictionaries like [{'dc': 'ny'}] to
+        read only from members whose ``dc`` tag has the value ``"ny"``.
+        To specify a priority-order for tag sets, provide a list of
+        tag sets: ``[{'dc': 'ny'}, {'dc': 'la'}, {}]``. A final, empty tag
+        set, ``{}``, means "read from any member that matches the mode,
+        ignoring tags." MongoClient tries each set of tags in turn
+        until it finds a set of tags with at least one matching member.
+
+           .. seealso:: `Data-Center Awareness
+               <https://www.mongodb.com/docs/manual/data-center-awareness/>`_
+        """
+        return list(self.__tag_sets) if self.__tag_sets else [{}]
+
+    @property
+    def max_staleness(self) -> int:
+        """The maximum estimated length of time (in seconds) a replica set
+        secondary can fall behind the primary in replication before it will
+        no longer be selected for operations, or -1 for no maximum."""
+        return self.__max_staleness
+
+    @property
+    def hedge(self) -> Optional[_Hedge]:
+        """The read preference ``hedge`` parameter.
+
+        A dictionary that configures how the server will perform hedged reads.
+        It consists of the following keys:
+
+        - ``enabled``: Enables or disables hedged reads in sharded clusters.
+
+        Hedged reads are automatically enabled in MongoDB 4.4+ when using a
+        ``nearest`` read preference. To explicitly enable hedged reads, set
+        the ``enabled`` key  to ``true``::
+
+            >>> Nearest(hedge={'enabled': True})
+
+        To explicitly disable hedged reads, set the ``enabled`` key  to
+        ``False``::
+
+            >>> Nearest(hedge={'enabled': False})
+
+        .. versionadded:: 3.11
+        """
+        return self.__hedge
+
+    @property
+    def min_wire_version(self) -> int:
+        """The wire protocol version the server must support.
+
+        Some read preferences impose version requirements on all servers (e.g.
+        maxStalenessSeconds requires MongoDB 3.4 / maxWireVersion 5).
+
+        All servers' maxWireVersion must be at least this read preference's
+        `min_wire_version`, or the driver raises
+        :exc:`~pymongo.errors.ConfigurationError`.
+        """
+        return 0 if self.__max_staleness == -1 else 5
+
+    def __repr__(self):
+        return "%s(tag_sets=%r, max_staleness=%r, hedge=%r)" % (
+            self.name,
+            self.__tag_sets,
+            self.__max_staleness,
+            self.__hedge,
+        )
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, _ServerMode):
+            return (
+                self.mode == other.mode
+                and self.tag_sets == other.tag_sets
+                and self.max_staleness == other.max_staleness
+                and self.hedge == other.hedge
+            )
+        return NotImplemented
+
+    def __ne__(self, other: Any) -> bool:
+        return not self == other
+
+    def __getstate__(self):
+        """Return value of object for pickling.
+
+        Needed explicitly because __slots__() defined.
+        """
+        return {
+            "mode": self.__mode,
+            "tag_sets": self.__tag_sets,
+            "max_staleness": self.__max_staleness,
+            "hedge": self.__hedge,
+        }
+
+    def __setstate__(self, value):
+        """Restore from pickling."""
+        self.__mode = value["mode"]
+        self.__mongos_mode = _MONGOS_MODES[self.__mode]
+        self.__tag_sets = _validate_tag_sets(value["tag_sets"])
+        self.__max_staleness = _validate_max_staleness(value["max_staleness"])
+        self.__hedge = _validate_hedge(value["hedge"])
+
+
+class Primary(_ServerMode):
+    """Primary read preference.
+
+    * When directly connected to one mongod queries are allowed if the server
+      is standalone or a replica set primary.
+    * When connected to a mongos queries are sent to the primary of a shard.
+    * When connected to a replica set queries are sent to the primary of
+      the replica set.
+    """
+
+    __slots__ = ()
+
+    def __init__(self) -> None:
+        super(Primary, self).__init__(_PRIMARY)
+
+    def __call__(self, selection: Any) -> Any:
+        """Apply this read preference to a Selection."""
+        return selection.primary_selection
+
+    def __repr__(self):
+        return "Primary()"
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, _ServerMode):
+            return other.mode == _PRIMARY
+        return NotImplemented
+
+
+class PrimaryPreferred(_ServerMode):
+    """PrimaryPreferred read preference.
+
+    * When directly connected to one mongod queries are allowed to standalone
+      servers, to a replica set primary, or to replica set secondaries.
+    * When connected to a mongos queries are sent to the primary of a shard if
+      available, otherwise a shard secondary.
+    * When connected to a replica set queries are sent to the primary if
+      available, otherwise a secondary.
+
+    .. note:: When a :class:`~pymongo.mongo_client.MongoClient` is first
+      created reads will be routed to an available secondary until the
+      primary of the replica set is discovered.
+
+    :Parameters:
+      - `tag_sets`: The :attr:`~tag_sets` to use if the primary is not
+        available.
+      - `max_staleness`: (integer, in seconds) The maximum estimated
+        length of time a replica set secondary can fall behind the primary in
+        replication before it will no longer be selected for operations.
+        Default -1, meaning no maximum. If it is set, it must be at least
+        90 seconds.
+      - `hedge`: The :attr:`~hedge` to use if the primary is not available.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
+    """
+
+    __slots__ = ()
+
+    def __init__(
+        self,
+        tag_sets: Optional[_TagSets] = None,
+        max_staleness: int = -1,
+        hedge: Optional[_Hedge] = None,
+    ) -> None:
+        super(PrimaryPreferred, self).__init__(_PRIMARY_PREFERRED, tag_sets, max_staleness, hedge)
+
+    def __call__(self, selection: Any) -> Any:
+        """Apply this read preference to Selection."""
+        if selection.primary:
+            return selection.primary_selection
+        else:
+            return secondary_with_tags_server_selector(
+                self.tag_sets, max_staleness_selectors.select(self.max_staleness, selection)
+            )
+
+
+class Secondary(_ServerMode):
+    """Secondary read preference.
+
+    * When directly connected to one mongod queries are allowed to standalone
+      servers, to a replica set primary, or to replica set secondaries.
+    * When connected to a mongos queries are distributed among shard
+      secondaries. An error is raised if no secondaries are available.
+    * When connected to a replica set queries are distributed among
+      secondaries. An error is raised if no secondaries are available.
+
+    :Parameters:
+      - `tag_sets`: The :attr:`~tag_sets` for this read preference.
+      - `max_staleness`: (integer, in seconds) The maximum estimated
+        length of time a replica set secondary can fall behind the primary in
+        replication before it will no longer be selected for operations.
+        Default -1, meaning no maximum. If it is set, it must be at least
+        90 seconds.
+      - `hedge`: The :attr:`~hedge` for this read preference.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
+    """
+
+    __slots__ = ()
+
+    def __init__(
+        self,
+        tag_sets: Optional[_TagSets] = None,
+        max_staleness: int = -1,
+        hedge: Optional[_Hedge] = None,
+    ) -> None:
+        super(Secondary, self).__init__(_SECONDARY, tag_sets, max_staleness, hedge)
+
+    def __call__(self, selection: Any) -> Any:
+        """Apply this read preference to Selection."""
+        return secondary_with_tags_server_selector(
+            self.tag_sets, max_staleness_selectors.select(self.max_staleness, selection)
+        )
+
+
+class SecondaryPreferred(_ServerMode):
+    """SecondaryPreferred read preference.
+
+    * When directly connected to one mongod queries are allowed to standalone
+      servers, to a replica set primary, or to replica set secondaries.
+    * When connected to a mongos queries are distributed among shard
+      secondaries, or the shard primary if no secondary is available.
+    * When connected to a replica set queries are distributed among
+      secondaries, or the primary if no secondary is available.
+
+    .. note:: When a :class:`~pymongo.mongo_client.MongoClient` is first
+      created reads will be routed to the primary of the replica set until
+      an available secondary is discovered.
+
+    :Parameters:
+      - `tag_sets`: The :attr:`~tag_sets` for this read preference.
+      - `max_staleness`: (integer, in seconds) The maximum estimated
+        length of time a replica set secondary can fall behind the primary in
+        replication before it will no longer be selected for operations.
+        Default -1, meaning no maximum. If it is set, it must be at least
+        90 seconds.
+      - `hedge`: The :attr:`~hedge` for this read preference.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
+    """
+
+    __slots__ = ()
+
+    def __init__(
+        self,
+        tag_sets: Optional[_TagSets] = None,
+        max_staleness: int = -1,
+        hedge: Optional[_Hedge] = None,
+    ) -> None:
+        super(SecondaryPreferred, self).__init__(
+            _SECONDARY_PREFERRED, tag_sets, max_staleness, hedge
+        )
+
+    def __call__(self, selection: Any) -> Any:
+        """Apply this read preference to Selection."""
+        secondaries = secondary_with_tags_server_selector(
+            self.tag_sets, max_staleness_selectors.select(self.max_staleness, selection)
+        )
+
+        if secondaries:
+            return secondaries
+        else:
+            return selection.primary_selection
+
+
+class Nearest(_ServerMode):
+    """Nearest read preference.
+
+    * When directly connected to one mongod queries are allowed to standalone
+      servers, to a replica set primary, or to replica set secondaries.
+    * When connected to a mongos queries are distributed among all members of
+      a shard.
+    * When connected to a replica set queries are distributed among all
+      members.
+
+    :Parameters:
+      - `tag_sets`: The :attr:`~tag_sets` for this read preference.
+      - `max_staleness`: (integer, in seconds) The maximum estimated
+        length of time a replica set secondary can fall behind the primary in
+        replication before it will no longer be selected for operations.
+        Default -1, meaning no maximum. If it is set, it must be at least
+        90 seconds.
+      - `hedge`: The :attr:`~hedge` for this read preference.
+
+    .. versionchanged:: 3.11
+       Added ``hedge`` parameter.
+    """
+
+    __slots__ = ()
+
+    def __init__(
+        self,
+        tag_sets: Optional[_TagSets] = None,
+        max_staleness: int = -1,
+        hedge: Optional[_Hedge] = None,
+    ) -> None:
+        super(Nearest, self).__init__(_NEAREST, tag_sets, max_staleness, hedge)
+
+    def __call__(self, selection: Any) -> Any:
+        """Apply this read preference to Selection."""
+        return member_with_tags_server_selector(
+            self.tag_sets, max_staleness_selectors.select(self.max_staleness, selection)
+        )
+
+
+class _AggWritePref:
+    """Agg $out/$merge write preference.
+
+    * If there are readable servers and there is any pre-5.0 server, use
+      primary read preference.
+    * Otherwise use `pref` read preference.
+
+    :Parameters:
+      - `pref`: The read preference to use on MongoDB 5.0+.
+    """
+
+    __slots__ = ("pref", "effective_pref")
+
+    def __init__(self, pref):
+        self.pref = pref
+        self.effective_pref = ReadPreference.PRIMARY
+
+    def selection_hook(self, topology_description):
+        common_wv = topology_description.common_wire_version
+        if (
+            topology_description.has_readable_server(ReadPreference.PRIMARY_PREFERRED)
+            and common_wv
+            and common_wv < 13
+        ):
+            self.effective_pref = ReadPreference.PRIMARY
+        else:
+            self.effective_pref = self.pref
+
+    def __call__(self, selection):
+        """Apply this read preference to a Selection."""
+        return self.effective_pref(selection)
+
+    def __repr__(self):
+        return "_AggWritePref(pref=%r)" % (self.pref,)
+
+    # Proxy other calls to the effective_pref so that _AggWritePref can be
+    # used in place of an actual read preference.
+    def __getattr__(self, name):
+        return getattr(self.effective_pref, name)
+
+
+_ALL_READ_PREFERENCES = (Primary, PrimaryPreferred, Secondary, SecondaryPreferred, Nearest)
+
+
+def make_read_preference(
+    mode: int, tag_sets: Optional[_TagSets], max_staleness: int = -1
+) -> _ServerMode:
+    if mode == _PRIMARY:
+        if tag_sets not in (None, [{}]):
+            raise ConfigurationError("Read preference primary cannot be combined with tags")
+        if max_staleness != -1:
+            raise ConfigurationError(
+                "Read preference primary cannot be combined with maxStalenessSeconds"
+            )
+        return Primary()
+    return _ALL_READ_PREFERENCES[mode](tag_sets, max_staleness)  # type: ignore
+
+
+_MODES = (
+    "PRIMARY",
+    "PRIMARY_PREFERRED",
+    "SECONDARY",
+    "SECONDARY_PREFERRED",
+    "NEAREST",
+)
+
+
+class ReadPreference(object):
+    """An enum that defines the read preference modes supported by PyMongo.
+
+    See :doc:`/examples/high_availability` for code examples.
+
+    A read preference is used in three cases:
+
+    :class:`~pymongo.mongo_client.MongoClient` connected to a single mongod:
+
+    - ``PRIMARY``: Queries are allowed if the server is standalone or a replica
+      set primary.
+    - All other modes allow queries to standalone servers, to a replica set
+      primary, or to replica set secondaries.
+
+    :class:`~pymongo.mongo_client.MongoClient` initialized with the
+    ``replicaSet`` option:
+
+    - ``PRIMARY``: Read from the primary. This is the default, and provides the
+      strongest consistency. If no primary is available, raise
+      :class:`~pymongo.errors.AutoReconnect`.
+
+    - ``PRIMARY_PREFERRED``: Read from the primary if available, or if there is
+      none, read from a secondary.
+
+    - ``SECONDARY``: Read from a secondary. If no secondary is available,
+      raise :class:`~pymongo.errors.AutoReconnect`.
+
+    - ``SECONDARY_PREFERRED``: Read from a secondary if available, otherwise
+      from the primary.
+
+    - ``NEAREST``: Read from any member.
+
+    :class:`~pymongo.mongo_client.MongoClient` connected to a mongos, with a
+    sharded cluster of replica sets:
+
+    - ``PRIMARY``: Read from the primary of the shard, or raise
+      :class:`~pymongo.errors.OperationFailure` if there is none.
+      This is the default.
+
+    - ``PRIMARY_PREFERRED``: Read from the primary of the shard, or if there is
+      none, read from a secondary of the shard.
+
+    - ``SECONDARY``: Read from a secondary of the shard, or raise
+      :class:`~pymongo.errors.OperationFailure` if there is none.
+
+    - ``SECONDARY_PREFERRED``: Read from a secondary of the shard if available,
+      otherwise from the shard primary.
+
+    - ``NEAREST``: Read from any shard member.
+    """
+
+    PRIMARY = Primary()
+    PRIMARY_PREFERRED = PrimaryPreferred()
+    SECONDARY = Secondary()
+    SECONDARY_PREFERRED = SecondaryPreferred()
+    NEAREST = Nearest()
+
+
+def read_pref_mode_from_name(name: str) -> int:
+    """Get the read preference mode from mongos/uri name."""
+    return _MONGOS_MODES.index(name)
+
+
+class MovingAverage(object):
+    """Tracks an exponentially-weighted moving average."""
+
+    average: Optional[float]
+
+    def __init__(self) -> None:
+        self.average = None
+
+    def add_sample(self, sample: float) -> None:
+        if sample < 0:
+            # Likely system time change while waiting for hello response
+            # and not using time.monotonic. Ignore it, the next one will
+            # probably be valid.
+            return
+        if self.average is None:
+            self.average = sample
+        else:
+            # The Server Selection Spec requires an exponentially weighted
+            # average with alpha = 0.2.
+            self.average = 0.8 * self.average + 0.2 * sample
+
+    def get(self) -> Optional[float]:
+        """Get the calculated average, or None if no samples yet."""
+        return self.average
+
+    def reset(self) -> None:
+        self.average = None