feat: End a stream with an error when neither an event nor a heartbeat arrives for 30 seconds.

README.md

@@ -292,6 +292,8 @@ async for row in client.run_eventql_query(
 
 *Note that each row returned by the generator matches the projection specified in your query.*
 
+While the query runs, the server sends a heartbeat every second whenever there is no row to send. If neither a row nor a heartbeat arrives for 30 seconds, the generator closes the connection and raises a `HeartbeatTimeoutError`, which is a subclass of `ServerError`. For details, see [Detecting a Stalled Connection](#detecting-a-stalled-connection).
+
 #### Aborting a Query
 
 If you need to abort a query use `break` or `return` within the `async for` loop. However, this only works if there is currently an iteration going on.
@@ -456,6 +458,26 @@ async for event in client.observe_events(
 
 *Note that `from_latest_event` and `lower_bound` can not be provided at the same time.*
 
+#### Detecting a Stalled Connection
+
+While observing, the server sends a heartbeat every second whenever there is no new event. If neither an event nor a heartbeat arrives for 30 seconds, for example because a proxy keeps the connection open but no longer passes anything on, the generator closes the connection and raises a `HeartbeatTimeoutError`. Since `HeartbeatTimeoutError` is a subclass of `ServerError`, you can handle it specifically or together with other server errors:
+
+```python
+from eventsourcingdb import HeartbeatTimeoutError, ObserveEventsOptions
+
+try:
+  async for event in client.observe_events(
+    subject = '/books/42',
+    options = ObserveEventsOptions(
+      recursive = False
+    ),
+  ):
+    pass
+except HeartbeatTimeoutError:
+  # Neither an event nor a heartbeat arrived for 30 seconds.
+  pass
+```
+
 #### Aborting Observing
 
 If you need to abort observing use `break` or `return` within the `async for` loop. However, this only works if there is currently an iteration going on.

eventsourcingdb/__init__.py

@@ -4,6 +4,7 @@
 from .errors import (
     ClientError,
     CustomError,
+    HeartbeatTimeoutError,
     InternalError,
     ServerError,
     ValidationError,
@@ -39,6 +40,7 @@
     "Event",
     "EventCandidate",
     "EventType",
+    "HeartbeatTimeoutError",
     "IfEventIsMissingDuringObserve",
     "IfEventIsMissingDuringRead",
     "InternalError",

eventsourcingdb/client.py

@@ -16,6 +16,7 @@
 from .parse_raw_message import parse_raw_message
 from .read_event_types import EventType, is_event_type
 from .read_events import ReadEventsOptions
+from .read_lines_with_heartbeat_timeout import read_lines_with_heartbeat_timeout
 from .read_subjects import is_subject
 from .write_events import Precondition
 
@@ -185,7 +186,7 @@ async def run_eventql_query(self, query: str) -> AsyncGenerator[Any]:
 
         async with response:
             self._validate_response(response)
-            async for raw_message in response.body:
+            async for raw_message in read_lines_with_heartbeat_timeout(response.body):
                 message = parse_raw_message(raw_message)
 
                 if is_heartbeat(message):
@@ -222,7 +223,7 @@ async def observe_events(
 
         async with response:
             self._validate_response(response)
-            async for raw_message in response.body:
+            async for raw_message in read_lines_with_heartbeat_timeout(response.body):
                 message = parse_raw_message(raw_message)
 
                 if is_heartbeat(message):

eventsourcingdb/errors/__init__.py

@@ -1,12 +1,14 @@
 from .client_error import ClientError
 from .custom_error import CustomError
+from .heartbeat_timeout_error import HeartbeatTimeoutError
 from .internal_error import InternalError
 from .server_error import ServerError
 from .validation_error import ValidationError
 
 __all__ = [
     "ClientError",
     "CustomError",
+    "HeartbeatTimeoutError",
     "InternalError",
     "ServerError",
     "ValidationError",

eventsourcingdb/errors/heartbeat_timeout_error.py

@@ -0,0 +1,6 @@
+from .server_error import ServerError
+
+
+class HeartbeatTimeoutError(ServerError):
+    def __init__(self) -> None:
+        super().__init__("No event and no heartbeat arrived for 30 seconds")

eventsourcingdb/read_lines_with_heartbeat_timeout.py

@@ -0,0 +1,27 @@
+import asyncio
+from collections.abc import AsyncGenerator
+
+from aiohttp import StreamReader
+
+from .errors import HeartbeatTimeoutError
+
+HEARTBEAT_TIMEOUT_SECONDS = 30
+
+
+async def read_lines_with_heartbeat_timeout(stream: StreamReader) -> AsyncGenerator[bytes, None]:
+    while True:
+        timeout = asyncio.timeout(HEARTBEAT_TIMEOUT_SECONDS)
+        try:
+            async with timeout:
+                line = await stream.readline()
+        except TimeoutError as error:
+            # Only this timeout means that the server went silent, any other
+            # timeout (for example the one of the HTTP session) is passed on.
+            if not timeout.expired():
+                raise
+            raise HeartbeatTimeoutError() from error
+
+        if line == b'':
+            return
+
+        yield line

tests/shared/stream_server.py

@@ -0,0 +1,75 @@
+import asyncio
+from collections.abc import Awaitable, Callable
+from types import TracebackType
+from typing import Self
+
+from aiohttp import web
+
+from eventsourcingdb import Client
+
+WriteStream = Callable[[web.StreamResponse], Awaitable[None]]
+
+
+class StreamServer:
+    """
+    A local HTTP server that stands in for EventSourcingDB. It answers
+    observe-events and run-eventql-query with whatever the given function
+    writes to the stream, so that tests can control every single line.
+    """
+
+    def __init__(self, write_stream: WriteStream) -> None:
+        self.__write_stream = write_stream
+        self.__runner: web.AppRunner | None = None
+        self.__client: Client | None = None
+        self.client_disconnected = asyncio.Event()
+
+    async def __aenter__(self) -> Self:
+        application = web.Application()
+        application.router.add_post('/api/v1/observe-events', self.__handle)
+        application.router.add_post('/api/v1/run-eventql-query', self.__handle)
+
+        # Cancelling the handler when the client disconnects is what lets
+        # tests notice that the client closed the connection.
+        self.__runner = web.AppRunner(application, handler_cancellation=True)
+        await self.__runner.setup()
+        await web.TCPSite(self.__runner, '127.0.0.1', 0).start()
+        host, port = self.__runner.addresses[0][:2]
+
+        self.__client = Client(base_url=f'http://{host}:{port}', api_token='secret')
+        await self.__client.__aenter__()
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_val: BaseException | None = None,
+        exc_tb: TracebackType | None = None,
+    ) -> None:
+        if self.__client is not None:
+            await self.__client.__aexit__(exc_type, exc_val, exc_tb)
+        if self.__runner is not None:
+            await self.__runner.cleanup()
+
+    def get_client(self) -> Client:
+        if self.__client is None:
+            raise RuntimeError('Stream server must be started before getting a client.')
+
+        return self.__client
+
+    async def __handle(self, request: web.Request) -> web.StreamResponse:
+        response = web.StreamResponse(
+            headers={
+                'Server': 'EventSourcingDB/test',
+                'Content-Type': 'application/x-ndjson',
+            },
+        )
+        await response.prepare(request)
+
+        try:
+            await self.__write_stream(response)
+        except asyncio.CancelledError:
+            self.client_disconnected.set()
+            raise
+
+        return response