PYTHON-4576 Allow update to supply sort option (#1881)

Author: blink1073

doc/changelog.rst

@@ -19,6 +19,10 @@ PyMongo 4.11 brings a number of changes including:
   until either connection succeeds or a server selection timeout error is raised.
 - Added :func:`repr` support to :class:`pymongo.operations.IndexModel`.
 - Added :func:`repr` support to :class:`pymongo.operations.SearchIndexModel`.
+- Added ``sort`` parameter to
+  :meth:`~pymongo.collection.Collection.update_one`, :meth:`~pymongo.collection.Collection.replace_one`,
+  :class:`~pymongo.operations.UpdateOne`, and
+  :class:`~pymongo.operations.UpdateMany`,
 
 Issues Resolved
 ...............

pymongo/asynchronous/bulk.py

@@ -109,6 +109,7 @@ def __init__(
         self.uses_array_filters = False
         self.uses_hint_update = False
         self.uses_hint_delete = False
+        self.uses_sort = False
         self.is_retryable = True
         self.retrying = False
         self.started_retryable_write = False
@@ -144,6 +145,7 @@ def add_update(
         collation: Optional[Mapping[str, Any]] = None,
         array_filters: Optional[list[Mapping[str, Any]]] = None,
         hint: Union[str, dict[str, Any], None] = None,
+        sort: Optional[Mapping[str, Any]] = None,
     ) -> None:
         """Create an update document and add it to the list of ops."""
         validate_ok_for_update(update)
@@ -159,6 +161,9 @@ def add_update(
         if hint is not None:
             self.uses_hint_update = True
             cmd["hint"] = hint
+        if sort is not None:
+            self.uses_sort = True
+            cmd["sort"] = sort
         if multi:
             # A bulk_write containing an update_many is not retryable.
             self.is_retryable = False
@@ -171,6 +176,7 @@ def add_replace(
         upsert: bool = False,
         collation: Optional[Mapping[str, Any]] = None,
         hint: Union[str, dict[str, Any], None] = None,
+        sort: Optional[Mapping[str, Any]] = None,
     ) -> None:
         """Create a replace document and add it to the list of ops."""
         validate_ok_for_replace(replacement)
@@ -181,6 +187,9 @@ def add_replace(
         if hint is not None:
             self.uses_hint_update = True
             cmd["hint"] = hint
+        if sort is not None:
+            self.uses_sort = True
+            cmd["sort"] = sort
         self.ops.append((_UPDATE, cmd))
 
     def add_delete(
@@ -699,6 +708,10 @@ async def execute_no_results(
             raise ConfigurationError(
                 "Must be connected to MongoDB 4.2+ to use hint on unacknowledged update commands."
             )
+        if unack and self.uses_sort and conn.max_wire_version < 25:
+            raise ConfigurationError(
+                "Must be connected to MongoDB 8.0+ to use sort on unacknowledged update commands."
+            )
         # Cannot have both unacknowledged writes and bypass document validation.
         if self.bypass_doc_val:
             raise OperationFailure(

pymongo/asynchronous/client_bulk.py

@@ -118,6 +118,7 @@ def __init__(
         self.uses_array_filters = False
         self.uses_hint_update = False
         self.uses_hint_delete = False
+        self.uses_sort = False
 
         self.is_retryable = self.client.options.retry_writes
         self.retrying = False
@@ -148,6 +149,7 @@ def add_update(
         collation: Optional[Mapping[str, Any]] = None,
         array_filters: Optional[list[Mapping[str, Any]]] = None,
         hint: Union[str, dict[str, Any], None] = None,
+        sort: Optional[Mapping[str, Any]] = None,
     ) -> None:
         """Create an update document and add it to the list of ops."""
         validate_ok_for_update(update)
@@ -169,6 +171,9 @@ def add_update(
         if collation is not None:
             self.uses_collation = True
             cmd["collation"] = collation
+        if sort is not None:
+            self.uses_sort = True
+            cmd["sort"] = sort
         if multi:
             # A bulk_write containing an update_many is not retryable.
             self.is_retryable = False
@@ -184,6 +189,7 @@ def add_replace(
         upsert: Optional[bool] = None,
         collation: Optional[Mapping[str, Any]] = None,
         hint: Union[str, dict[str, Any], None] = None,
+        sort: Optional[Mapping[str, Any]] = None,
     ) -> None:
         """Create a replace document and add it to the list of ops."""
         validate_ok_for_replace(replacement)
@@ -202,6 +208,9 @@ def add_replace(
         if collation is not None:
             self.uses_collation = True
             cmd["collation"] = collation
+        if sort is not None:
+            self.uses_sort = True
+            cmd["sort"] = sort
         self.ops.append(("replace", cmd))
         self.namespaces.append(namespace)
         self.total_ops += 1

pymongo/asynchronous/collection.py

@@ -993,6 +993,7 @@ async def _update(
         session: Optional[AsyncClientSession] = None,
         retryable_write: bool = False,
         let: Optional[Mapping[str, Any]] = None,
+        sort: Optional[Mapping[str, Any]] = None,
         comment: Optional[Any] = None,
     ) -> Optional[Mapping[str, Any]]:
         """Internal update / replace helper."""
@@ -1024,6 +1025,14 @@ async def _update(
             if not isinstance(hint, str):
                 hint = helpers_shared._index_document(hint)
             update_doc["hint"] = hint
+        if sort is not None:
+            if not acknowledged and conn.max_wire_version < 25:
+                raise ConfigurationError(
+                    "Must be connected to MongoDB 8.0+ to use sort on unacknowledged update commands."
+                )
+            common.validate_is_mapping("sort", sort)
+            update_doc["sort"] = sort
+
         command = {"update": self.name, "ordered": ordered, "updates": [update_doc]}
         if let is not None:
             common.validate_is_mapping("let", let)
@@ -1079,6 +1088,7 @@ async def _update_retryable(
         hint: Optional[_IndexKeyHint] = None,
         session: Optional[AsyncClientSession] = None,
         let: Optional[Mapping[str, Any]] = None,
+        sort: Optional[Mapping[str, Any]] = None,
         comment: Optional[Any] = None,
     ) -> Optional[Mapping[str, Any]]:
         """Internal update / replace helper."""
@@ -1102,6 +1112,7 @@ async def _update(
                 session=session,
                 retryable_write=retryable_write,
                 let=let,
+                sort=sort,
                 comment=comment,
             )
 
@@ -1122,6 +1133,7 @@ async def replace_one(
         hint: Optional[_IndexKeyHint] = None,
         session: Optional[AsyncClientSession] = None,
         let: Optional[Mapping[str, Any]] = None,
+        sort: Optional[Mapping[str, Any]] = None,
         comment: Optional[Any] = None,
     ) -> UpdateResult:
         """Replace a single document matching the filter.
@@ -1176,8 +1188,13 @@ async def replace_one(
             aggregate expression context (e.g. "$$var").
         :param comment: A user-provided comment to attach to this
             command.
+        :param sort: Specify which document the operation updates if the query matches
+            multiple documents. The first document matched by the sort order will be updated.
+            This option is only supported on MongoDB 8.0 and above.
         :return: - An instance of :class:`~pymongo.results.UpdateResult`.
 
+        .. versionchanged:: 4.11
+           Added ``sort`` parameter.
         .. versionchanged:: 4.1
            Added ``let`` parameter.
            Added ``comment`` parameter.
@@ -1209,6 +1226,7 @@ async def replace_one(
                 hint=hint,
                 session=session,
                 let=let,
+                sort=sort,
                 comment=comment,
             ),
             write_concern.acknowledged,
@@ -1225,6 +1243,7 @@ async def update_one(
         hint: Optional[_IndexKeyHint] = None,
         session: Optional[AsyncClientSession] = None,
         let: Optional[Mapping[str, Any]] = None,
+        sort: Optional[Mapping[str, Any]] = None,
         comment: Optional[Any] = None,
     ) -> UpdateResult:
         """Update a single document matching the filter.
@@ -1283,11 +1302,16 @@ async def update_one(
             constant or closed expressions that do not reference document
             fields. Parameters can then be accessed as variables in an
             aggregate expression context (e.g. "$$var").
+        :param sort: Specify which document the operation updates if the query matches
+            multiple documents. The first document matched by the sort order will be updated.
+            This option is only supported on MongoDB 8.0 and above.
         :param comment: A user-provided comment to attach to this
             command.
 
         :return: - An instance of :class:`~pymongo.results.UpdateResult`.
 
+        .. versionchanged:: 4.11
+           Added ``sort`` parameter.
         .. versionchanged:: 4.1
            Added ``let`` parameter.
            Added ``comment`` parameter.
@@ -1322,6 +1346,7 @@ async def update_one(
                 hint=hint,
                 session=session,
                 let=let,
+                sort=sort,
                 comment=comment,
             ),
             write_concern.acknowledged,