feat: add on_assign, on_revoke, on_lost callbacks for Confluent subscriber (#2789)

Add rebalance callback parameters to the Confluent subscriber, allowing
users to specify on_assign, on_revoke, and on_lost callbacks that are
passed through to confluent_kafka.Consumer.subscribe().

This enables monitoring partition assignment/revocation events without
monkey-patching AsyncConfluentConsumer.

Fixes #1676

Author: Br1an67

faststream/confluent/broker/registrator.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from typing import (
     TYPE_CHECKING,
     Annotated,
@@ -78,6 +78,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: Literal[False] = False,
         max_records: int | None = None,
         # broker args
@@ -138,6 +142,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: Literal[True] = ...,
         max_records: int | None = None,
         # broker args
@@ -198,6 +206,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: Literal[False] = False,
         max_records: int | None = None,
         # broker args
@@ -258,6 +270,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: bool = False,
         max_records: int | None = None,
         # broker args
@@ -322,6 +338,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: bool = False,
         max_records: int | None = None,
         # broker args
@@ -464,6 +484,12 @@ def subscriber(
                 return the ALSO. See method docs below.
             batch: Whether to consume messages in batches or not.
             max_records: Number of messages to consume as one batch.
+            on_assign: Callback called when partitions are assigned to the consumer
+                during a rebalance. Receives ``(consumer, partitions)`` arguments.
+            on_revoke: Callback called when partitions are revoked from the consumer
+                during a rebalance. Receives ``(consumer, partitions)`` arguments.
+            on_lost: Callback called when partitions are lost (e.g., due to session
+                timeout). Receives ``(consumer, partitions)`` arguments.
             dependencies: Dependencies list (`[Dependant(),]`) to apply to the subscriber.
             parser: Parser to map original **Message** object to FastStream one.
             decoder: Function to decode FastStream msg bytes body to python objects.
@@ -506,6 +532,9 @@ def subscriber(
                 "session_timeout_ms": session_timeout_ms,
                 "heartbeat_interval_ms": heartbeat_interval_ms,
                 "isolation_level": isolation_level,
+                "on_assign": on_assign,
+                "on_revoke": on_revoke,
+                "on_lost": on_lost,
             },
             auto_commit=auto_commit,
             # subscriber args

faststream/confluent/broker/router.py

@@ -135,6 +135,10 @@ def __init__(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: bool = False,
         max_records: int | None = None,
         # broker args
@@ -274,6 +278,9 @@ def __init__(
                 return the ALSO. See method docs below.
             batch: Whether to consume messages in batches or not.
             max_records: Number of messages to consume as one batch.
+            on_assign: Callback called when partitions are assigned to the consumer.
+            on_revoke: Callback called when partitions are revoked from the consumer.
+            on_lost: Callback called when partitions are lost.
             dependencies: Dependencies list (`[Dependant(),]`) to apply to the subscriber.
             parser: Parser to map original **Message** object to FastStream one.
             decoder: Function to decode FastStream msg bytes body to python objects.
@@ -309,6 +316,9 @@ def __init__(
             session_timeout_ms=session_timeout_ms,
             heartbeat_interval_ms=heartbeat_interval_ms,
             isolation_level=isolation_level,
+            on_assign=on_assign,
+            on_revoke=on_revoke,
+            on_lost=on_lost,
             max_records=max_records,
             batch=batch,
             # basic args

faststream/confluent/fastapi/fastapi.py

@@ -430,6 +430,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: Literal[False] = False,
         max_records: int | None = None,
         # broker args
@@ -497,6 +501,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: Literal[False] = False,
         max_records: int | None = None,
         # broker args
@@ -564,6 +572,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: Literal[True] = ...,
         max_records: int | None = None,
         # broker args
@@ -622,6 +634,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: bool = False,
         max_records: int | None = None,
         # broker args
@@ -693,6 +709,10 @@ def subscriber(
             "read_uncommitted",
             "read_committed",
         ] = "read_uncommitted",
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
         batch: bool = False,
         max_records: int | None = None,
         # broker args
@@ -838,6 +858,12 @@ def subscriber(
                 return the ALSO. See method docs below.
             batch: Whether to consume messages in batches or not.
             max_records: Number of messages to consume as one batch.
+            on_assign: Callback called when partitions are assigned to the consumer
+                during a rebalance. Receives ``(consumer, partitions)`` arguments.
+            on_revoke: Callback called when partitions are revoked from the consumer
+                during a rebalance. Receives ``(consumer, partitions)`` arguments.
+            on_lost: Callback called when partitions are lost (e.g., due to session
+                timeout). Receives ``(consumer, partitions)`` arguments.
             dependencies: Dependencies list (`[Dependant(),]`) to apply to the subscriber.
             parser: Parser to map original **Message** object to FastStream one.
             decoder: Function to decode FastStream msg bytes body to python objects.
@@ -950,6 +976,9 @@ def subscriber(
             session_timeout_ms=session_timeout_ms,
             heartbeat_interval_ms=heartbeat_interval_ms,
             isolation_level=isolation_level,
+            on_assign=on_assign,
+            on_revoke=on_revoke,
+            on_lost=on_lost,
             batch=batch,
             max_records=max_records,
             # broker args

faststream/confluent/helpers/client.py

@@ -215,10 +215,18 @@ def __init__(
         connections_max_idle_ms: int = 540000,
         isolation_level: str = "read_uncommitted",
         allow_auto_create_topics: bool = True,
+        # rebalance callbacks
+        on_assign: Callable[..., None] | None = None,
+        on_revoke: Callable[..., None] | None = None,
+        on_lost: Callable[..., None] | None = None,
     ) -> None:
         self.admin_client = admin_service
         self.logger_state = logger
 
+        self._on_assign = on_assign
+        self._on_revoke = on_revoke
+        self._on_lost = on_lost
+
         self.topics = list(topics)
         self.partitions = partitions
 
@@ -292,10 +300,17 @@ async def start(self) -> None:
             )
 
         if self.topics:
+            subscribe_kwargs: dict[str, Any] = {"topics": self.topics}
+            if self._on_assign is not None:
+                subscribe_kwargs["on_assign"] = self._on_assign
+            if self._on_revoke is not None:
+                subscribe_kwargs["on_revoke"] = self._on_revoke
+            if self._on_lost is not None:
+                subscribe_kwargs["on_lost"] = self._on_lost
             await run_in_executor(
                 self._thread_pool,
                 self.consumer.subscribe,
-                topics=self.topics,
+                **subscribe_kwargs,
             )
 
         elif self.partitions: