Add more tests

Author: NoahStapp

pymongo/asynchronous/collection.py

@@ -2946,6 +2946,7 @@ async def _aggregate(
             session,
             retryable=not cmd._performs_write,
             operation=_Op.AGGREGATE,
+            is_aggregate_write=cmd._performs_write,
         )
 
     async def aggregate(

pymongo/asynchronous/mongo_client.py

@@ -2010,6 +2010,7 @@ async def _retry_internal(
         retryable: bool = False,
         operation_id: Optional[int] = None,
         is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2022,6 +2023,7 @@ async def _retry_internal(
         :param read_pref: Topology of read operation, defaults to None
         :param retryable: If the operation should be retried once, defaults to None
         :param is_run_command: If this is a runCommand operation, defaults to False
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
 
         :return: Output of the calling func()
         """
@@ -2037,6 +2039,7 @@ async def _retry_internal(
             retryable=retryable,
             operation_id=operation_id,
             is_run_command=is_run_command,
+            is_aggregate_write=is_aggregate_write,
         ).run()
 
     async def _retryable_read(
@@ -2049,6 +2052,7 @@ async def _retryable_read(
         retryable: bool = True,
         operation_id: Optional[int] = None,
         is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2065,6 +2069,7 @@ async def _retryable_read(
         :param retryable: if we should attempt retries
             (may not always be supported even if supplied), defaults to False
         :param is_run_command: If this is a runCommand operation, defaults to False.
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2084,6 +2089,7 @@ async def _retryable_read(
                 retryable=retryable,
                 operation_id=operation_id,
                 is_run_command=is_run_command,
+                is_aggregate_write=is_aggregate_write,
             )
 
     async def _retryable_write(
@@ -2754,6 +2760,7 @@ def __init__(
         retryable: bool = False,
         operation_id: Optional[int] = None,
         is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
@@ -2777,6 +2784,7 @@ def __init__(
         self._operation_id = operation_id
         self._attempt_number = 0
         self._is_run_command = is_run_command
+        self._is_aggregate_write = is_aggregate_write
 
     async def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2822,6 +2830,9 @@ async def run(self) -> T:
                     self._client.options.retry_reads and self._client.options.retry_writes
                 ):
                     raise
+                if self._is_aggregate_write and not self._client.options.retry_writes:
+                    raise
+
                 # Execute specialized catch on read
                 if self._is_read:
                     if isinstance(exc, (ConnectionFailure, OperationFailure)):
@@ -2870,9 +2881,9 @@ async def run(self) -> T:
                     always_retryable = exc_to_check.has_error_label("RetryableError") and overloaded
 
                     # Always retry abortTransaction and commitTransaction up to once
-                    if not (self._client.options.retry_writes and self._retryable) and (
-                        not always_retryable
-                        and self._operation not in ["abortTransaction", "commitTransaction"]
+                    if self._operation not in ["abortTransaction", "commitTransaction"] and (
+                        not self._client.options.retry_writes
+                        or not (self._retryable or always_retryable)
                     ):
                         raise
                     if retryable_write_label or always_retryable:

pymongo/synchronous/collection.py

@@ -2939,6 +2939,7 @@ def _aggregate(
             session,
             retryable=not cmd._performs_write,
             operation=_Op.AGGREGATE,
+            is_aggregate_write=cmd._performs_write,
         )
 
     def aggregate(

pymongo/synchronous/mongo_client.py

@@ -2006,6 +2006,7 @@ def _retry_internal(
         retryable: bool = False,
         operation_id: Optional[int] = None,
         is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2018,6 +2019,7 @@ def _retry_internal(
         :param read_pref: Topology of read operation, defaults to None
         :param retryable: If the operation should be retried once, defaults to None
         :param is_run_command: If this is a runCommand operation, defaults to False
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
 
         :return: Output of the calling func()
         """
@@ -2033,6 +2035,7 @@ def _retry_internal(
             retryable=retryable,
             operation_id=operation_id,
             is_run_command=is_run_command,
+            is_aggregate_write=is_aggregate_write,
         ).run()
 
     def _retryable_read(
@@ -2045,6 +2048,7 @@ def _retryable_read(
         retryable: bool = True,
         operation_id: Optional[int] = None,
         is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2061,6 +2065,7 @@ def _retryable_read(
         :param retryable: if we should attempt retries
             (may not always be supported even if supplied), defaults to False
         :param is_run_command: If this is a runCommand operation, defaults to False.
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2080,6 +2085,7 @@ def _retryable_read(
                 retryable=retryable,
                 operation_id=operation_id,
                 is_run_command=is_run_command,
+                is_aggregate_write=is_aggregate_write,
             )
 
     def _retryable_write(
@@ -2744,6 +2750,7 @@ def __init__(
         retryable: bool = False,
         operation_id: Optional[int] = None,
         is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
@@ -2767,6 +2774,7 @@ def __init__(
         self._operation_id = operation_id
         self._attempt_number = 0
         self._is_run_command = is_run_command
+        self._is_aggregate_write = is_aggregate_write
 
     def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2812,6 +2820,9 @@ def run(self) -> T:
                     self._client.options.retry_reads and self._client.options.retry_writes
                 ):
                     raise
+                if self._is_aggregate_write and not self._client.options.retry_writes:
+                    raise
+
                 # Execute specialized catch on read
                 if self._is_read:
                     if isinstance(exc, (ConnectionFailure, OperationFailure)):
@@ -2860,9 +2871,9 @@ def run(self) -> T:
                     always_retryable = exc_to_check.has_error_label("RetryableError") and overloaded
 
                     # Always retry abortTransaction and commitTransaction up to once
-                    if not (self._client.options.retry_writes and self._retryable) and (
-                        not always_retryable
-                        and self._operation not in ["abortTransaction", "commitTransaction"]
+                    if self._operation not in ["abortTransaction", "commitTransaction"] and (
+                        not self._client.options.retry_writes
+                        or not (self._retryable or always_retryable)
                     ):
                         raise
                     if retryable_write_label or always_retryable: