Skip to content

BatchSubscriber

faststream.kafka.subscriber.usecase.BatchSubscriber #

BatchSubscriber(
    config: KafkaSubscriberConfig,
    specification: SubscriberSpecification[Any, Any],
    calls: CallsCollection[tuple[ConsumerRecord, ...]],
    batch_timeout_ms: int,
    max_records: int | None,
)

Bases: LogicSubscriber[tuple['ConsumerRecord', ...]]

Source code in faststream/kafka/subscriber/usecase.py
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
def __init__(
    self,
    config: "KafkaSubscriberConfig",
    specification: "SubscriberSpecification[Any, Any]",
    calls: "CallsCollection[tuple[ConsumerRecord, ...]]",
    batch_timeout_ms: int,
    max_records: int | None,
) -> None:
    reg = (
        Address(config.pattern, KAFKA_ADDRESS_SYNTAX).regex
        if config.pattern
        else None
    )

    self.parser = AioKafkaBatchParser(
        msg_class=KafkaMessage if config.ack_first else KafkaAckableMessage,
        regex=reg,
    )
    config.decoder = self.parser.decode_batch
    config.parser = self.parser.parse_batch
    super().__init__(config, specification, calls)

    self.batch_timeout_ms = batch_timeout_ms
    self.max_records = max_records

parser instance-attribute #

parser = AioKafkaBatchParser(
    msg_class=KafkaMessage
    if config.ack_first
    else KafkaAckableMessage,
    regex=reg,
)

batch_timeout_ms instance-attribute #

batch_timeout_ms = batch_timeout_ms

max_records instance-attribute #

max_records = max_records

lock instance-attribute #

lock: AbstractContextManager[Any] = nullcontext()

graceful_timeout instance-attribute #

graceful_timeout: float | None

calls instance-attribute #

calls = calls

specification instance-attribute #

specification = specification

ack_policy instance-attribute #

ack_policy = config.ack_policy

running instance-attribute #

running = False

tasks instance-attribute #

tasks: list[Task[Any]] = []

consumer instance-attribute #

consumer: Optional[AIOKafkaConsumer] = None

batch instance-attribute #

batch: bool

group_id instance-attribute #

group_id = config.group_id

pattern property #

pattern: Address | None

The pattern this Subscriber was declared with, and its Broker address.

topics property #

topics: list[Topic]

topic_names_for_subscribe property #

topic_names_for_subscribe: list[str]

partitions property #

partitions: list[TopicPartition]

The assignment as the consumer receives it: the client library's tuples, prefixed.

Declared with faststream.kafka.TopicPartition; handed to aiokafka as its own.

builder property #

builder: Callable[..., AIOKafkaConsumer]

client_id property #

client_id: str | None

topics_to_create property #

topics_to_create: list[Topic]

topic_names property #

topic_names: list[str]

get_msg async #

get_msg(
    consumer: AIOKafkaConsumer,
) -> tuple[ConsumerRecord, ...]
Source code in faststream/kafka/subscriber/usecase.py
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
async def get_msg(
    self,
    consumer: "AIOKafkaConsumer",
) -> tuple["ConsumerRecord", ...]:
    assert consumer, "You should setup subscriber at first."

    messages = await consumer.getmany(
        timeout_ms=self.batch_timeout_ms,
        max_records=self.max_records,
    )

    if not messages:  # pragma: no cover
        await anyio.sleep(self.batch_timeout_ms / 1000)
        return ()

    return tuple(chain(*messages.values()))

get_log_context #

get_log_context(
    message: Optional[
        StreamMessage[tuple[ConsumerRecord, ...]]
    ],
) -> dict[str, str]
Source code in faststream/kafka/subscriber/usecase.py
500
501
502
503
504
505
506
507
508
509
510
511
512
513
def get_log_context(
    self,
    message: Optional["StreamMessage[tuple[ConsumerRecord, ...]]"],
) -> dict[str, str]:
    if message is None:
        topic = ",".join(self.topic_names)
    else:
        topic = message.raw_message[0].topic

    return self.build_log_context(
        message=message,
        topic=topic,
        group_id=self.group_id,
    )

start async #

start() -> None

Start the consumer.

Source code in faststream/kafka/subscriber/usecase.py
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
async def start(self) -> None:
    """Start the consumer."""
    await super().start()
    await self._ensure_topics()

    self.consumer = consumer = self.builder(
        group_id=self.group_id,
        client_id=self.client_id,
        **self._connection_args,
    )

    self.parser._setup(consumer)

    pattern = self.pattern
    if self.topics or pattern:
        consumer.subscribe(
            topics=self.topic_names_for_subscribe,
            pattern=pattern.broker_address if pattern else None,
            listener=make_logging_listener(
                consumer=consumer,
                logger=self._outer_config.logger.logger.logger,
                log_extra=self.get_log_context(None),
                listener=self._listener,
            ),
        )

    elif self.partitions:
        consumer.assign(partitions=self.partitions)

    await consumer.start()

    self._post_start()

    if self.calls:
        self.add_task(self._run_consume_loop, (self.consumer,))

stop async #

stop() -> None
Source code in faststream/kafka/subscriber/usecase.py
263
264
265
266
267
268
async def stop(self) -> None:
    await super().stop()

    if self.consumer is not None:
        await self.consumer.stop()
        self.consumer = None

add_call #

add_call(
    *,
    parser_: Optional[CustomCallable],
    decoder_: Optional[CustomCallable],
    dependencies_: Sequence[Dependant],
    codec_: Optional[CodecProto] = None,
) -> Self
Source code in faststream/_internal/endpoint/subscriber/usecase.py
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def add_call(
    self,
    *,
    parser_: Optional["CustomCallable"],
    decoder_: Optional["CustomCallable"],
    dependencies_: Sequence["Dependant"],
    codec_: Optional["CodecProto"] = None,
) -> Self:
    self._call_options = _CallOptions(
        parser=parser_,
        decoder=decoder_,
        dependencies=dependencies_,
        codec=codec_,
    )
    return self

consume async #

consume(msg: MsgType) -> Any

Consume a message asynchronously.

Source code in faststream/_internal/endpoint/subscriber/usecase.py
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
async def consume(self, msg: MsgType) -> Any:
    """Consume a message asynchronously."""
    if not self.running:
        return None

    try:
        return await self.process_message(msg)

    except StopConsume:
        # Stop handler at StopConsume exception
        await self.stop()

    except SystemExit:
        # Stop handler at `exit()` call
        await self.stop()

        if app := self._outer_config.context.get("app"):
            app.exit()

    except Exception:  # nosec B110  # ruff: ignore[try-except-pass]
        # All other exceptions were logged by CriticalLogMiddleware
        pass

process_message async #

process_message(msg: MsgType) -> Response

Execute all message processing stages.

Source code in faststream/_internal/endpoint/subscriber/usecase.py
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
async def process_message(self, msg: MsgType) -> "Response":
    """Execute all message processing stages."""
    context = self._outer_config.context
    logger_state = self._outer_config.logger

    async with AsyncExitStack() as stack:
        stack.enter_context(self.lock)

        # Enter context before middlewares
        stack.enter_context(
            context.scopes(
                (
                    ("handler_", self),
                    ("logger", logger_state.logger.logger),
                    *self._outer_config.extra_context.items(),
                ),
            ),
        )

        # reserve place for context scope under the middlewares __aexit__
        # because middlewares should be exited before context scope release
        message_scope = stack.enter_context(ExitStack())

        # enter all middlewares
        middlewares: list[BaseMiddleware] = []
        for base_m in self.__build__middlewares_stack():
            middleware = base_m(msg, context=context)
            middlewares.append(middleware)
            await stack.enter_async_context(middleware)

        cache: dict[Any, Any] = {}
        parsing_error: Exception | None = None
        for h in self.calls:
            try:
                message = await h.is_suitable(msg, cache)
            except Exception as e:
                parsing_error = e
                break

            if message is not None:
                message_scope.enter_context(
                    context.scopes(
                        (
                            ("log_context", self.get_log_context(message)),
                            ("message", message),
                        ),
                    ),
                )

                result_msg = ensure_response(
                    await h.call(
                        message=message,
                        # consumer middlewares
                        _extra_middlewares=(
                            m.consume_scope for m in middlewares[::-1]
                        ),
                    ),
                )

                if not result_msg.correlation_id:
                    result_msg.correlation_id = message.correlation_id

                for p in chain(
                    self.__get_response_publisher(message),
                    h.handler._publishers,
                ):
                    await p._publish(
                        result_msg.as_publish_command(),
                        _extra_middlewares=(
                            m.publish_scope for m in middlewares[::-1]
                        ),
                    )

                # Return data for tests
                return result_msg

        # Suitable handler was not found or
        # parsing/decoding exception occurred.
        # Reraise it to catch in tests
        if parsing_error:
            raise parsing_error

        error_msg = f"There is no suitable handler for {msg=}"
        raise SubscriberNotFound(error_msg)

    # An error was raised and processed by some middleware
    return ensure_response(None)

get_one async #

get_one(*, timeout: float = 5.0) -> KafkaMessage | None
Source code in faststream/kafka/subscriber/usecase.py
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
@override
async def get_one(
    self,
    *,
    timeout: float = 5.0,
) -> "KafkaMessage | None":
    assert not self.calls, (
        "You can't use `get_one` method if subscriber has registered handlers."
    )

    assert self.consumer, "You should start subscriber at first."

    raw_messages = await self.consumer.getmany(
        timeout_ms=timeout * 1000,
        max_records=1,
    )

    if not raw_messages:
        return None

    ((raw_message,),) = raw_messages.values()

    context = self._outer_config.context

    async_parser, async_decoder = self._get_parser_and_decoder()

    msg: KafkaMessage | None = await process_msg(  # type: ignore[assignment]
        msg=raw_message,
        middlewares=(
            m(raw_message, context=context) for m in self._broker_middlewares
        ),
        parser=async_parser,
        decoder=async_decoder,
    )
    return msg

schema #

schema() -> dict[str, SubscriberSpec]
Source code in faststream/_internal/endpoint/subscriber/usecase.py
525
526
527
def schema(self) -> dict[str, "SubscriberSpec"]:
    self._build_fastdepends_model()
    return self.specification.get_schema()

add_task #

add_task(
    func: Callable[..., Coroutine[Any, Any, Any]],
    func_args: tuple[Any, ...] | None = None,
    func_kwargs: dict[str, Any] | None = None,
    *,
    restart_on_failure: bool = True,
) -> None
Source code in faststream/_internal/endpoint/subscriber/mixins.py
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def add_task(
    self,
    func: Callable[..., Coroutine[Any, Any, Any]],
    func_args: tuple[Any, ...] | None = None,
    func_kwargs: dict[str, Any] | None = None,
    *,
    restart_on_failure: bool = True,
) -> None:
    args = func_args or ()
    kwargs = func_kwargs or {}
    task = asyncio.create_task(func(*args, **kwargs))
    callback = TaskCallbackSupervisor(
        func,
        func_args,
        func_kwargs,
        self,
        restart_on_failure=restart_on_failure,
    )
    task.add_done_callback(callback)
    self.tasks.append(task)

consume_one async #

consume_one(msg: MsgType) -> None
Source code in faststream/kafka/subscriber/usecase.py
378
379
async def consume_one(self, msg: MsgType) -> None:
    await self.consume(msg)

build_log_context staticmethod #

build_log_context(
    message: Optional[StreamMessage[Any]],
    topic: str,
    group_id: str | None = None,
) -> dict[str, str]
Source code in faststream/kafka/subscriber/usecase.py
394
395
396
397
398
399
400
401
402
403
404
@staticmethod
def build_log_context(
    message: Optional["StreamMessage[Any]"],
    topic: str,
    group_id: str | None = None,
) -> dict[str, str]:
    return {
        "topic": topic,
        "group_id": group_id or "",
        "message_id": getattr(message, "message_id", ""),
    }