Feat: AMQP pub sub

async_processor/__init__.py

@@ -2,6 +2,8 @@
 from async_processor.function_service import FunctionAsyncExecutor
 from async_processor.processor import AsyncProcessor, Processor
 from async_processor.types import (
+    AMQPInputConfig,
+    AMQPOutputConfig,
     AWSAccessKeyAuth,
     CoreNATSOutputConfig,
     Input,
@@ -31,6 +33,8 @@
     "NATSOutputConfig",
     "OutputMessage",
     "WorkerConfig",
+    "AMQPInputConfig",
+    "AMQPOutputConfig",
     "ProcessStatus",
     "SQSInputConfig",
     "SQSOutputConfig",

async_processor/amqp_pub_sub.py

@@ -0,0 +1,106 @@
+from contextlib import asynccontextmanager
+from typing import AsyncIterator, Optional
+
+from aio_pika import Message, connect_robust
+from aio_pika.exceptions import AMQPError
+
+from async_processor.types import (
+    AMQPInputConfig,
+    AMQPOutputConfig,
+    Input,
+    InputFetchAckFailure,
+    InputMessageFetchFailure,
+    Output,
+)
+
+
+class AMQPInput(Input):
+    def __init__(self, config: AMQPInputConfig):
+        self._queue_url = config.queue_url
+        self._queue_name = config.queue_name
+        self._connection = None
+        self._channel = None
+        self._queue = None
+
+    async def _get_connect(self):
+        if not self._connection:
+            self._connection = await connect_robust(self._queue_url)
+        return self._connection
+
+    async def _get_channel(self):
+        if not  self._channel:
+            await self._get_connect()
+            self._channel = await self._connection.channel()
+        return self._channel
+
+    async def _get_queue(self):
+        if not self._queue:
+            await self._get_channel()
+            self._queue = await self._channel.declare_queue(self._queue_name)
+        return self._queue
+
+    @asynccontextmanager
+    async def get_input_message(
+        self,
+    ) -> AsyncIterator[Optional[str]]:
+        message = None
+        queue = await self._get_queue()
+        try:
+            message = await queue.get(fail=False, timeout=5)
+        except Exception as ex:
+            raise InputMessageFetchFailure(f"Error fetch input message: {ex}") from ex
+        if not message:
+            yield None
+            return
+        try:
+            yield message.body.decode()
+        finally:
+            if message:
+                try:
+                    await message.ack()
+                except Exception as ex:
+                    raise InputFetchAckFailure(
+                        f"Error publishing input message: {ex}"
+                    ) from ex
+
+    async def publish_input_message(
+        self, serialized_input_message: bytes, request_id: str
+    ):
+        try:
+            channel = await self._get_channel()
+            await channel.default_exchange.publish(
+                Message(body=serialized_input_message), routing_key=self._queue_name
+            )
+        except AMQPError as ex:
+            raise InputFetchAckFailure(f"Error publishing input message: {ex}") from ex
+
+
+class AMQPOutput(Output):
+    def __init__(self, config: AMQPOutputConfig):
+        self._queue_url = config.queue_url
+        self._queue_name = config.queue_name
+        self._connection = None
+        self._channel = None
+        self._queue = None
+
+    async def _get_connect(self):
+        if not self._connection:
+            self._connection = await connect_robust(self._queue_url)
+        return self._connection
+
+    async def _get_channel(self):
+        if not self._channel:
+            await self._get_connect()
+            self._channel = await self._connection.channel()
+        return self._channel
+
+    async def publish_output_message(
+        self, serialized_output_message: bytes, request_id: Optional[str]
+    ):
+        try:
+            queue = await self._get_channel()
+            await queue.default_exchange.publish(
+                Message(body=serialized_output_message), routing_key=self._queue_name
+            )
+        except AMQPError as ex:
+            raise InputFetchAckFailure(f"Error publishing output message: {ex}") from ex

async_processor/types.py

@@ -216,6 +216,17 @@ def to_input(self) -> Input:
         return SQSInput(self)
 
 
+class AMQPInputConfig(InputConfig):
+    type: constr(regex=r"^amqp$") = "amqp"
+    queue_url: str
+    queue_name: str
+
+    def to_input(self) -> Input:
+        from async_processor.amqp_pub_sub import AMQPInput
+
+        return AMQPInput(self)
+
+
 class NATSInputConfig(InputConfig):
     type: constr(regex=r"^nats$") = "nats"
 
@@ -288,6 +299,18 @@ def to_output(self) -> Output:
         return SQSOutput(self)
 
 
+class AMQPOutputConfig(OutputConfig):
+    type: constr(regex=r"^amqp$") = "amqp"
+
+    queue_url: str
+    queue_name: str
+
+    def to_output(self) -> Output:
+        from async_processor.amqp_pub_sub import AMQPOutput
+
+        return AMQPOutput(self)
+
+
 class NATSOutputConfig(OutputConfig):
     type: constr(regex=r"^nats$") = "nats"
 

examples/amqp.py

@@ -0,0 +1,25 @@
+from async_processor import (
+    AMQPInputConfig,
+    AMQPOutputConfig,
+    InputMessage,
+    Processor,
+    WorkerConfig,
+)
+
+
+class MultiplicationProcessor(Processor):
+    def process(self, input_message: InputMessage) -> int:
+        body = input_message.body
+        return body["x"] * body["y"]
+
+
+app = MultiplicationProcessor().build_app(
+    worker_config=WorkerConfig(
+        input_config=AMQPInputConfig(
+            queue_url="amqp://guest:guest@localhost:5672/", queue_name="home1"
+        ),
+        output_config=AMQPOutputConfig(
+            queue_url="amqp://guest:guest@localhost:5672/", queue_name="home2"
+        ),
+    ),
+)

examples/run.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+module_name="$1"
+
+if [ -z "$module_name" ]; then
+    echo "Usage: $0 <module_name>"
+    echo "Example: $0 [amqp]"
+    exit 1
+fi
+
+gunicorn "$module_name:app" --workers 1 --worker-class uvicorn.workers.UvicornWorker --bind 127.0.0.1:8000