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)