PYTHON-2767 Support snapshot reads on secondaries (#656)

Add the MongoClient.start_session snapshot option.

(cherry picked from commit 14160ae)

Author: ShaneHarvey

pymongo/client_session.py

@@ -125,14 +125,29 @@ class SessionOptions(object):
     """Options for a new :class:`ClientSession`.
 
     :Parameters:
-      - `causal_consistency` (optional): If True (the default), read
-        operations are causally ordered within the session.
+      - `causal_consistency` (optional): If True, read operations are causally
+        ordered within the session. Defaults to True when the ``snapshot``
+        option is ``False``.
       - `default_transaction_options` (optional): The default
         TransactionOptions to use for transactions started on this session.
+      - `snapshot` (optional): If True, then all reads performed using this
+        session will read from the same snapshot. This option is incompatible
+        with ``causal_consistency=True``. Defaults to ``False``.
+
+    .. versionchanged:: 3.12
+       Added the ``snapshot`` parameter.
     """
     def __init__(self,
-                 causal_consistency=True,
-                 default_transaction_options=None):
+                 causal_consistency=None,
+                 default_transaction_options=None,
+                 snapshot=False):
+        if snapshot:
+            if causal_consistency:
+                raise ConfigurationError('snapshot reads do not support '
+                                         'causal_consistency=True')
+            causal_consistency = False
+        elif causal_consistency is None:
+            causal_consistency = True
         self._causal_consistency = causal_consistency
         if default_transaction_options is not None:
             if not isinstance(default_transaction_options, TransactionOptions):
@@ -141,6 +156,7 @@ def __init__(self,
                     "pymongo.client_session.TransactionOptions, not: %r" %
                     (default_transaction_options,))
         self._default_transaction_options = default_transaction_options
+        self._snapshot = snapshot
 
     @property
     def causal_consistency(self):
@@ -156,6 +172,14 @@ def default_transaction_options(self):
         """
         return self._default_transaction_options
 
+    @property
+    def snapshot(self):
+        """Whether snapshot reads are configured.
+
+        .. versionadded:: 3.12
+        """
+        return self._snapshot
+
 
 class TransactionOptions(object):
     """Options for :meth:`ClientSession.start_transaction`.
@@ -388,6 +412,7 @@ def __init__(self, client, server_session, options, authset, implicit):
         self._authset = authset
         self._cluster_time = None
         self._operation_time = None
+        self._snapshot_time = None
         # Is this an implicitly created session?
         self._implicit = implicit
         self._transaction = _Transaction(None, client)
@@ -603,6 +628,10 @@ def start_transaction(self, read_concern=None, write_concern=None,
         """
         self._check_ended()
 
+        if self.options.snapshot:
+            raise InvalidOperation("Transactions are not supported in "
+                                   "snapshot sessions")
+
         if self.in_transaction:
             raise InvalidOperation("Transaction already in progress")
 
@@ -781,6 +810,12 @@ def _process_response(self, reply):
         """Process a response to a command that was run with this session."""
         self._advance_cluster_time(reply.get('$clusterTime'))
         self._advance_operation_time(reply.get('operationTime'))
+        if self._options.snapshot and self._snapshot_time is None:
+            if 'cursor' in reply:
+                ct = reply['cursor'].get('atClusterTime')
+            else:
+                ct = reply.get('atClusterTime')
+            self._snapshot_time = ct
         if self.in_transaction and self._transaction.sharded:
             recovery_token = reply.get('recoveryToken')
             if recovery_token:
@@ -854,15 +889,9 @@ def _apply_to(self, command, is_retryable, read_preference):
 
                 if self._transaction.opts.read_concern:
                     rc = self._transaction.opts.read_concern.document
-                else:
-                    rc = {}
-
-                if (self.options.causal_consistency
-                        and self.operation_time is not None):
-                    rc['afterClusterTime'] = self.operation_time
-
-                if rc:
-                    command['readConcern'] = rc
+                    if rc:
+                        command['readConcern'] = rc
+                self._update_read_concern(command)
 
             command['txnNumber'] = self._server_session.transaction_id
             command['autocommit'] = False
@@ -871,6 +900,17 @@ def _start_retryable_write(self):
         self._check_ended()
         self._server_session.inc_transaction_id()
 
+    def _update_read_concern(self, cmd):
+        if (self.options.causal_consistency
+                and self.operation_time is not None):
+            cmd.setdefault('readConcern', {})[
+                'afterClusterTime'] = self.operation_time
+        if self.options.snapshot:
+            rc = cmd.setdefault('readConcern', {})
+            rc['level'] = 'snapshot'
+            if self._snapshot_time is not None:
+                rc['atClusterTime'] = self._snapshot_time
+
 
 class _ServerSession(object):
     def __init__(self, generation):

pymongo/message.py

@@ -317,12 +317,8 @@ def as_command(self, sock_info):
         if session:
             session._apply_to(cmd, False, self.read_preference)
             # Explain does not support readConcern.
-            if (not explain and session.options.causal_consistency
-                    and session.operation_time is not None
-                    and not session.in_transaction):
-                cmd.setdefault(
-                    'readConcern', {})[
-                    'afterClusterTime'] = session.operation_time
+            if not explain and not session.in_transaction:
+                session._update_read_concern(cmd)
         sock_info.send_cluster_time(cmd, session, self.client)
         # Support auto encryption
         client = self.client

pymongo/mongo_client.py

@@ -1878,8 +1878,9 @@ def __start_session(self, implicit, **kwargs):
             self, server_session, opts, authset, implicit)
 
     def start_session(self,
-                      causal_consistency=True,
-                      default_transaction_options=None):
+                      causal_consistency=None,
+                      default_transaction_options=None,
+                      snapshot=False):
         """Start a logical session.
 
         This method takes the same parameters as
@@ -1904,7 +1905,8 @@ def start_session(self,
         return self.__start_session(
             False,
             causal_consistency=causal_consistency,
-            default_transaction_options=default_transaction_options)
+            default_transaction_options=default_transaction_options,
+            snapshot=snapshot)
 
     def _get_server_session(self):
         """Internal: start or resume a _ServerSession."""

pymongo/network.py

@@ -93,10 +93,8 @@ def command(sock_info, dbname, spec, slave_ok, is_mongos,
     if read_concern and not (session and session.in_transaction):
         if read_concern.level:
             spec['readConcern'] = read_concern.document
-        if (session and session.options.causal_consistency
-                and session.operation_time is not None):
-            spec.setdefault(
-                'readConcern', {})['afterClusterTime'] = session.operation_time
+        if session:
+            session._update_read_concern(spec)
     if collation is not None:
         spec['collation'] = collation
 

test/load_balancer/test_sessions_unified.py

@@ -0,0 +1,23 @@
+# Copyright 2021-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.
+
+import sys
+
+sys.path[0:0] = [""]
+
+from test import unittest
+from test.test_sessions_unified import *
+
+if __name__ == '__main__':
+    unittest.main()

test/sessions/dirty-session-errors.json renamed to test/sessions/legacy/dirty-session-errors.json

@@ -668,4 +668,4 @@
       }
     }
   ]
-}
+}

test/sessions/server-support.json renamed to test/sessions/legacy/server-support.json

@@ -0,0 +1,105 @@
+{
+  "description": "snapshot-sessions-not-supported-server-error",
+  "schemaVersion": "1.0",
+  "runOnRequirements": [
+    {
+      "minServerVersion": "3.6",
+      "maxServerVersion": "4.4.99",
+      "topologies": [
+        "replicaset, sharded-replicaset"
+      ]
+    }
+  ],
+  "createEntities": [
+    {
+      "client": {
+        "id": "client0",
+        "observeEvents": [
+          "commandStartedEvent",
+          "commandFailedEvent"
+        ]
+      }
+    },
+    {
+      "database": {
+        "id": "database0",
+        "client": "client0",
+        "databaseName": "database0"
+      }
+    },
+    {
+      "collection": {
+        "id": "collection0",
+        "database": "database0",
+        "collectionName": "collection0"
+      }
+    },
+    {
+      "session": {
+        "id": "session0",
+        "client": "client0",
+        "sessionOptions": {
+          "snapshot": true
+        }
+      }
+    }
+  ],
+  "initialData": [
+    {
+      "collectionName": "collection0",
+      "databaseName": "database0",
+      "documents": [
+        {
+          "_id": 1,
+          "x": 11
+        }
+      ]
+    }
+  ],
+  "tests": [
+    {
+      "description": "Server returns an error on find with snapshot",
+      "operations": [
+        {
+          "name": "find",
+          "object": "collection0",
+          "arguments": {
+            "session": "session0",
+            "filter": {}
+          },
+          "expectError": {
+            "isError": true,
+            "isClientError": false
+          }
+        }
+      ],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "events": [
+            {
+              "commandStartedEvent": {
+                "command": {
+                  "find": "collection0",
+                  "readConcern": {
+                    "level": "snapshot",
+                    "atClusterTime": {
+                      "$$exists": false
+                    }
+                  }
+                },
+                "commandName": "find",
+                "databaseName": "database0"
+              }
+            },
+            {
+              "commandFailedEvent": {
+                "commandName": "find"
+              }
+            }
+          ]
+        }
+      ]
+    }
+  ]
+}