PYTHON-5809 - Send afterClusterTime on writes in causally-consistent sessions (#2851)

Author: NoahStapp

.evergreen/resync-specs.sh

@@ -117,6 +117,9 @@ do
     apm|APM|command-monitoring|command_monitoring)
       cpjson command-logging-and-monitoring/tests/monitoring command_monitoring
       ;;
+    causal|causal-consistency|causal_consistency)
+      cpjson causal-consistency/tests/ causal_consistency
+      ;;
     command-logging|command_logging)
       cpjson command-logging-and-monitoring/tests/logging command_logging
       ;;

.evergreen/spec-patch/PYTHON-5809.patch

@@ -172,6 +172,7 @@
     WTimeoutError,
 )
 from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES
+from pymongo.operations import _WRITES_WITH_CLUSTER_TIME
 from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import ReadPreference, _ServerMode
 from pymongo.server_type import SERVER_TYPE
@@ -1111,7 +1112,12 @@ def _apply_to(
             return
         self._check_ended()
         self._materialize(conn.logical_session_timeout_minutes)
-        if self.options.snapshot:
+        # Add afterClusterTime on snapshot reads or writes in causally-consistent sessions
+        if self.options.snapshot or (
+            self.options.causal_consistency
+            and not self.in_transaction
+            and operation in _WRITES_WITH_CLUSTER_TIME
+        ):
             self._update_read_concern(command, conn)
 
         self._server_session.last_use = time.monotonic()

pymongo/asynchronous/client_session.py

@@ -82,6 +82,22 @@ class _Op(str, enum.Enum):
     TEST = "testOperation"
 
 
+_WRITES_WITH_CLUSTER_TIME = frozenset(
+    {
+        _Op.INSERT.value,
+        _Op.UPDATE.value,
+        _Op.FIND_AND_MODIFY.value,
+        _Op.DELETE.value,
+        _Op.BULK_WRITE.value,
+        _Op.CREATE.value,
+        _Op.CREATE_INDEXES.value,
+        _Op.DROP.value,
+        _Op.DROP_DATABASE.value,
+        _Op.DROP_INDEXES.value,
+    }
+)
+
+
 class InsertOne(Generic[_DocumentType]):
     """Represents an insert_one operation."""
 

pymongo/operations.py

@@ -169,6 +169,7 @@
     WTimeoutError,
 )
 from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES
+from pymongo.operations import _WRITES_WITH_CLUSTER_TIME
 from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import ReadPreference, _ServerMode
 from pymongo.server_type import SERVER_TYPE
@@ -1107,7 +1108,12 @@ def _apply_to(
             return
         self._check_ended()
         self._materialize(conn.logical_session_timeout_minutes)
-        if self.options.snapshot:
+        # Add afterClusterTime on snapshot reads or writes in causally-consistent sessions
+        if self.options.snapshot or (
+            self.options.causal_consistency
+            and not self.in_transaction
+            and operation in _WRITES_WITH_CLUSTER_TIME
+        ):
             self._update_read_concern(command, conn)
 
         self._server_session.last_use = time.monotonic()

pymongo/synchronous/client_session.py

@@ -0,0 +1,33 @@
+# 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.
+
+"""Test the causal consistency unified spec tests."""
+from __future__ import annotations
+
+import os
+import pathlib
+import sys
+
+sys.path[0:0] = [""]
+
+from test import unittest
+from test.asynchronous.unified_format import generate_test_classes, get_test_path
+
+_IS_SYNC = False
+
+# Generate unified tests.
+globals().update(generate_test_classes(get_test_path("causal_consistency"), module=__name__))
+
+if __name__ == "__main__":
+    unittest.main()

test/asynchronous/test_causal_consistency.py

@@ -1138,52 +1138,7 @@ async def _test_no_read_concern(self, op):
             self.assertIsNone(rc)
 
     @async_client_context.require_no_standalone
-    async def test_writes_do_not_include_read_concern(self):
-        await self._test_no_read_concern(
-            lambda coll, session: coll.bulk_write([InsertOne[dict]({})], session=session)
-        )
-        await self._test_no_read_concern(lambda coll, session: coll.insert_one({}, session=session))
-        await self._test_no_read_concern(
-            lambda coll, session: coll.insert_many([{}], session=session)
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.replace_one({"_id": 1}, {"x": 1}, session=session)
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.update_one({}, {"$set": {"X": 1}}, session=session)
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.update_many({}, {"$set": {"x": 1}}, session=session)
-        )
-        await self._test_no_read_concern(lambda coll, session: coll.delete_one({}, session=session))
-        await self._test_no_read_concern(
-            lambda coll, session: coll.delete_many({}, session=session)
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.find_one_and_replace({"x": 1}, {"y": 1}, session=session)
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.find_one_and_update(
-                {"y": 1}, {"$set": {"x": 1}}, session=session
-            )
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.find_one_and_delete({"x": 1}, session=session)
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.create_index("foo", session=session)
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.create_indexes(
-                [IndexModel([("bar", ASCENDING)])], session=session
-            )
-        )
-        await self._test_no_read_concern(
-            lambda coll, session: coll.drop_index("foo_1", session=session)
-        )
-        await self._test_no_read_concern(lambda coll, session: coll.drop_indexes(session=session))
-
-        # Not a write, but explain also doesn't support readConcern.
+    async def test_explain_does_not_include_read_concern(self):
         await self._test_no_read_concern(
             lambda coll, session: coll.find({}, session=session).explain()
         )

test/asynchronous/test_session.py

@@ -761,6 +761,12 @@ async def _databaseOperation_createChangeStream(self, target, *args, **kwargs):
     async def _collectionOperation_createChangeStream(self, target, *args, **kwargs):
         return await self.__entityOperation_createChangeStream(target, *args, **kwargs)
 
+    async def _clientOperation_dropDatabase(self, target, **kwargs):
+        self.__raise_if_unsupported("dropDatabase", target, AsyncMongoClient)
+        return await target.drop_database(
+            name_or_database=kwargs.pop("database"), session=kwargs.pop("session", None)
+        )
+
     async def _databaseOperation_runCommand(self, target, **kwargs):
         self.__raise_if_unsupported("runCommand", target, AsyncDatabase)
         # Ensure the first key is the command name.