rocketmq-sdk 0.1.0__py3-none-any.whl

rocketmq/__init__.py

@@ -0,0 +1,75 @@
+"""
+RocketMQ Python SDK — schema-aware AMQP client.
+
+Pydantic + aio-pika wrapper that mirrors the ``rocketmq.js`` API.
+
+Usage::
+
+    from rocketmq import connect
+    from rocketmq.schema import BaseSchema
+
+    class Order(BaseSchema):
+        id: str
+        customer_id: str
+        qty: int
+
+    async def main():
+        mq = await connect(url="amqp://localhost")
+        orders = await mq.queue("orders", Order)
+        await orders.send(Order(id="1", customer_id="c1", qty=5))
+        await orders.consume(lambda msg: print(msg))
+        await mq.close()
+"""
+
+__version__ = "0.1.0"
+
+# Error codes
+# Core API
+from rocketmq.client import QueueHandle, RocketMQ, connect
+from rocketmq.error_codes import BrokerErrorCode
+
+# Errors
+from rocketmq.errors import (
+    ConnectionError_,
+    ConsumeError,
+    PublishError,
+    QueueError,
+    RocketMQError,
+    SchemaError,
+    SchemaValidationError,
+    SerializationError,
+    TimeoutError_,
+)
+
+# Schema
+from rocketmq.schema import BaseSchema, Proto, ProtoType
+
+# Serializer
+from rocketmq.serializer import JsonSerializer, Serializer
+
+
+__all__ = [
+    # Schema
+    "BaseSchema",
+    # Error codes
+    "BrokerErrorCode",
+    "ConnectionError_",
+    "ConsumeError",
+    "JsonSerializer",
+    "Proto",
+    "ProtoType",
+    "PublishError",
+    "QueueError",
+    "QueueHandle",
+    "RocketMQ",
+    # Errors
+    "RocketMQError",
+    "SchemaError",
+    "SchemaValidationError",
+    "SerializationError",
+    # Serializer
+    "Serializer",
+    "TimeoutError_",
+    # Core
+    "connect",
+]

rocketmq/amqp.py

@@ -0,0 +1,222 @@
+"""
+Thin wrapper around aio-pika's connection and channel.
+
+Exposes only the operations the SDK needs without leaking the full
+aio-pika surface.  No schema or validation logic lives here.
+
+Usage::
+
+    conn = await AmqpConnection.connect("amqp://localhost")
+    ch = await conn.create_channel()
+    await ch.assert_queue("orders")
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
+from typing import Any, final
+
+import aio_pika
+from aio_pika.abc import AbstractChannel, AbstractIncomingMessage
+
+
+@dataclass(frozen=True, slots=True)
+class AssertQueueReply:
+    """Result of declaring a queue."""
+
+    queue: str
+    message_count: int
+    consumer_count: int
+
+
+@dataclass(frozen=True, slots=True)
+class AssertExchangeReply:
+    """Result of declaring an exchange."""
+
+    exchange: str
+
+
+ConsumeHandler = Callable[[AbstractIncomingMessage], Coroutine[Any, Any, None]]
+
+
+@final
+class AmqpChannel:
+    """Wraps an aio-pika channel with the operations the SDK needs.
+
+    No schema or validation logic.  Just queue/exchange/publish/consume.
+
+    Usage::
+
+        ch = await conn.create_channel()
+        await ch.assert_queue("orders", arguments={"x-schema": proto_str})
+    """
+
+    def __init__(self, channel: AbstractChannel) -> None:
+        self._ch = channel
+
+    @property
+    def raw(self) -> AbstractChannel:
+        """Exposes the raw aio-pika channel for event listeners."""
+        return self._ch
+
+    async def assert_queue(
+        self,
+        name: str,
+        *,
+        durable: bool = False,
+        arguments: dict[str, Any] | None = None,
+    ) -> AssertQueueReply:
+        """Declares a queue on the broker.
+
+        Usage::
+
+            reply = await ch.assert_queue("orders", durable=True)
+        """
+        queue = await self._ch.declare_queue(
+            name,
+            durable=durable,
+            arguments=arguments,
+        )
+        return AssertQueueReply(
+            queue=queue.name,
+            message_count=queue.declaration_result.message_count or 0,
+            consumer_count=queue.declaration_result.consumer_count or 0,
+        )
+
+    async def assert_exchange(
+        self,
+        name: str,
+        exchange_type: str = "direct",
+        *,
+        durable: bool = False,
+    ) -> AssertExchangeReply:
+        """Declares an exchange on the broker."""
+        exchange = await self._ch.declare_exchange(
+            name,
+            type=aio_pika.ExchangeType(exchange_type),
+            durable=durable,
+        )
+        return AssertExchangeReply(exchange=exchange.name)
+
+    async def bind_queue(self, queue: str, exchange: str, routing_key: str) -> None:
+        """Binds a queue to an exchange with a routing key."""
+        queue_obj = await self._ch.get_queue(queue)
+        _ = await queue_obj.bind(exchange, routing_key)
+
+    async def publish(
+        self,
+        exchange: str,
+        routing_key: str,
+        body: bytes,
+        *,
+        content_type: str = "application/json",
+        persistent: bool = True,
+        headers: dict[str, Any] | None = None,
+    ) -> None:
+        """Publishes a message to an exchange with a routing key."""
+        delivery_mode = (
+            aio_pika.DeliveryMode.PERSISTENT if persistent else aio_pika.DeliveryMode.NOT_PERSISTENT
+        )
+        exchange_obj = await self._ch.get_exchange(exchange)
+        message = aio_pika.Message(
+            body=body,
+            content_type=content_type,
+            delivery_mode=delivery_mode,
+            headers=headers,
+        )
+        _ = await exchange_obj.publish(message, routing_key=routing_key)
+
+    async def send_to_queue(
+        self,
+        queue: str,
+        body: bytes,
+        *,
+        content_type: str = "application/json",
+        persistent: bool = True,
+    ) -> None:
+        """Publishes a message directly to a queue via the default exchange.
+
+        Usage::
+
+            await ch.send_to_queue("orders", b'{"id": 1}')
+        """
+        delivery_mode = (
+            aio_pika.DeliveryMode.PERSISTENT if persistent else aio_pika.DeliveryMode.NOT_PERSISTENT
+        )
+        # WHY default exchange: publishing to "" with routing_key=queue_name
+        # is the AMQP standard way to send directly to a named queue.
+        exchange = await self._ch.get_exchange("")
+        message = aio_pika.Message(
+            body=body,
+            content_type=content_type,
+            delivery_mode=delivery_mode,
+        )
+        _ = await exchange.publish(message, routing_key=queue)
+
+    async def consume(
+        self,
+        queue: str,
+        handler: ConsumeHandler,
+        *,
+        no_ack: bool = True,
+        arguments: dict[str, Any] | None = None,
+    ) -> str:
+        """Subscribes to a queue and returns the consumer tag.
+
+        Usage::
+
+            tag = await ch.consume("orders", on_message, no_ack=True)
+            tag = await ch.consume("orders", on_message, arguments={"x-consumer-schema": proto})
+        """
+        queue_obj = await self._ch.get_queue(queue)
+        tag = await queue_obj.consume(handler, no_ack=no_ack, arguments=arguments)
+        return tag
+
+    async def ack(self, message: AbstractIncomingMessage) -> None:
+        """Acknowledges a message."""
+        await message.ack()
+
+    async def nack(self, message: AbstractIncomingMessage, *, requeue: bool = True) -> None:
+        """Negative-acknowledges a message."""
+        await message.nack(requeue=requeue)
+
+    async def set_qos(self, prefetch_count: int) -> None:
+        """Sets prefetch count (QoS) on the channel."""
+        _ = await self._ch.set_qos(prefetch_count=prefetch_count)
+
+    async def close(self) -> None:
+        """Closes the channel."""
+        await self._ch.close()
+
+
+@final
+class AmqpConnection:
+    """Manages the AMQP connection lifecycle.
+
+    Wraps aio-pika's ``connect_robust`` and exposes a single
+    ``create_channel()`` method.
+
+    Usage::
+
+        conn = await AmqpConnection.connect("amqp://localhost")
+        ch = await conn.create_channel()
+    """
+
+    def __init__(self, connection: aio_pika.abc.AbstractConnection) -> None:
+        self._conn = connection
+
+    @staticmethod
+    async def connect(url: str) -> AmqpConnection:
+        """Opens a new AMQP connection to the given URL."""
+        conn = await aio_pika.connect_robust(url)
+        return AmqpConnection(conn)
+
+    async def create_channel(self) -> AmqpChannel:
+        """Creates a new channel on this connection."""
+        ch = await self._conn.channel()
+        return AmqpChannel(ch)
+
+    async def close(self) -> None:
+        """Closes the underlying TCP connection."""
+        await self._conn.close()