async def process(self, stream: StreamT[bytes]) -> None:
self.agent_started.set()
try:
async for i, event in aenumerate(stream.events()):
self.agent_started_processing.set()
self.processed_total += 1
await self.on_agent_event(stream, event)
key = event.message.tp, event.message.offset
if key in self.seen_offsets:
print(f'EVENT PROCESSED TWICE: {key}')
await self.crash(Exception(f'Event processed twice: {key}')
)
self.seen_offsets.add(key)
assert key in self.seen_offsets
if self.processed_total >= self.num_messages:
self.agent_stopped_processing.set()
await self.sleep(0)
except asyncio.CancelledError:
if self.processed_total < self.num_messages:
print('I WAS CANCELLED?!?!?')
raise
except Exception as exc:
print(f'AGENT RAISED ERROR: {exc!r}')
await self.crash(exc)
raise
def enumerate(self, start: int = 0) -> AsyncIterable[Tuple[int, T_co]]:
"""Enumerate values received on this stream.
Unlike Python's built-in ``enumerate``, this works with
async generators.
"""
return aenumerate(self, start)
async def kvjoin(
self,
items: Union[AsyncIterable[Tuple[K, V]], Iterable[Tuple[K, V]]],
reply_to: ReplyToArg = None,
) -> List[Any]: # pragma: no cover
"""RPC map operation on list of ``(key, value)`` pairs.
A join returns the results in order, and only returns once
all values have been processed.
"""
reply_to = self._get_strtopic(reply_to or self.app.conf.reply_to)
barrier = BarrierState(reply_to)
# Map correlation_id -> index
posindex: MutableMapping[str, int] = {
cid: i
async for i, cid in aenumerate(
self._barrier_send(barrier, items, reply_to))
}
# All the messages have been sent so finalize the barrier.
barrier.finalize()
# wait until all replies received
await barrier
# then construct a list in the correct order.
values: List = [None] * barrier.total
async for correlation_id, value in barrier.iterate():
values[posindex[correlation_id]] = value
return values
async def _slurp_stream(self) -> None:
buf: List[EventT] = []
can_log_done = True
try:
async for i, event in aenumerate(self._read_changelog()):
buf.append(event)
await self.table.on_changelog_event(event)
if len(buf) >= self._buffer_size:
self.table.apply_changelog_batch(buf)
buf.clear()
if self._should_stop_reading():
break
remaining = self._remaining_total()
if remaining and not i % 10_000:
can_log_done = True
self.log.info('Waiting for %s records...', remaining)
elif not remaining and can_log_done:
can_log_done = False
self.log.info('All up to date')
except StopAsyncIteration:
self.log.info('Got stop iteration')
pass
finally:
self.log.info('Stopped reading!')
if buf:
self.table.apply_changelog_batch(buf)
buf.clear()
async def test_aenumerate():
it = (a async for a in aenumerate(aiter([1, 2, 3, 4, 5])))
assert await anext(it) == (0, 1)
assert await anext(it) == (1, 2)
assert await anext(it) == (2, 3)
assert await anext(it) == (3, 4)
assert await anext(it) == (4, 5)
with pytest.raises(StopAsyncIteration):
await anext(it)
sentinel = object()
assert await anext(it, sentinel) is sentinel
assert repr(aiter([1, 2, 3, 4]))