From 47656502de63925fe9404e700b8d035f6626dbf6 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Fri, 14 Aug 2026 16:04:22 +0100 Subject: [PATCH] feat: add asyncio client in orbita.aio Mirror the synchronous API method-for-method on grpc.aio, sharing the result dataclasses, validation helpers, and error types so an application can move between the two without translating anything. The transport grows an aio channel opener and an asyncio retry loop with the same read-only-retry and ambiguous-mutation posture. The one structural difference is connection setup: __init__ cannot await, so the data client records configuration and performs the limits handshake in connect(), with async-with connecting on entry. Admin has no handshake, so its channel still opens eagerly. Unit tests reuse the in-process sync gRPC server through pytest-asyncio, and the integration suite gains an async pass verified against the pinned server image. --- README.md | 28 ++- pyproject.toml | 2 + src/orbita/__init__.py | 6 +- src/orbita/_transport.py | 50 +++- src/orbita/aio.py | 490 ++++++++++++++++++++++++++++++++++++++ tests/test_aio.py | 174 ++++++++++++++ tests/test_integration.py | 77 ++++++ 7 files changed, 822 insertions(+), 5 deletions(-) create mode 100644 src/orbita/aio.py create mode 100644 tests/test_aio.py diff --git a/README.md b/README.md index 0c5227b..183def2 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,9 @@ # Orbita Python -`orbita-python` is a synchronous Python client for Orbita. It connects to one -load-balanced endpoint and leaves partition routing to the server, which is the -same complete non-transactional posture supported by Orbita's generated stubs. +`orbita-python` is a Python client for Orbita, with synchronous and asyncio +APIs. It connects to one load-balanced endpoint and leaves partition routing to +the server, which is the same complete non-transactional posture supported by +Orbita's generated stubs. The package is pinned to the wire contract at `orbita 0.2.0-dev@cb11cf03575e2c6e370fb83500038de45bd55e69`. Generated code @@ -36,6 +37,27 @@ with orbita.Client( print(f"already exists at version {created.current_version}") ``` +The asyncio API lives in `orbita.aio` and mirrors the synchronous one +method-for-method, sharing the same result and error types. The one difference +is that connection setup is explicit, because `__init__` cannot await: use +`async with`, which connects on entry, or call `await client.connect()`. + +```python +from orbita import aio + +async with aio.Client( + "http://127.0.0.1:7100", + credential="tenant-secret", +) as client: + catalog = client.keyspace("default") + entry = await catalog.get_entry(b"catalog/current") + async for page in catalog.pages(b"catalog/", include_values=True): + ... +``` + +`orbita.aio.AdminClient` is the asynchronous counterpart of +`orbita.admin.AdminClient`. + Use an explicit `http://` endpoint for plaintext h2c. Use `https://` for TLS; pass `tls_credentials=grpc.ssl_channel_credentials(...)` when the endpoint needs custom trust roots or a client certificate. diff --git a/pyproject.toml b/pyproject.toml index c112b4a..92062c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dev = [ "grpcio-tools==1.75.1", "mypy>=1.14,<2", "pytest>=8.3,<9", + "pytest-asyncio>=0.25,<2", "pytest-cov>=6,<8", "ruff>=0.9,<1", "types-protobuf>=6.31.1,<7", @@ -51,6 +52,7 @@ packages = ["src/orbita"] [tool.pytest.ini_options] testpaths = ["tests"] addopts = "-q -ra" +asyncio_mode = "auto" markers = ["integration: requires ORBITA_ENDPOINT and ORBITA_ROOT_CREDENTIAL"] [tool.ruff] diff --git a/src/orbita/__init__.py b/src/orbita/__init__.py index 03bd09d..7f2268d 100644 --- a/src/orbita/__init__.py +++ b/src/orbita/__init__.py @@ -1,4 +1,8 @@ -"""Synchronous Python client for Orbita.""" +"""Python client for Orbita. + +The synchronous API is exported here; the asyncio API lives in ``orbita.aio`` +and is not imported eagerly so synchronous users never touch ``grpc.aio``. +""" from orbita._transport import AmbiguousMutationError, CredentialError, RetryPolicy from orbita.client import ( diff --git a/src/orbita/_transport.py b/src/orbita/_transport.py index 4ab1a0d..5ed8315 100644 --- a/src/orbita/_transport.py +++ b/src/orbita/_transport.py @@ -2,12 +2,14 @@ from __future__ import annotations +import asyncio import time from dataclasses import dataclass -from typing import Any, Callable, Optional, Sequence, Tuple, TypeVar, cast +from typing import Any, Awaitable, Callable, Optional, Sequence, Tuple, TypeVar, cast from urllib.parse import SplitResult, urlsplit import grpc +import grpc.aio T = TypeVar("T") Metadata = Sequence[Tuple[str, str]] @@ -102,6 +104,21 @@ def open_channel(self, max_message_bytes: int = 0) -> grpc.Channel: credentials = self.tls_credentials or grpc.ssl_channel_credentials() return grpc.secure_channel(target, credentials, options=options) + def open_aio_channel(self, max_message_bytes: int = 0) -> grpc.aio.Channel: + options = list(self.channel_options) + if max_message_bytes: + options.extend( + ( + ("grpc.max_receive_message_length", max_message_bytes), + ("grpc.max_send_message_length", max_message_bytes), + ) + ) + target = self.endpoint.netloc + if self.endpoint.scheme == "http": + return grpc.aio.insecure_channel(target, options=options) + credentials = self.tls_credentials or grpc.ssl_channel_credentials() + return grpc.aio.secure_channel(target, credentials, options=options) + def metadata(self) -> Metadata: if self.token_source is None: return () @@ -147,6 +164,37 @@ def retry_read( raise RuntimeError("orbita: retry loop exhausted") +async def retry_read_aio( + call: Callable[..., Awaitable[T]], + request: Any, + config: TransportConfig, + timeout: Optional[float], +) -> T: + deadline = None if timeout is None else time.monotonic() + _validate_timeout(timeout) + backoff = config.retry_policy.initial_backoff + for attempt in range(1, config.retry_policy.max_attempts + 1): + metadata = config.metadata() + remaining = _remaining(deadline) + if remaining == 0: + raise TimeoutError("orbita: operation timed out") + try: + return await call(request, timeout=remaining, metadata=metadata) + except grpc.RpcError as error: + if error.code() != grpc.StatusCode.UNAVAILABLE: + raise + if attempt == config.retry_policy.max_attempts: + raise + sleep_for = backoff + if deadline is not None: + sleep_for = min(sleep_for, max(0.0, deadline - time.monotonic())) + if sleep_for > 0: + await asyncio.sleep(sleep_for) + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError("orbita: operation timed out") + backoff = min(backoff * 2, config.retry_policy.max_backoff) + raise RuntimeError("orbita: retry loop exhausted") + + def mutation_error(operation: str, error: grpc.RpcError, *, aborted: bool = False) -> Exception: ambiguous = { grpc.StatusCode.UNAVAILABLE, diff --git a/src/orbita/aio.py b/src/orbita/aio.py new file mode 100644 index 0000000..6257d60 --- /dev/null +++ b/src/orbita/aio.py @@ -0,0 +1,490 @@ +"""Asynchronous client for Orbita, mirroring the synchronous API on grpc.aio. + +The result and error types are shared with the synchronous client so an +application can move between the two without translating anything: the same +``Entry``, ``SetResult``, and ``AmbiguousMutationError`` mean the same things. +The one structural difference is connection setup. ``__init__`` cannot await, +so the constructor only records configuration and the limits handshake happens +in ``connect`` (or ``async with``), which is why the async client is explicit +about connecting where the synchronous one is not. +""" + +from __future__ import annotations + +from datetime import timedelta +from typing import Any, AsyncIterator, Callable, Optional, Sequence, Tuple, TypeVar, cast + +import grpc +import grpc.aio +from google.protobuf.message import Message + +from orbita._transport import ( + RetryPolicy, + TokenSource, + TransportConfig, + mutation_error, + retry_read_aio, +) +from orbita.client import ( + DeleteResult, + Entry, + Page, + Readiness, + ReadinessCondition, + SetResult, + _bytes, + _timeout, + _ttl_millis, + _uint32, + _uint64, +) +from orbita.v1 import ( + admin_pb2, + admin_pb2_grpc, + health_pb2, + health_pb2_grpc, + kv_pb2, + kv_pb2_grpc, +) + +MessageT = TypeVar("MessageT", bound=Message) +_ADMIN_MAX_MESSAGE_BYTES = 64 << 20 + + +class Client: + """An asynchronous connection to Orbita's data and health services. + + The constructor performs no I/O so it can run outside an event loop; + ``connect`` performs the limits handshake and opens the channel. Prefer + ``async with Client(...)``, which connects on entry and closes on exit. + """ + + def __init__( + self, + endpoint: str, + *, + credential: Optional[str] = None, + token_source: Optional[TokenSource] = None, + retry_policy: Optional[RetryPolicy] = None, + tls_credentials: Optional[grpc.ChannelCredentials] = None, + channel_options: Optional[Sequence[Tuple[str, Any]]] = None, + timeout: Optional[float] = None, + ) -> None: + self._config = TransportConfig.build( + endpoint, + credential=credential, + token_source=token_source, + retry_policy=retry_policy, + tls_credentials=tls_credentials, + channel_options=channel_options, + ) + self._connect_timeout = _timeout(timeout) + self._channel: Optional[grpc.aio.Channel] = None + self._kv: Any = None + self._health: Any = None + + async def connect(self) -> Client: + """Discover server limits and open the channel; idempotent.""" + + if self._channel is not None: + return self + bootstrap = self._config.open_aio_channel() + try: + limits = await retry_read_aio( + kv_pb2_grpc.KvStub(bootstrap).GetLimits, # type: ignore[no-untyped-call] + kv_pb2.GetLimitsRequest(), + self._config, + self._connect_timeout, + ) + finally: + await bootstrap.close(grace=None) + if limits.max_message_bytes < 1 or limits.max_message_bytes > 2**31 - 1: + raise ValueError( + f"orbita: invalid advertised max message size {limits.max_message_bytes}" + ) + self._channel = self._config.open_aio_channel(limits.max_message_bytes) + self._kv = kv_pb2_grpc.KvStub(self._channel) # type: ignore[no-untyped-call] + self._health = health_pb2_grpc.HealthStub(self._channel) # type: ignore[no-untyped-call] + return self + + async def close(self) -> None: + """Close the underlying gRPC channel.""" + + if self._channel is not None: + channel, self._channel = self._channel, None + self._kv = None + self._health = None + await channel.close(grace=None) + + async def __aenter__(self) -> Client: + return await self.connect() + + async def __aexit__(self, *_: object) -> None: + await self.close() + + def keyspace(self, name: str) -> Keyspace: + """Create access relative to one immutable keyspace name.""" + + return Keyspace(self, name) + + async def readiness(self, *, timeout: Optional[float] = None) -> Readiness: + """Check whether the connected node is ready to serve.""" + + response = await retry_read_aio( + self._require_connected(self._health).CheckReadiness, + health_pb2.CheckReadinessRequest(), + self._config, + timeout, + ) + return Readiness( + ready=response.ready, + conditions=tuple( + ReadinessCondition(name=item.name, met=item.met) for item in response.conditions + ), + ) + + def _require_connected(self, stub: Any) -> Any: + if stub is None: + raise RuntimeError("orbita: client is not connected; call connect() first") + return stub + + +class Keyspace: + """Byte-valued asynchronous access relative to one immutable keyspace name.""" + + def __init__(self, client: Client, name: str) -> None: + self._client = client + self._name = name + + @property + def name(self) -> str: + """The immutable keyspace name.""" + + return self._name + + async def get(self, key: bytes, *, timeout: Optional[float] = None) -> Optional[bytes]: + """Read a value, returning None when the key is absent.""" + + entry = await self.get_entry(key, timeout=timeout) + return None if entry is None else entry.value + + async def get_entry(self, key: bytes, *, timeout: Optional[float] = None) -> Optional[Entry]: + """Read a value with its version and expiry metadata.""" + + key = _bytes("key", key) + response = await retry_read_aio( + self._kv().Get, + kv_pb2.GetRequest(keyspace=self._name, key=key), + self._client._config, + timeout, + ) + if not response.found: + return None + expires = response.expires_at_millis if response.HasField("expires_at_millis") else None + return Entry( + key=key, + value=bytes(response.value), + version=response.version, + expires_at_millis=expires, + ) + + async def set( + self, + key: bytes, + value: bytes, + *, + if_not_present: bool = False, + if_version: Optional[int] = None, + ttl: Optional[timedelta] = None, + timeout: Optional[float] = None, + ) -> SetResult: + """Write one value without retrying an ambiguous mutation.""" + + if if_not_present and if_version is not None: + raise ValueError("orbita: set accepts only one condition") + request = kv_pb2.SetRequest( + keyspace=self._name, + key=_bytes("key", key), + value=_bytes("value", value), + ) + if if_not_present: + request.condition.if_not_present = True + elif if_version is not None: + request.condition.if_version = _uint64("if_version", if_version) + if ttl is not None: + request.ttl_millis = _ttl_millis(ttl) + try: + response = await self._kv().Set( + request, + timeout=_timeout(timeout), + metadata=self._client._config.metadata(), + ) + except grpc.RpcError as error: + raise mutation_error("set", error) from error + current = response.current_version if response.HasField("current_version") else None + return SetResult(response.applied, response.version, current) + + async def delete( + self, + key: bytes, + *, + if_version: Optional[int] = None, + timeout: Optional[float] = None, + ) -> DeleteResult: + """Delete one key without retrying an ambiguous mutation.""" + + request = kv_pb2.DeleteRequest(keyspace=self._name, key=_bytes("key", key)) + if if_version is not None: + request.condition.if_version = _uint64("if_version", if_version) + try: + response = await self._kv().Delete( + request, + timeout=_timeout(timeout), + metadata=self._client._config.metadata(), + ) + except grpc.RpcError as error: + raise mutation_error("delete", error) from error + current = response.current_version if response.HasField("current_version") else None + return DeleteResult(response.applied, response.existed, current) + + async def list_page( + self, + prefix: bytes = b"", + *, + cursor: bytes = b"", + limit: int = 0, + include_values: bool = False, + timeout: Optional[float] = None, + ) -> Page: + """Read one ordered page of a prefix scan.""" + + response = await retry_read_aio( + self._kv().List, + kv_pb2.ListRequest( + keyspace=self._name, + prefix=_bytes("prefix", prefix), + cursor=_bytes("cursor", cursor), + limit=_uint32("limit", limit), + include_values=include_values, + ), + self._client._config, + timeout, + ) + entries = tuple( + Entry( + key=bytes(item.key), + value=bytes(item.value) if include_values else None, + version=item.version, + expires_at_millis=( + item.expires_at_millis if item.HasField("expires_at_millis") else None + ), + ) + for item in response.entries + ) + return Page(entries=entries, next_cursor=bytes(response.next_cursor)) + + async def pages( + self, + prefix: bytes = b"", + *, + cursor: bytes = b"", + limit: int = 0, + include_values: bool = False, + timeout: Optional[float] = None, + ) -> AsyncIterator[Page]: + """Iterate pages until the server returns an empty cursor.""" + + next_cursor = _bytes("cursor", cursor) + while True: + page = await self.list_page( + prefix, + cursor=next_cursor, + limit=limit, + include_values=include_values, + timeout=timeout, + ) + yield page + if not page.next_cursor: + return + next_cursor = page.next_cursor + + def _kv(self) -> Any: + return self._client._require_connected(self._client._kv) + + +class AdminClient: + """An asynchronous connection to Orbita's operator-facing API. + + Admin has no limits handshake, so unlike ``Client`` the channel opens + eagerly and ``connect`` does not exist; ``async with`` only manages close. + """ + + def __init__( + self, + endpoint: str, + *, + root_credential: Optional[str] = None, + token_source: Optional[TokenSource] = None, + retry_policy: Optional[RetryPolicy] = None, + tls_credentials: Optional[grpc.ChannelCredentials] = None, + channel_options: Optional[Sequence[Tuple[str, Any]]] = None, + ) -> None: + self._config = TransportConfig.build( + endpoint, + credential=root_credential, + token_source=token_source, + retry_policy=retry_policy, + tls_credentials=tls_credentials, + channel_options=channel_options, + ) + self._channel = self._config.open_aio_channel(_ADMIN_MAX_MESSAGE_BYTES) + self._admin = admin_pb2_grpc.AdminStub(self._channel) # type: ignore[no-untyped-call] + + async def close(self) -> None: + """Close the underlying gRPC channel.""" + + await self._channel.close(grace=None) + + async def __aenter__(self) -> AdminClient: + return self + + async def __aexit__(self, *_: object) -> None: + await self.close() + + async def create_keyspace( + self, request: admin_pb2.CreateKeyspaceRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.Keyspace: + return cast( + admin_pb2.Keyspace, + await self._mutation("create keyspace", self._admin.CreateKeyspace, request, timeout), + ) + + async def update_keyspace( + self, request: admin_pb2.UpdateKeyspaceRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.Keyspace: + return cast( + admin_pb2.Keyspace, + await self._mutation("update keyspace", self._admin.UpdateKeyspace, request, timeout), + ) + + async def delete_keyspace( + self, request: admin_pb2.DeleteKeyspaceRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.DeleteKeyspaceResponse: + return cast( + admin_pb2.DeleteKeyspaceResponse, + await self._mutation("delete keyspace", self._admin.DeleteKeyspace, request, timeout), + ) + + async def list_keyspaces( + self, request: admin_pb2.ListKeyspacesRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.ListKeyspacesResponse: + return cast( + admin_pb2.ListKeyspacesResponse, + await self._read(self._admin.ListKeyspaces, request, timeout), + ) + + async def create_credential( + self, request: admin_pb2.CreateCredentialRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.CreateCredentialResponse: + return cast( + admin_pb2.CreateCredentialResponse, + await self._mutation( + "create credential", self._admin.CreateCredential, request, timeout, aborted=True + ), + ) + + async def revoke_credential( + self, request: admin_pb2.RevokeCredentialRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.RevokeCredentialResponse: + return cast( + admin_pb2.RevokeCredentialResponse, + await self._mutation( + "revoke credential", self._admin.RevokeCredential, request, timeout + ), + ) + + async def describe_cluster( + self, request: admin_pb2.DescribeClusterRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.DescribeClusterResponse: + return cast( + admin_pb2.DescribeClusterResponse, + await self._read(self._admin.DescribeCluster, request, timeout), + ) + + async def split_partition( + self, request: admin_pb2.SplitPartitionRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.SplitPartitionResponse: + return cast( + admin_pb2.SplitPartitionResponse, + await self._mutation("split partition", self._admin.SplitPartition, request, timeout), + ) + + async def merge_partitions( + self, request: admin_pb2.MergePartitionsRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.MergePartitionsResponse: + return cast( + admin_pb2.MergePartitionsResponse, + await self._mutation("merge partitions", self._admin.MergePartitions, request, timeout), + ) + + async def transfer_ownership( + self, request: admin_pb2.TransferOwnershipRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.TransferOwnershipResponse: + return cast( + admin_pb2.TransferOwnershipResponse, + await self._mutation( + "transfer ownership", self._admin.TransferOwnership, request, timeout + ), + ) + + async def finalize_upgrade( + self, request: admin_pb2.FinalizeUpgradeRequest, *, timeout: Optional[float] = None + ) -> admin_pb2.FinalizeUpgradeResponse: + return cast( + admin_pb2.FinalizeUpgradeResponse, + await self._mutation("finalize upgrade", self._admin.FinalizeUpgrade, request, timeout), + ) + + async def _read( + self, call: Callable[..., Any], request: Message, timeout: Optional[float] + ) -> Message: + response = await retry_read_aio(call, _clone(request), self._config, timeout) + return _clone(cast(Message, response)) + + async def _mutation( + self, + operation: str, + call: Callable[..., Any], + request: Message, + timeout: Optional[float], + *, + aborted: bool = False, + ) -> Message: + try: + response = await call( + _clone(request), + timeout=_timeout(timeout), + metadata=self._config.metadata(), + ) + except grpc.RpcError as error: + raise mutation_error(operation, error, aborted=aborted) from error + return _clone(cast(Message, response)) + + +def _clone(message: MessageT) -> MessageT: + clone = message.__class__() + clone.CopyFrom(message) + return clone + + +__all__ = [ + "AdminClient", + "Client", + "DeleteResult", + "Entry", + "Keyspace", + "Page", + "Readiness", + "ReadinessCondition", + "SetResult", +] diff --git a/tests/test_aio.py b/tests/test_aio.py new file mode 100644 index 0000000..ac79f94 --- /dev/null +++ b/tests/test_aio.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +from datetime import timedelta + +import grpc +import pytest + +import orbita +from orbita import aio +from orbita.v1 import admin_pb2 + + +async def test_async_client_discovers_limits_and_resolves_a_token_for_each_rpc(services) -> None: + tokens: list[str] = [] + + def token_source() -> str: + token = f"token-{len(tokens) + 1}" + tokens.append(token) + return token + + async with aio.Client(services.endpoint, token_source=token_source) as client: + assert await client.keyspace("default").get(b"key") == b"value" + + assert tokens == ["token-1", "token-2"] + assert services.kv.metadata[0]["authorization"] == "Bearer token-1" + assert services.kv.metadata[1]["authorization"] == "Bearer token-2" + + +async def test_async_get_distinguishes_missing_and_empty_values(services) -> None: + async with aio.Client(services.endpoint) as client: + keyspace = client.keyspace("default") + assert keyspace.name == "default" + assert await keyspace.get(b"missing") is None + assert await keyspace.get(b"empty") == b"" + assert await keyspace.get_entry(b"key") == orbita.Entry( + key=b"key", + value=b"value", + version=7, + expires_at_millis=123, + ) + + +async def test_async_read_retries_only_unavailable(services) -> None: + services.kv.get_failures = 2 + policy = orbita.RetryPolicy(max_attempts=3, initial_backoff=0, max_backoff=0) + async with aio.Client(services.endpoint, retry_policy=policy) as client: + assert await client.keyspace("default").get(b"key") == b"value" + assert services.kv.get_calls == 3 + + +async def test_async_read_deadline_expires_during_backoff(services) -> None: + services.kv.get_failures = 2 + policy = orbita.RetryPolicy(max_attempts=3, initial_backoff=0.05, max_backoff=0.05) + async with aio.Client(services.endpoint, retry_policy=policy) as client: + with pytest.raises(TimeoutError, match="operation timed out"): + await client.keyspace("default").get(b"key", timeout=0.001) + assert services.kv.get_calls == 1 + + +async def test_async_set_encodes_conditions_and_ttl(services) -> None: + async with aio.Client(services.endpoint) as client: + result = await client.keyspace("default").set( + b"key", + b"value", + if_not_present=True, + ttl=timedelta(microseconds=1), + ) + assert result == orbita.SetResult(applied=True, version=8, current_version=None) + assert services.kv.last_set is not None + assert services.kv.last_set.condition.WhichOneof("kind") == "if_not_present" + assert services.kv.last_set.ttl_millis == 1 + + +async def test_async_set_rejects_invalid_options_before_sending(services) -> None: + async with aio.Client(services.endpoint) as client: + keyspace = client.keyspace("default") + with pytest.raises(ValueError, match="only one condition"): + await keyspace.set(b"key", b"value", if_not_present=True, if_version=1) + with pytest.raises(ValueError, match="must not be negative"): + await keyspace.set(b"key", b"value", ttl=timedelta(microseconds=-1)) + assert services.kv.set_calls == 0 + + +async def test_async_failed_condition_is_a_successful_result(services) -> None: + async with aio.Client(services.endpoint) as client: + result = await client.keyspace("default").set(b"condition-failed", b"value") + assert result == orbita.SetResult(applied=False, version=0, current_version=9) + + +async def test_async_mutation_is_not_retried_and_reports_ambiguity(services) -> None: + services.kv.set_status = grpc.StatusCode.UNAVAILABLE + async with aio.Client(services.endpoint) as client: + with pytest.raises(orbita.AmbiguousMutationError) as caught: + await client.keyspace("default").set(b"key", b"value") + assert services.kv.set_calls == 1 + assert caught.value.operation == "set" + assert caught.value.rpc_error.code() == grpc.StatusCode.UNAVAILABLE + + +async def test_async_delete_encodes_a_version_condition(services) -> None: + async with aio.Client(services.endpoint) as client: + result = await client.keyspace("default").delete(b"key", if_version=7) + assert result == orbita.DeleteResult(applied=True, existed=True, current_version=None) + assert services.kv.last_delete is not None + assert services.kv.last_delete.condition.if_version == 7 + + +async def test_async_pages_carry_opaque_cursors_and_preserve_keys_only_entries(services) -> None: + async with aio.Client(services.endpoint) as client: + keyspace = client.keyspace("default") + keys_only = [page async for page in keyspace.pages(b"prefix/")] + with_values = [page async for page in keyspace.pages(b"prefix/", include_values=True)] + + assert [page.next_cursor for page in keys_only] == [b"opaque", b""] + assert keys_only[0].entries[0].value is None + assert with_values[1].entries[0].value == b"" + + +async def test_async_readiness_returns_unmet_conditions_without_an_rpc_error(services) -> None: + async with aio.Client(services.endpoint) as client: + readiness = await client.readiness() + assert readiness == orbita.Readiness( + ready=False, + conditions=(orbita.ReadinessCondition(name="wal-recovered", met=False),), + ) + + +async def test_async_calls_before_connect_fail_with_a_clear_error(services) -> None: + client = aio.Client(services.endpoint) + with pytest.raises(RuntimeError, match="not connected"): + await client.keyspace("default").get(b"key") + with pytest.raises(RuntimeError, match="not connected"): + await client.readiness() + await client.close() + assert services.kv.limits_calls == 0 + + +async def test_async_connect_is_idempotent(services) -> None: + client = aio.Client(services.endpoint) + try: + assert await client.connect() is client + assert await client.connect() is client + finally: + await client.close() + assert services.kv.limits_calls == 1 + + +async def test_async_invalid_advertised_message_size_is_rejected(services) -> None: + services.kv.max_message_bytes = 0 + client = aio.Client(services.endpoint) + with pytest.raises(ValueError, match="invalid advertised max message size"): + await client.connect() + await client.close() + + +async def test_async_admin_reads_retry_and_mutations_report_ambiguity(services) -> None: + services.admin.list_failures = 1 + policy = orbita.RetryPolicy(max_attempts=2, initial_backoff=0, max_backoff=0) + async with aio.AdminClient(services.endpoint, retry_policy=policy) as admin: + listed = await admin.list_keyspaces(admin_pb2.ListKeyspacesRequest()) + assert [keyspace.name for keyspace in listed.keyspaces] == ["default"] + assert services.admin.list_calls == 2 + + services.admin.credential_status = grpc.StatusCode.ABORTED + with pytest.raises(orbita.AmbiguousMutationError) as caught: + await admin.create_credential(admin_pb2.CreateCredentialRequest()) + assert caught.value.operation == "create credential" + + +async def test_async_admin_sends_the_root_credential(services) -> None: + async with aio.AdminClient(services.endpoint, root_credential="root-secret") as admin: + created = await admin.create_keyspace(admin_pb2.CreateKeyspaceRequest(name="payments")) + assert created.name == "payments" + assert services.admin.metadata[0]["authorization"] == "Bearer root-secret" diff --git a/tests/test_integration.py b/tests/test_integration.py index 944c4dc..b1c36e3 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import os import time import uuid @@ -8,6 +9,7 @@ import pytest import orbita +from orbita import aio from orbita.admin import AdminClient from orbita.v1 import admin_pb2 @@ -85,6 +87,63 @@ def test_client_against_orbita() -> None: ) +async def test_async_client_against_orbita() -> None: + assert ENDPOINT is not None + assert ROOT_CREDENTIAL is not None + name = f"python-integration-{uuid.uuid4().hex[:12]}" + + async with aio.AdminClient(ENDPOINT, root_credential=ROOT_CREDENTIAL) as admin: + await admin.create_keyspace(admin_pb2.CreateKeyspaceRequest(name=name), timeout=10) + try: + credential = await admin.create_credential( + admin_pb2.CreateCredentialRequest( + keyspaces=[name], + permissions=[admin_pb2.PERMISSION_READ, admin_pb2.PERMISSION_WRITE], + description="orbita-python async integration test", + ), + timeout=10, + ) + assert credential.secret + + async with aio.Client(ENDPOINT, credential=credential.secret, timeout=10) as client: + readiness = await client.readiness(timeout=10) + assert readiness.ready + keyspace = client.keyspace(name) + await _wait_for_credential_async(keyspace) + + created = await keyspace.set(b"items/a", b"one", if_not_present=True, timeout=10) + assert created.applied + + updated = await keyspace.set( + b"items/a", b"two", if_version=created.version, timeout=10 + ) + assert updated.applied + await keyspace.set(b"items/b", b"three", timeout=10) + + entry = await keyspace.get_entry(b"items/a", timeout=10) + assert entry is not None + assert entry.value == b"two" + + pages = [ + page + async for page in keyspace.pages( + b"items/", limit=1, include_values=True, timeout=10 + ) + ] + assert sorted(entry.key for page in pages for entry in page.entries) == [ + b"items/a", + b"items/b", + ] + + deleted = await keyspace.delete(b"items/a", if_version=updated.version, timeout=10) + assert deleted.applied and deleted.existed + assert await keyspace.get(b"items/a", timeout=10) is None + finally: + await admin.delete_keyspace( + admin_pb2.DeleteKeyspaceRequest(name=name, confirm_name=name), timeout=10 + ) + + def _wait_for_credential(keyspace: orbita.Keyspace) -> None: deadline = time.monotonic() + 10 while True: @@ -101,3 +160,21 @@ def _wait_for_credential(keyspace: orbita.Keyspace) -> None: if time.monotonic() >= deadline: raise time.sleep(0.1) + + +async def _wait_for_credential_async(keyspace: aio.Keyspace) -> None: + deadline = time.monotonic() + 10 + while True: + try: + await keyspace.get(b"credential-probe", timeout=2) + return + except grpc.RpcError as error: + if error.code() not in { + grpc.StatusCode.UNAUTHENTICATED, + grpc.StatusCode.UNAVAILABLE, + grpc.StatusCode.NOT_FOUND, + }: + raise + if time.monotonic() >= deadline: + raise + await asyncio.sleep(0.1)