diff --git a/docs/docs/SUMMARY.md b/docs/docs/SUMMARY.md index bf061a321ab..5c5b4db5b4d 100644 --- a/docs/docs/SUMMARY.md +++ b/docs/docs/SUMMARY.md @@ -119,6 +119,7 @@ search: - [Reference - Code API](api/index.md) - Public API - faststream + - [AckPolicy](public_api/faststream/AckPolicy.md) - [BaseMiddleware](public_api/faststream/BaseMiddleware.md) - [Context](public_api/faststream/Context.md) - [Depends](public_api/faststream/Depends.md) @@ -205,6 +206,7 @@ search: - [TestRedisBroker](public_api/faststream/redis/TestRedisBroker.md) - All API - faststream + - [AckPolicy](api/faststream/AckPolicy.md) - [BaseMiddleware](api/faststream/BaseMiddleware.md) - [Context](api/faststream/Context.md) - [Depends](api/faststream/Depends.md) @@ -472,8 +474,15 @@ search: - [encode_message](api/faststream/message/utils/encode_message.md) - [gen_cor_id](api/faststream/message/utils/gen_cor_id.md) - middlewares + - [AckPolicy](api/faststream/middlewares/AckPolicy.md) + - [AcknowledgementMiddleware](api/faststream/middlewares/AcknowledgementMiddleware.md) - [BaseMiddleware](api/faststream/middlewares/BaseMiddleware.md) - [ExceptionMiddleware](api/faststream/middlewares/ExceptionMiddleware.md) + - acknowledgement + - conf + - [AckPolicy](api/faststream/middlewares/acknowledgement/conf/AckPolicy.md) + - middleware + - [AcknowledgementMiddleware](api/faststream/middlewares/acknowledgement/middleware/AcknowledgementMiddleware.md) - base - [BaseMiddleware](api/faststream/middlewares/base/BaseMiddleware.md) - exception diff --git a/docs/docs/en/api/faststream/AckPolicy.md b/docs/docs/en/api/faststream/AckPolicy.md new file mode 100644 index 00000000000..4d7218c81be --- /dev/null +++ b/docs/docs/en/api/faststream/AckPolicy.md @@ -0,0 +1,11 @@ +--- +# 0.5 - API +# 2 - Release +# 3 - Contributing +# 5 - Template Page +# 10 - Default +search: + boost: 0.5 +--- + +::: faststream.AckPolicy diff --git a/docs/docs/en/api/faststream/middlewares/AckPolicy.md b/docs/docs/en/api/faststream/middlewares/AckPolicy.md new file mode 100644 index 00000000000..82d0033dfb4 --- /dev/null +++ b/docs/docs/en/api/faststream/middlewares/AckPolicy.md @@ -0,0 +1,11 @@ +--- +# 0.5 - API +# 2 - Release +# 3 - Contributing +# 5 - Template Page +# 10 - Default +search: + boost: 0.5 +--- + +::: faststream.middlewares.AckPolicy diff --git a/docs/docs/en/api/faststream/middlewares/AcknowledgementMiddleware.md b/docs/docs/en/api/faststream/middlewares/AcknowledgementMiddleware.md new file mode 100644 index 00000000000..d3e7d6a7639 --- /dev/null +++ b/docs/docs/en/api/faststream/middlewares/AcknowledgementMiddleware.md @@ -0,0 +1,11 @@ +--- +# 0.5 - API +# 2 - Release +# 3 - Contributing +# 5 - Template Page +# 10 - Default +search: + boost: 0.5 +--- + +::: faststream.middlewares.AcknowledgementMiddleware diff --git a/docs/docs/en/api/faststream/middlewares/acknowledgement/conf/AckPolicy.md b/docs/docs/en/api/faststream/middlewares/acknowledgement/conf/AckPolicy.md new file mode 100644 index 00000000000..8a92ec0a549 --- /dev/null +++ b/docs/docs/en/api/faststream/middlewares/acknowledgement/conf/AckPolicy.md @@ -0,0 +1,11 @@ +--- +# 0.5 - API +# 2 - Release +# 3 - Contributing +# 5 - Template Page +# 10 - Default +search: + boost: 0.5 +--- + +::: faststream.middlewares.acknowledgement.conf.AckPolicy diff --git a/docs/docs/en/api/faststream/middlewares/acknowledgement/middleware/AcknowledgementMiddleware.md b/docs/docs/en/api/faststream/middlewares/acknowledgement/middleware/AcknowledgementMiddleware.md new file mode 100644 index 00000000000..79b2956eb47 --- /dev/null +++ b/docs/docs/en/api/faststream/middlewares/acknowledgement/middleware/AcknowledgementMiddleware.md @@ -0,0 +1,11 @@ +--- +# 0.5 - API +# 2 - Release +# 3 - Contributing +# 5 - Template Page +# 10 - Default +search: + boost: 0.5 +--- + +::: faststream.middlewares.acknowledgement.middleware.AcknowledgementMiddleware diff --git a/faststream/__init__.py b/faststream/__init__.py index 09514567a8a..cad7e628bfd 100644 --- a/faststream/__init__.py +++ b/faststream/__init__.py @@ -4,7 +4,7 @@ from faststream._internal.utils import apply_types from faststream.annotations import ContextRepo, Logger from faststream.app import FastStream -from faststream.middlewares import BaseMiddleware, ExceptionMiddleware +from faststream.middlewares import AckPolicy, BaseMiddleware, ExceptionMiddleware from faststream.params import ( Context, Depends, @@ -16,6 +16,7 @@ __all__ = ( # middlewares + "AckPolicy", "BaseMiddleware", # params "Context", diff --git a/faststream/_internal/broker/broker.py b/faststream/_internal/broker/broker.py index 162ea5fbc6f..f4229be5c69 100644 --- a/faststream/_internal/broker/broker.py +++ b/faststream/_internal/broker/broker.py @@ -16,7 +16,6 @@ from fast_depends.pydantic import PydanticSerializer from typing_extensions import Doc, Self -from faststream._internal._compat import is_test_env from faststream._internal.constants import EMPTY from faststream._internal.context.repository import ContextRepo from faststream._internal.setup import ( @@ -164,12 +163,10 @@ def __init__( self._connection = None self._producer = None - # TODO: remove useless middleware filter - if not is_test_env(): - self._middlewares = ( - CriticalLogMiddleware(logger_state), - *self._middlewares, - ) + self._middlewares = ( + CriticalLogMiddleware(logger_state), + *self._middlewares, + ) self._state = EmptyState( depends_params=FastDependsData( diff --git a/faststream/_internal/subscriber/acknowledgement_watcher.py b/faststream/_internal/subscriber/acknowledgement_watcher.py deleted file mode 100644 index c86e59baf66..00000000000 --- a/faststream/_internal/subscriber/acknowledgement_watcher.py +++ /dev/null @@ -1,220 +0,0 @@ -import logging -from abc import ABC, abstractmethod -from collections import ( - Counter, - Counter as CounterType, -) -from typing import TYPE_CHECKING, Any, Optional, Union - -from faststream.exceptions import ( - AckMessage, - HandlerException, - NackMessage, - RejectMessage, - SkipMessage, -) - -if TYPE_CHECKING: - from types import TracebackType - - from faststream._internal.basic_types import LoggerProto - from faststream._internal.types import MsgType - from faststream.message import StreamMessage - - -class BaseWatcher(ABC): - """A base class for a watcher.""" - - max_tries: int - - def __init__( - self, - max_tries: int = 0, - logger: Optional["LoggerProto"] = None, - ) -> None: - self.logger = logger - self.max_tries = max_tries - - @abstractmethod - def add(self, message_id: str) -> None: - """Add a message.""" - raise NotImplementedError - - @abstractmethod - def is_max(self, message_id: str) -> bool: - """Check if the given message ID is the maximum attempt.""" - raise NotImplementedError - - @abstractmethod - def remove(self, message_id: str) -> None: - """Remove a message.""" - raise NotImplementedError - - -class EndlessWatcher(BaseWatcher): - """A class to watch and track messages.""" - - def add(self, message_id: str) -> None: - """Add a message to the list.""" - - def is_max(self, message_id: str) -> bool: - """Check if the given message ID is the maximum attempt.""" - return False - - def remove(self, message_id: str) -> None: - """Remove a message.""" - - -class OneTryWatcher(BaseWatcher): - """A class to watch and track messages.""" - - def add(self, message_id: str) -> None: - """Add a message.""" - - def is_max(self, message_id: str) -> bool: - """Check if the given message ID is the maximum attempt.""" - return True - - def remove(self, message_id: str) -> None: - """Remove a message.""" - - -class CounterWatcher(BaseWatcher): - """A class to watch and track the count of messages.""" - - memory: CounterType[str] - - def __init__( - self, - max_tries: int = 3, - logger: Optional["LoggerProto"] = None, - ) -> None: - super().__init__(logger=logger, max_tries=max_tries) - self.memory = Counter() - - def add(self, message_id: str) -> None: - """Check if the given message ID is the maximum attempt.""" - self.memory[message_id] += 1 - - def is_max(self, message_id: str) -> bool: - """Check if the number of tries for a message has exceeded the maximum allowed tries.""" - is_max = self.memory[message_id] > self.max_tries - if self.logger is not None: - if is_max: - self.logger.log( - logging.ERROR, - f"Already retried {self.max_tries} times. Skipped.", - ) - else: - self.logger.log( - logging.ERROR, - "Error is occurred. Pushing back to queue.", - ) - return is_max - - def remove(self, message_id: str) -> None: - """Remove a message from memory.""" - self.memory[message_id] = 0 - self.memory += Counter() - - -class WatcherContext: - """A class representing a context for a watcher.""" - - def __init__( - self, - message: "StreamMessage[MsgType]", - watcher: BaseWatcher, - logger: Optional["LoggerProto"] = None, - **extra_options: Any, - ) -> None: - self.watcher = watcher - self.message = message - self.extra_options = extra_options - self.logger = logger - - async def __aenter__(self) -> None: - self.watcher.add(self.message.message_id) - - async def __aexit__( - self, - exc_type: Optional[type[BaseException]], - exc_val: Optional[BaseException], - exc_tb: Optional["TracebackType"], - ) -> bool: - """Exit the asynchronous context manager.""" - if not exc_type: - await self.__ack() - - elif isinstance(exc_val, HandlerException): - if isinstance(exc_val, SkipMessage): - self.watcher.remove(self.message.message_id) - - elif isinstance(exc_val, AckMessage): - await self.__ack(**exc_val.extra_options) - - elif isinstance(exc_val, NackMessage): - await self.__nack(**exc_val.extra_options) - - elif isinstance(exc_val, RejectMessage): # pragma: no branch - await self.__reject(**exc_val.extra_options) - - # Exception was processed and suppressed - return True - - elif self.watcher.is_max(self.message.message_id): - await self.__reject() - - else: - await self.__nack() - - # Exception was not processed - return False - - async def __ack(self, **exc_extra_options: Any) -> None: - try: - await self.message.ack(**self.extra_options, **exc_extra_options) - except Exception as er: - if self.logger is not None: - self.logger.log(logging.ERROR, er, exc_info=er) - else: - self.watcher.remove(self.message.message_id) - - async def __nack(self, **exc_extra_options: Any) -> None: - try: - await self.message.nack(**self.extra_options, **exc_extra_options) - except Exception as er: - if self.logger is not None: - self.logger.log(logging.ERROR, er, exc_info=er) - - async def __reject(self, **exc_extra_options: Any) -> None: - try: - await self.message.reject(**self.extra_options, **exc_extra_options) - except Exception as er: - if self.logger is not None: - self.logger.log(logging.ERROR, er, exc_info=er) - else: - self.watcher.remove(self.message.message_id) - - -def get_watcher( - logger: Optional["LoggerProto"], - try_number: Union[bool, int], -) -> BaseWatcher: - """Get a watcher object based on the provided parameters. - - Args: - logger: Optional logger object for logging messages. - try_number: Optional parameter to specify the type of watcher. - - If set to True, an EndlessWatcher object will be returned. - - If set to False, a OneTryWatcher object will be returned. - - If set to an integer, a CounterWatcher object with the specified maximum number of tries will be returned. - """ - watcher: Optional[BaseWatcher] - if try_number is True: - watcher = EndlessWatcher() - elif try_number is False: - watcher = OneTryWatcher() - else: - watcher = CounterWatcher(logger=logger, max_tries=try_number) - return watcher diff --git a/faststream/_internal/subscriber/usecase.py b/faststream/_internal/subscriber/usecase.py index 0133be34931..58ed1a8e2ab 100644 --- a/faststream/_internal/subscriber/usecase.py +++ b/faststream/_internal/subscriber/usecase.py @@ -7,7 +7,6 @@ Any, Callable, Optional, - Union, overload, ) @@ -19,7 +18,6 @@ from faststream._internal.subscriber.utils import ( MultiLock, default_filter, - get_watcher_context, resolve_custom_func, ) from faststream._internal.types import ( @@ -29,6 +27,7 @@ ) from faststream._internal.utils.functions import sync_fake_context, to_async from faststream.exceptions import SetupError, StopConsume, SubscriberNotFound +from faststream.middlewares import AckPolicy, AcknowledgementMiddleware from faststream.response import ensure_response from faststream.specification.asyncapi.message import parse_handler_params from faststream.specification.asyncapi.utils import to_camelcase @@ -92,13 +91,12 @@ class SubscriberUsecase(SubscriberProto[MsgType]): def __init__( self, *, - no_ack: bool, no_reply: bool, - retry: Union[bool, int], broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[MsgType]"], default_parser: "AsyncCallable", default_decoder: "AsyncCallable", + ack_policy: AckPolicy, # AsyncAPI information title_: Optional[str], description_: Optional[str], @@ -110,9 +108,7 @@ def __init__( self._parser = default_parser self._decoder = default_decoder self._no_reply = no_reply - # Watcher args - self._no_ack = no_ack - self._retry = retry + self.ack_policy = ack_policy self._call_options = None self._call_decorators = () @@ -134,6 +130,15 @@ def __init__( self.description_ = description_ self.include_in_schema = include_in_schema + if self.ack_policy is not AckPolicy.DO_NOTHING: + self._broker_middlewares = ( + AcknowledgementMiddleware( + self.ack_policy, + self.extra_watcher_options, + ), + *self._broker_middlewares, + ) + def add_middleware(self, middleware: "BrokerMiddleware[MsgType]") -> None: self._broker_middlewares = (*self._broker_middlewares, middleware) @@ -157,8 +162,6 @@ def _setup( # type: ignore[override] self.graceful_timeout = graceful_timeout self.extra_context = extra_context - self.watcher = get_watcher_context(logger, self._no_ack, self._retry) - for call in self.calls: if parser := call.item_parser or broker_parser: async_parser = resolve_custom_func(to_async(parser), self._parser) @@ -345,15 +348,6 @@ async def process_message(self, msg: MsgType) -> "Response": break if message is not None: - # Acknowledgement scope - # TODO: move it to scope enter at `retry` option deprecation - await stack.enter_async_context( - self.watcher( - message, - **self.extra_watcher_options, - ), - ) - stack.enter_context( context.scope("log_context", self.get_log_context(message)), ) diff --git a/faststream/_internal/subscriber/utils.py b/faststream/_internal/subscriber/utils.py index 4dc615a9c04..fb9002c3a3f 100644 --- a/faststream/_internal/subscriber/utils.py +++ b/faststream/_internal/subscriber/utils.py @@ -1,7 +1,7 @@ import asyncio import inspect from collections.abc import Awaitable, Iterable -from contextlib import AbstractAsyncContextManager, AsyncExitStack, suppress +from contextlib import AsyncExitStack, suppress from functools import partial from typing import ( TYPE_CHECKING, @@ -15,18 +15,13 @@ import anyio from typing_extensions import Literal, Self, overload -from faststream._internal.subscriber.acknowledgement_watcher import ( - WatcherContext, - get_watcher, -) from faststream._internal.types import MsgType -from faststream._internal.utils.functions import fake_context, return_input, to_async +from faststream._internal.utils.functions import return_input, to_async from faststream.message.source_type import SourceType if TYPE_CHECKING: from types import TracebackType - from faststream._internal.basic_types import LoggerProto from faststream._internal.types import ( AsyncCallable, CustomCallable, @@ -90,24 +85,6 @@ async def default_filter(msg: "StreamMessage[Any]") -> bool: return not msg.processed -def get_watcher_context( - logger: Optional["LoggerProto"], - no_ack: bool, - retry: Union[bool, int], - **extra_options: Any, -) -> Callable[..., "AbstractAsyncContextManager[None]"]: - """Create Acknowledgement scope.""" - if no_ack: - return fake_context - - return partial( - WatcherContext, - watcher=get_watcher(logger, retry), - logger=logger, - **extra_options, - ) - - class MultiLock: """A class representing a multi lock.""" diff --git a/faststream/confluent/broker/registrator.py b/faststream/confluent/broker/registrator.py index 8eb5c583269..7187c56a03c 100644 --- a/faststream/confluent/broker/registrator.py +++ b/faststream/confluent/broker/registrator.py @@ -16,6 +16,7 @@ from faststream.confluent.publisher.factory import create_publisher from faststream.confluent.subscriber.factory import create_subscriber from faststream.exceptions import SetupError +from faststream.middlewares import AckPolicy if TYPE_CHECKING: from confluent_kafka import Message @@ -294,14 +295,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -565,14 +562,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -836,14 +829,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1110,14 +1099,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1174,9 +1159,8 @@ def subscriber( }, is_manual=not auto_commit, # subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=self._middlewares, broker_dependencies=self._dependencies, # Specification diff --git a/faststream/confluent/fastapi/fastapi.py b/faststream/confluent/fastapi/fastapi.py index 73f6972a408..ac24669769d 100644 --- a/faststream/confluent/fastapi/fastapi.py +++ b/faststream/confluent/fastapi/fastapi.py @@ -25,6 +25,7 @@ from faststream._internal.constants import EMPTY from faststream._internal.fastapi.router import StreamRouter from faststream.confluent.broker.broker import KafkaBroker as KB +from faststream.middlewares import AckPolicy if TYPE_CHECKING: from enum import Enum @@ -833,14 +834,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1607,14 +1604,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -2004,14 +1997,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -2187,8 +2176,7 @@ def subscriber( parser=parser, decoder=decoder, middlewares=middlewares, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/confluent/router.py b/faststream/confluent/router.py index cd1fc9509c5..c4d36fd7e0d 100644 --- a/faststream/confluent/router.py +++ b/faststream/confluent/router.py @@ -17,6 +17,7 @@ SubscriberRoute, ) from faststream.confluent.broker.registrator import KafkaRegistrator +from faststream.middlewares import AckPolicy if TYPE_CHECKING: from confluent_kafka import Message @@ -380,14 +381,10 @@ def __init__( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -445,8 +442,7 @@ def __init__( description=description, include_in_schema=include_in_schema, # FastDepends args - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, ) diff --git a/faststream/confluent/subscriber/factory.py b/faststream/confluent/subscriber/factory.py index 24107e56be5..744a47f744c 100644 --- a/faststream/confluent/subscriber/factory.py +++ b/faststream/confluent/subscriber/factory.py @@ -19,6 +19,7 @@ from faststream._internal.basic_types import AnyDict from faststream._internal.types import BrokerMiddleware from faststream.confluent.schemas import TopicPartition + from faststream.middlewares import AckPolicy @overload @@ -33,9 +34,8 @@ def create_subscriber( connection_data: "AnyDict", is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[tuple[ConfluentMsg, ...]]"], # Specification args @@ -57,9 +57,8 @@ def create_subscriber( connection_data: "AnyDict", is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[ConfluentMsg]"], # Specification args @@ -81,9 +80,8 @@ def create_subscriber( connection_data: "AnyDict", is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable[ "BrokerMiddleware[Union[ConfluentMsg, tuple[ConfluentMsg, ...]]]" @@ -109,9 +107,8 @@ def create_subscriber( connection_data: "AnyDict", is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable[ "BrokerMiddleware[Union[ConfluentMsg, tuple[ConfluentMsg, ...]]]" @@ -133,9 +130,8 @@ def create_subscriber( group_id=group_id, connection_data=connection_data, is_manual=is_manual, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, title_=title_, @@ -149,9 +145,8 @@ def create_subscriber( group_id=group_id, connection_data=connection_data, is_manual=is_manual, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, title_=title_, diff --git a/faststream/confluent/subscriber/usecase.py b/faststream/confluent/subscriber/usecase.py index 470fb578404..e14e2903d29 100644 --- a/faststream/confluent/subscriber/usecase.py +++ b/faststream/confluent/subscriber/usecase.py @@ -32,6 +32,7 @@ ) from faststream.confluent.client import AsyncConfluentConsumer from faststream.message import StreamMessage + from faststream.middlewares import AckPolicy class LogicSubscriber(ABC, SubscriberUsecase[MsgType]): @@ -58,9 +59,8 @@ def __init__( # Subscriber args default_parser: "AsyncCallable", default_decoder: "AsyncCallable", - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[MsgType]"], # AsyncAPI args @@ -72,9 +72,8 @@ def __init__( default_parser=default_parser, default_decoder=default_decoder, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -263,9 +262,8 @@ def __init__( connection_data: "AnyDict", is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Message]"], # AsyncAPI args @@ -285,9 +283,8 @@ def __init__( default_parser=self.parser.parse_message, default_decoder=self.parser.decode_message, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -328,9 +325,8 @@ def __init__( connection_data: "AnyDict", is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[tuple[Message, ...]]"], # AsyncAPI args @@ -352,9 +348,8 @@ def __init__( default_parser=self.parser.parse_message_batch, default_decoder=self.parser.decode_message_batch, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args diff --git a/faststream/kafka/broker/registrator.py b/faststream/kafka/broker/registrator.py index b5b7c4b0d9e..76cb4ec77b6 100644 --- a/faststream/kafka/broker/registrator.py +++ b/faststream/kafka/broker/registrator.py @@ -18,6 +18,7 @@ from faststream._internal.broker.abc_broker import ABCBroker from faststream.kafka.publisher.factory import create_publisher from faststream.kafka.subscriber.factory import create_subscriber +from faststream.middlewares import AckPolicy if TYPE_CHECKING: from aiokafka import TopicPartition @@ -396,14 +397,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -766,14 +763,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1136,14 +1129,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1509,14 +1498,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1576,9 +1561,8 @@ def subscriber( partitions=partitions, is_manual=not auto_commit, # subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=self._middlewares, broker_dependencies=self._dependencies, # Specification diff --git a/faststream/kafka/fastapi/fastapi.py b/faststream/kafka/fastapi/fastapi.py index 1b92a1029c2..3cd73426db2 100644 --- a/faststream/kafka/fastapi/fastapi.py +++ b/faststream/kafka/fastapi/fastapi.py @@ -28,6 +28,7 @@ from faststream._internal.constants import EMPTY from faststream._internal.fastapi.router import StreamRouter from faststream.kafka.broker.broker import KafkaBroker as KB +from faststream.middlewares import AckPolicy if TYPE_CHECKING: from asyncio import AbstractEventLoop @@ -944,14 +945,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1434,14 +1431,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -1924,14 +1917,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -2417,14 +2406,10 @@ def subscriber( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -2607,8 +2592,7 @@ def subscriber( parser=parser, decoder=decoder, middlewares=middlewares, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/kafka/router.py b/faststream/kafka/router.py index 1652e52901a..416d23df200 100644 --- a/faststream/kafka/router.py +++ b/faststream/kafka/router.py @@ -18,6 +18,7 @@ SubscriberRoute, ) from faststream.kafka.broker.registrator import KafkaRegistrator +from faststream.middlewares import AckPolicy if TYPE_CHECKING: from aiokafka import ConsumerRecord, TopicPartition @@ -483,14 +484,10 @@ def __init__( Iterable["SubscriberMiddleware[KafkaMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -555,8 +552,7 @@ def __init__( description=description, include_in_schema=include_in_schema, # FastDepends args - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, ) diff --git a/faststream/kafka/subscriber/factory.py b/faststream/kafka/subscriber/factory.py index e53654283eb..7542a11cd62 100644 --- a/faststream/kafka/subscriber/factory.py +++ b/faststream/kafka/subscriber/factory.py @@ -20,6 +20,7 @@ from faststream._internal.basic_types import AnyDict from faststream._internal.types import BrokerMiddleware + from faststream.middlewares import AckPolicy @overload @@ -36,9 +37,8 @@ def create_subscriber( partitions: Iterable["TopicPartition"], is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[tuple[ConsumerRecord, ...]]"], # Specification args @@ -62,9 +62,8 @@ def create_subscriber( partitions: Iterable["TopicPartition"], is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[ConsumerRecord]"], # Specification args @@ -88,9 +87,8 @@ def create_subscriber( partitions: Iterable["TopicPartition"], is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable[ "BrokerMiddleware[Union[ConsumerRecord, tuple[ConsumerRecord, ...]]]" @@ -118,9 +116,8 @@ def create_subscriber( partitions: Iterable["TopicPartition"], is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable[ "BrokerMiddleware[Union[ConsumerRecord, tuple[ConsumerRecord, ...]]]" @@ -163,9 +160,8 @@ def create_subscriber( connection_args=connection_args, partitions=partitions, is_manual=is_manual, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, title_=title_, @@ -181,9 +177,8 @@ def create_subscriber( connection_args=connection_args, partitions=partitions, is_manual=is_manual, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, title_=title_, diff --git a/faststream/kafka/subscriber/usecase.py b/faststream/kafka/subscriber/usecase.py index 03dae216871..f3cf9c2122b 100644 --- a/faststream/kafka/subscriber/usecase.py +++ b/faststream/kafka/subscriber/usecase.py @@ -36,6 +36,7 @@ from faststream._internal.publisher.proto import BasePublisherProto, ProducerProto from faststream._internal.setup import SetupState from faststream.message import StreamMessage + from faststream.middlewares import AckPolicy class LogicSubscriber(SubscriberUsecase[MsgType]): @@ -64,9 +65,8 @@ def __init__( # Subscriber args default_parser: "AsyncCallable", default_decoder: "AsyncCallable", - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[MsgType]"], # AsyncAPI args @@ -78,9 +78,8 @@ def __init__( default_parser=default_parser, default_decoder=default_decoder, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -292,9 +291,8 @@ def __init__( partitions: Iterable["TopicPartition"], is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[ConsumerRecord]"], # AsyncAPI args @@ -328,9 +326,8 @@ def __init__( default_parser=self.parser.parse_message, default_decoder=self.parser.decode_message, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -373,9 +370,8 @@ def __init__( partitions: Iterable["TopicPartition"], is_manual: bool, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable[ "BrokerMiddleware[Sequence[tuple[ConsumerRecord, ...]]]" @@ -414,9 +410,8 @@ def __init__( default_parser=self.parser.parse_message, default_decoder=self.parser.decode_message, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args diff --git a/faststream/middlewares/__init__.py b/faststream/middlewares/__init__.py index 0615c88194c..f8d57bdf50d 100644 --- a/faststream/middlewares/__init__.py +++ b/faststream/middlewares/__init__.py @@ -1,4 +1,11 @@ +from faststream.middlewares.acknowledgement.conf import AckPolicy +from faststream.middlewares.acknowledgement.middleware import AcknowledgementMiddleware from faststream.middlewares.base import BaseMiddleware from faststream.middlewares.exception import ExceptionMiddleware -__all__ = ("BaseMiddleware", "ExceptionMiddleware") +__all__ = ( + "AckPolicy", + "AcknowledgementMiddleware", + "BaseMiddleware", + "ExceptionMiddleware", +) diff --git a/faststream/middlewares/acknowledgement/__init__.py b/faststream/middlewares/acknowledgement/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/faststream/middlewares/acknowledgement/conf.py b/faststream/middlewares/acknowledgement/conf.py new file mode 100644 index 00000000000..60b910264dd --- /dev/null +++ b/faststream/middlewares/acknowledgement/conf.py @@ -0,0 +1,8 @@ +from enum import Enum + + +class AckPolicy(str, Enum): + ACK = "ack" + REJECT_ON_ERROR = "reject_on_error" + NACK_ON_ERROR = "nack_on_error" + DO_NOTHING = "do_nothing" diff --git a/faststream/middlewares/acknowledgement/middleware.py b/faststream/middlewares/acknowledgement/middleware.py new file mode 100644 index 00000000000..25137686495 --- /dev/null +++ b/faststream/middlewares/acknowledgement/middleware.py @@ -0,0 +1,111 @@ +import logging +from typing import TYPE_CHECKING, Any, Optional + +from faststream.exceptions import ( + AckMessage, + HandlerException, + NackMessage, + RejectMessage, +) +from faststream.middlewares.acknowledgement.conf import AckPolicy +from faststream.middlewares.base import BaseMiddleware + +if TYPE_CHECKING: + from types import TracebackType + + from faststream._internal.basic_types import AnyDict, AsyncFuncAny + from faststream._internal.context.repository import ContextRepo + from faststream.message import StreamMessage + + +class AcknowledgementMiddleware: + def __init__(self, ack_policy: AckPolicy, extra_options: "AnyDict") -> None: + self.ack_policy = ack_policy + self.extra_options = extra_options + + def __call__(self, msg: Optional[Any], context: "ContextRepo") -> "_AcknowledgementMiddleware": + return _AcknowledgementMiddleware( + msg, + ack_policy=self.ack_policy, + extra_options=self.extra_options, + context=context, + ) + + +class _AcknowledgementMiddleware(BaseMiddleware): + def __init__( + self, + msg: Optional[Any], + /, + *, + context: "ContextRepo", + ack_policy: AckPolicy, + extra_options: "AnyDict", + ) -> None: + super().__init__(msg, context=context) + self.ack_policy = ack_policy + self.extra_options = extra_options + self.logger = context.get_local("logger") + + async def consume_scope( + self, + call_next: "AsyncFuncAny", + msg: "StreamMessage[Any]", + ) -> Any: + self.message = msg + return await call_next(msg) + + async def __aexit__( + self, + exc_type: Optional[type[BaseException]] = None, + exc_val: Optional[BaseException] = None, + exc_tb: Optional["TracebackType"] = None, + ) -> Optional[bool]: + if self.ack_policy is AckPolicy.DO_NOTHING: + return False + + if not exc_type: + await self.__ack() + + elif isinstance(exc_val, HandlerException): + if isinstance(exc_val, AckMessage): + await self.__ack(**exc_val.extra_options) + + elif isinstance(exc_val, NackMessage): + await self.__nack(**exc_val.extra_options) + + elif isinstance(exc_val, RejectMessage): # pragma: no branch + await self.__reject(**exc_val.extra_options) + + # Exception was processed and suppressed + return True + + elif self.ack_policy is AckPolicy.REJECT_ON_ERROR: + await self.__reject() + + elif self.ack_policy is AckPolicy.NACK_ON_ERROR: + await self.__nack() + + # Exception was not processed + return False + + async def __ack(self, **exc_extra_options: Any) -> None: + try: + await self.message.ack(**exc_extra_options, **self.extra_options) + except Exception as er: + if self.logger is not None: + self.logger.log(logging.ERROR, er, exc_info=er) + + async def __nack(self, **exc_extra_options: Any) -> None: + try: + await self.message.nack(**exc_extra_options, **self.extra_options) + except Exception as er: + if self.logger is not None: + self.logger.log(logging.ERROR, er, exc_info=er) + + async def __reject(self, **exc_extra_options: Any) -> None: + try: + await self.message.reject(**exc_extra_options, **self.extra_options) + except Exception as er: + if self.logger is not None: + self.logger.log(logging.ERROR, er, exc_info=er) diff --git a/faststream/nats/broker/registrator.py b/faststream/nats/broker/registrator.py index 4ff8f0fbda8..f8feacb294a 100644 --- a/faststream/nats/broker/registrator.py +++ b/faststream/nats/broker/registrator.py @@ -5,6 +5,7 @@ from typing_extensions import Doc, override from faststream._internal.broker.abc_broker import ABCBroker +from faststream.middlewares import AckPolicy from faststream.nats.helpers import StreamBuilder from faststream.nats.publisher.factory import create_publisher from faststream.nats.publisher.specified import SpecificationPublisher @@ -160,14 +161,10 @@ def subscriber( # type: ignore[override] int, Doc("Number of workers to process messages concurrently."), ] = 1, - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -222,9 +219,8 @@ def subscriber( # type: ignore[override] inbox_prefix=inbox_prefix, ack_first=ack_first, # subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=self._middlewares, broker_dependencies=self._dependencies, # AsyncAPI diff --git a/faststream/nats/fastapi/fastapi.py b/faststream/nats/fastapi/fastapi.py index e2c5ae3d985..aa449acde77 100644 --- a/faststream/nats/fastapi/fastapi.py +++ b/faststream/nats/fastapi/fastapi.py @@ -32,6 +32,7 @@ from faststream.__about__ import SERVICE_NAME from faststream._internal.constants import EMPTY from faststream._internal.fastapi.router import StreamRouter +from faststream.middlewares import AckPolicy from faststream.nats.broker import NatsBroker from faststream.nats.subscriber.specified import SpecificationSubscriber @@ -692,14 +693,10 @@ def subscriber( # type: ignore[override] int, Doc("Number of workers to process messages concurrently."), ] = 1, - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -871,8 +868,7 @@ def subscriber( # type: ignore[override] decoder=decoder, middlewares=middlewares, max_workers=max_workers, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/nats/parser.py b/faststream/nats/parser.py index 0f3b2f4c8ce..3cb2c695a81 100644 --- a/faststream/nats/parser.py +++ b/faststream/nats/parser.py @@ -4,6 +4,7 @@ StreamMessage, decode_message, ) +from faststream.middlewares import AckPolicy from faststream.nats.message import ( NatsBatchMessage, NatsKvMessage, @@ -54,9 +55,9 @@ async def decode_message( class NatsParser(NatsBaseParser): """A class to parse NATS core messages.""" - def __init__(self, *, pattern: str, no_ack: bool) -> None: + def __init__(self, *, pattern: str, ack_policy: AckPolicy) -> None: super().__init__(pattern=pattern) - self.no_ack = no_ack + self.ack_policy = ack_policy async def parse_message( self, @@ -69,7 +70,7 @@ async def parse_message( headers = message.header or {} - if not self.no_ack: + if self.ack_policy is not AckPolicy.DO_NOTHING: message._ackd = True # prevent message from acking return NatsMessage( diff --git a/faststream/nats/publisher/producer.py b/faststream/nats/publisher/producer.py index f2ed4715e0b..62245c2c7c8 100644 --- a/faststream/nats/publisher/producer.py +++ b/faststream/nats/publisher/producer.py @@ -5,6 +5,7 @@ import nats from typing_extensions import override +from faststream import AckPolicy from faststream._internal.publisher.proto import ProducerProto from faststream._internal.subscriber.utils import resolve_custom_func from faststream.exceptions import FeatureNotSupportedException @@ -40,7 +41,7 @@ def __init__( parser: Optional["CustomCallable"], decoder: Optional["CustomCallable"], ) -> None: - default = NatsParser(pattern="", no_ack=False) + default = NatsParser(pattern="", ack_policy=AckPolicy.REJECT_ON_ERROR) self._parser = resolve_custom_func(parser, default.parse_message) self._decoder = resolve_custom_func(decoder, default.decode_message) @@ -111,7 +112,7 @@ def __init__( parser: Optional["CustomCallable"], decoder: Optional["CustomCallable"], ) -> None: - default = NatsParser(pattern="", no_ack=False) + default = NatsParser(pattern="", ack_policy=AckPolicy.REJECT_ON_ERROR) self._parser = resolve_custom_func(parser, default.parse_message) self._decoder = resolve_custom_func(decoder, default.decode_message) diff --git a/faststream/nats/router.py b/faststream/nats/router.py index 3b980730500..5d4f87b03be 100644 --- a/faststream/nats/router.py +++ b/faststream/nats/router.py @@ -16,6 +16,7 @@ BrokerRouter, SubscriberRoute, ) +from faststream.middlewares import AckPolicy from faststream.nats.broker.registrator import NatsRegistrator if TYPE_CHECKING: @@ -250,14 +251,10 @@ def __init__( int, Doc("Number of workers to process messages concurrently."), ] = 1, - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -307,8 +304,7 @@ def __init__( parser=parser, decoder=decoder, middlewares=middlewares, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/nats/schemas/js_stream.py b/faststream/nats/schemas/js_stream.py index 0d9b23ac4fd..3ad4fc2e4f5 100644 --- a/faststream/nats/schemas/js_stream.py +++ b/faststream/nats/schemas/js_stream.py @@ -6,6 +6,7 @@ from faststream._internal.proto import NameRequired from faststream._internal.utils.path import compile_path +from faststream.middlewares import AckPolicy if TYPE_CHECKING: from re import Pattern @@ -120,13 +121,10 @@ def __init__( "cluster may be available but for reads only.", ), ] = None, - no_ack: Annotated[ - bool, - Doc( - "Should stream acknowledge writes or not. Without acks publisher can't determine, does message " - "received by stream or not.", - ), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, template_owner: Optional[str] = None, duplicate_window: Annotated[ float, @@ -191,6 +189,7 @@ def __init__( super().__init__(name) subjects = subjects or [] + no_ack = ack_policy is AckPolicy.DO_NOTHING self.subjects = subjects self.declare = declare diff --git a/faststream/nats/subscriber/factory.py b/faststream/nats/subscriber/factory.py index 613a76a535c..586f0db870c 100644 --- a/faststream/nats/subscriber/factory.py +++ b/faststream/nats/subscriber/factory.py @@ -30,6 +30,7 @@ from faststream._internal.basic_types import AnyDict from faststream._internal.types import BrokerMiddleware + from faststream.middlewares import AckPolicy from faststream.nats.schemas import JStream, KvWatch, ObjWatch, PullSub @@ -59,9 +60,8 @@ def create_subscriber( max_workers: int, stream: Optional["JStream"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Any]"], # Specification information @@ -159,9 +159,8 @@ def create_subscriber( # basic args extra_options=extra_options, # Subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # Specification @@ -177,9 +176,8 @@ def create_subscriber( # basic args extra_options=extra_options, # Subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # Specification @@ -199,9 +197,8 @@ def create_subscriber( # basic args extra_options=extra_options, # Subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # Specification @@ -219,9 +216,8 @@ def create_subscriber( # basic args extra_options=extra_options, # Subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # Specification @@ -240,9 +236,8 @@ def create_subscriber( # basic args extra_options=extra_options, # Subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # Specification @@ -259,9 +254,8 @@ def create_subscriber( # basic args extra_options=extra_options, # Subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # Specification @@ -278,9 +272,8 @@ def create_subscriber( # basic args extra_options=extra_options, # Subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # Specification information diff --git a/faststream/nats/subscriber/usecase.py b/faststream/nats/subscriber/usecase.py index 9f15f004569..2c1e35cd56c 100644 --- a/faststream/nats/subscriber/usecase.py +++ b/faststream/nats/subscriber/usecase.py @@ -8,12 +8,10 @@ Callable, Generic, Optional, - Union, cast, ) import anyio -from fast_depends.dependencies import Dependant from nats.errors import ConnectionClosedError, TimeoutError from nats.js.api import ConsumerConfig, ObjectInfo from typing_extensions import Doc, override @@ -22,6 +20,7 @@ from faststream._internal.subscriber.usecase import SubscriberUsecase from faststream._internal.subscriber.utils import process_msg from faststream._internal.types import MsgType +from faststream.middlewares import AckPolicy from faststream.nats.helpers import KVBucketDeclarer, OSBucketDeclarer from faststream.nats.message import NatsMessage from faststream.nats.parser import ( @@ -41,6 +40,7 @@ from .state import ConnectedSubscriberState, EmptySubscriberState, SubscriberState if TYPE_CHECKING: + from fast_depends.dependencies import Dependant from nats.aio.msg import Msg from nats.aio.subscription import Subscription from nats.js import JetStreamContext @@ -82,10 +82,9 @@ def __init__( # Subscriber args default_parser: "AsyncCallable", default_decoder: "AsyncCallable", - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[MsgType]"], # AsyncAPI args title_: Optional[str], @@ -101,9 +100,8 @@ def __init__( default_parser=default_parser, default_decoder=default_decoder, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -235,10 +233,9 @@ def __init__( # Subscriber args default_parser: "AsyncCallable", default_decoder: "AsyncCallable", - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[MsgType]"], # AsyncAPI args title_: Optional[str], @@ -253,9 +250,8 @@ def __init__( default_parser=default_parser, default_decoder=default_decoder, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -306,17 +302,16 @@ def __init__( queue: str, extra_options: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Msg]"], # AsyncAPI args title_: Optional[str], description_: Optional[str], include_in_schema: bool, ) -> None: - parser_ = NatsParser(pattern=subject, no_ack=no_ack) + parser_ = NatsParser(pattern=subject, ack_policy=ack_policy) self.queue = queue @@ -328,9 +323,8 @@ def __init__( default_parser=parser_.parse_message, default_decoder=parser_.decode_message, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -416,10 +410,9 @@ def __init__( queue: str, extra_options: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Msg]"], # AsyncAPI args title_: Optional[str], @@ -434,9 +427,8 @@ def __init__( queue=queue, extra_options=extra_options, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -474,10 +466,9 @@ def __init__( queue: str, extra_options: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Msg]"], # AsyncAPI args title_: Optional[str], @@ -497,9 +488,8 @@ def __init__( default_parser=parser_.parse_message, default_decoder=parser_.decode_message, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -606,10 +596,9 @@ def __init__( queue: str, extra_options: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Msg]"], # AsyncAPI args title_: Optional[str], @@ -625,9 +614,8 @@ def __init__( queue=queue, extra_options=extra_options, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -669,10 +657,9 @@ def __init__( config: "ConsumerConfig", extra_options: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Msg]"], # AsyncAPI args title_: Optional[str], @@ -689,9 +676,8 @@ def __init__( extra_options=extra_options, queue="", # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -749,10 +735,9 @@ def __init__( config: "ConsumerConfig", extra_options: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[Msg]"], # AsyncAPI args title_: Optional[str], @@ -768,9 +753,8 @@ def __init__( config=config, extra_options=extra_options, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -814,10 +798,9 @@ def __init__( pull_sub: "PullSub", extra_options: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[list[Msg]]"], # AsyncAPI args title_: Optional[str], @@ -837,9 +820,8 @@ def __init__( default_parser=parser.parse_batch, default_decoder=parser.decode_batch, # Propagated args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI args @@ -931,7 +913,7 @@ def __init__( subject: str, config: "ConsumerConfig", kv_watch: "KvWatch", - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[KeyValue.Entry]"], # AsyncAPI args title_: Optional[str], @@ -945,9 +927,8 @@ def __init__( subject=subject, config=config, extra_options=None, - no_ack=True, + ack_policy=AckPolicy.DO_NOTHING, no_reply=True, - retry=False, default_parser=parser.parse_message, default_decoder=parser.decode_message, broker_middlewares=broker_middlewares, @@ -1085,7 +1066,7 @@ def __init__( subject: str, config: "ConsumerConfig", obj_watch: "ObjWatch", - broker_dependencies: Iterable[Dependant], + broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[list[Msg]]"], # AsyncAPI args title_: Optional[str], @@ -1101,9 +1082,8 @@ def __init__( subject=subject, config=config, extra_options=None, - no_ack=True, + ack_policy=AckPolicy.DO_NOTHING, no_reply=True, - retry=False, default_parser=parser.parse_message, default_decoder=parser.decode_message, broker_middlewares=broker_middlewares, diff --git a/faststream/nats/testing.py b/faststream/nats/testing.py index 4ee0d1602c0..d7095549936 100644 --- a/faststream/nats/testing.py +++ b/faststream/nats/testing.py @@ -11,6 +11,7 @@ from nats.aio.msg import Msg from typing_extensions import override +from faststream import AckPolicy from faststream._internal.subscriber.utils import resolve_custom_func from faststream._internal.testing.broker import TestBroker from faststream.exceptions import SubscriberNotFound @@ -70,7 +71,7 @@ class FakeProducer(NatsFastProducer): def __init__(self, broker: NatsBroker) -> None: self.broker = broker - default = NatsParser(pattern="", no_ack=False) + default = NatsParser(pattern="", ack_policy=AckPolicy.REJECT_ON_ERROR) self._parser = resolve_custom_func(broker._parser, default.parse_message) self._decoder = resolve_custom_func(broker._decoder, default.decode_message) diff --git a/faststream/rabbit/broker/registrator.py b/faststream/rabbit/broker/registrator.py index 20893edbd05..117aafc1de6 100644 --- a/faststream/rabbit/broker/registrator.py +++ b/faststream/rabbit/broker/registrator.py @@ -4,6 +4,7 @@ from typing_extensions import Doc, override from faststream._internal.broker.abc_broker import ABCBroker +from faststream.middlewares import AckPolicy from faststream.rabbit.publisher.factory import create_publisher from faststream.rabbit.publisher.specified import SpecificationPublisher from faststream.rabbit.publisher.usecase import PublishKwargs @@ -57,6 +58,10 @@ def subscriber( # type: ignore[override] Optional["AnyDict"], Doc("Extra consumer arguments to use in `queue.consume(...)` method."), ] = None, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, # broker arguments dependencies: Annotated[ Iterable["Dependant"], @@ -74,14 +79,6 @@ def subscriber( # type: ignore[override] Iterable["SubscriberMiddleware[RabbitMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - Union[bool, int], - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, no_reply: Annotated[ bool, Doc( @@ -113,9 +110,8 @@ def subscriber( # type: ignore[override] exchange=RabbitExchange.validate(exchange), consume_args=consume_args, # subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=self._middlewares, broker_dependencies=self._dependencies, # AsyncAPI diff --git a/faststream/rabbit/fastapi/fastapi.py b/faststream/rabbit/fastapi/fastapi.py index 92403800798..6411dbc9495 100644 --- a/faststream/rabbit/fastapi/fastapi.py +++ b/faststream/rabbit/fastapi/fastapi.py @@ -20,6 +20,7 @@ from faststream.__about__ import SERVICE_NAME from faststream._internal.constants import EMPTY from faststream._internal.fastapi.router import StreamRouter +from faststream.middlewares import AckPolicy from faststream.rabbit.broker.broker import RabbitBroker as RB from faststream.rabbit.schemas import ( RabbitExchange, @@ -511,14 +512,10 @@ def subscriber( # type: ignore[override] Iterable["SubscriberMiddleware[RabbitMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - Union[bool, int], - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -675,8 +672,7 @@ def subscriber( # type: ignore[override] parser=parser, decoder=decoder, middlewares=middlewares, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/rabbit/router.py b/faststream/rabbit/router.py index a11b5ddbb60..3edbf447b74 100644 --- a/faststream/rabbit/router.py +++ b/faststream/rabbit/router.py @@ -8,6 +8,7 @@ BrokerRouter, SubscriberRoute, ) +from faststream.middlewares import AckPolicy from faststream.rabbit.broker.registrator import RabbitRegistrator if TYPE_CHECKING: @@ -229,14 +230,10 @@ def __init__( Iterable["SubscriberMiddleware[RabbitMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - Union[bool, int], - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -270,8 +267,7 @@ def __init__( parser=parser, decoder=decoder, middlewares=middlewares, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/rabbit/subscriber/factory.py b/faststream/rabbit/subscriber/factory.py index b82210bd3d5..4554f9c9c5a 100644 --- a/faststream/rabbit/subscriber/factory.py +++ b/faststream/rabbit/subscriber/factory.py @@ -1,5 +1,5 @@ from collections.abc import Iterable -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Optional from faststream.rabbit.subscriber.specified import SpecificationSubscriber @@ -9,6 +9,7 @@ from faststream._internal.basic_types import AnyDict from faststream._internal.types import BrokerMiddleware + from faststream.middlewares import AckPolicy from faststream.rabbit.schemas import RabbitExchange, RabbitQueue @@ -18,11 +19,10 @@ def create_subscriber( exchange: "RabbitExchange", consume_args: Optional["AnyDict"], # Subscriber args - no_ack: bool, no_reply: bool, - retry: Union[bool, int], broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[IncomingMessage]"], + ack_policy: "AckPolicy", # AsyncAPI args title_: Optional[str], description_: Optional[str], @@ -32,9 +32,8 @@ def create_subscriber( queue=queue, exchange=exchange, consume_args=consume_args, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, title_=title_, diff --git a/faststream/rabbit/subscriber/usecase.py b/faststream/rabbit/subscriber/usecase.py index aeae29d970e..119938afb53 100644 --- a/faststream/rabbit/subscriber/usecase.py +++ b/faststream/rabbit/subscriber/usecase.py @@ -3,7 +3,6 @@ TYPE_CHECKING, Any, Optional, - Union, ) import anyio @@ -12,6 +11,7 @@ from faststream._internal.subscriber.usecase import SubscriberUsecase from faststream._internal.subscriber.utils import process_msg from faststream.exceptions import SetupError +from faststream.middlewares import AckPolicy from faststream.rabbit.parser import AioPikaParser from faststream.rabbit.publisher.fake import RabbitFakePublisher from faststream.rabbit.schemas import BaseRMQInformation @@ -54,9 +54,8 @@ def __init__( exchange: "RabbitExchange", consume_args: Optional["AnyDict"], # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: Union[bool, int], broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[IncomingMessage]"], # AsyncAPI args @@ -70,9 +69,8 @@ def __init__( default_parser=parser.parse_message, default_decoder=parser.decode_message, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -175,7 +173,7 @@ async def get_one( self, *, timeout: float = 5.0, - no_ack: bool = True, + ack_policy: AckPolicy = AckPolicy.REJECT_ON_ERROR, ) -> "Optional[RabbitMessage]": assert self._queue_obj, "You should start subscriber at first." # nosec B101 assert ( # nosec B101 @@ -185,6 +183,7 @@ async def get_one( sleep_interval = timeout / 10 raw_message: Optional[IncomingMessage] = None + no_ack = self.ack_policy is AckPolicy.DO_NOTHING with anyio.move_on_after(timeout): while ( # noqa: ASYNC110 raw_message := await self._queue_obj.get( diff --git a/faststream/redis/broker/registrator.py b/faststream/redis/broker/registrator.py index 9b10aae10e8..1599a10d46d 100644 --- a/faststream/redis/broker/registrator.py +++ b/faststream/redis/broker/registrator.py @@ -4,6 +4,7 @@ from typing_extensions import Doc, override from faststream._internal.broker.abc_broker import ABCBroker +from faststream.middlewares import AckPolicy from faststream.redis.message import UnifyRedisDict from faststream.redis.publisher.factory import create_publisher from faststream.redis.publisher.specified import SpecificationPublisher @@ -65,14 +66,10 @@ def subscriber( # type: ignore[override] Iterable["SubscriberMiddleware[UnifyRedisMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -104,9 +101,8 @@ def subscriber( # type: ignore[override] list=list, stream=stream, # subscriber args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=self._middlewares, broker_dependencies=self._dependencies, # AsyncAPI diff --git a/faststream/redis/fastapi/fastapi.py b/faststream/redis/fastapi/fastapi.py index ed97300b879..01a5432f5e5 100644 --- a/faststream/redis/fastapi/fastapi.py +++ b/faststream/redis/fastapi/fastapi.py @@ -25,6 +25,7 @@ from faststream.__about__ import SERVICE_NAME from faststream._internal.constants import EMPTY from faststream._internal.fastapi.router import StreamRouter +from faststream.middlewares import AckPolicy from faststream.redis.broker.broker import RedisBroker as RB from faststream.redis.message import UnifyRedisDict from faststream.redis.schemas import ListSub, PubSub, StreamSub @@ -461,14 +462,10 @@ def subscriber( # type: ignore[override] Iterable["SubscriberMiddleware[UnifyRedisMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -625,8 +622,7 @@ def subscriber( # type: ignore[override] parser=parser, decoder=decoder, middlewares=middlewares, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/redis/router.py b/faststream/redis/router.py index f7e60051fe8..32d937dcce0 100644 --- a/faststream/redis/router.py +++ b/faststream/redis/router.py @@ -8,6 +8,7 @@ BrokerRouter, SubscriberRoute, ) +from faststream.middlewares import AckPolicy from faststream.redis.broker.registrator import RedisRegistrator from faststream.redis.message import BaseMessage @@ -147,14 +148,10 @@ def __init__( Iterable["SubscriberMiddleware[UnifyRedisMessage]"], Doc("Subscriber middlewares to wrap incoming message processing."), ] = (), - retry: Annotated[ - bool, - Doc("Whether to `nack` message at processing exception."), - ] = False, - no_ack: Annotated[ - bool, - Doc("Whether to disable **FastStream** autoacknowledgement logic or not."), - ] = False, + ack_policy: Annotated[ + AckPolicy, + Doc("Whether to disable **FastStream** auto acknowledgement logic or not."), + ] = AckPolicy.REJECT_ON_ERROR, no_reply: Annotated[ bool, Doc( @@ -188,8 +185,7 @@ def __init__( parser=parser, decoder=decoder, middlewares=middlewares, - retry=retry, - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, title=title, description=description, diff --git a/faststream/redis/schemas/stream_sub.py b/faststream/redis/schemas/stream_sub.py index 50a0b6d606f..07488d5f863 100644 --- a/faststream/redis/schemas/stream_sub.py +++ b/faststream/redis/schemas/stream_sub.py @@ -3,6 +3,7 @@ from faststream._internal.proto import NameRequired from faststream.exceptions import SetupError +from faststream.middlewares import AckPolicy class StreamSub(NameRequired): @@ -27,11 +28,13 @@ def __init__( group: Optional[str] = None, consumer: Optional[str] = None, batch: bool = False, - no_ack: bool = False, + ack_policy: AckPolicy = AckPolicy.REJECT_ON_ERROR, last_id: Optional[str] = None, maxlen: Optional[int] = None, max_records: Optional[int] = None, ) -> None: + no_ack = ack_policy is AckPolicy.DO_NOTHING + if (group and not consumer) or (not group and consumer): msg = "You should specify `group` and `consumer` both" raise SetupError(msg) diff --git a/faststream/redis/subscriber/factory.py b/faststream/redis/subscriber/factory.py index 9238628332c..e46d3d6646c 100644 --- a/faststream/redis/subscriber/factory.py +++ b/faststream/redis/subscriber/factory.py @@ -18,6 +18,7 @@ from fast_depends.dependencies import Dependant from faststream._internal.types import BrokerMiddleware + from faststream.middlewares import AckPolicy from faststream.redis.message import UnifyRedisDict SubsciberType: TypeAlias = Union[ @@ -35,9 +36,8 @@ def create_subscriber( list: Union["ListSub", str, None], stream: Union["StreamSub", str, None], # Subscriber args - no_ack: bool = False, + ack_policy: "AckPolicy", no_reply: bool = False, - retry: bool = False, broker_dependencies: Iterable["Dependant"] = (), broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"] = (), # AsyncAPI args @@ -51,9 +51,8 @@ def create_subscriber( return AsyncAPIChannelSubscriber( channel=channel_sub, # basic args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # AsyncAPI args @@ -67,9 +66,8 @@ def create_subscriber( return AsyncAPIStreamBatchSubscriber( stream=stream_sub, # basic args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # AsyncAPI args @@ -80,9 +78,8 @@ def create_subscriber( return AsyncAPIStreamSubscriber( stream=stream_sub, # basic args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # AsyncAPI args @@ -96,9 +93,8 @@ def create_subscriber( return AsyncAPIListBatchSubscriber( list=list_sub, # basic args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # AsyncAPI args @@ -109,9 +105,8 @@ def create_subscriber( return AsyncAPIListSubscriber( list=list_sub, # basic args - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_dependencies=broker_dependencies, broker_middlewares=broker_middlewares, # AsyncAPI args diff --git a/faststream/redis/subscriber/usecase.py b/faststream/redis/subscriber/usecase.py index 70ccaf73e9d..995fc704547 100644 --- a/faststream/redis/subscriber/usecase.py +++ b/faststream/redis/subscriber/usecase.py @@ -54,6 +54,7 @@ CustomCallable, ) from faststream.message import StreamMessage as BrokerStreamMessage + from faststream.middlewares import AckPolicy TopicName: TypeAlias = bytes @@ -71,9 +72,8 @@ def __init__( default_parser: "AsyncCallable", default_decoder: "AsyncCallable", # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -85,9 +85,8 @@ def __init__( default_parser=default_parser, default_decoder=default_decoder, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -215,9 +214,8 @@ def __init__( *, channel: "PubSub", # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -230,9 +228,8 @@ def __init__( default_parser=parser.parse_message, default_decoder=parser.decode_message, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -341,9 +338,8 @@ def __init__( default_parser: "AsyncCallable", default_decoder: "AsyncCallable", # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -355,9 +351,8 @@ def __init__( default_parser=default_parser, default_decoder=default_decoder, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -448,9 +443,8 @@ def __init__( *, list: ListSub, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -464,9 +458,8 @@ def __init__( default_parser=parser.parse_message, default_decoder=parser.decode_message, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -497,9 +490,8 @@ def __init__( *, list: ListSub, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -513,9 +505,8 @@ def __init__( default_parser=parser.parse_message, default_decoder=parser.decode_message, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -551,9 +542,8 @@ def __init__( default_parser: "AsyncCallable", default_decoder: "AsyncCallable", # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -565,9 +555,8 @@ def __init__( default_parser=default_parser, default_decoder=default_decoder, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -742,9 +731,8 @@ def __init__( *, stream: StreamSub, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -758,9 +746,8 @@ def __init__( default_parser=parser.parse_message, default_decoder=parser.decode_message, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI @@ -811,9 +798,8 @@ def __init__( *, stream: StreamSub, # Subscriber args - no_ack: bool, + ack_policy: "AckPolicy", no_reply: bool, - retry: bool, broker_dependencies: Iterable["Dependant"], broker_middlewares: Iterable["BrokerMiddleware[UnifyRedisDict]"], # AsyncAPI args @@ -827,9 +813,8 @@ def __init__( default_parser=parser.parse_message, default_decoder=parser.decode_message, # Propagated options - no_ack=no_ack, + ack_policy=ack_policy, no_reply=no_reply, - retry=retry, broker_middlewares=broker_middlewares, broker_dependencies=broker_dependencies, # AsyncAPI diff --git a/tests/brokers/base/fastapi.py b/tests/brokers/base/fastapi.py index 38c9475ae20..e1e9c159540 100644 --- a/tests/brokers/base/fastapi.py +++ b/tests/brokers/base/fastapi.py @@ -534,7 +534,6 @@ async def hello_router2() -> str: async def test_dependency_overrides(self, mock: Mock, queue: str) -> None: router = self.router_class() - router2 = self.router_class() def dep1() -> None: mock.not_call() @@ -547,10 +546,39 @@ def dep2() -> None: args, kwargs = self.get_subscriber_params(queue) - @router2.subscriber(*args, **kwargs) + @router.subscriber(*args, **kwargs) async def hello_router2(dep: None = Depends(dep1)) -> str: return "hi" + app.include_router(router) + + async with self.patch_broker(router.broker) as br: + with TestClient(app) as client: + assert client.app_state["broker"] is br + + r = await br.request( + "hi", + queue, + timeout=0.5, + ) + assert await r.decode() == "hi", r + + mock.assert_called_once() + assert not mock.not_call.called + + @pytest.mark.xfail(reason="https://github.com/airtai/faststream/issues/1742") + async def test_nested_router(self, mock: Mock, queue: str) -> None: + router = self.router_class() + router2 = self.router_class() + + app = FastAPI() + + args, kwargs = self.get_subscriber_params(queue) + + @router2.subscriber(*args, **kwargs) + async def hello_router2() -> str: + return "hi" + router.include_router(router2) app.include_router(router) diff --git a/tests/brokers/base/router.py b/tests/brokers/base/router.py index 68c8c8552c0..1f50340fb96 100644 --- a/tests/brokers/base/router.py +++ b/tests/brokers/base/router.py @@ -447,9 +447,8 @@ def subscriber() -> None: ... sub = next(iter(pub_broker._subscribers)) publisher = next(iter(pub_broker._publishers)) - - assert len((*sub._broker_middlewares, *sub.calls[0].item_middlewares)) == 3 - assert len((*publisher._broker_middlewares, *publisher._middlewares)) == 3 + assert len((*sub._broker_middlewares, *sub.calls[0].item_middlewares)) == 5 + assert len((*publisher._broker_middlewares, *publisher._middlewares)) == 4 async def test_router_include_with_middlewares( self, @@ -473,8 +472,8 @@ def subscriber() -> None: ... publisher = next(iter(pub_broker._publishers)) sub_middlewares = (*sub._broker_middlewares, *sub.calls[0].item_middlewares) - assert len(sub_middlewares) == 3, sub_middlewares - assert len((*publisher._broker_middlewares, *publisher._middlewares)) == 3 + assert len(sub_middlewares) == 5, sub_middlewares + assert len((*publisher._broker_middlewares, *publisher._middlewares)) == 4 async def test_router_parser( self, diff --git a/tests/brokers/confluent/test_consume.py b/tests/brokers/confluent/test_consume.py index 6f61f01d5fc..9bb954ee2f2 100644 --- a/tests/brokers/confluent/test_consume.py +++ b/tests/brokers/confluent/test_consume.py @@ -4,6 +4,7 @@ import pytest +from faststream import AckPolicy from faststream.confluent import KafkaBroker from faststream.confluent.annotations import KafkaMessage from faststream.confluent.client import AsyncConfluentConsumer @@ -251,7 +252,9 @@ async def test_consume_no_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - args, kwargs = self.get_subscriber_params(queue, group_id="test", no_ack=True) + args, kwargs = self.get_subscriber_params( + queue, group_id="test", ack_policy=AckPolicy.DO_NOTHING + ) @consume_broker.subscriber(*args, **kwargs) async def handler(msg: KafkaMessage) -> None: diff --git a/tests/brokers/kafka/test_consume.py b/tests/brokers/kafka/test_consume.py index f5e5421fbd6..cb0f32db436 100644 --- a/tests/brokers/kafka/test_consume.py +++ b/tests/brokers/kafka/test_consume.py @@ -5,6 +5,7 @@ import pytest from aiokafka import AIOKafkaConsumer +from faststream import AckPolicy from faststream.exceptions import AckMessage from faststream.kafka import KafkaBroker, TopicPartition from faststream.kafka.annotations import KafkaMessage @@ -296,7 +297,9 @@ async def test_consume_no_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue, group_id="test", no_ack=True) + @consume_broker.subscriber( + queue, group_id="test", ack_policy=AckPolicy.DO_NOTHING + ) async def handler(msg: KafkaMessage) -> None: event.set() diff --git a/tests/brokers/nats/test_consume.py b/tests/brokers/nats/test_consume.py index d7200e627f4..d253bcb716a 100644 --- a/tests/brokers/nats/test_consume.py +++ b/tests/brokers/nats/test_consume.py @@ -5,6 +5,7 @@ import pytest from nats.aio.msg import Msg +from faststream import AckPolicy from faststream.exceptions import AckMessage from faststream.nats import ConsumerConfig, JStream, NatsBroker, PullSub from faststream.nats.annotations import NatsMessage @@ -169,7 +170,7 @@ async def test_core_consume_no_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue, no_ack=True) + @consume_broker.subscriber(queue, ack_policy=AckPolicy.DO_NOTHING) async def handler(msg: NatsMessage) -> None: if not msg.raw_message._ackd: event.set() @@ -278,7 +279,7 @@ async def test_consume_no_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue, no_ack=True) + @consume_broker.subscriber(queue, ack_policy=AckPolicy.DO_NOTHING) async def handler(msg: NatsMessage) -> None: event.set() diff --git a/tests/brokers/rabbit/test_consume.py b/tests/brokers/rabbit/test_consume.py index 742daac225d..4f51bb48999 100644 --- a/tests/brokers/rabbit/test_consume.py +++ b/tests/brokers/rabbit/test_consume.py @@ -5,6 +5,7 @@ import pytest from aio_pika import IncomingMessage, Message +from faststream import AckPolicy from faststream.exceptions import AckMessage, NackMessage, RejectMessage, SkipMessage from faststream.rabbit import RabbitBroker, RabbitExchange, RabbitQueue from faststream.rabbit.annotations import RabbitMessage @@ -26,7 +27,7 @@ async def test_consume_from_exchange( ) -> None: consume_broker = self.get_broker() - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) def h(m) -> None: event.set() @@ -56,7 +57,6 @@ async def test_consume_with_get_old( @consume_broker.subscriber( queue=RabbitQueue(name=queue, passive=True), exchange=RabbitExchange(name=exchange.name, passive=True), - retry=True, ) def h(m) -> None: event.set() @@ -92,7 +92,7 @@ async def test_consume_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage) -> None: event.set() @@ -126,7 +126,7 @@ async def test_consume_manual_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage) -> None: await msg.ack() event.set() @@ -160,7 +160,7 @@ async def test_consume_exception_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage) -> None: try: raise AckMessage @@ -196,7 +196,7 @@ async def test_consume_manual_nack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage): await msg.nack() event.set() @@ -231,7 +231,7 @@ async def test_consume_exception_nack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage) -> None: try: raise NackMessage @@ -267,7 +267,7 @@ async def test_consume_manual_reject( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage): await msg.reject() event.set() @@ -302,7 +302,7 @@ async def test_consume_exception_reject( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue=queue, exchange=exchange, retry=1) + @consume_broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage) -> None: try: raise RejectMessage @@ -386,7 +386,9 @@ async def test_consume_no_ack( ) -> None: consume_broker = self.get_broker(apply_types=True) - @consume_broker.subscriber(queue, exchange=exchange, retry=1, no_ack=True) + @consume_broker.subscriber( + queue, exchange=exchange, ack_policy=AckPolicy.DO_NOTHING + ) async def handler(msg: RabbitMessage) -> None: event.set() diff --git a/tests/brokers/rabbit/test_test_client.py b/tests/brokers/rabbit/test_test_client.py index 856ed80668c..cf766697164 100644 --- a/tests/brokers/rabbit/test_test_client.py +++ b/tests/brokers/rabbit/test_test_client.py @@ -207,18 +207,18 @@ async def test_consume_manual_ack( consume2 = asyncio.Event() consume3 = asyncio.Event() - @broker.subscriber(queue=queue, exchange=exchange, retry=1) + @broker.subscriber(queue=queue, exchange=exchange) async def handler(msg: RabbitMessage) -> None: await msg.raw_message.ack() consume.set() - @broker.subscriber(queue=queue + "1", exchange=exchange, retry=1) + @broker.subscriber(queue=queue + "1", exchange=exchange) async def handler2(msg: RabbitMessage): await msg.raw_message.nack() consume2.set() raise ValueError - @broker.subscriber(queue=queue + "2", exchange=exchange, retry=1) + @broker.subscriber(queue=queue + "2", exchange=exchange) async def handler3(msg: RabbitMessage): await msg.raw_message.reject() consume3.set() diff --git a/tests/brokers/test_pushback.py b/tests/brokers/test_pushback.py deleted file mode 100644 index afb064ff69d..00000000000 --- a/tests/brokers/test_pushback.py +++ /dev/null @@ -1,124 +0,0 @@ -from unittest.mock import AsyncMock - -import pytest - -from faststream._internal.subscriber.acknowledgement_watcher import ( - CounterWatcher, - EndlessWatcher, - WatcherContext, -) -from faststream.exceptions import NackMessage, SkipMessage - - -@pytest.fixture() -def message(): - return AsyncMock(message_id=1, committed=None) - - -@pytest.mark.asyncio() -async def test_push_back_correct(async_mock: AsyncMock, message) -> None: - watcher = CounterWatcher(3) - - context = WatcherContext( - message=message, - watcher=watcher, - ) - - async with context: - await async_mock() - - async_mock.assert_awaited_once() - message.ack.assert_awaited_once() - assert not watcher.memory.get(message.message_id) - - -@pytest.mark.asyncio() -async def test_push_back_endless_correct(async_mock: AsyncMock, message) -> None: - watcher = EndlessWatcher() - - context = WatcherContext( - message=message, - watcher=watcher, - ) - - async with context: - await async_mock() - - async_mock.assert_awaited_once() - message.ack.assert_awaited_once() - - -@pytest.mark.asyncio() -async def test_push_back_watcher(async_mock: AsyncMock, message) -> None: - watcher = CounterWatcher(3) - - context = WatcherContext( - message=message, - watcher=watcher, - ) - - async_mock.side_effect = ValueError("Ooops!") - - while not message.reject.called: - with pytest.raises(ValueError): # noqa: PT011 - async with context: - await async_mock() - - assert not message.ack.await_count - assert message.nack.await_count == 3 - message.reject.assert_awaited_once() - - -@pytest.mark.asyncio() -async def test_push_endless_back_watcher(async_mock: AsyncMock, message) -> None: - watcher = EndlessWatcher() - - context = WatcherContext( - message=message, - watcher=watcher, - ) - - async_mock.side_effect = ValueError("Ooops!") - - while message.nack.await_count < 10: - with pytest.raises(ValueError): # noqa: PT011 - async with context: - await async_mock() - - assert not message.ack.called - assert not message.reject.called - assert message.nack.await_count == 10 - - -@pytest.mark.asyncio() -async def test_ignore_skip(async_mock: AsyncMock, message) -> None: - watcher = CounterWatcher(3) - - context = WatcherContext( - message=message, - watcher=watcher, - ) - - async with context: - raise SkipMessage - - assert not message.nack.called - assert not message.reject.called - assert not message.ack.called - - -@pytest.mark.asyncio() -async def test_additional_params_with_handler_exception( - async_mock: AsyncMock, message -) -> None: - watcher = EndlessWatcher() - - context = WatcherContext( - message=message, - watcher=watcher, - ) - - async with context: - raise NackMessage(delay=5) - - message.nack.assert_called_with(delay=5)