Feat: AMQP module

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,161 @@
+from contextlib import asynccontextmanager
+from typing import AsyncIterator, Optional
+
+from aio_pika import Message, connect_robust
+from aio_pika.abc import AbstractChannel, AbstractConnection, AbstractQueue
+from aio_pika.exceptions import ChannelNotFoundEntity, QueueEmpty
+
+from async_processor.logger import logger
+from async_processor.types import (
+    AMQPInputConfig,
+    AMQPOutputConfig,
+    Input,
+    InputFetchAckFailure,
+    InputMessageFetchFailure,
+    Output,
+)
+
+
+class AMQPInput(Input):
+    def __init__(self, config: AMQPInputConfig):
+        self._url = config.url
+        self._queue_name = config.queue_name
+        self._wait_time_seconds = config.wait_time_seconds
+        self._nc = None
+        self._ch = None
+        self._queue = None
+
+    async def _validate_queue_exists(self):
+        channel = await self._get_channel()
+        try:
+            self._queue = await channel.declare_queue(self._queue_name, passive=True)
+        except ChannelNotFoundEntity as ex:
+            raise Exception(
+                f"Queue {self._queue_name!r} does not exist."
+                " Please create the queue before running the async processor."
+            ) from ex
+
+    async def __aenter__(self):
+        await self._validate_queue_exists()
+        return self
+
+    async def _get_connect(self) -> AbstractConnection:
+        if self._nc:
+            return self._nc
+        self._nc = await connect_robust(self._url)
+        return self._nc
+
+    async def _get_channel(self) -> AbstractChannel:
+        if self._ch:
+            return self._ch
+        connection = await self._get_connect()
+        self._ch = await connection.channel()
+        return self._ch
+
+    async def _get_queue(self) -> AbstractQueue:
+        if self._queue:
+            return self._queue
+        channel = await self._get_channel()
+        self._queue = await channel.declare_queue(self._queue_name, passive=True)
+        return self._queue
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        if not self._nc:
+            return
+        try:
+            await self._nc.close()
+        except Exception:
+            logger.exception("Failed to drain and close nats connection")
+
+    @asynccontextmanager
+    async def get_input_message(
+        self,
+    ) -> AsyncIterator[Optional[bytes]]:
+        message = None
+        queue = await self._get_queue()
+        try:
+            message = await queue.get(fail=False, timeout=self._wait_time_seconds)
+        except QueueEmpty:
+            logger.debug("No message in queue")
+        except Exception as ex:
+            raise InputMessageFetchFailure(f"Error fetch input message: {ex}") from ex
+        if not message:
+            yield None
+            return
+        try:
+            yield message.body
+        finally:
+            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
+    ):
+        channel = await self._get_channel()
+        await channel.default_exchange.publish(
+            Message(body=serialized_input_message), routing_key=self._queue_name
+        )
+
+
+class AMQPOutput(Output):
+    def __init__(self, config: AMQPOutputConfig):
+        self._url = config.url
+        self._queue_name = config.queue_name
+        self._queue = None
+        self._nc = None
+        self._ch = None
+
+    async def _validate_queue_exists(self):
+        channel = await self._get_channel()
+        try:
+            self._queue = await channel.declare_queue(self._queue_name, passive=True)
+        except ChannelNotFoundEntity as ex:
+            raise Exception(
+                f"Queue {self._queue_name!r} does not exist."
+                " Please create the queue before running the async processor."
+            ) from ex
+
+    async def __aenter__(self):
+        await self._validate_queue_exists()
+        return self
+
+    async def _get_connect(self) -> AbstractConnection:
+        if self._nc:
+            return self._nc
+        self._nc = await connect_robust(self._url)
+        return self._nc
+
+    async def _get_channel(self) -> AbstractChannel:
+        if self._ch:
+            return self._ch
+        connection = await self._get_connect()
+        self._ch = await connection.channel()
+        return self._ch
+
+    async def _get_queue(self) -> AbstractQueue:
+        if self._queue:
+            return self._queue
+        channel = await self._get_channel()
+        self._queue = await channel.declare_queue(self._queue_name, passive=True)
+        return self._queue
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        if not self._nc:
+            return
+        try:
+            await self._nc.close()
+        except Exception:
+            logger.exception("Failed to drain and close nats connection")
+
+    async def publish_output_message(
+        self, serialized_output_message: bytes, request_id: Optional[str]
+    ):
+        channel = await self._get_channel()
+        await self._get_queue()
+        await channel.default_exchange.publish(
+            Message(body=serialized_output_message), routing_key=self._queue_name
+        )

async_processor/types.py

@@ -216,6 +216,20 @@ def to_input(self) -> Input:
         return SQSInput(self)
 
 
+class AMQPInputConfig(InputConfig):
+    type: constr(regex=r"^amqp$") = "amqp"
+    url: constr(
+        regex=r"^(?:amqp|amqps):\/\/(?:([^:/?#\s]+)(?::([^@/?#\s]+))?@)?([^/?#\s]+)(?::(\d+))?\/?([^?#\s]*)?(?:\?(.*))?$"
+    )
+    queue_name: str
+    wait_time_seconds: conint(ge=1, le=20) = 5
+
+    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 +302,21 @@ def to_output(self) -> Output:
         return SQSOutput(self)
 
 
+class AMQPOutputConfig(OutputConfig):
+    type: constr(regex=r"^amqp$") = "amqp"
+
+    url: constr(
+        regex=r"^(?:amqp|amqps):\/\/(?:([^:/?#\s]+)(?::([^@/?#\s]+))?@)?([^/?#\s]+)(?::(\d+))?\/?([^?#\s]*)?(?:\?(.*))?$"
+    )
+    queue_name: str
+    wait_time_seconds: confloat(ge=1) = 5
+
+    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,63 @@
+import asyncio
+import json
+import random
+import uuid
+
+import aio_pika
+
+from async_processor import (
+    AMQPInputConfig,
+    AMQPOutputConfig,
+    InputMessage,
+    Processor,
+    WorkerConfig,
+)
+
+# change this config
+input_url = "amqp://guest:guest@localhost:5672/"
+input_queue_name = "home1"
+output_url = "amqp://guest:guest@localhost:5672/"
+output_queue_name = "home2"
+
+
+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(url=input_url, queue_name=input_queue_name),
+        output_config=AMQPOutputConfig(url=output_url, queue_name=output_queue_name),
+    ),
+)
+
+
+async def send_request(url: str, routing_key: str):
+    connection = await aio_pika.connect_robust(url)
+
+    async with connection:
+        request_id = str(uuid.uuid4())
+
+        channel = await connection.channel()
+
+        payload = json.dumps(
+            {
+                "request_id": request_id,
+                "body": {"x": random.randint(1, 100), "y": random.randint(1, 100)},
+            }
+        )
+        print(payload)
+        await channel.default_exchange.publish(
+            aio_pika.Message(payload.encode()), routing_key=routing_key
+        )
+
+
+async def test():
+    for _ in range(100):
+        await send_request(url=input_url, routing_key=input_queue_name)
+
+
+if __name__ == "__main__":
+    asyncio.run(test())

examples/run.sh

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

examples/test.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+module_name="$1"
+
+if [ -z "$module_name" ]; then
+    echo "Usage: $0 <module_name>"
+    echo "Example: $0 [amqp]"
+    exit 1
+fi
+
+if [ -f "${module_name}" ]; then
+    module_name="${module_name}"
+else
+    module_name="$module_name.py"
+fi
+
+python $module_name