KafkaBroker

async def close(
    self,
    exc_type: Optional[Type[BaseException]] = None,
    exc_val: Optional[BaseException] = None,
    exec_tb: Optional[TracebackType] = None,
) -> None:
    """Closes the object.

    Args:
        exc_type: The type of the exception being handled, if any.
        exc_val: The exception instance being handled, if any.
        exec_tb: The traceback of the exception being handled, if any.

    Returns:
        None

    Raises:
        NotImplementedError: If the method is not implemented.

    """
    super()._abc_close(exc_type, exc_val, exec_tb)

    for h in self.handlers.values():
        await h.close()

    if self._connection is not None:
        await self._close(exc_type, exc_val, exec_tb)

async def connect(
    self,
    *args: Any,
    **kwargs: Any,
) -> ConsumerConnectionParams:
    """Establishes a connection to Kafka and returns connection parameters.

    Args:
        *args: Additional arguments.
        **kwargs: Additional keyword arguments.

    Returns:
        ConsumerConnectionParams: The connection parameters.
    """
    connection = await super().connect(*args, **kwargs)
    for p in self._publishers.values():
        p._producer = self._producer
    return connection

def include_router(self, router: BrokerRouter[Any, MsgType]) -> None:
    """Includes a router in the current object.

    Args:
        router: The router to be included.

    Returns:
        None

    """
    for r in router._handlers:
        self.subscriber(*r.args, **r.kwargs)(r.call)

    self._publishers = {**self._publishers, **router._publishers}

def include_routers(self, *routers: BrokerRouter[Any, MsgType]) -> None:
    """Includes routers in the current object.

    Args:
        *routers: Variable length argument list of routers to include.

    Returns:
        None

    """
    for r in routers:
        self.include_router(r)

@override
async def publish(  # type: ignore[override]
    self,
    *args: Any,
    **kwargs: Any,
) -> None:
    """Publish a message to Kafka.

    Args:
        *args: Positional arguments for message publishing.
        **kwargs: Keyword arguments for message publishing.

    Raises:
        RuntimeError: If KafkaBroker is not started yet.
    """
    assert self._producer, NOT_CONNECTED_YET  # nosec B101
    return await self._producer.publish(*args, **kwargs)

async def publish_batch(
    self,
    *args: Any,
    **kwargs: Any,
) -> None:
    """Publish a batch of messages to Kafka.

    Args:
        *args: Positional arguments for message publishing.
        **kwargs: Keyword arguments for message publishing.

    Raises:
        RuntimeError: If KafkaBroker is not started yet.
    """
    assert self._producer, NOT_CONNECTED_YET  # nosec B101
    await self._producer.publish_batch(*args, **kwargs)

@override
def publisher(  # type: ignore[override]
    self,
    topic: str,
    key: Optional[bytes] = None,
    partition: Optional[int] = None,
    timestamp_ms: Optional[int] = None,
    headers: Optional[Dict[str, str]] = None,
    reply_to: str = "",
    batch: bool = False,
    # AsyncAPI information
    title: Optional[str] = None,
    description: Optional[str] = None,
    schema: Optional[Any] = None,
    include_in_schema: bool = True,
) -> Publisher:
    """Create a message publisher for the specified topic.

    Args:
        topic (str): The topic to publish messages to.
        key (Optional[bytes]): Message key.
        partition (Optional[int]): Partition to send the message to.
        timestamp_ms (Optional[int]): Message timestamp in milliseconds.
        headers (Optional[Dict[str, str]]): Message headers.
        reply_to (str): The topic to which responses should be sent.
        batch (bool): Whether to publish messages in batches.
        title (Optional[str]): AsyncAPI title.
        description (Optional[str]): AsyncAPI description.
        schema (Optional[Any]): AsyncAPI schema.
        include_in_schema (bool): Whether to include the message publisher in the AsyncAPI schema.

    Returns:
        Publisher: A message publisher.
    """
    publisher = self._publishers.get(
        topic,
        Publisher(
            topic=topic,
            client_id=self.client_id,
            key=key,
            batch=batch,
            partition=partition,
            timestamp_ms=timestamp_ms,
            headers=headers,
            reply_to=reply_to,
            title=title,
            _description=description,
            _schema=schema,
            include_in_schema=include_in_schema,
        ),
    )
    super().publisher(topic, publisher)
    if self._producer is not None:
        publisher._producer = self._producer
    return publisher

async def start(self) -> None:
    """Start the KafkaBroker and message handlers."""
    context.set_global(
        "default_log_context",
        self._get_log_context(None, ""),
    )

    await super().start()

    for handler in self.handlers.values():
        c = self._get_log_context(None, handler.topics, handler.group_id)
        self._log(f"`{handler.call_name}` waiting for messages", extra=c)
        await handler.start(**(self._connection or {}))

@override
def subscriber(  # type: ignore[override]
    self,
    *topics: str,
    group_id: Optional[str] = None,
    key_deserializer: Optional[Callable[[bytes], Any]] = None,
    value_deserializer: Optional[Callable[[bytes], Any]] = None,
    fetch_max_wait_ms: int = 500,
    fetch_max_bytes: int = 52428800,
    fetch_min_bytes: int = 1,
    max_partition_fetch_bytes: int = 1 * 1024 * 1024,
    auto_offset_reset: Literal[
        "latest",
        "earliest",
        "none",
    ] = "latest",
    auto_commit: bool = True,
    auto_commit_interval_ms: int = 5000,
    check_crcs: bool = True,
    partition_assignment_strategy: Sequence[AbstractPartitionAssignor] = (
        RoundRobinPartitionAssignor,
    ),
    max_poll_interval_ms: int = 300000,
    rebalance_timeout_ms: Optional[int] = None,
    session_timeout_ms: int = 10000,
    heartbeat_interval_ms: int = 3000,
    consumer_timeout_ms: int = 200,
    max_poll_records: Optional[int] = None,
    exclude_internal_topics: bool = True,
    isolation_level: Literal[
        "read_uncommitted",
        "read_committed",
    ] = "read_uncommitted",
    # broker arguments
    dependencies: Sequence[Depends] = (),
    parser: Optional[
        Union[
            CustomParser[aiokafka.ConsumerRecord, KafkaMessage],
            CustomParser[Tuple[aiokafka.ConsumerRecord, ...], KafkaMessage],
        ]
    ] = None,
    decoder: Optional[CustomDecoder] = None,
    middlewares: Optional[
        Sequence[
            Callable[
                [aiokafka.ConsumerRecord],
                BaseMiddleware,
            ]
        ]
    ] = None,
    filter: Union[
        Filter[KafkaMessage],
        Filter[StreamMessage[Tuple[aiokafka.ConsumerRecord, ...]]],
    ] = default_filter,
    batch: bool = False,
    max_records: Optional[int] = None,
    batch_timeout_ms: int = 200,
    no_ack: bool = False,
    # AsyncAPI information
    title: Optional[str] = None,
    description: Optional[str] = None,
    include_in_schema: bool = True,
    **original_kwargs: Any,
) -> Callable[
    [Callable[P_HandlerParams, T_HandlerReturn]],
    Union[
        HandlerCallWrapper[
            aiokafka.ConsumerRecord, P_HandlerParams, T_HandlerReturn
        ],
        HandlerCallWrapper[
            Tuple[aiokafka.ConsumerRecord, ...], P_HandlerParams, T_HandlerReturn
        ],
    ],
]:
    """Create a message subscriber for the specified topics.

    Args:
        *topics (str): The topics to subscribe to.
        group_id (Optional[str]): The Kafka consumer group ID.
        key_deserializer (Optional[Callable[[bytes], Any]]): Key deserializer function.
        value_deserializer (Optional[Callable[[bytes], Any]]): Value deserializer function.
        fetch_max_wait_ms (int): The maximum time to wait for data.
        fetch_max_bytes (int): The maximum number of bytes to fetch.
        fetch_min_bytes (int): The minimum number of bytes to fetch.
        max_partition_fetch_bytes (int): The maximum bytes to fetch for a partition.
        auto_offset_reset (Literal["latest", "earliest", "none"]): Auto offset reset policy.
        auto_commit (bool): Whether to enable auto-commit.
        auto_commit_interval_ms (int): Auto-commit interval in milliseconds.
        check_crcs (bool): Whether to check CRCs.
        partition_assignment_strategy (Sequence[AbstractPartitionAssignor]): Partition assignment strategy.
        max_poll_interval_ms (int): Maximum poll interval in milliseconds.
        rebalance_timeout_ms (Optional[int]): Rebalance timeout in milliseconds.
        session_timeout_ms (int): Session timeout in milliseconds.
        heartbeat_interval_ms (int): Heartbeat interval in milliseconds.
        consumer_timeout_ms (int): Consumer timeout in milliseconds.
        max_poll_records (Optional[int]): Maximum number of records to poll.
        exclude_internal_topics (bool): Whether to exclude internal topics.
        isolation_level (Literal["read_uncommitted", "read_committed"]): Isolation level.
        dependencies (Sequence[Depends]): Additional dependencies for message handling.
        parser (Optional[Union[CustomParser[aiokafka.ConsumerRecord], CustomParser[Tuple[aiokafka.ConsumerRecord, ...]]]]): Message parser.
        decoder (Optional[CustomDecoder]): Message decoder.
        middlewares (Optional[Sequence[Callable[[aiokafka.ConsumerRecord], BaseMiddleware]]]): Message middlewares.
        filter (Union[Filter[KafkaMessage], Filter[StreamMessage[Tuple[aiokafka.ConsumerRecord, ...]]]]): Message filter.
        batch (bool): Whether to process messages in batches.
        max_records (Optional[int]): Maximum number of records to process in each batch.
        batch_timeout_ms (int): Batch timeout in milliseconds.
        no_ack (bool): Whether not to ack/nack/reject messages.
        title (Optional[str]): AsyncAPI title.
        description (Optional[str]): AsyncAPI description.
        include_in_schema (bool): Whether to include the message handler in the AsyncAPI schema.
        **original_kwargs: Additional keyword arguments.

    Returns:
        Callable: A decorator that wraps a message handler function.
    """
    super().subscriber()

    self._setup_log_context(topics, group_id)

    if not auto_commit and not group_id:
        raise ValueError("You should install `group_id` with manual commit mode")

    key = Handler.get_routing_hash(topics, group_id)
    builder = partial(
        aiokafka.AIOKafkaConsumer,
        key_deserializer=key_deserializer,
        value_deserializer=value_deserializer,
        fetch_max_wait_ms=fetch_max_wait_ms,
        fetch_max_bytes=fetch_max_bytes,
        fetch_min_bytes=fetch_min_bytes,
        max_partition_fetch_bytes=max_partition_fetch_bytes,
        auto_offset_reset=auto_offset_reset,
        enable_auto_commit=auto_commit,
        auto_commit_interval_ms=auto_commit_interval_ms,
        check_crcs=check_crcs,
        partition_assignment_strategy=partition_assignment_strategy,
        max_poll_interval_ms=max_poll_interval_ms,
        rebalance_timeout_ms=rebalance_timeout_ms,
        session_timeout_ms=session_timeout_ms,
        heartbeat_interval_ms=heartbeat_interval_ms,
        consumer_timeout_ms=consumer_timeout_ms,
        max_poll_records=max_poll_records,
        exclude_internal_topics=exclude_internal_topics,
        isolation_level=isolation_level,
    )
    handler = self.handlers.get(
        key,
        Handler(
            *topics,
            log_context_builder=partial(
                self._get_log_context,
                topics=topics,
                group_id=group_id,
            ),
            is_manual=not auto_commit,
            group_id=group_id,
            client_id=self.client_id,
            builder=builder,
            description=description,
            title=title,
            batch=batch,
            batch_timeout_ms=batch_timeout_ms,
            max_records=max_records,
            include_in_schema=include_in_schema,
            graceful_timeout=self.graceful_timeout,
        ),
    )

    self.handlers[key] = handler

    def consumer_wrapper(
        func: Callable[P_HandlerParams, T_HandlerReturn],
    ) -> HandlerCallWrapper[
        aiokafka.ConsumerRecord, P_HandlerParams, T_HandlerReturn
    ]:
        """A wrapper function for a consumer handler.

        Args:
            func : The consumer handler function to be wrapped.

        Returns:
            The wrapped handler call.

        Raises:
            NotImplementedError: If silent animals are not supported.
        !!! note

            The above docstring is autogenerated by docstring-gen library (https://docstring-gen.airt.ai)
        """
        handler_call, dependant = self._wrap_handler(
            func=func,
            extra_dependencies=dependencies,
            no_ack=no_ack,
            **original_kwargs,
        )

        handler.add_call(
            handler=handler_call,
            filter=filter,
            middlewares=middlewares,
            parser=parser or self._global_parser,
            decoder=decoder or self._global_decoder,
            dependant=dependant,
        )

        return handler_call

    return consumer_wrapper