fix(batch): surface background stream failures on exit instead of returning partial results

test/test_batch_stream_async.py

@@ -0,0 +1,127 @@
+"""Unit tests for the async batch-stream failure handling (no cluster needed)."""
+
+import asyncio
+import logging
+from typing import Optional
+
+import pytest
+
+from weaviate.collections.batch.async_ import _BatchBaseAsync, _BgTasks
+from weaviate.collections.batch.base import _BatchDataWrapper
+from weaviate.collections.batch.batch_wrapper import _ContextManagerAsync
+from weaviate.exceptions import WeaviateBatchStreamError
+
+
+class _NotAnException(BaseException):
+    """Escapes the wrappers' `except Exception` like grpc.aio's CancelledError does."""
+
+
+def _bare_batch(**mangled) -> _BatchBaseAsync:
+    batch = object.__new__(_BatchBaseAsync)
+    for name, value in mangled.items():
+        setattr(batch, f"_BatchBaseAsync__{name}", value)
+    return batch
+
+
+class _FakeTimeouts:
+    insert = 1
+
+
+class _FakeConnection:
+    timeout_config = _FakeTimeouts()
+
+
+def test_wait_raises_the_background_exception_and_keeps_partial_results() -> None:
+    # a background failure must not come back as a success, and a user catching it
+    # must still see what failed
+    class FakeBgTasks:
+        async def gather(self, timeout=None) -> None:
+            return None
+
+    partial = _BatchDataWrapper()
+    partial.failed_objects = ["sentinel-failure"]  # type: ignore[list-item]
+    backup = _BatchDataWrapper()
+    batch = _bare_batch(
+        bg_exception=RuntimeError("boom"),
+        bg_tasks=FakeBgTasks(),
+        connection=_FakeConnection(),
+        results_for_wrapper=partial,
+        results_for_wrapper_backup=backup,
+    )
+
+    with pytest.raises(RuntimeError, match="boom"):
+        asyncio.run(batch._wait())
+    assert backup.failed_objects == ["sentinel-failure"]
+
+
+def test_flush_raises_when_a_task_dies_without_setting_bg_exception() -> None:
+    # loop_wrapper/recv_wrapper only catch Exception, so a BaseException kills a task with
+    # __bg_exception unset; flush() must notice the dead task instead of spinning forever
+    async def run() -> None:
+        async def dies() -> None:
+            raise _NotAnException("boom")
+
+        loop_task = asyncio.get_running_loop().create_task(asyncio.sleep(3600))
+        recv_task = asyncio.get_running_loop().create_task(dies())
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        assert recv_task.done()
+
+        batch = _bare_batch(
+            bg_exception=None,
+            bg_tasks=_BgTasks(recv=recv_task, loop=loop_task),
+            batch_objects=[object()],
+            batch_references=[],
+        )
+        try:
+            await asyncio.wait_for(batch.flush(), timeout=2)  # a regression hangs here
+        finally:
+            loop_task.cancel()
+
+    with pytest.raises(WeaviateBatchStreamError, match="stream has ended"):
+        asyncio.run(run())
+
+
+class _FakeBatch:
+    """Stands in for _BatchBaseAsync behind the context manager."""
+
+    def __init__(self, wait_error: Optional[BaseException] = None) -> None:
+        self.wait_error = wait_error
+        self.wait_called = False
+
+    async def _start(self) -> None:
+        pass
+
+    async def _shutdown(self) -> None:
+        pass
+
+    async def _wait(self) -> None:
+        self.wait_called = True
+        if self.wait_error is not None:
+            raise self.wait_error
+
+
+def test_aexit_raises_a_background_failure_on_a_clean_block() -> None:
+    fake = _FakeBatch(WeaviateBatchStreamError("bg died"))
+
+    async def run() -> None:
+        async with _ContextManagerAsync(fake):  # type: ignore[arg-type]
+            pass
+
+    with pytest.raises(WeaviateBatchStreamError, match="bg died"):
+        asyncio.run(run())
+
+
+def test_aexit_keeps_the_users_exception_over_a_background_failure(caplog) -> None:
+    # the block's own exception must not be replaced by the background failure
+    fake = _FakeBatch(WeaviateBatchStreamError("bg died"))
+
+    async def run() -> None:
+        async with _ContextManagerAsync(fake):  # type: ignore[arg-type]
+            raise ValueError("user code")
+
+    with caplog.at_level(logging.WARNING, logger="weaviate-client"):
+        with pytest.raises(ValueError, match="user code"):
+            asyncio.run(run())
+    assert fake.wait_called  # still drained/awaited
+    assert "bg died" in caplog.text

test/test_batch_stream_sync.py

@@ -0,0 +1,126 @@
+"""Unit tests for the sync batch-stream failure handling (no cluster needed)."""
+
+import logging
+import time
+from typing import Optional
+
+import pytest
+
+from weaviate.collections.batch.base import _BatchDataWrapper
+from weaviate.collections.batch.batch_wrapper import _ContextManagerSync
+from weaviate.collections.batch.sync import _BatchBaseSync
+from weaviate.exceptions import WeaviateBatchStreamError
+
+
+def _bare_batch(**mangled) -> _BatchBaseSync:
+    batch = object.__new__(_BatchBaseSync)
+    for name, value in mangled.items():
+        setattr(batch, f"_BatchBaseSync__{name}", value)
+    return batch
+
+
+class _FakeThreads:
+    def __init__(self, alive: bool = False) -> None:
+        self.alive = alive
+
+    def join(self, timeout=None) -> None:
+        return None
+
+    def is_alive(self) -> bool:
+        return self.alive
+
+
+class _FakeTimeouts:
+    insert = 1
+
+
+class _FakeConnection:
+    timeout_config = _FakeTimeouts()
+
+
+def _batch_for_wait(**mangled) -> _BatchBaseSync:
+    defaults = {
+        "bg_exception": None,
+        "bg_threads": _FakeThreads(alive=False),
+        "connection": _FakeConnection(),
+        "results_for_wrapper": _BatchDataWrapper(),
+        "results_for_wrapper_backup": _BatchDataWrapper(),
+        "batch_objects": [],
+        "batch_references": [],
+    }
+    defaults.update(mangled)
+    return _bare_batch(**defaults)
+
+
+def test_wait_raises_the_background_exception_and_keeps_partial_results() -> None:
+    partial = _BatchDataWrapper()
+    partial.failed_objects = ["sentinel-failure"]  # type: ignore[list-item]
+    backup = _BatchDataWrapper()
+    batch = _batch_for_wait(
+        bg_exception=RuntimeError("boom"),
+        results_for_wrapper=partial,
+        results_for_wrapper_backup=backup,
+    )
+
+    with pytest.raises(RuntimeError, match="boom"):
+        batch._wait()
+    assert backup.failed_objects == ["sentinel-failure"]
+
+
+def test_wait_names_unsent_data_when_the_threads_are_gone() -> None:
+    batch = _batch_for_wait(batch_objects=[object(), object()], batch_references=[object()])
+    with pytest.raises(
+        WeaviateBatchStreamError, match="ended with 2 objects and 1 references unsent"
+    ):
+        batch._wait()
+
+
+def test_start_raises_the_stored_background_exception_at_once() -> None:
+    # a background thread that died right away (e.g. on a closed connection) used to leave
+    # _start polling for 60 s and then blaming the network (#2139)
+    batch = _bare_batch(bg_exception=None, bg_threads=None)
+
+    def start_dead_threads() -> None:
+        setattr(batch, "_BatchBaseSync__bg_threads", _FakeThreads(alive=False))
+        setattr(batch, "_BatchBaseSync__bg_exception", RuntimeError("connection closed"))
+
+    setattr(batch, "_BatchBaseSync__start_bg_threads", start_dead_threads)
+    started = time.monotonic()
+    with pytest.raises(RuntimeError, match="connection closed"):
+        batch._start()
+    assert time.monotonic() - started < 5
+
+
+class _FakeBatch:
+    def __init__(self, wait_error: Optional[BaseException] = None) -> None:
+        self.wait_error = wait_error
+        self.wait_called = False
+
+    def _start(self) -> None:
+        pass
+
+    def _shutdown(self) -> None:
+        pass
+
+    def _wait(self) -> None:
+        self.wait_called = True
+        if self.wait_error is not None:
+            raise self.wait_error
+
+
+def test_exit_raises_a_background_failure_on_a_clean_block() -> None:
+    # the sync colour used to swallow this entirely
+    fake = _FakeBatch(WeaviateBatchStreamError("bg died"))
+    with pytest.raises(WeaviateBatchStreamError, match="bg died"):
+        with _ContextManagerSync(fake):  # type: ignore[type-var]
+            pass
+
+
+def test_exit_keeps_the_users_exception_over_a_background_failure(caplog) -> None:
+    fake = _FakeBatch(WeaviateBatchStreamError("bg died"))
+    with caplog.at_level(logging.WARNING, logger="weaviate-client"):
+        with pytest.raises(ValueError, match="user code"):
+            with _ContextManagerSync(fake):  # type: ignore[type-var]
+                raise ValueError("user code")
+    assert fake.wait_called
+    assert "bg died" in caplog.text

weaviate/collections/batch/async_.py

@@ -202,6 +202,16 @@ async def _wait(self) -> None:
             self.__results_for_wrapper.imported_shards
         )
 
+        if self.__bg_exception is not None:
+            # surface the background failure instead of returning partial results as a success
+            raise self.__bg_exception
+        n_objs, n_refs = len(self.__batch_objects), len(self.__batch_references)
+        if n_objs + n_refs > 0 and not self.__all_tasks_alive():
+            # the tasks are gone with data still queued: the batch did NOT complete
+            raise WeaviateBatchStreamError(
+                f"batch stream ended with {n_objs} objects and {n_refs} references unsent"
+            )
+
     async def _shutdown(self) -> None:
         self.__is_stopped.set()
 
@@ -532,6 +542,8 @@ async def flush(self) -> None:
         # bg thread is sending objs+refs automatically, so simply wait for everything to be done
         while len(self.__batch_objects) > 0 or len(self.__batch_references) > 0:
             await asyncio.sleep(0.01)
+            # a dead task means nothing drains the queues: raise, don't spin forever
+            self.__check_bg_tasks_alive()
 
     async def _add_object(
         self,
@@ -628,4 +640,4 @@ def __check_bg_tasks_alive(self) -> None:
         if self.__all_tasks_alive():
             return
 
-        raise self.__bg_exception or Exception("Batch tasks died unexpectedly")
+        raise self.__bg_exception or WeaviateBatchStreamError("the batch stream has ended")

weaviate/collections/batch/batch_wrapper.py

@@ -508,7 +508,13 @@ def __init__(self, current_batch: T):
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         self.__current_batch._shutdown()
-        self.__current_batch._wait()
+        try:
+            self.__current_batch._wait()
+        except Exception as e:
+            if exc_type is None:
+                raise
+            # the exception leaving the block wins; the background failure is only logged
+            logger.warning(f"batch stream failed while the block raised {exc_type.__name__}: {e}")
 
     def __enter__(self) -> P:
         self.__current_batch._start()
@@ -521,7 +527,13 @@ def __init__(self, current_batch: _BatchBaseAsync):
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await self.__current_batch._shutdown()
-        await self.__current_batch._wait()
+        try:
+            await self.__current_batch._wait()
+        except Exception as e:
+            if exc_type is None:
+                raise
+            # the exception leaving the block wins; the background failure is only logged
+            logger.warning(f"batch stream failed while the block raised {exc_type.__name__}: {e}")
 
     async def __aenter__(self) -> Q:
         await self.__current_batch._start()

weaviate/collections/batch/sync.py

@@ -125,7 +125,10 @@ def _start(self) -> None:
         logger.info("Provisioned stream to the server for batch processing")
         now = time.time()
         while not self.__all_threads_alive():
-            # wait for the recv threads to be started
+            if self.__bg_exception is not None:
+                # a background thread already died (e.g. on a closed connection): raise its
+                # error now instead of polling for 60 s and blaming the network
+                raise self.__bg_exception
             time.sleep(0.01)
             if time.time() - now > 60:
                 raise WeaviateBatchStreamError(
@@ -152,6 +155,16 @@ def _wait(self) -> None:
             self.__results_for_wrapper.imported_shards
         )
 
+        if self.__bg_exception is not None:
+            # surface the background failure instead of returning partial results as a success
+            raise self.__bg_exception
+        n_objs, n_refs = len(self.__batch_objects), len(self.__batch_references)
+        if n_objs + n_refs > 0 and not self.__all_threads_alive():
+            # the threads are gone with data still queued: the batch did NOT complete
+            raise WeaviateBatchStreamError(
+                f"batch stream ended with {n_objs} objects and {n_refs} references unsent"
+            )
+
     def _shutdown(self) -> None:
         # Shutdown the current batch and wait for all requests to be finished
         self.__is_stopped.set()
@@ -644,4 +657,4 @@ def __check_bg_threads_alive(self) -> None:
         if self.__all_threads_alive():
             return
 
-        raise self.__bg_exception or Exception("Batch thread died unexpectedly")
+        raise self.__bg_exception or WeaviateBatchStreamError("the batch stream has ended")

weaviate/exceptions.py

@@ -418,5 +418,8 @@ def __init__(self, pb: version.Version, grpc: version.Version) -> None:
         )
 
 
-class _BatchStreamShutdownError(Exception):
-    """Internal exception to signal that the batch stream was shutdown."""
+class _BatchStreamShutdownError(WeaviateBatchStreamError):
+    """Internal exception to signal that the batch stream was shutdown (gRPC ABORTED)."""
+
+    def __init__(self, message: str = "the server aborted the batch stream") -> None:
+        super().__init__(message)