Make async iterator fully async

Fixes #804

Author: mariusconjeaud

neomodel/async_/database.py

@@ -7,7 +7,7 @@
 import sys
 import time
 from contextvars import ContextVar
-from typing import TYPE_CHECKING, Any, Callable, TextIO
+from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, TextIO
 from urllib.parse import quote, unquote, urlparse
 
 from neo4j import (
@@ -753,6 +753,73 @@ async def _run_cypher_query(
 
         return results, meta
 
+    async def _stream_cypher_query(
+        self,
+        session: AsyncSession | AsyncTransaction,
+        query: str,
+        params: dict[str, Any],
+        handle_unique: bool,
+        resolve_objects: bool,
+    ) -> AsyncIterator[tuple[list, tuple[str, ...]]]:
+        """
+        Stream query results one record at a time without loading all into memory.
+
+        This is an internal method used for async iteration. It yields results
+        as they arrive from the database instead of collecting them all first.
+
+        :param session: Neo4j session or transaction
+        :param query: Cypher query string
+        :param params: Query parameters
+        :param handle_unique: Whether to raise UniqueProperty on constraint violations
+        :param resolve_objects: Whether to resolve nodes to neomodel objects
+        :yields: Tuple of (values_list, keys_tuple) for each record
+        """
+        try:
+            start = time.time()
+            if self._parallel_runtime:
+                query = "CYPHER runtime=parallel " + query
+
+            response: AsyncResult = await session.run(query=query, parameters=params)
+            keys = response.keys()
+
+            # Stream results one record at a time
+            async for record in response:
+                values = list(record.values())
+
+                if resolve_objects:
+                    # Resolve objects for this single record
+                    for idx, value in enumerate(values):
+                        values[idx] = self._object_resolution(value)
+
+                yield values, keys
+
+            end = time.time()
+            tte = end - start
+            if os.environ.get("NEOMODEL_CYPHER_DEBUG", False) and tte > float(
+                os.environ.get("NEOMODEL_SLOW_QUERIES", 0)
+            ):
+                logger.debug(
+                    "query: "
+                    + query
+                    + "\nparams: "
+                    + repr(params)
+                    + f"\ntook: {tte:.2g}s\n"
+                )
+
+        except ClientError as e:
+            if e.code == "Neo.ClientError.Schema.ConstraintValidationFailed":
+                if hasattr(e, "message") and e.message is not None:
+                    if "already exists with label" in e.message and handle_unique:
+                        raise UniqueProperty(e.message) from e
+                    raise ConstraintValidationFailed(e.message) from e
+                raise ConstraintValidationFailed(
+                    "A constraint validation failed"
+                ) from e
+
+            exc_info = sys.exc_info()
+            if exc_info[1] is not None and exc_info[2] is not None:
+                raise exc_info[1].with_traceback(exc_info[2])
+
     async def get_id_method(self) -> str:
         db_version = await self.database_version
         if db_version is None:

neomodel/async_/match.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass
 from typing import Any, AsyncIterator, Optional, Union
 
+from neomodel._async_compat.util import AsyncUtil
 from neomodel.async_ import relationship_manager
 from neomodel.async_.database import adb
 from neomodel.async_.node import AsyncStructuredNode
@@ -1198,25 +1199,77 @@ async def _execute(self, lazy: bool = False, dict_output: bool = False) -> Any:
                         for item in self._ast.additional_return
                     ]
         query = self.build_query()
-        results, prop_names = await adb.cypher_query(
-            query,
-            self._query_params,
-            resolve_objects=True,
-        )
-        if dict_output:
-            for item in results:
-                yield dict(zip(prop_names, item))
-            return
-        # The following is not as elegant as it could be but had to be copied from the
-        # version prior to cypher_query with the resolve_objects capability.
-        # It seems that certain calls are only supposed to be focusing to the first
-        # result item returned (?)
-        if results and len(results[0]) == 1:
-            for n in results:
-                yield n[0]
+
+        # Use streaming for async code to avoid loading all results into memory
+        if AsyncUtil.is_async_code:
+            # Helper to process streaming results
+            async def process_stream(stream_iterator):
+                first_result = True
+                result_has_single_column = False
+                async for values, prop_names in stream_iterator:
+                    if first_result:
+                        # Determine format on first result
+                        result_has_single_column = len(values) == 1
+                        first_result = False
+
+                    if dict_output:
+                        yield dict(zip(prop_names, values))
+                    elif result_has_single_column:
+                        yield values[0]
+                    else:
+                        yield values
+
+            # Stream results one by one from the database
+            if adb._active_transaction:
+                # Use current transaction if active
+                stream = adb._stream_cypher_query(
+                    adb._active_transaction,
+                    query,
+                    self._query_params,
+                    handle_unique=True,
+                    resolve_objects=True,
+                )
+                async for item in process_stream(stream):
+                    yield item
+                return
+            else:
+                # Create a session for streaming
+                # Note: We need to keep the session open during iteration
+                async with adb.driver.session(
+                    database=adb._database_name,
+                    impersonated_user=adb.impersonated_user,
+                ) as session:
+                    stream = adb._stream_cypher_query(
+                        session,
+                        query,
+                        self._query_params,
+                        handle_unique=True,
+                        resolve_objects=True,
+                    )
+                    async for item in process_stream(stream):
+                        yield item
+                return
         else:
-            for result in results:
-                yield result
+            # Sync code path: use traditional approach (fetch all results)
+            results, prop_names = await adb.cypher_query(
+                query,
+                self._query_params,
+                resolve_objects=True,
+            )
+            if dict_output:
+                for item in results:
+                    yield dict(zip(prop_names, item))
+                return
+            # The following is not as elegant as it could be but had to be copied from the
+            # version prior to cypher_query with the resolve_objects capability.
+            # It seems that certain calls are only supposed to be focusing to the first
+            # result item returned (?)
+            if results and len(results[0]) == 1:
+                for n in results:
+                    yield n[0]
+            else:
+                for result in results:
+                    yield result
 
 
 class AsyncBaseSet:
@@ -1243,6 +1296,16 @@ async def all(self, lazy: bool = False) -> list:
         return results
 
     async def __aiter__(self) -> AsyncIterator:
+        """
+        Async iterator that streams results from the database one at a time.
+
+        This provides true async iteration without loading all results into memory first.
+        For large result sets, this is much more memory efficient than using all().
+
+        Example:
+            async for node in Coffee.nodes:
+                print(node.name)  # Process each node as it arrives
+        """
         ast = await self.query_cls(self).build_ast()
         async for item in ast._execute():
             yield item

neomodel/sync_/database.py

@@ -7,7 +7,7 @@
 import sys
 import time
 from contextvars import ContextVar
-from typing import TYPE_CHECKING, Any, Callable, TextIO
+from typing import TYPE_CHECKING, Any, Callable, Iterator, TextIO
 from urllib.parse import quote, unquote, urlparse
 
 from neo4j import (
@@ -749,6 +749,73 @@ def _run_cypher_query(
 
         return results, meta
 
+    def _stream_cypher_query(
+        self,
+        session: Session | Transaction,
+        query: str,
+        params: dict[str, Any],
+        handle_unique: bool,
+        resolve_objects: bool,
+    ) -> Iterator[tuple[list, tuple[str, ...]]]:
+        """
+        Stream query results one record at a time without loading all into memory.
+
+        This is an internal method used for iteration. It yields results
+        as they arrive from the database instead of collecting them all first.
+
+        :param session: Neo4j session or transaction
+        :param query: Cypher query string
+        :param params: Query parameters
+        :param handle_unique: Whether to raise UniqueProperty on constraint violations
+        :param resolve_objects: Whether to resolve nodes to neomodel objects
+        :yields: Tuple of (values_list, keys_tuple) for each record
+        """
+        try:
+            start = time.time()
+            if self._parallel_runtime:
+                query = "CYPHER runtime=parallel " + query
+
+            response: Result = session.run(query=query, parameters=params)
+            keys = response.keys()
+
+            # Stream results one record at a time
+            for record in response:
+                values = list(record.values())
+
+                if resolve_objects:
+                    # Resolve objects for this single record
+                    for idx, value in enumerate(values):
+                        values[idx] = self._object_resolution(value)
+
+                yield values, keys
+
+            end = time.time()
+            tte = end - start
+            if os.environ.get("NEOMODEL_CYPHER_DEBUG", False) and tte > float(
+                os.environ.get("NEOMODEL_SLOW_QUERIES", 0)
+            ):
+                logger.debug(
+                    "query: "
+                    + query
+                    + "\nparams: "
+                    + repr(params)
+                    + f"\ntook: {tte:.2g}s\n"
+                )
+
+        except ClientError as e:
+            if e.code == "Neo.ClientError.Schema.ConstraintValidationFailed":
+                if hasattr(e, "message") and e.message is not None:
+                    if "already exists with label" in e.message and handle_unique:
+                        raise UniqueProperty(e.message) from e
+                    raise ConstraintValidationFailed(e.message) from e
+                raise ConstraintValidationFailed(
+                    "A constraint validation failed"
+                ) from e
+
+            exc_info = sys.exc_info()
+            if exc_info[1] is not None and exc_info[2] is not None:
+                raise exc_info[1].with_traceback(exc_info[2])
+
     def get_id_method(self) -> str:
         db_version = self.database_version
         if db_version is None: