diff --git a/src/s2_sdk/_batching.py b/src/s2_sdk/_batching.py index e57f7d3..d66090f 100644 --- a/src/s2_sdk/_batching.py +++ b/src/s2_sdk/_batching.py @@ -45,7 +45,13 @@ async def append_record_batches( *, batching: Batching | None = None, ) -> AsyncIterable[list[Record]]: - """Group records into batches based on count, bytes, and linger time.""" + """Group records into batches based on count, bytes, and linger time. + + Note: + If iteration over ``records`` raises, or batching is cancelled or stopped + before completion, records already consumed but not yet yielded as a batch + may be discarded. + """ if batching is None: batching = Batching() validate_batching(batching.max_records, batching.max_bytes) @@ -89,10 +95,6 @@ async def append_record_batches( acc.add(record) yield acc.take() - except Exception: - if not acc.is_empty(): - yield acc.take() - raise finally: if next_record_task is not None and not next_record_task.done(): next_record_task.cancel() diff --git a/tests/test_batching.py b/tests/test_batching.py index 900e76f..0a09b14 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -84,28 +84,6 @@ async def delayed_records(): assert len(batches[1]) == 2 # r2, r3 -@pytest.mark.asyncio -@pytest.mark.parametrize( - "linger", - [timedelta(0), timedelta(milliseconds=10)], - ids=["zero_linger", "non_zero_linger"], -) -async def test_batch_accumulator_flushes_when_source_iter_raises(linger): - async def records(): - yield Record(body=b"r1") - yield Record(body=b"r2") - raise RuntimeError("err") - - batches = [] - with pytest.raises(RuntimeError, match="err"): - async for batch in append_record_batches( - records(), batching=Batching(linger=linger) - ): - batches.append(batch) - - assert batches == [[Record(body=b"r1"), Record(body=b"r2")]] - - @pytest.mark.asyncio async def test_append_inputs_skips_empty_batches(): inputs = []