fix (#1748): add Kafka publish no_confirm option

faststream/confluent/broker/broker.py

@@ -481,6 +481,7 @@ async def publish(  # type: ignore[override]
         correlation_id: Optional[str] = None,
         *,
         reply_to: str = "",
+        no_confirm: bool = False,
         # extra options to be compatible with test client
         **kwargs: Any,
     ) -> Optional[Any]:
@@ -496,6 +497,7 @@ async def publish(  # type: ignore[override]
             headers=headers,
             correlation_id=correlation_id,
             reply_to=reply_to,
+            no_confirm=no_confirm,
             **kwargs,
         )
 
@@ -534,6 +536,7 @@ async def publish_batch(
         headers: Optional[Dict[str, str]] = None,
         reply_to: str = "",
         correlation_id: Optional[str] = None,
+        no_confirm: bool = False,
     ) -> None:
         assert self._producer, NOT_CONNECTED_YET  # nosec B101
 
@@ -551,6 +554,7 @@ async def publish_batch(
             headers=headers,
             reply_to=reply_to,
             correlation_id=correlation_id,
+            no_confirm=no_confirm,
         )
 
     @override

faststream/confluent/client.py

@@ -1,4 +1,7 @@
+import asyncio
 import logging
+from contextlib import suppress
+from threading import Thread
 from time import time
 from typing import (
     TYPE_CHECKING,
@@ -24,8 +27,17 @@
 from faststream.utils.functions import call_or_await
 
 if TYPE_CHECKING:
+    from typing_extensions import NotRequired, TypedDict
+
     from faststream.types import AnyDict, LoggerProto
 
+    class _SendKwargs(TypedDict):
+        value: Optional[Union[str, bytes]]
+        key: Optional[Union[str, bytes]]
+        headers: Optional[List[Tuple[str, Union[str, bytes]]]]
+        partition: NotRequired[int]
+        timestamp: NotRequired[int]
+
 
 class AsyncConfluentProducer:
     """An asynchronous Python Kafka client using the "confluent-kafka" package."""
@@ -102,9 +114,21 @@ def __init__(
 
         self.producer = Producer(self.config, logger=self.logger)
 
+        self.__running = True
+        self._poll_thread = Thread(target=self._poll_loop)
+        self._poll_thread.start()
+
+    def _poll_loop(self) -> None:
+        while self.__running:
+            with suppress(Exception):
+                self.producer.poll(0.1)
+
     async def stop(self) -> None:
         """Stop the Kafka producer and flush remaining messages."""
-        await call_or_await(self.producer.flush)
+        if self.__running:
+            await call_or_await(self.producer.flush)
+            self.__running = False
+            self._poll_thread.join()
 
     async def send(
         self,
@@ -114,9 +138,10 @@ async def send(
         partition: Optional[int] = None,
         timestamp_ms: Optional[int] = None,
         headers: Optional[List[Tuple[str, Union[str, bytes]]]] = None,
+        no_confirm: bool = False,
     ) -> None:
         """Sends a single message to a Kafka topic."""
-        kwargs: AnyDict = {
+        kwargs: _SendKwargs = {
             "value": value,
             "key": key,
             "headers": headers,
@@ -128,9 +153,18 @@ async def send(
         if timestamp_ms is not None:
             kwargs["timestamp"] = timestamp_ms
 
+        result_future: asyncio.Future[Optional[Message]] = asyncio.Future()
+
+        def ack_callback(err: Any, msg: Optional[Message]) -> None:
+            if err or (msg is not None and (err := msg.error())):
+                result_future.set_exception(KafkaException(err))
+            else:
+                result_future.set_result(msg)
+
         # should be sync to prevent segfault
-        self.producer.produce(topic, **kwargs)
-        self.producer.poll(0)
+        self.producer.produce(topic, **kwargs, on_delivery=ack_callback)
+        if not no_confirm:
+            await result_future
 
     def create_batch(self) -> "BatchBuilder":
         """Creates a batch for sending multiple messages."""

faststream/confluent/publisher/usecase.py

@@ -176,6 +176,7 @@ async def publish(
         headers: Optional[Dict[str, str]] = None,
         correlation_id: Optional[str] = None,
         reply_to: str = "",
+        no_confirm: bool = False,
         # publisher specific
         _extra_middlewares: Iterable["PublisherMiddleware"] = (),
     ) -> Optional[Any]:
@@ -184,6 +185,7 @@ async def publish(
         kwargs: AnyDict = {
             "key": key or self.key,
             # basic args
+            "no_confirm": no_confirm,
             "topic": topic or self.topic,
             "partition": partition or self.partition,
             "timestamp_ms": timestamp_ms,
@@ -243,8 +245,9 @@ async def publish(
         partition: Optional[int] = None,
         timestamp_ms: Optional[int] = None,
         headers: Optional[Dict[str, str]] = None,
-        reply_to: str = "",
         correlation_id: Optional[str] = None,
+        reply_to: str = "",
+        no_confirm: bool = False,
         # publisher specific
         _extra_middlewares: Iterable["PublisherMiddleware"] = (),
     ) -> None:
@@ -258,6 +261,7 @@ async def publish(
 
         kwargs: AnyDict = {
             "topic": topic or self.topic,
+            "no_confirm": no_confirm,
             "partition": partition or self.partition,
             "timestamp_ms": timestamp_ms,
             "headers": headers or self.headers,

faststream/confluent/schemas/partition.py

@@ -1,15 +1,16 @@
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from confluent_kafka import TopicPartition as ConfluentPartition
-from typing_extensions import NotRequired, TypedDict
 
+if TYPE_CHECKING:
+    from typing_extensions import NotRequired, TypedDict
 
-class _TopicKwargs(TypedDict):
-    topic: str
-    partition: int
-    offset: int
-    metadata: NotRequired[str]
-    leader_epoch: NotRequired[int]
+    class _TopicKwargs(TypedDict):
+        topic: str
+        partition: int
+        offset: int
+        metadata: NotRequired[str]
+        leader_epoch: NotRequired[int]
 
 
 class TopicPartition:

faststream/confluent/testing.py

@@ -100,6 +100,7 @@ async def publish(  # type: ignore[override]
         headers: Optional[Dict[str, str]] = None,
         correlation_id: Optional[str] = None,
         *,
+        no_confirm: bool = False,
         reply_to: str = "",
         rpc: bool = False,
         rpc_timeout: Optional[float] = None,
@@ -147,6 +148,7 @@ async def publish_batch(
         headers: Optional[Dict[str, str]] = None,
         reply_to: str = "",
         correlation_id: Optional[str] = None,
+        no_confirm: bool = False,
     ) -> None:
         """Publish a batch of messages to the Kafka broker."""
         for handler in self.broker._subscribers.values():  # pragma: no branch

faststream/kafka/broker/broker.py

@@ -719,6 +719,10 @@ async def publish(  # type: ignore[override]
             str,
             Doc("Reply message topic name to send response."),
         ] = "",
+        no_confirm: Annotated[
+            bool,
+            Doc("Do not wait for Kafka publish confirmation."),
+        ] = False,
         # extra options to be compatible with test client
         **kwargs: Any,
     ) -> Optional[Any]:
@@ -741,6 +745,7 @@ async def publish(  # type: ignore[override]
             headers=headers,
             correlation_id=correlation_id,
             reply_to=reply_to,
+            no_confirm=no_confirm,
             **kwargs,
         )
 
@@ -861,6 +866,10 @@ async def publish_batch(
                 "**correlation_id** is a useful option to trace messages."
             ),
         ] = None,
+        no_confirm: Annotated[
+            bool,
+            Doc("Do not wait for Kafka publish confirmation."),
+        ] = False,
     ) -> None:
         assert self._producer, NOT_CONNECTED_YET  # nosec B101
 
@@ -879,6 +888,7 @@ async def publish_batch(
             headers=headers,
             reply_to=reply_to,
             correlation_id=correlation_id,
+            no_confirm=no_confirm,
         )
 
     @override

faststream/kafka/publisher/producer.py

@@ -47,6 +47,7 @@ async def publish(  # type: ignore[override]
         timestamp_ms: Optional[int] = None,
         headers: Optional[Dict[str, str]] = None,
         reply_to: str = "",
+        no_confirm: bool = False,
     ) -> None:
         """Publish a message to a topic."""
         message, content_type = encode_message(message)
@@ -63,14 +64,16 @@ async def publish(  # type: ignore[override]
                 reply_to,
             )
 
-        await self._producer.send(
+        send_future = await self._producer.send(
             topic=topic,
             value=message,
             key=key,
             partition=partition,
             timestamp_ms=timestamp_ms,
             headers=[(i, (j or "").encode()) for i, j in headers_to_send.items()],
         )
+        if not no_confirm:
+            await send_future
 
     async def stop(self) -> None:
         await self._producer.stop()
@@ -84,6 +87,7 @@ async def publish_batch(
         timestamp_ms: Optional[int] = None,
         headers: Optional[Dict[str, str]] = None,
         reply_to: str = "",
+        no_confirm: bool = False,
     ) -> None:
         """Publish a batch of messages to a topic."""
         batch = self._producer.create_batch()
@@ -114,7 +118,9 @@ async def publish_batch(
                 headers=[(i, j.encode()) for i, j in final_headers.items()],
             )
 
-        await self._producer.send_batch(batch, topic, partition=partition)
+        send_future = await self._producer.send_batch(batch, topic, partition=partition)
+        if not no_confirm:
+            await send_future
 
     @override
     async def request(self, *args: Any, **kwargs: Any) -> Optional[Any]:

faststream/kafka/publisher/usecase.py

@@ -276,6 +276,10 @@ async def publish(
             str,
             Doc("Reply message topic name to send response."),
         ] = "",
+        no_confirm: Annotated[
+            bool,
+            Doc("Do not wait for Kafka publish confirmation."),
+        ] = False,
         # publisher specific
         _extra_middlewares: Annotated[
             Iterable["PublisherMiddleware"],
@@ -311,6 +315,7 @@ async def publish(
             reply_to=reply_to,
             correlation_id=correlation_id,
             timestamp_ms=timestamp_ms,
+            no_confirm=no_confirm,
         )
 
     @override
@@ -440,6 +445,10 @@ async def publish(
                 "**correlation_id** is a useful option to trace messages."
             ),
         ] = None,
+        no_confirm: Annotated[
+            bool,
+            Doc("Do not wait for Kafka publish confirmation."),
+        ] = False,
         # publisher specific
         _extra_middlewares: Annotated[
             Iterable["PublisherMiddleware"],
@@ -479,4 +488,5 @@ async def publish(
             reply_to=reply_to,
             correlation_id=correlation_id,
             timestamp_ms=timestamp_ms,
+            no_confirm=no_confirm,
         )

faststream/kafka/testing.py

@@ -105,6 +105,7 @@ async def publish(  # type: ignore[override]
         rpc: bool = False,
         rpc_timeout: Optional[float] = None,
         raise_timeout: bool = False,
+        no_confirm: bool = False,
     ) -> Optional[Any]:
         """Publish a message to the Kafka broker."""
         incoming = build_message(
@@ -184,6 +185,7 @@ async def publish_batch(
         headers: Optional[Dict[str, str]] = None,
         reply_to: str = "",
         correlation_id: Optional[str] = None,
+        no_confirm: bool = False,
     ) -> None:
         """Publish a batch of messages to the Kafka broker."""
         for handler in self.broker._subscribers.values():  # pragma: no branch

pyproject.toml

@@ -291,6 +291,7 @@ minversion = "7.0"
 addopts = "-q -m 'not slow'"
 testpaths = ["tests"]
 markers = ["rabbit", "kafka", "confluent", "nats", "redis", "slow", "all"]
+asyncio_default_fixture_loop_scope = "function"
 
 [tool.coverage.run]
 parallel = true

tests/brokers/base/middlewares.py

@@ -34,7 +34,7 @@ async def test_subscriber_middleware(
         raw_broker,
     ):
         async def mid(call_next, msg):
-            mock.start(msg._decoded_body)
+            mock.start(await msg.decode())
             result = await call_next(msg)
             mock.end()
             event.set()
@@ -468,7 +468,7 @@ async def subscriber1(m):
 
         @broker.subscriber(*args, **kwargs)
         async def subscriber2(msg=Context("message")):
-            mock(msg.decoded_body)
+            mock(await msg.decode())
             event.set()
 
         broker = self.patch_broker(raw_broker, broker)
@@ -509,7 +509,7 @@ async def subscriber1(m):
 
         @broker.subscriber(*args2, **kwargs2)
         async def subscriber2(msg=Context("message")):
-            mock(msg.decoded_body)
+            mock(await msg.decode())
 
         broker = self.patch_broker(raw_broker, broker)
 
@@ -581,7 +581,7 @@ async def subscriber1(m):
 
         @broker.subscriber(*args2, **kwargs2)
         async def subscriber2(msg=Context("message")):
-            mock(msg.decoded_body)
+            mock(await msg.decode())
 
         broker = self.patch_broker(raw_broker, broker)
 
@@ -632,7 +632,7 @@ async def subscriber2(m):
 
         @broker.subscriber(*args3, **kwargs3)
         async def subscriber3(msg=Context("message")):
-            mock(msg.decoded_body)
+            mock(await msg.decode())
             if mock.call_count > 1:
                 event.set()