55
66from aio_pika import DeliveryMode , ExchangeType , Message , connect_robust
77from aio_pika .abc import AbstractChannel , AbstractQueue , AbstractRobustConnection
8- from taskiq .abc .broker import AsyncBroker
9- from taskiq .abc .result_backend import AsyncResultBackend
10- from taskiq .message import BrokerMessage
8+ from taskiq import AckableMessage , AsyncBroker , AsyncResultBackend , BrokerMessage
119
1210_T = TypeVar ("_T" ) # noqa: WPS111
1311
@@ -228,21 +226,20 @@ async def kick(self, message: BrokerMessage) -> None:
228226 if self .write_channel is None :
229227 raise ValueError ("Please run startup before kicking." )
230228
231- message_base_params : dict [str , Any ] = {
229+ message_base_params : Dict [str , Any ] = {
232230 "body" : message .message ,
233231 "headers" : {
234232 "task_id" : message .task_id ,
235233 "task_name" : message .task_name ,
236234 ** message .labels ,
237235 },
238236 "delivery_mode" : DeliveryMode .PERSISTENT ,
237+ "priority" : parse_val (
238+ int ,
239+ message .labels .get ("priority" ),
240+ ),
239241 }
240242
241- message_base_params ["priority" ] = parse_val (
242- int ,
243- message .labels .get ("priority" ),
244- )
245-
246243 delay : Optional [int ] = parse_val (int , message .labels .get ("delay" ))
247244 rmq_message : Message = Message (** message_base_params )
248245
@@ -268,7 +265,7 @@ async def kick(self, message: BrokerMessage) -> None:
268265 routing_key = self ._delay_queue_name ,
269266 )
270267
271- async def listen (self ) -> AsyncGenerator [bytes , None ]:
268+ async def listen (self ) -> AsyncGenerator [AckableMessage , None ]:
272269 """
273270 Listen to queue.
274271
@@ -284,5 +281,8 @@ async def listen(self) -> AsyncGenerator[bytes, None]:
284281 queue = await self .declare_queues (self .read_channel )
285282 async with queue .iterator () as iterator :
286283 async for message in iterator :
287- async with message .process ():
288- yield message .body
284+ yield AckableMessage (
285+ data = message .body ,
286+ ack = message .ack ,
287+ reject = message .reject ,
288+ )
0 commit comments