"""Streams."""
import asyncio
import os
import reprlib
import typing
import weakref
from asyncio import CancelledError
from contextvars import ContextVar
from typing import (
Any,
AsyncIterable,
AsyncIterator,
Callable,
Deque,
Dict,
Iterable,
Iterator,
List,
Mapping,
MutableSequence,
NamedTuple,
Optional,
Sequence,
Set,
Tuple,
Union,
cast,
)
from mode import Seconds, Service, get_logger, shortlabel, want_seconds
from mode.utils.aiter import aenumerate, aiter
from mode.utils.futures import current_task, maybe_async, notify
from mode.utils.queues import ThrowableQueue
from mode.utils.types.trees import NodeT
from . import joins
from .exceptions import ImproperlyConfigured, Skip
from .types import TP, AppT, ConsumerT, EventT, K, ModelArg, ModelT, TopicT
from .types.joins import JoinT
from .types.models import FieldDescriptorT
from .types.serializers import SchemaT
from .types.streams import (
GroupByKeyArg,
JoinableT,
Processor,
StreamT,
T,
T_co,
T_contra,
)
from .types.topics import ChannelT
from .types.tuples import Message
NO_CYTHON = bool(os.environ.get("NO_CYTHON", False))
if not NO_CYTHON: # pragma: no cover
try:
from ._cython.streams import StreamIterator as _CStreamIterator
except ImportError:
_CStreamIterator = None
else: # pragma: no cover
_CStreamIterator = None
__all__ = [
"Stream",
"current_event",
]
logger = get_logger(__name__)
if typing.TYPE_CHECKING: # pragma: no cover
_current_event: ContextVar[Optional[weakref.ReferenceType[EventT]]]
_current_event = ContextVar("current_event")
[docs]def current_event() -> Optional[EventT]:
"""Return the event currently being processed, or None."""
eventref = _current_event.get(None)
return eventref() if eventref is not None else None
async def maybe_forward(value: Any, channel: ChannelT) -> Any:
if isinstance(value, EventT):
await value.forward(channel)
else:
await channel.send(value=value)
return value
class _LinkedListDirection(NamedTuple):
attr: str
getter: Callable[[StreamT], Optional[StreamT]]
_LinkedListDirectionFwd = _LinkedListDirection("_next", lambda n: n._next)
_LinkedListDirectionBwd = _LinkedListDirection("_prev", lambda n: n._prev)
[docs]class Stream(StreamT[T_co], Service):
"""A stream: async iterator processing events in channels/topics."""
logger = logger
# Service starting/stopping logs use severity DEBUG in this class.
mundane_level = "debug"
#: Number of events processed by this instance so far.
events_total: int = 0
_processors: MutableSequence[Processor]
_anext_started = False
_passive = False
_finalized = False
_passive_started: asyncio.Event
def __init__(
self,
channel: AsyncIterator[T_co],
*,
app: AppT,
processors: Iterable[Processor[T]] = None,
combined: List[JoinableT] = None,
on_start: Optional[Callable] = None,
join_strategy: Optional[JoinT] = None,
beacon: Optional[NodeT] = None,
concurrency_index: Optional[int] = None,
prev: Optional[StreamT] = None,
active_partitions: Optional[Set[TP]] = None,
enable_acks: bool = True,
prefix: str = "",
loop: Optional[asyncio.AbstractEventLoop] = None,
) -> None:
Service.__init__(self, loop=loop, beacon=beacon)
self.app = app
self.channel = channel
self.outbox = self.app.FlowControlQueue(
maxsize=self.app.conf.stream_buffer_maxsize,
clear_on_resume=True,
)
self._passive_started = asyncio.Event()
self.join_strategy = join_strategy
self.combined = combined if combined is not None else []
self.concurrency_index = concurrency_index
self._prev = prev
self.active_partitions = active_partitions
self.enable_acks = enable_acks
self.prefix = prefix
self._processors = list(processors) if processors else []
self._on_start = on_start
# attach beacon to channel, or if iterable attach to current task.
task = current_task(loop=self.loop)
if task is not None:
self.task_owner = task
# Generate message handler
self._on_stream_event_in = self.app.sensors.on_stream_event_in
self._on_stream_event_out = self.app.sensors.on_stream_event_out
self._on_message_in = self.app.sensors.on_message_in
self._on_message_out = self.app.sensors.on_message_out
self._skipped_value = object()
[docs] def get_active_stream(self) -> StreamT:
"""Return the currently active stream.
A stream can be derived using ``Stream.group_by`` etc,
so if this stream was used to create another derived
stream, this function will return the stream being actively
consumed from. E.g. in the example::
>>> @app.agent()
... async def agent(a):
.. a = a
... b = a.group_by(Withdrawal.account_id)
... c = b.through('backup_topic')
... async for value in c:
... ...
The return value of ``a.get_active_stream()`` would be ``c``.
Notes:
The chain of streams that leads to the active stream
is decided by the :attr:`_next` attribute. To get
to the active stream we just traverse this linked-list::
>>> def get_active_stream(self):
... node = self
... while node._next:
... node = node._next
"""
return list(self._iter_ll_forwards())[-1]
[docs] def get_root_stream(self) -> StreamT:
"""Get the root stream that this stream was derived from."""
return list(self._iter_ll_backwards())[-1]
def _iter_ll_forwards(self) -> Iterator[StreamT]:
return self._iter_ll(_LinkedListDirectionFwd)
def _iter_ll_backwards(self) -> Iterator[StreamT]:
return self._iter_ll(_LinkedListDirectionBwd)
def _iter_ll(self, dir_: _LinkedListDirection) -> Iterator[StreamT]:
node: Optional[StreamT] = self
seen: Set[StreamT] = set()
while node:
if node in seen:
raise RuntimeError(f"Loop in Stream.{dir_.attr}: Call support!")
seen.add(node)
yield node
node = dir_.getter(node)
[docs] def add_processor(self, processor: Processor[T]) -> None:
"""Add processor callback executed whenever a new event is received.
Processor functions can be async or non-async, must accept
a single argument, and should return the value, mutated or not.
For example a processor handling a stream of numbers may modify
the value::
def double(value: int) -> int:
return value * 2
stream.add_processor(double)
"""
self._processors.append(processor)
[docs] def info(self) -> Mapping[str, Any]:
"""Return stream settings as a dictionary."""
# used by e.g. .clone to reconstruct keyword arguments
# needed to create a clone of the stream.
return {
"app": self.app,
"channel": self.channel,
"processors": self._processors,
"on_start": self._on_start,
"loop": self.loop,
"combined": self.combined,
"beacon": self.beacon,
"concurrency_index": self.concurrency_index,
"prev": self._prev,
"active_partitions": self.active_partitions,
}
[docs] def clone(self, **kwargs: Any) -> StreamT:
"""Create a clone of this stream.
Notes:
If the cloned stream is supposed to supersede this stream,
like in ``group_by``/``through``/etc., you should use
:meth:`_chain` instead so `stream._next = cloned_stream`
is set and :meth:`get_active_stream` returns the cloned stream.
"""
return self.__class__(**{**self.info(), **kwargs})
def _chain(self, **kwargs: Any) -> StreamT:
assert not self._finalized
self._next = new_stream = self.clone(
on_start=self.maybe_start,
prev=self,
# move processors to active stream
processors=list(self._processors),
**kwargs,
)
# delete moved processors from self
self._processors.clear()
return new_stream
[docs] def noack(self) -> "StreamT":
"""Create new stream where acks are manual."""
self._next = new_stream = self.clone(
enable_acks=False,
)
return new_stream
[docs] async def items(self) -> AsyncIterator[Tuple[K, T_co]]:
"""Iterate over the stream as ``key, value`` pairs.
Examples:
.. sourcecode:: python
@app.agent(topic)
async def mytask(stream):
async for key, value in stream.items():
print(key, value)
"""
async for event in self.events():
yield event.key, cast(T_co, event.value)
[docs] async def events(self) -> AsyncIterable[EventT]:
"""Iterate over the stream as events exclusively.
This means the stream must be iterating over a channel,
or at least an iterable of event objects.
"""
async for _ in self: # noqa: F841
if self.current_event is not None:
yield self.current_event
[docs] async def take(self, max_: int, within: Seconds) -> AsyncIterable[Sequence[T_co]]:
"""Buffer n values at a time and yield a list of buffered values.
Arguments:
max_: Max number of messages to receive. When more than this
number of messages are received within the specified number of
seconds then we flush the buffer immediately.
within: Timeout for when we give up waiting for another value,
and process the values we have.
Warning: If there's no timeout (i.e. `timeout=None`),
the agent is likely to stall and block buffered events for an
unreasonable length of time(!).
"""
buffer: List[T_co] = []
events: List[EventT] = []
buffer_add = buffer.append
event_add = events.append
buffer_size = buffer.__len__
buffer_full = asyncio.Event()
buffer_consumed = asyncio.Event()
timeout = want_seconds(within) if within else None
stream_enable_acks: bool = self.enable_acks
buffer_consuming: Optional[asyncio.Future] = None
channel_it = aiter(self.channel)
# We add this processor to populate the buffer, and the stream
# is passively consumed in the background (enable_passive below).
async def add_to_buffer(value: T) -> T:
try:
# buffer_consuming is set when consuming buffer after timeout.
nonlocal buffer_consuming
if buffer_consuming is not None:
try:
await buffer_consuming
finally:
buffer_consuming = None
buffer_add(cast(T_co, value))
event = self.current_event
if event is None:
raise RuntimeError("Take buffer found current_event is None")
event_add(event)
if buffer_size() >= max_:
# signal that the buffer is full and should be emptied.
buffer_full.set()
# strict wait for buffer to be consumed after buffer full.
# If max is 1000, we are not allowed to return 1001 values.
buffer_consumed.clear()
await self.wait(buffer_consumed)
except CancelledError: # pragma: no cover
raise
except Exception as exc:
self.log.exception("Error adding to take buffer: %r", exc)
await self.crash(exc)
return value
# Disable acks to ensure this method acks manually
# events only after they are consumed by the user
self.enable_acks = False
self.add_processor(add_to_buffer)
self._enable_passive(cast(ChannelT, channel_it))
try:
while not self.should_stop:
# wait until buffer full, or timeout
await self.wait_for_stopped(buffer_full, timeout=timeout)
if buffer:
# make sure background thread does not add new items to
# buffer while we read.
buffer_consuming = self.loop.create_future()
try:
yield list(buffer)
finally:
buffer.clear()
for event in events:
await self.ack(event)
events.clear()
# allow writing to buffer again
notify(buffer_consuming)
buffer_full.clear()
buffer_consumed.set()
else: # pragma: no cover
pass
else: # pragma: no cover
pass
finally:
# Restore last behaviour of "enable_acks"
self.enable_acks = stream_enable_acks
self._processors.remove(add_to_buffer)
[docs] async def take_events(
self, max_: int, within: Seconds
) -> AsyncIterable[Sequence[EventT]]:
"""Buffer n events at a time and yield a list of buffered events.
Arguments:
max_: Max number of messages to receive. When more than this
number of messages are received within the specified number of
seconds then we flush the buffer immediately.
within: Timeout for when we give up waiting for another value,
and process the values we have.
Warning: If there's no timeout (i.e. `timeout=None`),
the agent is likely to stall and block buffered events for an
unreasonable length of time(!).
"""
buffer: List[T_co] = []
events: List[EventT] = []
buffer_add = buffer.append
event_add = events.append
buffer_size = buffer.__len__
buffer_full = asyncio.Event()
buffer_consumed = asyncio.Event()
timeout = want_seconds(within) if within else None
stream_enable_acks: bool = self.enable_acks
buffer_consuming: Optional[asyncio.Future] = None
channel_it = aiter(self.channel)
# We add this processor to populate the buffer, and the stream
# is passively consumed in the background (enable_passive below).
async def add_to_buffer(value: T) -> T:
try:
# buffer_consuming is set when consuming buffer after timeout.
nonlocal buffer_consuming
if buffer_consuming is not None:
try:
await buffer_consuming
finally:
buffer_consuming = None
buffer_add(cast(T_co, value))
event = self.current_event
if event is None:
raise RuntimeError("Take buffer found current_event is None")
event_add(event)
if buffer_size() >= max_:
# signal that the buffer is full and should be emptied.
buffer_full.set()
# strict wait for buffer to be consumed after buffer full.
# If max is 1000, we are not allowed to return 1001 values.
buffer_consumed.clear()
await self.wait(buffer_consumed)
except CancelledError: # pragma: no cover
raise
except Exception as exc:
self.log.exception("Error adding to take buffer: %r", exc)
await self.crash(exc)
return value
# Disable acks to ensure this method acks manually
# events only after they are consumed by the user
self.enable_acks = False
self.add_processor(add_to_buffer)
self._enable_passive(cast(ChannelT, channel_it))
try:
while not self.should_stop:
# wait until buffer full, or timeout
await self.wait_for_stopped(buffer_full, timeout=timeout)
if buffer:
# make sure background thread does not add new items to
# buffer while we read.
buffer_consuming = self.loop.create_future()
try:
yield list(events)
finally:
buffer.clear()
for event in events:
await self.ack(event)
events.clear()
# allow writing to buffer again
notify(buffer_consuming)
buffer_full.clear()
buffer_consumed.set()
else: # pragma: no cover
pass
else: # pragma: no cover
pass
finally:
# Restore last behaviour of "enable_acks"
self.enable_acks = stream_enable_acks
self._processors.remove(add_to_buffer)
[docs] async def take_with_timestamp(
self, max_: int, within: Seconds, timestamp_field_name: str
) -> AsyncIterable[Sequence[T_co]]:
"""Buffer n values at a time and yield a list of buffered values with the
timestamp when the message was added to kafka.
Arguments:
max_: Max number of messages to receive. When more than this
number of messages are received within the specified number of
seconds then we flush the buffer immediately.
within: Timeout for when we give up waiting for another value,
and process the values we have.
Warning: If there's no timeout (i.e. `timeout=None`),
the agent is likely to stall and block buffered events for an
unreasonable length of time(!).
timestamp_field_name: the name of the field containing kafka timestamp,
that is going to be added to the value
"""
buffer: List[T_co] = []
events: List[EventT] = []
buffer_add = buffer.append
event_add = events.append
buffer_size = buffer.__len__
buffer_full = asyncio.Event()
buffer_consumed = asyncio.Event()
timeout = want_seconds(within) if within else None
stream_enable_acks: bool = self.enable_acks
buffer_consuming: Optional[asyncio.Future] = None
channel_it = aiter(self.channel)
# We add this processor to populate the buffer, and the stream
# is passively consumed in the background (enable_passive below).
async def add_to_buffer(value: T) -> T:
try:
# buffer_consuming is set when consuming buffer after timeout.
nonlocal buffer_consuming
if buffer_consuming is not None:
try:
await buffer_consuming
finally:
buffer_consuming = None
event = self.current_event
if isinstance(value, dict) and timestamp_field_name:
value[timestamp_field_name] = event.message.timestamp
buffer_add(value)
if event is None:
raise RuntimeError("Take buffer found current_event is None")
event_add(event)
if buffer_size() >= max_:
# signal that the buffer is full and should be emptied.
buffer_full.set()
# strict wait for buffer to be consumed after buffer full.
# If max is 1000, we are not allowed to return 1001 values.
buffer_consumed.clear()
await self.wait(buffer_consumed)
except CancelledError: # pragma: no cover
raise
except Exception as exc:
self.log.exception("Error adding to take buffer: %r", exc)
await self.crash(exc)
return value
# Disable acks to ensure this method acks manually
# events only after they are consumed by the user
self.enable_acks = False
self.add_processor(add_to_buffer)
self._enable_passive(cast(ChannelT, channel_it))
try:
while not self.should_stop:
# wait until buffer full, or timeout
await self.wait_for_stopped(buffer_full, timeout=timeout)
if buffer:
# make sure background thread does not add new items to
# buffer while we read.
buffer_consuming = self.loop.create_future()
try:
yield list(buffer)
finally:
buffer.clear()
for event in events:
await self.ack(event)
events.clear()
# allow writing to buffer again
notify(buffer_consuming)
buffer_full.clear()
buffer_consumed.set()
else: # pragma: no cover
pass
else: # pragma: no cover
pass
finally:
# Restore last behaviour of "enable_acks"
self.enable_acks = stream_enable_acks
self._processors.remove(add_to_buffer)
[docs] 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)
[docs] async def noack_take(
self, max_: int, within: Seconds
) -> AsyncIterable[Sequence[T_co]]:
"""
Buffer n values at a time and yield a list of buffered values.
:param max_: Max number of messages to receive. When more than this
number of messages are received within the specified number of
seconds then we flush the buffer immediately.
:param within: Timeout for when we give up waiting for another value,
and process the values we have.
Warning: If there's no timeout (i.e. `timeout=None`),
the agent is likely to stall and block buffered events for an
unreasonable length of time(!).
"""
buffer: List[T_co] = []
events: List[EventT] = []
buffer_add = buffer.append
event_add = events.append
buffer_size = buffer.__len__
buffer_full = asyncio.Event()
buffer_consumed = asyncio.Event()
timeout = want_seconds(within) if within else None
stream_enable_acks: bool = self.enable_acks
buffer_consuming: Optional[asyncio.Future] = None
channel_it = aiter(self.channel)
# We add this processor to populate the buffer, and the stream
# is passively consumed in the background (enable_passive below).
async def add_to_buffer(value: T) -> T:
try:
# buffer_consuming is set when consuming buffer
# after timeout.
nonlocal buffer_consuming
if buffer_consuming is not None:
try:
await buffer_consuming
finally:
buffer_consuming = None
# We want to save events instead of values to allow for manual ack
event = self.current_event
buffer_add(cast(T_co, event))
if event is None:
raise RuntimeError("Take buffer found current_event is None")
event_add(event)
if buffer_size() >= max_:
# signal that the buffer is full and should be emptied.
buffer_full.set()
# strict wait for buffer to be consumed after buffer
# full.
# If max is 1000, we are not allowed to return 1001
# values.
buffer_consumed.clear()
await self.wait(buffer_consumed)
except CancelledError: # pragma: no cover
raise
except Exception as exc:
self.log.exception("Error adding to take buffer: %r", exc)
await self.crash(exc)
return value
# Disable acks to ensure this method acks manually
# events only after they are consumed by the user
self.enable_acks = False
self.add_processor(add_to_buffer)
self._enable_passive(cast(ChannelT, channel_it))
try:
while not self.should_stop:
# wait until buffer full, or timeout
await self.wait_for_stopped(buffer_full, timeout=timeout)
if buffer:
# make sure background thread does not add new items to
# buffer while we read.
buffer_consuming = self.loop.create_future()
try:
yield list(buffer)
finally:
buffer.clear()
# code change: We want to manually ack
# for event in events:
# await self.ack(event)
events.clear()
# allow writing to buffer again
notify(buffer_consuming)
buffer_full.clear()
buffer_consumed.set()
else: # pragma: no cover
pass
else: # pragma: no cover
pass
finally:
# Restore last behaviour of "enable_acks"
self.enable_acks = stream_enable_acks
self._processors.remove(add_to_buffer)
[docs] def through(self, channel: Union[str, ChannelT]) -> StreamT:
"""Forward values to in this stream to channel.
Send messages received on this stream to another channel,
and return a new stream that consumes from that channel.
Notes:
The messages are forwarded after any processors have been
applied.
Example:
.. sourcecode:: python
topic = app.topic('foo')
@app.agent(topic)
async def mytask(stream):
async for value in stream.through(app.topic('bar')):
# value was first received in topic 'foo',
# then forwarded and consumed from topic 'bar'
print(value)
"""
if self._finalized:
# if agent restart we reuse the same stream object
# which already have done the stream.through()
# so on iteration we set the finalized flag
# and make this through() a noop.
return self
if self.concurrency_index is not None:
raise ImproperlyConfigured(
"Agent with concurrency>1 cannot use stream.through!"
)
# ridiculous mypy
if isinstance(channel, str):
channelchannel = cast(ChannelT, self.derive_topic(channel))
else:
channelchannel = channel
channel_it = aiter(channelchannel)
if self._next is not None:
raise ImproperlyConfigured("Stream is already using group_by/through")
through = self._chain(channel=channel_it)
async def forward(value: T) -> T:
event = self.current_event
return await maybe_forward(event, channelchannel)
self.add_processor(forward)
self._enable_passive(cast(ChannelT, channel_it), declare=True)
return through
def _enable_passive(self, channel: ChannelT, *, declare: bool = False) -> None:
if not self._passive:
self._passive = True
self.add_future(self._passive_drainer(channel, declare))
async def _passive_drainer(self, channel: ChannelT, declare: bool = False) -> None:
try:
if declare:
await channel.maybe_declare()
self._passive_started.set()
try:
async for item in self: # pragma: no cover
...
except BaseException as exc:
# forward the exception to the final destination channel,
# e.g. in through/group_by/etc.
await channel.throw(exc)
finally:
self._channel_stop_iteration(channel)
self._passive = False
def _channel_stop_iteration(self, channel: Any) -> None:
try:
on_stop_iteration = channel.on_stop_iteration
except AttributeError:
pass
else:
on_stop_iteration()
[docs] def echo(self, *channels: Union[str, ChannelT]) -> StreamT:
"""Forward values to one or more channels.
Unlike :meth:`through`, we don't consume from these channels.
"""
_channels = [
self.derive_topic(c) if isinstance(c, str) else c for c in channels
]
async def echoing(value: T) -> T:
await asyncio.wait(
[
asyncio.ensure_future(maybe_forward(value, channel))
for channel in _channels
],
return_when=asyncio.ALL_COMPLETED,
)
return value
self.add_processor(echoing)
return self
[docs] def group_by(
self,
key: GroupByKeyArg,
*,
name: Optional[str] = None,
topic: Optional[TopicT] = None,
partitions: Optional[int] = None,
) -> StreamT:
"""Create new stream that repartitions the stream using a new key.
Arguments:
key: The key argument decides how the new key is generated,
it can be a field descriptor, a callable, or an async
callable.
Note: The ``name`` argument must be provided if the key
argument is a callable.
name: Suffix to use for repartitioned topics.
This argument is required if `key` is a callable.
Examples:
Using a field descriptor to use a field in the event as the new
key:
.. sourcecode:: python
s = withdrawals_topic.stream()
# values in this stream are of type Withdrawal
async for event in s.group_by(Withdrawal.account_id):
...
Using an async callable to extract a new key:
.. sourcecode:: python
s = withdrawals_topic.stream()
async def get_key(withdrawal):
return await aiohttp.get(
f'http://e.com/resolve_account/{withdrawal.account_id}')
async for event in s.group_by(get_key):
...
Using a regular callable to extract a new key:
.. sourcecode:: python
s = withdrawals_topic.stream()
def get_key(withdrawal):
return withdrawal.account_id.upper()
async for event in s.group_by(get_key):
...
"""
if self._finalized:
# see note in self.through()
return self
channel: ChannelT
if self.concurrency_index is not None:
raise ImproperlyConfigured(
"Agent with concurrency>1 cannot use stream.group_by!"
)
if not name:
if isinstance(key, FieldDescriptorT):
name = key.ident
else:
raise TypeError("group_by with callback must set name=topic_suffix")
if topic is not None:
channel = topic
else:
prefix = ""
if self.prefix and not cast(TopicT, self.channel).has_prefix:
prefix = self.prefix + "-"
suffix = f"-{name}-repartition"
p = partitions if partitions else self.app.conf.topic_partitions
channel = cast(ChannelT, self.channel).derive(
prefix=prefix, suffix=suffix, partitions=p, internal=True
)
format_key = self._format_key
channel_it = aiter(channel)
if self._next is not None:
raise ImproperlyConfigured("Stream already uses group_by/through")
grouped = self._chain(channel=channel_it)
async def repartition(value: T) -> T:
event = self.current_event
if event is None:
raise RuntimeError("Cannot repartition stream with non-topic channel")
new_key = await format_key(key, value)
await event.forward(channel, key=new_key)
return value
self.add_processor(repartition)
self._enable_passive(cast(ChannelT, channel_it), declare=True)
return grouped
[docs] def filter(self, fun: Processor[T]) -> StreamT:
"""Filter values from stream using callback.
The callback may be a traditional function, lambda function,
or an `async def` function.
This method is useful for filtering events before repartitioning
a stream.
Examples:
>>> async for v in stream.filter(lambda: v > 1000).group_by(...):
... # do something
"""
async def on_value(value: T) -> T:
if not await maybe_async(fun(value)):
raise Skip()
else:
return value
self.add_processor(on_value)
return self
async def _format_key(self, key: GroupByKeyArg, value: T_contra) -> str:
try:
if isinstance(key, FieldDescriptorT):
return key.getattr(cast(ModelT, value))
return await maybe_async(cast(Callable, key)(value))
except BaseException as exc:
self.log.exception("Error in grouping key : %r", exc)
raise Skip() from exc
[docs] def derive_topic(
self,
name: str,
*,
schema: Optional[SchemaT] = None,
key_type: ModelArg = None,
value_type: ModelArg = None,
prefix: str = "",
suffix: str = "",
) -> TopicT:
"""Create Topic description derived from the K/V type of this stream.
Arguments:
name: Topic name.
key_type: Specific key type to use for this topic.
If not set, the key type of this stream will be used.
value_type: Specific value type to use for this topic.
If not set, the value type of this stream will be used.
Raises:
ValueError: if the stream channel is not a topic.
"""
if isinstance(self.channel, TopicT):
return cast(TopicT, self.channel).derive_topic(
topics=[name],
schema=schema,
key_type=key_type,
value_type=value_type,
prefix=prefix,
suffix=suffix,
)
raise ValueError("Cannot derive topic from non-topic channel.")
[docs] async def throw(self, exc: BaseException) -> None:
"""Send exception to stream iteration."""
await cast(ChannelT, self.channel).throw(exc)
[docs] def combine(self, *nodes: JoinableT, **kwargs: Any) -> StreamT:
"""Combine streams and tables into joined stream."""
# A combined stream is composed of multiple streams that
# all share the same outbox.
# The resulting stream's `on_merge` callback can be used to
# process values from all the combined streams, and e.g.
# joins uses this to consolidate multiple values into one.
if self._finalized:
# see note in self.through()
return self
stream = self._chain(combined=self.combined + list(nodes))
for node in stream.combined:
node.contribute_to_stream(stream)
return stream
[docs] def contribute_to_stream(self, active: StreamT) -> None:
"""Add stream as node in joined stream."""
self.outbox = active.outbox
[docs] async def remove_from_stream(self, stream: StreamT) -> None:
"""Remove as node in a joined stream."""
await self.stop()
[docs] def join(self, *fields: FieldDescriptorT) -> StreamT:
"""Create stream where events are joined."""
return self._join(joins.RightJoin(stream=self, fields=fields))
[docs] def left_join(self, *fields: FieldDescriptorT) -> StreamT:
"""Create stream where events are joined by LEFT JOIN."""
return self._join(joins.LeftJoin(stream=self, fields=fields))
[docs] def inner_join(self, *fields: FieldDescriptorT) -> StreamT:
"""Create stream where events are joined by INNER JOIN."""
return self._join(joins.InnerJoin(stream=self, fields=fields))
[docs] def outer_join(self, *fields: FieldDescriptorT) -> StreamT:
"""Create stream where events are joined by OUTER JOIN."""
return self._join(joins.OuterJoin(stream=self, fields=fields))
def _join(self, join_strategy: JoinT) -> StreamT:
return self.clone(join_strategy=join_strategy)
[docs] async def on_merge(self, value: T = None) -> Optional[T]:
"""Signal called when an event is to be joined."""
# TODO for joining streams
# The join strategy.process method can return None
# to eat the value, and on the next event create a merged
# event out of the previous event and new event.
join_strategy = self.join_strategy
if join_strategy:
value = await join_strategy.process(value)
return value
[docs] async def on_start(self) -> None:
"""Signal called when the stream starts."""
if self._on_start:
await self._on_start()
if self._passive:
await self._passive_started.wait()
[docs] async def stop(self) -> None:
"""Stop this stream."""
# Stop all related streams (created by .through/.group_by/etc.)
for s in cast(Stream, self.get_root_stream())._iter_ll_forwards():
await Service.stop(cast(Service, s))
[docs] async def on_stop(self) -> None:
"""Signal that the stream is stopping."""
self._passive = False
self._passive_started.clear()
for table_or_stream in self.combined:
await table_or_stream.remove_from_stream(self)
def __iter__(self) -> Any:
return self
def __next__(self) -> T:
raise NotImplementedError("Streams are asynchronous: use `async for`")
def __aiter__(self) -> AsyncIterator[T_co]: # pragma: no cover
if _CStreamIterator is not None:
return self._c_aiter()
else:
return self._py_aiter()
async def _c_aiter(self) -> AsyncIterator[T_co]: # pragma: no cover
self.log.dev("Using Cython optimized __aiter__")
skipped_value = self._skipped_value
self._finalized = True
started_by_aiter = await self.maybe_start()
it = _CStreamIterator(self)
try:
while not self.should_stop:
do_ack = self.enable_acks
value, sensor_state = await it.next() # noqa: B305
try:
if value is not skipped_value:
self.events_total += 1
yield value
finally:
event, self.current_event = self.current_event, None
it.after(event, do_ack or value is skipped_value, sensor_state)
except StopAsyncIteration:
# We are not allowed to propagate StopAsyncIteration in __aiter__
# (if we do, it'll be converted to RuntimeError by CPython).
# It can be raised when streaming over a list:
# async for value in app.stream([1, 2, 3, 4]):
# ...
# To support that, we just return here and that will stop
# the iteration.
return
finally:
self._channel_stop_iteration(self.channel)
if started_by_aiter:
await self.stop()
self.service_reset()
def _set_current_event(self, event: Optional[EventT] = None) -> None:
if event is None:
_current_event.set(None)
else:
_current_event.set(weakref.ref(event))
self.current_event = event
async def _py_aiter(self) -> AsyncIterator[T_co]:
self._finalized = True
started_by_aiter = await self.maybe_start()
on_merge = self.on_merge
on_stream_event_out = self._on_stream_event_out
on_message_out = self._on_message_out
# get from channel
channel = self.channel
if isinstance(channel, ChannelT):
chan_is_channel = True
chan = cast(ChannelT, self.channel)
chan_queue = chan.queue
chan_queue_empty = chan_queue.empty
chan_errors = chan_queue._errors
chan_quick_get = chan_queue.get_nowait
else:
chan_is_channel = False
chan_queue = cast(ThrowableQueue, None)
chan_queue_empty = cast(Callable, None)
chan_errors = cast(Deque, None)
chan_quick_get = cast(Callable, None)
chan_slow_get = channel.__anext__
# Topic description -> processors
processors = self._processors
# Sensor: on_stream_event_in
on_stream_event_in = self._on_stream_event_in
# localize global variables
create_ref = weakref.ref
_maybe_async = maybe_async
event_cls = EventT
_current_event_contextvar = _current_event
consumer: ConsumerT = self.app.consumer
unacked: Set[Message] = consumer.unacked
add_unacked: Callable[[Message], None] = unacked.add
acking_topics: Set[str] = self.app.topics.acking_topics
on_message_in = self._on_message_in
sleep = asyncio.sleep
trace = self.app.trace
_shortlabel = shortlabel
sensor_state: Optional[Dict] = None
skipped_value = self._skipped_value
try:
while not self.should_stop:
event = None
do_ack = self.enable_acks # set to False to not ack event.
# wait for next message
value: Any = None
# we iterate until on_merge gives value.
while value is None and event is None:
await sleep(0)
# get message from channel
# This inlines ThrowableQueue.get for performance:
# We selectively call `await Q.put`/`Q.put_nowait`,
# and prefer the latter if the queue is non-empty.
channel_value: Any
if chan_is_channel:
if chan_errors:
raise chan_errors.popleft()
if chan_queue_empty():
channel_value = await chan_slow_get()
else:
channel_value = chan_quick_get()
else:
# chan is an AsyncIterable
channel_value = await chan_slow_get()
if isinstance(channel_value, event_cls):
event = channel_value
message = event.message
topic = message.topic
tp = message.tp
offset = message.offset
if (
not self.app.flow_control.is_active()
or message.generation_id != self.app.consumer_generation_id
):
value = skipped_value
self.log.dev(
"Skipping message %r with generation_id %r because "
"app generation_id is %r flow_control.is_active %r",
message,
message.generation_id,
self.app.consumer_generation_id,
self.app.flow_control.is_active(),
)
break
if topic in acking_topics and not message.tracked:
message.tracked = True
# This inlines Consumer.track_message(message)
add_unacked(message)
on_message_in(message.tp, message.offset, message)
# XXX ugh this should be in the consumer somehow
# call Sensors
sensor_state = on_stream_event_in(tp, offset, self, event)
# set task-local current_event
_current_event_contextvar.set(create_ref(event))
# set Stream._current_event
self.current_event = event
# Stream yields Event.value
value = event.value
else:
value = channel_value
self.current_event = None
sensor_state = None
# reduce using processors
try:
for processor in processors:
with trace(f"processor-{_shortlabel(processor)}"):
value = await _maybe_async(processor(value))
value = await on_merge(value)
except Skip:
# We want to ack the filtered message
# otherwise the lag would increase
value = skipped_value
try:
if value is not skipped_value:
self.events_total += 1
yield value
finally:
self.current_event = None
# We want to ack the filtered out message
# otherwise the lag would increase
if event is not None and (do_ack or value is skipped_value):
# This inlines self.ack
last_stream_to_ack = event.ack()
message = event.message
tp = event.message.tp
offset = event.message.offset
on_stream_event_out(tp, offset, self, event, sensor_state)
if last_stream_to_ack:
on_message_out(tp, offset, message)
except StopAsyncIteration:
# We are not allowed to propagate StopAsyncIteration in __aiter__
# (if we do, it'll be converted to RuntimeError by CPython).
# It can be raised when streaming over a list:
# async for value in app.stream([1, 2, 3, 4]):
# ...
# To support that, we just return here and that will stop
# the iteration.
return
finally:
self._channel_stop_iteration(channel)
if started_by_aiter:
# if `async for` called Stream.start()
# we need to stop it before ending iteration.
await self.stop()
# reset to allow calling .start again on next `async for`
self.service_reset()
async def __anext__(self) -> T: # pragma: no cover
...
[docs] async def ack(self, event: EventT) -> bool:
"""Ack event.
This will decrease the reference count of the event message by one,
and when the reference count reaches zero, the worker will
commit the offset so that the message will not be seen by a worker
again.
Arguments:
event: Event to ack.
"""
# WARNING: This function is duplicated in __aiter__
last_stream_to_ack = event.ack()
message = event.message
tp = message.tp
offset = message.offset
self._on_stream_event_out(tp, offset, self, event)
if last_stream_to_ack:
self._on_message_out(tp, offset, message)
return last_stream_to_ack
def __and__(self, other: Any) -> Any:
return self.combine(self, other)
def __copy__(self) -> Any:
return self.clone()
def _repr_info(self) -> str:
if self.combined:
return reprlib.repr(self.combined)
return reprlib.repr(self.channel)
@property
def label(self) -> str:
"""Return description of stream, used in graphs and logs."""
return f"{type(self).__name__}: {self._repr_channel()}"
def _repr_channel(self) -> str:
return reprlib.repr(self.channel)
@property
def shortlabel(self) -> str:
"""Return short description of stream."""
# used for shortlabel(stream), which is used by statsd to generate ids
# note: str(channel) returns topic name when it's a topic, so
# this will be:
# "Channel: <ANON>", for channel or
# "Topic: withdrawals", for a topic.
# statsd then uses that as part of the id.
return f"Stream: {self._human_channel()}"
def _human_channel(self) -> str:
if self.combined:
return "&".join(s._human_channel() for s in self.combined)
return f"{type(self.channel).__name__}: {self.channel}"
