Handling cursors separately

Author: apetenchea

arangoasync/aql.py

@@ -23,7 +23,7 @@
     AQLQueryTrackingSetError,
     AQLQueryValidateError,
 )
-from arangoasync.executor import ApiExecutor
+from arangoasync.executor import ApiExecutor, DefaultApiExecutor, NonAsyncExecutor
 from arangoasync.request import Method, Request
 from arangoasync.response import Response
 from arangoasync.result import Result
@@ -326,7 +326,14 @@ async def execute(
         def response_handler(resp: Response) -> Cursor:
             if not resp.is_success:
                 raise AQLQueryExecuteError(resp, request)
-            return Cursor(self._executor, self.deserializer.loads(resp.raw_body))
+            if self._executor.context == "async":
+                # We cannot have a cursor getting back async jobs
+                executor: NonAsyncExecutor = DefaultApiExecutor(
+                    self._executor.connection
+                )
+            else:
+                executor = cast(NonAsyncExecutor, self._executor)
+            return Cursor(executor, self.deserializer.loads(resp.raw_body))
 
         return await self._executor.execute(request, response_handler)
 

arangoasync/cursor.py

@@ -12,7 +12,7 @@
     CursorNextError,
     CursorStateError,
 )
-from arangoasync.executor import ApiExecutor
+from arangoasync.executor import NonAsyncExecutor
 from arangoasync.request import Method, Request
 from arangoasync.response import Response
 from arangoasync.serialization import Deserializer, Serializer
@@ -39,7 +39,7 @@ class Cursor:
             is created.
     """
 
-    def __init__(self, executor: ApiExecutor, data: Json) -> None:
+    def __init__(self, executor: NonAsyncExecutor, data: Json) -> None:
         self._executor = executor
         self._cached: Optional[bool] = None
         self._count: Optional[int] = None

arangoasync/database.py

@@ -2,6 +2,7 @@
     "Database",
     "StandardDatabase",
     "TransactionDatabase",
+    "AsyncDatabase",
 ]
 
 

arangoasync/executor.py

@@ -1,3 +1,11 @@
+__all__ = [
+    "ApiExecutor",
+    "DefaultApiExecutor",
+    "NonAsyncExecutor",
+    "TransactionApiExecutor",
+    "AsyncApiExecutor",
+]
+
 from typing import Callable, Optional, TypeVar
 
 from arangoasync.connection import Connection
@@ -11,10 +19,10 @@
 T = TypeVar("T")
 
 
-class DefaultApiExecutor:
-    """Default API executor.
+class ExecutorContext:
+    """Base class for API executors.
 
-    Responsible for executing requests and handling responses.
+    Not to be exported publicly.
 
     Args:
         connection: HTTP connection.
@@ -27,10 +35,6 @@ def __init__(self, connection: Connection) -> None:
     def connection(self) -> Connection:
         return self._conn
 
-    @property
-    def context(self) -> str:
-        return "default"
-
     @property
     def db_name(self) -> str:
         return self._conn.db_name
@@ -49,6 +53,23 @@ def serialize(self, data: Json) -> str:
     def deserialize(self, data: bytes) -> Json:
         return self.deserializer.loads(data)
 
+
+class DefaultApiExecutor(ExecutorContext):
+    """Default API executor.
+
+    Responsible for executing requests and handling responses.
+
+    Args:
+        connection: HTTP connection.
+    """
+
+    def __init__(self, connection: Connection) -> None:
+        super().__init__(connection)
+
+    @property
+    def context(self) -> str:
+        return "default"
+
     async def execute(
         self, request: Request, response_handler: Callable[[Response], T]
     ) -> T:
@@ -62,7 +83,7 @@ async def execute(
         return response_handler(response)
 
 
-class TransactionApiExecutor(DefaultApiExecutor):
+class TransactionApiExecutor(ExecutorContext):
     """Executes transaction API requests.
 
     Args:
@@ -97,7 +118,7 @@ async def execute(
         return response_handler(response)
 
 
-class AsyncApiExecutor(DefaultApiExecutor):
+class AsyncApiExecutor(ExecutorContext):
     """Executes asynchronous API requests (jobs).
 
     Args:
@@ -144,3 +165,4 @@ async def execute(
 
 
 ApiExecutor = DefaultApiExecutor | TransactionApiExecutor | AsyncApiExecutor
+NonAsyncExecutor = DefaultApiExecutor | TransactionApiExecutor

arangoasync/job.py

@@ -97,7 +97,8 @@ async def result(self) -> T:
         Raises:
             ArangoError: If the job raised an exception or there was a problem with
                 the request.
-            AsyncJobResultError: If retrieval fails.
+            AsyncJobResultError: If retrieval fails, because job no longer exists or
+                is still pending.
 
         References:
             - `get-the-results-of-an-async-job <https://docs.arangodb.com/stable/develop/http-api/jobs/#get-the-results-of-an-async-job>`__
@@ -112,10 +113,11 @@ async def result(self) -> T:
             # The job result is available on the server
             return self._response_handler(response)
 
-        # The job is not known (anymore).
-        # We can tell the status from the HTTP status code.
         if response.status_code == 204:
+            # The job is still in the pending queue or not yet finished.
             raise AsyncJobResultError(response, request, self._not_done())
+        # The job is not known (anymore).
+        # We can tell the status from the HTTP status code.
         if response.error_code == HTTP_NOT_FOUND:
             raise AsyncJobResultError(response, request, self._not_found())
         raise AsyncJobResultError(response, request)

tests/test_async.py

@@ -7,6 +7,7 @@
     AQLQueryExecuteError,
     AsyncJobCancelError,
     AsyncJobListError,
+    AsyncJobResultError,
 )
 
 
@@ -101,3 +102,34 @@ async def test_async_result(db, bad_db, doc_col, docs):
     job4 = await aql.execute("RETURN 1")
     await job4.wait()
     await job4.clear()
+
+    # Attempt to get the result of a pending job
+    job5 = await aql.execute("RETURN SLEEP(5)")
+    time.sleep(1)
+    with pytest.raises(AsyncJobResultError):
+        _ = await job5.result()
+    await job5.wait()
+
+
+@pytest.mark.asyncio
+async def test_async_cursor(db, doc_col, docs):
+    # Insert some documents first
+    await asyncio.gather(*(doc_col.insert(doc) for doc in docs))
+
+    async_db = db.begin_async_execution()
+    aql = async_db.aql
+    job = await aql.execute(
+        f"FOR d IN {doc_col.name} SORT d._key RETURN d",
+        count=True,
+        batch_size=1,
+        ttl=1000,
+    )
+    await job.wait()
+
+    # Get the cursor. Bear in mind that its underlying executor is async.
+    doc_cnt = 0
+    cursor = await job.result()
+    async with cursor as ctx:
+        async for _ in ctx:
+            doc_cnt += 1
+    assert doc_cnt == len(docs)