asyncapi-python 0.1.0__py3-none-any.whl

asyncapi_python/amqp/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from .connection import channel_pool, AmqpPool
+from .consumer import Consumer
+from .producer import Producer
+from .base_application import BaseApplication

asyncapi_python/amqp/base_application.py

@@ -0,0 +1,61 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from .connection import channel_pool, AmqpPool
+from .consumer import Consumer
+from .producer import Producer
+from abc import ABC, abstractmethod
+from typing import Literal, TypedDict
+
+
+class Queue(TypedDict):
+    name: str | None
+    durable: bool
+    exclusive: bool
+    auto_delete: bool
+
+
+class Exchange(TypedDict):
+    name: str | None
+    type: Literal["topic", "direct", "fanout", "default", "headers"]
+    durable: bool
+    auto_delete: bool
+
+
+class BaseApplication(ABC):
+    def __init__(self, amqp_uri: str):
+        self._uri = amqp_uri
+        self._has_started = False
+        self._pool = channel_pool(self._uri)
+        self._consumer = Consumer(self._pool)
+
+    def _assert_started(self):
+        if not self._has_started:
+            cls_name = self.__class__.__name__
+            raise AssertionError(
+                f"Invoke of {cls_name}::request or {cls_name}::publish "
+                + "occurred before {cls_name}::start"
+            )
+
+    async def start(self, blocking: bool = True):
+        async with self._pool.acquire() as ch:
+            reply_queue = await ch.declare_queue(exclusive=True)
+            self._producer = Producer(self._pool, reply_queue)
+            await self._producer.run()
+            self._has_started = True
+        if blocking:
+            await self._consumer.run_blocking(timeout=None)
+        else:
+            await self._consumer.run()

asyncapi_python/amqp/connection.py

@@ -0,0 +1,42 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from functools import cache
+from aio_pika.robust_connection import (
+    AbstractRobustConnection,
+    AbstractRobustChannel,
+    connect_robust,
+)
+from aio_pika.pool import Pool
+
+
+@cache
+def connection_pool(amqp_uri: str) -> Pool[AbstractRobustConnection]:
+    async def get_connection():
+        return await connect_robust(amqp_uri)
+
+    return Pool(get_connection, max_size=2)
+
+
+@cache
+def channel_pool(amqp_uri: str) -> Pool[AbstractRobustChannel]:
+    async def get_channel():
+        async with connection_pool(amqp_uri).acquire() as connection:
+            return await connection.channel()
+
+    return Pool(get_channel, max_size=10)
+
+
+AmqpPool = Pool[AbstractRobustChannel]

asyncapi_python/amqp/consumer.py

@@ -0,0 +1,80 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from .message_handler import AbstractMessageHandler, MessageHandler, RpcMessageHandler
+from .message_handler_params import MessageHandlerParams
+from .utils import encode_message, decode_message
+
+import asyncio
+from aio_pika import Message
+from aio_pika.pool import Pool
+from aio_pika.abc import AbstractRobustChannel
+from asyncio import Future
+from typing import Callable, TypeVar
+from pydantic import BaseModel
+from logging import getLogger
+
+T = TypeVar("T", bound=BaseModel)
+U = TypeVar("U", bound=BaseModel)
+
+
+class Consumer:
+    def __init__(self, channel_pool: Pool[AbstractRobustChannel]):
+        self._handlers: dict[MessageHandlerParams, AbstractMessageHandler] = {}
+        self._logger = getLogger(__name__)
+        self._pool = channel_pool
+
+    async def run_blocking(self, timeout: int | float | None):
+        await self.run()
+        if timeout is not None:
+            await asyncio.sleep(timeout)
+        else:
+            await Future()
+
+    async def run(self):
+        async with self._pool.acquire() as channel:
+            for params, handler in self._handlers.items():
+                await params.setup_consume(handler, channel)
+
+    async def _reply_callback(self, message: Message, routing_key: str):
+        async with self._pool.acquire() as channel:
+            await channel.default_exchange.publish(message, routing_key)
+
+    def on(
+        self,
+        *,
+        params: MessageHandlerParams,
+        input_type: type[T],
+        output_type: type[U] | None,
+        callback: Callable,
+    ):
+        handler: AbstractMessageHandler
+        if params in self._handlers:
+            raise AssertionError(f"Only one handler for `{params}` is allowed")
+        if output_type is None:
+            handler = MessageHandler(
+                name=params.root.name,
+                callback=callback,
+                decode_message=lambda x: decode_message(x, input_type),
+            )
+        else:
+            handler = RpcMessageHandler(
+                name=params.root.name,
+                callback=callback,
+                reply_callback=self._reply_callback,
+                encode_message=encode_message,
+                decode_message=lambda x: decode_message(x, input_type),
+            )
+        self._handlers[params] = handler

asyncapi_python/amqp/message_handler.py

@@ -0,0 +1,89 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from abc import ABC, abstractmethod
+from aio_pika.message import AbstractIncomingMessage, Message
+from typing import Awaitable, Callable, Generic, TypeVar
+from pydantic import BaseModel
+from logging import getLogger
+
+
+T = TypeVar("T", bound=BaseModel)
+U = TypeVar("U", bound=BaseModel | None)
+V = TypeVar("V", bound=BaseModel)
+
+
+class AbstractMessageHandler(ABC, Generic[T, U]):
+    def __init__(self, name: str, callback: Callable[[T], Awaitable[U]]):
+        self._logger = getLogger(__name__)
+        self._name = name
+        self._callback = callback
+
+    async def __call__(self, message: AbstractIncomingMessage) -> None:
+        cls = f"{self.__class__.__name__}#{self._name}"
+        self._logger.info(f"{cls}: got message: {message.info()}")
+        self._logger.debug(f"content: {message.body!r}")
+        await self.on_call(message)
+        await message.ack()
+        self._logger.info(f"{cls} finished message: {message.info()}")
+
+    @abstractmethod
+    async def on_call(self, message: AbstractIncomingMessage) -> None:
+        raise NotImplementedError
+
+
+class MessageHandler(AbstractMessageHandler[T, None]):
+    def __init__(
+        self,
+        name: str,
+        callback: Callable[[T], Awaitable[None]],
+        decode_message: Callable[[bytes], T],
+    ):
+        super().__init__(name, callback)
+        self._decode_message = decode_message
+
+    async def on_call(self, message: AbstractIncomingMessage) -> None:
+        message_body = self._decode_message(message.body)
+        await self._callback(message_body)
+
+
+class RpcMessageHandler(AbstractMessageHandler[T, V]):
+    def __init__(
+        self,
+        name: str,
+        callback: Callable[[T], Awaitable[V]],
+        reply_callback: Callable[[Message, str], Awaitable[None]],
+        decode_message: Callable[[bytes], T],
+        encode_message: Callable[[V], bytes],
+    ):
+        super().__init__(name, callback)
+        self._reply_callback = reply_callback
+        self._decode_message = decode_message
+        self._encode_message = encode_message
+
+    async def on_call(self, message: AbstractIncomingMessage) -> None:
+        if message.correlation_id is None:
+            raise AssertionError("RPC Call got empty correlation_id")
+        if message.reply_to is None:
+            raise AssertionError("RPC Call got empty reply_to header")
+        message_body = self._decode_message(message.body)
+        result = await self._callback(message_body)
+        await self._reply_callback(
+            Message(
+                self._encode_message(result),
+                correlation_id=message.correlation_id,
+            ),
+            message.reply_to,
+        )

asyncapi_python/amqp/message_handler_params.py

@@ -0,0 +1,62 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pydantic import RootModel, BaseModel, ConfigDict, computed_field
+from typing import Literal
+from aio_pika.abc import AbstractRobustChannel
+from .message_handler import AbstractMessageHandler
+
+
+class ExchangeHandlerParams(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    kind: Literal["exchange"] = "exchange"
+    type: Literal["direct", "fanout", "topic", "headers"]
+    name: str
+    routing_key: str | None
+    auto_delete: bool = False
+
+
+class QueueHandlerParams(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    kind: Literal["queue"] = "queue"
+    name: str
+    exclusive: bool = False
+    auto_delete: bool = False
+    durable: bool = False
+
+
+class MessageHandlerParams(RootModel):
+    model_config = ConfigDict(frozen=True)
+    root: QueueHandlerParams | ExchangeHandlerParams
+
+    async def setup_consume(
+        self,
+        handler: AbstractMessageHandler,
+        channel: AbstractRobustChannel,
+    ):
+        match self.root:
+            case ExchangeHandlerParams(
+                name=name, routing_key=rk, type=et, auto_delete=ad
+            ):
+                exchange = await channel.declare_exchange(name, type=et, auto_delete=ad)
+                queue = await channel.declare_queue(exclusive=True)
+                await queue.bind(exchange, rk)
+            case QueueHandlerParams(name=name, exclusive=ex, auto_delete=ad, durable=d):
+                queue = await channel.declare_queue(
+                    name, exclusive=ex, auto_delete=ad, durable=d
+                )
+            case _:
+                raise NotImplementedError
+        await queue.consume(handler)

asyncapi_python/amqp/producer.py

@@ -0,0 +1,103 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from aio_pika import Message
+from aio_pika.pool import Pool
+from aio_pika.abc import (
+    AbstractRobustChannel,
+    AbstractRobustQueue,
+    AbstractIncomingMessage,
+)
+from logging import getLogger
+from pydantic import BaseModel
+from typing import TypeVar
+from asyncio import Future
+from uuid import uuid4
+from .utils import encode_message, decode_message
+
+T = TypeVar("T", bound=BaseModel)
+U = TypeVar("U", bound=BaseModel)
+
+
+class Producer:
+    def __init__(
+        self,
+        channel_pool: Pool[AbstractRobustChannel],
+        reply_queue: AbstractRobustQueue,
+    ):
+        if not reply_queue.exclusive:
+            raise AssertionError("Reply queue must be exclusive")
+
+        self._logger = getLogger(__name__)
+        self._pool = channel_pool
+        self._replies: dict[str, Future[AbstractIncomingMessage]] = {}
+        self._reply_queue = reply_queue
+        self._reply_consumer_tag: str | None = None
+
+    async def _on_reply(self, msg: AbstractIncomingMessage):
+        await msg.ack()
+        if msg.correlation_id is None or msg.correlation_id not in self._replies:
+            return
+        future = self._replies.pop(msg.correlation_id)
+        future.set_result(msg)
+
+    async def run(self):
+        self._reply_consumer_tag = await self._reply_queue.consume(self._on_reply)
+
+    async def publish(
+        self,
+        message: T,
+        exchange: str | None,
+        routing_key: str | None,
+    ):
+        outbound_message = Message(
+            body=encode_message(message),
+        )
+        async with self._pool.acquire() as channel:
+            await (
+                await channel.get_exchange(exchange)
+                if exchange is not None
+                else channel.default_exchange
+            ).publish(outbound_message, routing_key or "")
+
+    async def request(
+        self,
+        message: T,
+        exchange: str | None,
+        routing_key: str | None,
+        output_type: type[U],
+    ) -> U:
+        if not self._reply_consumer_tag:
+            raise AssertionError(
+                "Cannot make requests that expect replies before Consumer::start is called"
+            )
+        corr_id = str(uuid4())
+        outbound_message = Message(
+            body=encode_message(message),
+            correlation_id=corr_id,
+            reply_to=self._reply_queue.name,
+        )
+        reply_future = Future[AbstractIncomingMessage]()
+        async with self._pool.acquire() as channel:
+            await (
+                await channel.get_exchange(exchange)
+                if exchange is not None
+                else channel.default_exchange
+            ).publish(outbound_message, routing_key or "")
+            self._logger.info(f"Sent request {message}")
+            self._replies[corr_id] = reply_future
+        res = decode_message((await reply_future).body, output_type)
+        self._logger.info(f"Got response {res}")
+        return res

asyncapi_python/amqp/utils.py

@@ -0,0 +1,27 @@
+# Copyright 2024 Yaroslav Petrov <yaroslav.v.petrov@gmail.com>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pydantic import BaseModel
+from typing import TypeVar
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def encode_message(message: T) -> bytes:
+    return message.model_dump_json().encode()
+
+
+def decode_message(message: bytes, schema: type[T]) -> T:
+    return schema.model_validate_json(message)