Merge branch 'release/0.6.0'

Author: s3rius

docs/examples/extending/broker.py

@@ -1,6 +1,6 @@
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Union
 
-from taskiq import AsyncBroker, BrokerMessage
+from taskiq import AckableMessage, AsyncBroker, BrokerMessage
 
 
 class MyBroker(AsyncBroker):
@@ -23,7 +23,7 @@ async def kick(self, message: BrokerMessage) -> None:
         # Send a message.message.
         pass
 
-    async def listen(self) -> AsyncGenerator[bytes, None]:
+    async def listen(self) -> AsyncGenerator[Union[bytes, AckableMessage], None]:
         while True:
             # Get new message.
             new_message: bytes = ...  # type: ignore

docs/extending-taskiq/broker.md

@@ -23,7 +23,28 @@ As a broker developer, please send only raw bytes from the `message` field of a
 :::
 
 
-The `listen` method should yield raw bytes that were sent over the network.
+## Acknowledgement
+
+The `listen` method should yield raw bytes of a message.
+But if your broker supports acking or rejecting messages, the broker should return `taskiq.AckableMessage`
+with required fields.
+
+For example:
+
+```python
+
+async def listen(self) -> AsyncGenerator[AckableMessage, None]:
+   for message in self.my_channel:
+      yield AckableMessage(
+         data=message.bytes,
+         # Ack is a function that takes no parameters.
+         # So you either set here method of a message,
+         # or you can make a closure.
+         ack=message.ack
+         # Can be set to None if broker doesn't support it.
+         reject=message.reject
+      )
+```
 
 ## Conventions
 

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq"
-version = "0.5.0"
+version = "0.6.0"
 description = "Distributed task queue with full async support"
 authors = ["Pavel Kirilin <win10@list.ru>"]
 maintainers = ["Pavel Kirilin <win10@list.ru>"]

pyproject.toml

@@ -6,12 +6,22 @@
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.abc.schedule_source import ScheduleSource
+from taskiq.acks import AckableMessage
 from taskiq.brokers.inmemory_broker import InMemoryBroker
 from taskiq.brokers.shared_broker import async_shared_broker
 from taskiq.brokers.zmq_broker import ZeroMQBroker
 from taskiq.context import Context
 from taskiq.events import TaskiqEvents
-from taskiq.exceptions import TaskiqError
+from taskiq.exceptions import (
+    NoResultError,
+    RejectError,
+    ResultGetError,
+    ResultIsReadyError,
+    SecurityError,
+    SendTaskError,
+    TaskiqError,
+    TaskiqResultTimeoutError,
+)
 from taskiq.funcs import gather
 from taskiq.message import BrokerMessage, TaskiqMessage
 from taskiq.middlewares.prometheus_middleware import PrometheusMiddleware
@@ -21,28 +31,45 @@
 from taskiq.state import TaskiqState
 from taskiq.task import AsyncTaskiqTask
 
+try:
+    # Python 3.8+
+    from importlib.metadata import version  # noqa: WPS433
+except ImportError:
+    # Python 3.7
+    from importlib_metadata import version  # noqa: WPS433
+
+__version__ = version("taskiq")
 __all__ = [
+    "__version__",
     "gather",
     "Context",
     "AsyncBroker",
     "TaskiqError",
+    "RejectError",
     "TaskiqState",
     "TaskiqResult",
     "ZeroMQBroker",
     "TaskiqEvents",
+    "SecurityError",
     "TaskiqMessage",
     "BrokerMessage",
+    "ResultGetError",
     "ScheduledTask",
     "TaskiqDepends",
+    "NoResultError",
+    "SendTaskError",
+    "AckableMessage",
     "InMemoryBroker",
     "ScheduleSource",
     "TaskiqScheduler",
     "TaskiqFormatter",
     "AsyncTaskiqTask",
     "TaskiqMiddleware",
+    "ResultIsReadyError",
     "AsyncResultBackend",
     "async_shared_broker",
-    "AsyncTaskiqDecoratedTask",
-    "SimpleRetryMiddleware",
     "PrometheusMiddleware",
+    "SimpleRetryMiddleware",
+    "AsyncTaskiqDecoratedTask",
+    "TaskiqResultTimeoutError",
 ]

taskiq/__init__.py

@@ -2,8 +2,9 @@
 import sys
 from typing import Dict
 
-from importlib_metadata import entry_points, version
+from importlib_metadata import entry_points
 
+from taskiq import __version__
 from taskiq.abc.cmd import TaskiqCMD
 
 
@@ -60,7 +61,7 @@ def main() -> None:  # noqa: WPS210  # pragma: no cover
     args, _ = parser.parse_known_args()
 
     if args.version:
-        print(version("taskiq"))  # noqa: WPS421
+        print(__version__)  # noqa: WPS421
         return
 
     if args.subcommand is None:

taskiq/__main__.py

@@ -24,6 +24,7 @@
 from typing_extensions import ParamSpec, Self, TypeAlias
 
 from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.acks import AckableMessage
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.events import TaskiqEvents
 from taskiq.formatters.json_formatter import JSONFormatter
@@ -68,7 +69,10 @@ class AsyncBroker(ABC):
     """
 
     available_tasks: Dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
+    # True only if broker runs in worker process.
     is_worker_process: bool = False
+    # True only if broker runs in scheduler process.
+    is_scheduler_process: bool = False
 
     def __init__(
         self,
@@ -182,13 +186,19 @@ async def kick(
         """
 
     @abstractmethod
-    def listen(self) -> AsyncGenerator[bytes, None]:
+    def listen(self) -> AsyncGenerator[Union[bytes, AckableMessage], None]:
         """
         This function listens to new messages and yields them.
 
         This it the main point for workers.
         This function is used to get new tasks from the network.
 
+        If your broker support acknowledgement, then you
+        should wrap your message in AckableMessage dataclass.
+
+        If your messages was wrapped in AckableMessage dataclass,
+        taskiq will call ack when finish processing message.
+
         :yield: incoming messages.
         :return: nothing.
         """

taskiq/abc/broker.py

@@ -0,0 +1,21 @@
+import dataclasses
+from typing import Awaitable, Callable, Optional, Union
+
+
+@dataclasses.dataclass
+class AckableMessage:
+    """
+    Message that can be acknowledged.
+
+    If your broker support message acknowledgement,
+    please return this type of message, so we'll be
+    able to mark this message as acknowledged after
+    the function will be executed.
+
+    It adds more reliability to brokers and system
+    as a whole.
+    """
+
+    data: bytes
+    ack: Callable[[], Union[None, Awaitable[None]]]
+    reject: Optional[Callable[[], Union[None, Awaitable[None]]]] = None