diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..19bcd4f --- /dev/null +++ b/.gitignore @@ -0,0 +1,49 @@ +# Python bytecode / caches +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python + +# Virtualenvs +.venv/ +venv/ +env/ +ENV/ + +# Build / dist artifacts +build/ +dist/ +*.egg-info/ +.eggs/ +wheels/ +pip-wheel-metadata/ + +# Test / type-check / lint caches +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.tox/ +.nox/ +.coverage +.coverage.* +htmlcov/ +coverage.xml + +# Local-only dev / smoke-test scripts (kept out of source control) +scripts/ + +# Editor / IDE +.vscode/ +.idea/ +*.swp +*.swo + +# OS +.DS_Store +Thumbs.db + +# Secrets / env +.env +.env.* +!.env.example diff --git a/cloudrift/cache/base.py b/cloudrift/cache/base.py index eda45e8..e588156 100644 --- a/cloudrift/cache/base.py +++ b/cloudrift/cache/base.py @@ -27,8 +27,25 @@ async def exists(self, key: str) -> bool: """Return ``True`` if *key* exists.""" @abstractmethod - async def expire(self, key: str, seconds: int) -> bool: - """Set a timeout on *key*. Returns ``True`` if the timeout was set.""" + async def expire( + self, + key: str, + seconds: int, + nx: bool = False, + xx: bool = False, + ) -> bool: + """Set a timeout on *key*. Returns ``True`` if the timeout was set. + + Args: + key: Target key. + seconds: TTL in seconds. + nx: Only set the TTL if the key has no existing TTL. + xx: Only set the TTL if the key already has a TTL. + + ``nx`` and ``xx`` are mutually exclusive. Backends that don't support + these flags natively should emulate them via ``ttl(key)`` and document + that the operation is not atomic. + """ @abstractmethod async def ttl(self, key: str) -> int: @@ -54,6 +71,39 @@ async def hgetall(self, key: str) -> dict[bytes, bytes]: async def hdel(self, key: str, *fields: str) -> int: """Delete fields from the hash at *key*. Returns number of fields removed.""" + @abstractmethod + async def sadd(self, key: str, *members: bytes | str) -> int: + """Add one or more *members* to the set at *key*. + + Returns the number of members that were newly added (i.e. not already + present). This "was-new" signal is the foundation of unique-element + deduplication patterns (e.g. DAU/MAU tracking). + """ + + @abstractmethod + async def srem(self, key: str, *members: bytes | str) -> int: + """Remove one or more *members* from the set at *key*. Returns the number removed.""" + + @abstractmethod + async def scard(self, key: str) -> int: + """Return the number of elements in the set at *key*.""" + + @abstractmethod + async def sismember(self, key: str, member: bytes | str) -> bool: + """Return ``True`` if *member* is in the set at *key*.""" + + @abstractmethod + async def smembers(self, key: str) -> "set[bytes]": + """Return all members of the set at *key*.""" + + @abstractmethod + async def sinter(self, *keys: str) -> "set[bytes]": + """Return the members common to all sets at *keys* (set intersection). + + With a single key this is equivalent to :meth:`smembers`. A missing key + is treated as an empty set, so any missing key yields an empty result. + """ + @abstractmethod async def lpush(self, key: str, *values: bytes | str) -> int: """Prepend values to the list at *key*. Returns new list length.""" @@ -102,6 +152,27 @@ async def setex(self, key: str, value: bytes | str, ttl: int) -> None: """Atomic set-with-TTL. Default delegates to ``set(key, value, ttl=ttl)``.""" await self.set(key, value, ttl=ttl) + @asynccontextmanager + async def pipeline(self): + """Batch multiple commands. + + Usage: + async with cache.pipeline() as pipe: + pipe.sadd("k", "m") + pipe.expire("k", 60) + # commands execute on context exit + + The default implementation queues calls and replays them sequentially + on exit — it provides no atomicity and no round-trip savings. Redis + backends override this with a true server-side pipeline. Callers that + depend on atomicity or batching performance must check the backend. + """ + pipe = _SequentialPipeline(self) + try: + yield pipe + finally: + await pipe.execute() + async def health_check(self) -> bool: """Return True if the cache server is reachable.""" try: @@ -116,6 +187,33 @@ async def __aexit__(self, exc_type, exc, tb) -> None: await self.close() +class _SequentialPipeline: + """Default pipeline that records calls and replays them sequentially on execute. + + Provides no atomicity and no round-trip savings — exists so the + ``pipeline()`` API works on every backend. Redis backends bypass this with + a true server-side pipeline. + """ + + def __init__(self, backend: "CacheBackend") -> None: + self._backend = backend + self._ops: list[tuple[str, tuple, dict]] = [] + + def __getattr__(self, name: str): + def queue(*args, **kwargs): + self._ops.append((name, args, kwargs)) + return self + return queue + + async def execute(self) -> list: + results = [] + ops, self._ops = self._ops, [] + for name, args, kwargs in ops: + method = getattr(self._backend, name) + results.append(await method(*args, **kwargs)) + return results + + class _RedisMixin: """Concrete Redis implementation shared by all Redis-backed cache backends. @@ -148,9 +246,55 @@ async def exists(self, key: str) -> bool: except RedisError as e: raise CacheError(str(e)) from e - async def expire(self, key: str, seconds: int) -> bool: + async def expire( + self, + key: str, + seconds: int, + nx: bool = False, + xx: bool = False, + ) -> bool: + if nx and xx: + raise ValueError("expire() flags `nx` and `xx` are mutually exclusive") + try: + return bool(await self._client.expire(key, seconds, nx=nx, xx=xx)) + except RedisError as e: + raise CacheError(str(e)) from e + + async def sadd(self, key: str, *members: bytes | str) -> int: + try: + return await self._client.sadd(key, *members) + except RedisError as e: + raise CacheError(str(e)) from e + + async def srem(self, key: str, *members: bytes | str) -> int: + try: + return await self._client.srem(key, *members) + except RedisError as e: + raise CacheError(str(e)) from e + + async def scard(self, key: str) -> int: + try: + return await self._client.scard(key) + except RedisError as e: + raise CacheError(str(e)) from e + + async def sismember(self, key: str, member: bytes | str) -> bool: + try: + return bool(await self._client.sismember(key, member)) + except RedisError as e: + raise CacheError(str(e)) from e + + async def smembers(self, key: str) -> "set[bytes]": + try: + return await self._client.smembers(key) + except RedisError as e: + raise CacheError(str(e)) from e + + async def sinter(self, *keys: str) -> "set[bytes]": + if not keys: + raise ValueError("sinter() requires at least one key") try: - return bool(await self._client.expire(key, seconds)) + return set(await self._client.sinter(*keys)) except RedisError as e: raise CacheError(str(e)) from e diff --git a/cloudrift/messaging/azure_bus.py b/cloudrift/messaging/azure_bus.py index 1cc2243..36eeb6d 100644 --- a/cloudrift/messaging/azure_bus.py +++ b/cloudrift/messaging/azure_bus.py @@ -235,6 +235,55 @@ async def delete(self, receipt_handle: str) -> None: pass del self._receiver_tokens[rid] + async def dead_letter(self, receipt_handle: str, reason: str) -> None: + entry = self._pending.pop(receipt_handle, None) + if entry is None: + raise MessagingError( + f"No pending message for receipt handle: {receipt_handle!r}. " + "Call receive() first and use the returned receipt_handle." + ) + receiver, message = entry + try: + await receiver.dead_letter_message( + message, reason=reason, error_description=reason + ) + except ResourceNotFoundError as e: + raise QueueNotFoundError(str(e)) from e + except HttpResponseError as e: + raise MessagingError(str(e)) from e + finally: + rid = id(receiver) + if rid in self._receiver_tokens: + _, token_set = self._receiver_tokens[rid] + token_set.discard(receipt_handle) + if not token_set: + try: + await receiver.__aexit__(None, None, None) + except Exception: + pass + del self._receiver_tokens[rid] + + async def get_queue_depth(self) -> int: + # Message counts live on the management plane, not the data plane. The + # async admin client lives at azure.servicebus.aio.management (note: NOT + # azure.servicebus.management.aio) and accepts the same async credential. + from azure.servicebus.aio.management import ServiceBusAdministrationClient + + if self._connection_string: + admin = ServiceBusAdministrationClient.from_connection_string( + self._connection_string + ) + else: + admin = ServiceBusAdministrationClient(self._namespace, credential=self._credential) + try: + async with admin: + props = await admin.get_queue_runtime_properties(self.queue_name) + return props.active_message_count + except ResourceNotFoundError as e: + raise QueueNotFoundError(f"Queue not found: {self.queue_name}") from e + except HttpResponseError as e: + raise MessagingError(str(e)) from e + async def health_check(self) -> bool: try: client = await self._ensure() diff --git a/cloudrift/messaging/base.py b/cloudrift/messaging/base.py index a028f75..91a44bc 100644 --- a/cloudrift/messaging/base.py +++ b/cloudrift/messaging/base.py @@ -33,6 +33,28 @@ async def receive(self, max_messages: int = 1, wait_time: int = 0) -> list[Messa async def delete(self, receipt_handle: str) -> None: """Delete/acknowledge a message by its receipt handle.""" + @abstractmethod + async def dead_letter(self, receipt_handle: str, reason: str) -> None: + """Move a received message to the dead-letter queue and acknowledge it. + + Args: + receipt_handle: The receipt handle from a previously received message. + reason: A human-readable reason recorded with the dead-lettered message. + + Azure Service Bus implements this natively via ``dead_letter_message``. + SQS has no native per-message dead-letter API, so backends emulate it by + sending the message body to a configured dead-letter queue and then + deleting the original from the source queue. + """ + + @abstractmethod + async def get_queue_depth(self) -> int: + """Return the approximate number of messages waiting in the queue. + + This is an estimate: cloud queues report it asynchronously and it may + lag in-flight (received-but-not-yet-deleted) messages. + """ + @abstractmethod async def purge(self) -> None: """Delete all messages in the queue.""" diff --git a/cloudrift/messaging/sqs.py b/cloudrift/messaging/sqs.py index d0e3f72..49c19ad 100644 --- a/cloudrift/messaging/sqs.py +++ b/cloudrift/messaging/sqs.py @@ -28,6 +28,7 @@ def __init__( session: aioboto3.Session, *, endpoint_url: str | None = None, + dlq_url: str | None = None, max_pool_connections: int = 50, connect_timeout: float = 10.0, read_timeout: float = 60.0, @@ -36,6 +37,9 @@ def __init__( self.queue_url = queue_url self._session = session self._endpoint_url = endpoint_url + # Explicit DLQ URL; if None it is resolved lazily from the source queue's + # RedrivePolicy the first time dead_letter() is called. + self._dlq_url = dlq_url self._config = Config( max_pool_connections=max_pool_connections, connect_timeout=connect_timeout, @@ -45,6 +49,10 @@ def __init__( self._client_cm = None self._client = None self._lock = asyncio.Lock() + # receipt_handle → raw message body (JSON string), retained between + # receive() and delete()/dead_letter() so emulated dead-lettering can + # re-send the original payload to the DLQ. + self._pending: dict[str, str] = {} # ------------------------------------------------------------------ # Factory constructors @@ -158,15 +166,18 @@ async def receive(self, max_messages: int = 1, wait_time: int = 0) -> list[Messa WaitTimeSeconds=wait_time, AttributeNames=["All"], ) - return [ - Message( - id=m["MessageId"], - body=json.loads(m["Body"]), - receipt_handle=m["ReceiptHandle"], - attributes=m.get("Attributes", {}), + messages = [] + for m in response.get("Messages", []): + self._pending[m["ReceiptHandle"]] = m["Body"] + messages.append( + Message( + id=m["MessageId"], + body=json.loads(m["Body"]), + receipt_handle=m["ReceiptHandle"], + attributes=m.get("Attributes", {}), + ) ) - for m in response.get("Messages", []) - ] + return messages except ClientError as e: self._raise(e) @@ -176,6 +187,75 @@ async def delete(self, receipt_handle: str) -> None: await client.delete_message(QueueUrl=self.queue_url, ReceiptHandle=receipt_handle) except ClientError as e: self._raise(e) + finally: + self._pending.pop(receipt_handle, None) + + async def dead_letter(self, receipt_handle: str, reason: str) -> None: + """Emulated dead-letter for SQS: sends to DLQ then deletes from source. + + Warning: these are two separate API calls with no cross-queue + transaction. If the process dies between them (or the DLQ send + succeeds but the delete fails), the message may appear in both + queues (double-processed) or in neither (lost). For strict + dead-lettering, prefer the native SQS redrive policy and let the + service move the message after maxReceiveCount is reached. + """ + client = await self._ensure() + body = self._pending.get(receipt_handle) + if body is None: + raise MessagingError( + f"No pending message for receipt handle: {receipt_handle!r}. " + "Call receive() first and use the returned receipt_handle." + ) + dlq_url = await self._resolve_dlq_url(client) + try: + await client.send_message( + QueueUrl=dlq_url, + MessageBody=body, + MessageAttributes={ + "DeadLetterReason": {"DataType": "String", "StringValue": reason} + }, + ) + await client.delete_message(QueueUrl=self.queue_url, ReceiptHandle=receipt_handle) + except ClientError as e: + self._raise(e) + finally: + self._pending.pop(receipt_handle, None) + + async def get_queue_depth(self) -> int: + client = await self._ensure() + try: + response = await client.get_queue_attributes( + QueueUrl=self.queue_url, + AttributeNames=["ApproximateNumberOfMessages"], + ) + return int(response["Attributes"]["ApproximateNumberOfMessages"]) + except ClientError as e: + self._raise(e) + + async def _resolve_dlq_url(self, client) -> str: + """Return the configured DLQ URL, deriving it from RedrivePolicy if needed.""" + if self._dlq_url is not None: + return self._dlq_url + try: + response = await client.get_queue_attributes( + QueueUrl=self.queue_url, AttributeNames=["RedrivePolicy"] + ) + except ClientError as e: + self._raise(e) + redrive = response.get("Attributes", {}).get("RedrivePolicy") + if not redrive: + raise MessagingError( + f"No dead-letter queue configured for {self.queue_url}. Pass dlq_url= " + "when constructing the backend, or set a RedrivePolicy on the queue." + ) + target_arn = json.loads(redrive)["deadLetterTargetArn"] + dlq_name = target_arn.rsplit(":", 1)[-1] + try: + self._dlq_url = (await client.get_queue_url(QueueName=dlq_name))["QueueUrl"] + except ClientError as e: + self._raise(e) + return self._dlq_url async def purge(self) -> None: client = await self._ensure() diff --git a/tests/test_cache.py b/tests/test_cache.py index 08c36c3..8297cc3 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -117,6 +117,52 @@ async def test_hdel(cache): assert await cache.hget("h3", "x") is None +# --------------------------------------------------------------------------- +# Set commands +# --------------------------------------------------------------------------- + +async def test_sadd_returns_newly_added(cache): + assert await cache.sadd("s", b"a", b"b") == 2 + # adding an existing member does not count as new + assert await cache.sadd("s", b"a", b"c") == 1 + + +async def test_smembers_and_scard(cache): + await cache.sadd("s2", b"x", b"y", b"z") + assert await cache.smembers("s2") == {b"x", b"y", b"z"} + assert await cache.scard("s2") == 3 + + +async def test_srem_and_sismember(cache): + await cache.sadd("s3", b"a", b"b") + assert await cache.sismember("s3", b"a") + assert await cache.srem("s3", b"a") == 1 + assert not await cache.sismember("s3", b"a") + + +async def test_sinter(cache): + await cache.sadd("set_a", b"1", b"2", b"3") + await cache.sadd("set_b", b"2", b"3", b"4") + await cache.sadd("set_c", b"3", b"4", b"5") + assert await cache.sinter("set_a", "set_b") == {b"2", b"3"} + assert await cache.sinter("set_a", "set_b", "set_c") == {b"3"} + + +async def test_sinter_single_key_equals_smembers(cache): + await cache.sadd("only", b"a", b"b") + assert await cache.sinter("only") == await cache.smembers("only") + + +async def test_sinter_missing_key_is_empty(cache): + await cache.sadd("present", b"a", b"b") + assert await cache.sinter("present", "absent") == set() + + +async def test_sinter_no_keys_raises(cache): + with pytest.raises(ValueError, match="at least one key"): + await cache.sinter() + + # --------------------------------------------------------------------------- # List commands # --------------------------------------------------------------------------- diff --git a/tests/test_messaging.py b/tests/test_messaging.py index 2c5f6f5..8bfd9f9 100644 --- a/tests/test_messaging.py +++ b/tests/test_messaging.py @@ -1,3 +1,5 @@ +import json + import boto3 import pytest from moto.server import ThreadedMotoServer @@ -6,6 +8,7 @@ REGION = "us-east-1" QUEUE_NAME = "test-queue" +DLQ_NAME = "test-queue-dlq" @pytest.fixture(scope="module") @@ -18,15 +21,32 @@ def moto_server(): @pytest.fixture -async def sqs_backend(moto_server): - sqs = boto3.client( +def sqs_client(moto_server): + return boto3.client( "sqs", region_name=REGION, endpoint_url=moto_server, aws_access_key_id="test", aws_secret_access_key="test", ) - queue_url = sqs.create_queue(QueueName=QUEUE_NAME)["QueueUrl"] + + +@pytest.fixture +async def sqs_backend(moto_server, sqs_client): + # Create a DLQ and wire the source queue to it via RedrivePolicy so + # dead_letter() can resolve the target from the queue itself. + dlq_url = sqs_client.create_queue(QueueName=DLQ_NAME)["QueueUrl"] + dlq_arn = sqs_client.get_queue_attributes( + QueueUrl=dlq_url, AttributeNames=["QueueArn"] + )["Attributes"]["QueueArn"] + queue_url = sqs_client.create_queue( + QueueName=QUEUE_NAME, + Attributes={ + "RedrivePolicy": json.dumps( + {"deadLetterTargetArn": dlq_arn, "maxReceiveCount": "5"} + ) + }, + )["QueueUrl"] backend = get_queue( "sqs", queue_url=queue_url, @@ -35,9 +55,11 @@ async def sqs_backend(moto_server): region=REGION, endpoint_url=moto_server, ) + backend._dlq_test_url = dlq_url # expose for assertions yield backend await backend.close() - sqs.delete_queue(QueueUrl=queue_url) + sqs_client.delete_queue(QueueUrl=queue_url) + sqs_client.delete_queue(QueueUrl=dlq_url) async def test_send_and_receive(sqs_backend): @@ -79,3 +101,38 @@ def test_invalid_provider(): async def test_health_check(sqs_backend): result = await sqs_backend.health_check() assert result is True + + +async def test_get_queue_depth(sqs_backend): + assert await sqs_backend.get_queue_depth() == 0 + await sqs_backend.send({"a": 1}) + await sqs_backend.send({"b": 2}) + assert await sqs_backend.get_queue_depth() == 2 + + +async def test_dead_letter(sqs_backend, sqs_client): + await sqs_backend.send({"task": "boom"}) + messages = await sqs_backend.receive(max_messages=1) + assert messages + + await sqs_backend.dead_letter(messages[0].receipt_handle, reason="poison message") + + # Removed from the source queue... + assert await sqs_backend.receive(max_messages=1) == [] + # ...and delivered to the DLQ with the reason recorded. + dlq = sqs_client.receive_message( + QueueUrl=sqs_backend._dlq_test_url, + MaxNumberOfMessages=1, + MessageAttributeNames=["All"], + ) + received = dlq["Messages"][0] + assert json.loads(received["Body"]) == {"task": "boom"} + assert ( + received["MessageAttributes"]["DeadLetterReason"]["StringValue"] + == "poison message" + ) + + +async def test_dead_letter_unknown_handle(sqs_backend): + with pytest.raises(Exception, match="No pending message"): + await sqs_backend.dead_letter("bogus-handle", reason="x") diff --git a/tests/test_messaging_azure.py b/tests/test_messaging_azure.py new file mode 100644 index 0000000..903d1b8 --- /dev/null +++ b/tests/test_messaging_azure.py @@ -0,0 +1,113 @@ +"""Unit tests for the Azure Service Bus backend. + +Azure Service Bus has no local emulator (unlike SQS via moto), so these tests +mock the data-plane receiver and the management-plane admin client. They verify +the cloudrift glue logic — receipt-handle bookkeeping, dead-letter dispatch, and +queue-depth extraction — not the Azure SDK itself. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# Azure is an optional dependency and has no local emulator; skip the whole +# module when the SDK isn't installed (e.g. the default `.[dev]` CI install). +pytest.importorskip("azure.servicebus") + +from cloudrift.core.exceptions import MessagingError # noqa: E402 +from cloudrift.messaging.azure_bus import AzureServiceBusBackend # noqa: E402 + + +@pytest.fixture +def backend(): + return AzureServiceBusBackend.from_connection_string( + connection_string="Endpoint=sb://fake.servicebus.windows.net/;" + "SharedAccessKeyName=k;SharedAccessKey=secret", + queue_name="test-queue", + ) + + +# --------------------------------------------------------------------------- +# dead_letter (native, data plane) +# --------------------------------------------------------------------------- + +async def test_dead_letter_calls_native_api(backend): + receiver = AsyncMock() + message = MagicMock() + handle = "lock-token-1" + # Simulate the state receive() would have left behind. + backend._pending[handle] = (receiver, message) + backend._receiver_tokens[id(receiver)] = (receiver, {handle}) + + await backend.dead_letter(handle, reason="poison message") + + receiver.dead_letter_message.assert_awaited_once_with( + message, reason="poison message", error_description="poison message" + ) + # Bookkeeping is cleared and the now-empty receiver is closed. + assert handle not in backend._pending + assert id(receiver) not in backend._receiver_tokens + receiver.__aexit__.assert_awaited() + + +async def test_dead_letter_unknown_handle_raises(backend): + with pytest.raises(MessagingError, match="No pending message"): + await backend.dead_letter("nope", reason="x") + + +async def test_dead_letter_keeps_receiver_open_for_other_messages(backend): + receiver = AsyncMock() + msg_a, msg_b = MagicMock(), MagicMock() + backend._pending["a"] = (receiver, msg_a) + backend._pending["b"] = (receiver, msg_b) + backend._receiver_tokens[id(receiver)] = (receiver, {"a", "b"}) + + await backend.dead_letter("a", reason="bad") + + # "b" is still pending, so the shared receiver must stay open. + assert "b" in backend._pending + assert id(receiver) in backend._receiver_tokens + receiver.__aexit__.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# get_queue_depth (management plane) +# --------------------------------------------------------------------------- + +async def test_get_queue_depth(backend): + props = MagicMock(active_message_count=7) + admin = AsyncMock() + admin.get_queue_runtime_properties = AsyncMock(return_value=props) + admin.__aenter__ = AsyncMock(return_value=admin) + admin.__aexit__ = AsyncMock(return_value=None) + + with patch( + "azure.servicebus.aio.management.ServiceBusAdministrationClient" + ) as Admin: + Admin.from_connection_string.return_value = admin + depth = await backend.get_queue_depth() + + assert depth == 7 + admin.get_queue_runtime_properties.assert_awaited_once_with("test-queue") + + +async def test_get_queue_depth_uses_namespace_and_credential(): + credential = object() + backend = AzureServiceBusBackend( + "q", fully_qualified_namespace="ns.servicebus.windows.net", credential=credential + ) + props = MagicMock(active_message_count=0) + admin = AsyncMock() + admin.get_queue_runtime_properties = AsyncMock(return_value=props) + admin.__aenter__ = AsyncMock(return_value=admin) + admin.__aexit__ = AsyncMock(return_value=None) + + with patch( + "azure.servicebus.aio.management.ServiceBusAdministrationClient" + ) as Admin: + Admin.return_value = admin + depth = await backend.get_queue_depth() + + assert depth == 0 + # Namespace path constructs the client directly with the async credential. + Admin.assert_called_once_with("ns.servicebus.windows.net", credential=credential)