Refactoring allow_retry implementation

Author: apetenchea

arango/aql.py

@@ -313,7 +313,7 @@ def execute(
             development to catch issues early. If set to False, warnings are
             returned with the query result. There is a server configuration
             option "--query.fail-on-warning" for setting the default value for
-            this behaviour so it does not need to be set per-query.
+            this behaviour, so it does not need to be set per-query.
         :type fail_on_warning: bool
         :param profile: Return additional profiling details in the cursor,
             unless the query cache is used.
@@ -437,7 +437,7 @@ def execute(
         def response_handler(resp: Response) -> Cursor:
             if not resp.is_success:
                 raise AQLQueryExecuteError(resp, request)
-            return Cursor(self._conn, resp.body)
+            return Cursor(self._conn, resp.body, allow_retry=allow_retry)
 
         return self._execute(request, response_handler)
 

arango/cursor.py

@@ -27,6 +27,10 @@ class Cursor:
     :type init_data: dict
     :param cursor_type: Cursor type ("cursor" or "export").
     :type cursor_type: str
+    :param allow_retry: If set to True, the cursor will always attempt to fetch
+            the latest batch from server even if the previous attempt failed.
+            This option is only available for server versions 3.11 and above.
+    :type allow_retry: bool
     """
 
     __slots__ = [
@@ -41,16 +45,19 @@ class Cursor:
         "_has_more",
         "_batch",
         "_next_batch_id",
+        "_allow_retry",
     ]
 
     def __init__(
         self,
         connection: BaseConnection,
         init_data: Json,
         cursor_type: str = "cursor",
+        allow_retry: bool = False,
     ) -> None:
         self._conn = connection
         self._type = cursor_type
+        self._allow_retry = allow_retry
         self._batch: Deque[Any] = deque()
         self._id = None
         self._count: Optional[int] = None
@@ -280,33 +287,12 @@ def fetch(self) -> Json:
         """
         if self._id is None:
             raise CursorStateError("cursor ID not set")
-        request = Request(method="post", endpoint=f"/_api/{self._type}/{self._id}")
-        resp = self._conn.send_request(request)
-
-        if not resp.is_success:
-            raise CursorNextError(resp, request)
 
-        return self._update(resp.body)
-
-    def retry(self) -> Json:
-        """Retry fetching the next batch from server and update the cursor.
+        endpoint = f"/_api/{self._type}/{self._id}"
+        if self._allow_retry and self._next_batch_id is not None:
+            endpoint += f"/{self._next_batch_id}"
 
-        Available only if the ``allowRetry`` query options is enabled.
-        Introduced in 3.11.
-
-        :return: New batch details.
-        :rtype: dict
-        :raise arango.exceptions.CursorNextError: If batch retrieval fails.
-        :raise arango.exceptions.CursorStateError: If cursor ID is not set.
-        """
-        if self._id is None:
-            raise CursorStateError("cursor ID not set")
-        if self._id is None:
-            raise CursorStateError("nextBatchId not set")
-        request = Request(
-            method="post",
-            endpoint=f"/_api/{self._type}/{self._id}/{self._next_batch_id}",
-        )
+        request = Request(method="post", endpoint=endpoint)
         resp = self._conn.send_request(request)
 
         if not resp.is_success:

docs/cursor.rst

@@ -33,8 +33,7 @@ number of items in the result set may or may not be known in advance.
         'FOR doc IN students FILTER doc.age > @val RETURN doc',
         bind_vars={'val': 17},
         batch_size=2,
-        count=True,
-        allow_retry=True
+        count=True
     )
 
     # Get the cursor ID.
@@ -73,11 +72,7 @@ number of items in the result set may or may not be known in advance.
     cursor.pop()
 
     # Fetch the next batch and add them to the cursor object.
-    try:
-        cursor.fetch()
-    except CursorNextError:
-        # Retry fetching the latest batch from the cursor.
-        cursor.retry()
+    cursor.fetch()
 
     # Delete the cursor from the server.
     cursor.close()
@@ -121,3 +116,54 @@ instead.
         cursor.fetch()
     while not cursor.empty(): # Pop until nothing is left on the cursor.
         cursor.pop()
+
+With ArangoDB 3.11.0 or higher, you can also use the `allow_retry`
+parameter of :func:`arango.aql.AQL.execute` to automatically retry
+the request if the cursor encountered any issues during the previous
+fetch operation. Note that this feature causes the server to cache the
+last batch. To allow re-fetching of the very last batch of the query,
+the server cannot automatically delete the cursor. Once you have successfully
+received the last batch, you should call :func:`arango.cursor.Cursor.close`.
+
+**Example:**
+
+.. code-block:: python
+
+    from arango import ArangoClient
+
+    # Initialize the ArangoDB client.
+    client = ArangoClient()
+
+    # Connect to "test" database as root user.
+    db = client.db('test', username='root', password='passwd')
+
+    # Set up some test data to query against.
+    db.collection('students').insert_many([
+        {'_key': 'Abby', 'age': 22},
+        {'_key': 'John', 'age': 18},
+        {'_key': 'Mary', 'age': 21},
+        {'_key': 'Suzy', 'age': 23},
+        {'_key': 'Dave', 'age': 20}
+    ])
+
+    # Execute an AQL query which returns a cursor object.
+    cursor = db.aql.execute(
+        'FOR doc IN students FILTER doc.age > @val RETURN doc',
+        bind_vars={'val': 17},
+        batch_size=2,
+        count=True,
+        allow_retry=True
+    )
+
+    while cursor.has_more():
+        try:
+            cursor.fetch()
+        except ConnectionError:
+            # Retry the request.
+            continue
+
+    while not cursor.empty():
+        cursor.pop()
+
+    # Delete the cursor from the server.
+    cursor.close()

tests/test_cursor.py

@@ -263,7 +263,7 @@ def test_cursor_manual_fetch_and_pop(db, col, docs):
     assert err.value.message == "current batch is empty"
 
 
-def test_cursor_retry_disabled(db, col, db_version):
+def test_cursor_retry_disabled(db, col, docs, db_version):
     cursor = db.aql.execute(
         f"FOR d IN {col.name} SORT d._key RETURN d",
         count=True,
@@ -275,12 +275,17 @@ def test_cursor_retry_disabled(db, col, db_version):
     )
     result = cursor.fetch()
     assert result["id"] == cursor.id
-    cursor._next_batch_id = "2"
 
-    if db_version >= version.parse("3.11.0"):
-        with pytest.raises(CursorNextError) as err:
-            cursor.retry()
-            assert err.value.message == "batch id not found"
+    while not cursor.empty():
+        cursor.pop()
+
+    # The next batch ID should have no effect
+    cursor._next_batch_id = "2"
+    result = cursor.fetch()
+    if db_version >= version.parse("3.11.1"):
+        assert result["next_batch_id"] == "4"
+    doc = cursor.pop()
+    assert clean_doc(doc) == docs[2]
 
     assert cursor.close(ignore_missing=True)
 
@@ -312,7 +317,7 @@ def test_cursor_retry(db, col, docs, db_version):
     # Decrease the next batch ID as if the previous fetch failed
     if db_version >= version.parse("3.11.0"):
         cursor._next_batch_id = "2"
-        result = cursor.retry()
+        result = cursor.fetch()
         assert result["id"] == cursor.id
         assert result["next_batch_id"] == "3"
         doc = cursor.pop()
@@ -335,6 +340,12 @@ def test_cursor_retry(db, col, docs, db_version):
     doc = cursor.pop()
     assert clean_doc(doc) == docs[-1]
 
+    if db_version >= version.parse("3.11.0"):
+        # We should be able to fetch the last batch again
+        cursor.fetch()
+        doc = cursor.pop()
+        assert clean_doc(doc) == docs[-1]
+
     if db_version >= version.parse("3.11.0"):
         assert cursor.close()
     else: