Updated aio-pika broker

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

poetry.lock

@@ -15,6 +15,23 @@ black = "^22.6.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.0"
+isort = "^5.10.1"
+mypy = "^0.971"
+
+[tool.mypy]
+strict = true
+ignore_missing_imports = true
+allow_subclassing_any = true
+allow_untyped_calls = true
+pretty = true
+show_error_codes = true
+implicit_reexport = true
+allow_untyped_decorators = true
+warn_return_any = false
+
+[tool.isort]
+profile = "black"
+multi_line_output = 3
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

pyproject.toml

@@ -1 +1 @@
-__version__ = '0.1.0'
+__version__ = "0.1.0"

taskiq_aio_pika/__init__.py

@@ -1,11 +1,12 @@
-from typing import Any, AsyncGenerator, Dict, Optional, TypeVar
+from asyncio import AbstractEventLoop
+from typing import Any, AsyncGenerator, Optional, TypeVar
+
+from aio_pika import Channel, ExchangeType, Message, connect_robust
+from aio_pika.abc import AbstractChannel, AbstractRobustConnection
+from aio_pika.pool import Pool
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
-from taskiq.message import TaskiqMessage
-from aio_pika.abc import AbstractRobustConnection
-from aio_pika.pool import Pool
-from asyncio import AbstractEventLoop
-from aio_pika import connect_robust, Channel, Message, ExchangeType
+from taskiq.message import BrokerMessage
 
 _T = TypeVar("_T")
 
@@ -35,7 +36,7 @@ async def _get_rmq_connection() -> AbstractRobustConnection:
             loop=loop,
         )
 
-        async def get_channel() -> Channel:
+        async def get_channel() -> AbstractChannel:
             async with self.connection_pool.acquire() as connection:
                 return await connection.channel()
 
@@ -62,30 +63,30 @@ async def startup(self) -> None:
             queue = await channel.declare_queue(self.queue_name)
             await queue.bind(exchange=exchange, routing_key="*")
 
-    async def kick(self, message: TaskiqMessage) -> None:
+    async def kick(self, message: BrokerMessage) -> None:
         rmq_msg = Message(
-            body=message.json().encode(),
-            content_type="application/json",
+            body=message.message.encode(),
             headers={
                 "task_id": message.task_id,
                 "task_name": message.task_name,
+                **message.headers,
             },
         )
         async with self.channel_pool.acquire() as channel:
             exchange = await channel.get_exchange(self.exchange_name, ensure=False)
             await exchange.publish(rmq_msg, routing_key=message.task_id)
 
-    async def listen(self) -> AsyncGenerator[TaskiqMessage, None]:
+    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
         async with self.channel_pool.acquire() as channel:
             await channel.set_qos(prefetch_count=self.qos)
             queue = await channel.get_queue(self.queue_name, ensure=False)
             async with queue.iterator() as queue_iter:
                 async for rmq_message in queue_iter:
                     async with rmq_message.process():
                         try:
-                            yield TaskiqMessage.parse_raw(
+                            yield BrokerMessage.parse_raw(
                                 rmq_message.body,
-                                content_type=rmq_message.content_type,
+                                content_type=rmq_message.content_type or "",
                             )
                         except ValueError:
                             continue