python-etcd3-async 0.1.0__py3-none-any.whl

etcd3/async_client_extensions.py

@@ -0,0 +1,878 @@
+"""Extensions for async etcd3 client.
+
+This module provides auxiliary classes to split the functionality of AsyncMultiEndpointEtcd3Client
+into smaller, more manageable components.
+"""
+
+import asyncio
+import contextvars
+import logging
+import random
+from typing import Any
+
+import grpc
+
+import etcdrpc
+
+from etcd3 import exceptions
+from etcd3 import transactions
+from etcd3 import utils
+from etcd3.client import KVMetadata
+from etcd3.leases import AsyncLease
+from etcd3.locks import AsyncLock
+from etcd3.shared import (
+    _make_async_handle_errors_with_tracing,
+    _make_async_handle_generator_errors_with_tracing,
+    RequestBuilder,
+)
+from etcd3.trace import get_current_span, is_tracing_available
+from etcd3.trace_attributes import (
+    build_response_attributes,
+    build_lease_attributes,
+    TraceAttributes,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# Context variables for request tracking
+_request_context: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar(
+    "request_context", default=None
+)
+
+
+def get_request_context() -> dict[str, Any] | None:
+    """Get the current request context.
+
+    :returns: The current request context dictionary, or None if not set
+    """
+    return _request_context.get()
+
+
+def set_request_context(**kwargs) -> contextvars.Token:
+    """Set values in the request context.
+
+    :param kwargs: Key-value pairs to set in the context
+    :returns: A token that can be used to reset the context
+    """
+    current = _request_context.get() or {}
+    new_context = {**current, **kwargs}
+    return _request_context.set(new_context)
+
+
+def _manage_client_errors(self, exc):
+    """Error handler for extensions that routes to client's error handler."""
+    self.client._manage_grpc_errors(exc)
+
+
+def _set_span_attributes_from_response(response, endpoint: str = None):
+    """Set span attributes from etcd response.
+    
+    :param response: etcd response object
+    :param endpoint: Optional endpoint string
+    """
+    if not is_tracing_available():
+        return
+    
+    span = get_current_span()
+    if span is None:
+        return
+    
+    # Build attributes from response
+    attrs = build_response_attributes(header=getattr(response, 'header', None))
+    
+    # Add count if available
+    if hasattr(response, 'count'):
+        attrs[TraceAttributes.ETCD_COUNT] = response.count
+    
+    # Add deleted if available
+    if hasattr(response, 'deleted'):
+        attrs[TraceAttributes.ETCD_DELETED] = response.deleted
+    
+    # Add endpoint if provided
+    if endpoint:
+        attrs[TraceAttributes.ETCD_ENDPOINT] = endpoint
+    
+    # Set attributes on span
+    for key, value in attrs.items():
+        if value is not None:
+            span.set_attribute(key, value)
+
+
+# Create trace-aware error handlers
+_async_handle_errors = _make_async_handle_errors_with_tracing(
+    _manage_client_errors, "etcd3.operation"
+)
+_async_handle_generator_errors = _make_async_handle_generator_errors_with_tracing(
+    _manage_client_errors, "etcd3.operation"
+)
+# Error handler for lease keepalive that propagates errors to allow detection of connection loss
+_async_handle_keepalive_errors = _make_async_handle_generator_errors_with_tracing(
+    _manage_client_errors, "etcd3.keepalive", propagate_errors=True
+)
+
+
+class AsyncKVOperations:
+    """Handles key-value operations for AsyncMultiEndpointEtcd3Client."""
+
+    def __init__(self, client, max_concurrent=100):
+        self.client = client
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+        self._batch_semaphore = asyncio.Semaphore(10)  # Limit concurrent batches
+
+    async def _get_kv_stub(self):
+        """Get the KV stub using the current channel with auto-reconnection.
+
+        Creates a new stub each time to ensure it uses the latest channel,
+        which supports automatic reconnection after etcd server restart.
+        """
+        channel = await self.client._get_channel()
+        return etcdrpc.KVStub(channel)
+
+    @_async_handle_errors
+    async def get_response(self, key, **kwargs):
+        """Get the value of a key from etcd.
+
+        :param key: key in etcd to get
+        :type key: str or bytes
+        :param kwargs: Additional arguments (revision, limit, sort_order, etc.)
+        :returns: RangeResponse from etcd
+        """
+        range_request = self.client._build_get_range_request(key, **kwargs)
+
+        kv_stub = await self._get_kv_stub()
+        response = await kv_stub.Range(
+            range_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+        
+        # Set trace attributes from response
+        endpoint = str(self.client.endpoint_in_use) if self.client.endpoint_in_use else None
+        _set_span_attributes_from_response(response, endpoint)
+        
+        return response
+
+    async def get(self, key, **kwargs):
+        """Get the value of a key from etcd."""
+        range_response = await self.get_response(key, **kwargs)
+        if range_response.count < 1:
+            return None, None
+        kv = range_response.kvs[0]
+        return kv.value, KVMetadata(kv, range_response.header)
+
+    async def get_prefix_response(self, key_prefix, **kwargs):
+        """Get a range of keys with a prefix."""
+        if any(kwarg in kwargs for kwarg in ("key", "range_end")):
+            raise TypeError("Don't use key or range_end with prefix")
+        kwargs["range_end"] = utils.prefix_range_end(utils.to_bytes(key_prefix))
+        return await self.get_response(key_prefix, **kwargs)
+
+    @_async_handle_generator_errors
+    async def get_prefix(self, key_prefix, **kwargs):
+        """Get a range of keys with a prefix."""
+        range_response = await self.get_prefix_response(key_prefix, **kwargs)
+        for kv in range_response.kvs:
+            yield kv.value, KVMetadata(kv, range_response.header)
+
+    async def get_range_response(
+        self, range_start, range_end, sort_order=None, sort_target="key", **kwargs
+    ):
+        """Get a range of keys."""
+        if any(kwarg in kwargs for kwarg in ("key", "range_end")):
+            raise TypeError("Don't use key or range_end with get_range")
+        kwargs["range_end"] = range_end
+        if sort_order is not None:
+            kwargs["sort_order"] = sort_order
+            kwargs["sort_target"] = sort_target
+        return await self.get_response(range_start, **kwargs)
+
+    @_async_handle_generator_errors
+    async def get_range(self, range_start, range_end, **kwargs):
+        """Get a range of keys."""
+        range_response = await self.get_range_response(range_start, range_end, **kwargs)
+        for kv in range_response.kvs:
+            yield kv.value, KVMetadata(kv, range_response.header)
+
+    async def get_all_response(
+        self, sort_order=None, sort_target="key", keys_only=False
+    ):
+        """Get all keys currently stored in etcd."""
+        kwargs = {}
+        if sort_order is not None:
+            kwargs["sort_order"] = sort_order
+            kwargs["sort_target"] = sort_target
+        if keys_only:
+            kwargs["keys_only"] = True
+        # Use b'\x00' as the minimum key and b'\x00' as range_end to get all keys
+        # This works because etcd stores keys in sorted order by key bytes
+        # b'\x00' is the smallest non-empty key, and range_end=b'\x00' means
+        # include keys >= b'\x00' up to (but not including) b'\x00'
+        # Actually, we need range_end > key. Using high bytes works.
+        return await self.get_response(b"\x00", range_end=b"\xff" * 16, **kwargs)
+
+    @_async_handle_generator_errors
+    async def get_all(self, **kwargs):
+        """Get all keys currently stored in etcd."""
+        range_response = await self.get_all_response(**kwargs)
+        for kv in range_response.kvs:
+            yield kv.value, KVMetadata(kv, range_response.header)
+
+    @_async_handle_errors
+    async def put(
+        self,
+        key,
+        value=None,
+        lease=None,
+        prev_kv=False,
+        ignore_value=False,
+        ignore_lease=False,
+    ):
+        """Save a value to etcd.
+
+        :param key: key in etcd to set
+        :param value: value to set key to (optional if ignore_value=True)
+        :type value: bytes
+        :param lease: Lease to associate with this key.
+        :type lease: either :class:`.AsyncLease`, or int (ID of lease)
+        :param prev_kv: return the previous key-value pair
+        :type prev_kv: bool
+        :param ignore_value: if True, update the key using the current value
+        :type ignore_value: bool
+        :param ignore_lease: if True, update the key using the current lease
+        :type ignore_lease: bool
+        """
+        put_request = RequestBuilder.build_put_request(
+            key=key,
+            value=value,
+            lease=lease,
+            prev_kv=prev_kv,
+            ignore_value=ignore_value,
+            ignore_lease=ignore_lease,
+        )
+
+        kv_stub = await self._get_kv_stub()
+        put_response = await kv_stub.Put(
+            put_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+        
+        # Set trace attributes from response
+        endpoint = str(self.client.endpoint_in_use) if self.client.endpoint_in_use else None
+        _set_span_attributes_from_response(put_response, endpoint)
+
+        if prev_kv:
+            return KVMetadata(put_response.prev_kv, put_response.header)
+        return put_response.header
+
+    @_async_handle_errors
+    async def put_if_not_exists(self, key, value, lease=None):
+        """Save a value to etcd only if the key doesn't exist."""
+        status, _ = await self.client.transaction(
+            compare=[transactions.Create(key) == 0],
+            success=[transactions.Put(key, value, lease=lease)],
+            failure=[],
+        )
+
+        return status
+
+    @_async_handle_errors
+    async def replace(self, key, initial_value, new_value):
+        """Replace the value of a key in etcd."""
+        status, _ = await self.client.transaction(
+            compare=[transactions.Value(key) == initial_value],
+            success=[transactions.Put(key, new_value)],
+            failure=[],
+        )
+
+        return status
+
+    @_async_handle_errors
+    async def delete(self, key, prev_kv=False, return_response=False, range_end=None):
+        """Delete a key or range of keys from etcd.
+
+        :param key: key in etcd to delete
+        :type key: str or bytes
+        :param prev_kv: return the deleted key-value pair
+        :type prev_kv: bool
+        :param return_response: return the full response
+        :type return_response: bool
+        :param range_end: End of the key range (for range deletes)
+        :type range_end: str or bytes, optional
+        :returns: Number of deleted keys or full response
+        """
+        delete_request = RequestBuilder.build_delete_request(
+            key=key,
+            prev_kv=prev_kv,
+            range_end=range_end,
+        )
+
+        kv_stub = await self._get_kv_stub()
+        delete_response = await kv_stub.DeleteRange(
+            delete_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+        
+        # Set trace attributes from response
+        endpoint = str(self.client.endpoint_in_use) if self.client.endpoint_in_use else None
+        _set_span_attributes_from_response(delete_response, endpoint)
+
+        if return_response:
+            return delete_response
+
+        return delete_response.deleted
+
+    @_async_handle_errors
+    async def delete_prefix(self, prefix):
+        """Delete a range of keys with a prefix."""
+        delete_request = RequestBuilder.build_delete_request(
+            key=prefix,
+            range_end=utils.prefix_range_end(utils.to_bytes(prefix)),
+            prev_kv=False,
+        )
+
+        kv_stub = await self._get_kv_stub()
+        delete_response = await kv_stub.DeleteRange(
+            delete_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+
+        return delete_response.deleted
+
+    async def batch_put(self, items, concurrency=10, atomic=True):
+        """Batch put multiple key-value pairs.
+
+        By default, this method executes all puts as a single atomic transaction,
+        ensuring either all operations succeed or none do. If atomic=False,
+        operations are executed in parallel with concurrency control.
+
+        :param items: Iterable of (key, value) tuples to put
+        :type items: Iterable[tuple[str | bytes, str | bytes]]
+        :param concurrency: Maximum number of concurrent put operations (only used if atomic=False)
+        :type concurrency: int
+        :param atomic: If True (default), execute as atomic transaction. If False, parallel execution.
+        :type atomic: bool
+        :returns: Number of successful operations (atomic=True: either all or none)
+        :rtype: int
+        :raises BatchOperationError: If atomic=True and some operations fail
+        """
+        # Convert to list to allow multiple iterations
+        items_list = list(items)
+        if not items_list:
+            return 0
+
+        if atomic:
+            # Use atomic transaction for all operations
+            # This ensures either all succeed or none do
+            success_ops = [transactions.Put(key, value) for key, value in items_list]
+            # Use a dummy compare that always succeeds (create_revision > -1 is always true)
+            compare = [transactions.Create(key) > -1 for key, _ in items_list]
+
+            try:
+                success, _ = await self.client.transaction(
+                    compare=compare,
+                    success=success_ops,
+                    failure=[],
+                )
+                if success:
+                    return len(items_list)
+                else:
+                    # Transaction failed - this shouldn't happen with our compare
+                    raise exceptions.BatchOperationError(
+                        f"Atomic batch put failed for {len(items_list)} items"
+                    )
+            except Exception as exc:
+                if isinstance(exc, exceptions.BatchOperationError):
+                    raise
+                raise exceptions.BatchOperationError(f"Batch put failed: {exc}") from exc
+        else:
+            # Non-atomic: parallel execution with concurrency control
+            semaphore = asyncio.Semaphore(concurrency)
+            results = []
+
+            async def put_with_semaphore(key, value):
+                async with semaphore:
+                    try:
+                        await self.put(key, value)
+                        results.append((key, 1))
+                    except Exception:
+                        results.append((key, 0))
+
+            async with asyncio.TaskGroup() as tg:
+                for k, v in items_list:
+                    tg.create_task(put_with_semaphore(k, v))
+
+            # Count successful operations
+            return sum(1 for _, success in results if success == 1)
+
+    async def batch_delete(self, keys, concurrency=10, atomic=True):
+        """Batch delete multiple keys.
+
+        By default, this method executes all deletes as a single atomic transaction,
+        ensuring either all operations succeed or none do. If atomic=False,
+        operations are executed in parallel with concurrency control.
+
+        :param keys: Iterable of keys to delete
+        :type keys: Iterable[str | bytes]
+        :param concurrency: Maximum number of concurrent delete operations (only used if atomic=False)
+        :type concurrency: int
+        :param atomic: If True (default), execute as atomic transaction. If False, parallel execution.
+        :type atomic: bool
+        :returns: Number of successful operations (atomic=True: either all or none)
+        :rtype: int
+        :raises BatchOperationError: If atomic=True and some operations fail
+        """
+        # Convert to list to allow multiple iterations
+        keys_list = list(keys)
+        if not keys_list:
+            return 0
+
+        if atomic:
+            # Use atomic transaction for all operations
+            # This ensures either all succeed or none do
+            delete_ops = [transactions.Delete(key) for key in keys_list]
+            # Use a dummy compare that always succeeds (create_revision > -1 is always true)
+            compare = [transactions.Create(key) > -1 for key in keys_list]
+
+            try:
+                success, _ = await self.client.transaction(
+                    compare=compare,
+                    success=delete_ops,
+                    failure=[],
+                )
+                if success:
+                    return len(keys_list)
+                else:
+                    raise exceptions.BatchOperationError(
+                        f"Atomic batch delete failed for {len(keys_list)} keys"
+                    )
+            except Exception as exc:
+                if isinstance(exc, exceptions.BatchOperationError):
+                    raise
+                raise exceptions.BatchOperationError(f"Batch delete failed: {exc}") from exc
+        else:
+            # Non-atomic: parallel execution with concurrency control
+            semaphore = asyncio.Semaphore(concurrency)
+            results = []
+
+            async def delete_with_semaphore(key):
+                async with semaphore:
+                    try:
+                        result = await self.delete(key)
+                        # delete returns True/False or deleted count
+                        if isinstance(result, bool):
+                            results.append((key, 1 if result else 0))
+                        else:
+                            results.append((key, 1 if result > 0 else 0))
+                    except Exception:
+                        results.append((key, 0))
+
+            async with asyncio.TaskGroup() as tg:
+                for k in keys_list:
+                    tg.create_task(delete_with_semaphore(k))
+
+            # Count successful operations
+            return sum(success for _, success in results)
+
+    async def put_many(self, items, timeout=None):
+        """Put multiple key-value pairs with controlled concurrency.
+
+        This is a simpler interface for batch puts that handles all
+        the asyncio complexity internally.
+
+        :param items: Dict of {key: value} pairs or iterable of (key, value) tuples
+        :type items: dict or Iterable[tuple]
+        :param timeout: Optional timeout for each individual put operation
+        :type timeout: float
+        :returns: Tuple of (success_count, failed_keys)
+        :rtype: tuple[int, list]
+        """
+        if isinstance(items, dict):
+            items = list(items.items())
+
+        success_count = 0
+        failed_keys = []
+
+        for key, value in items:
+            try:
+                async with self._semaphore:
+                    await asyncio.wait_for(self.put(key, value), timeout=timeout)
+                    success_count += 1
+            except Exception as exc:
+                failed_keys.append((key, exc))
+
+        return success_count, failed_keys
+
+
+class AsyncLeaseOperations:
+    """Handles lease operations for AsyncMultiEndpointEtcd3Client."""
+
+    def __init__(self, client):
+        self.client = client
+
+    async def _get_lease_stub(self):
+        """Get the lease stub using the current channel with auto-reconnection.
+
+        Creates a new stub each time to ensure it uses the latest channel,
+        which supports automatic reconnection after etcd server restart.
+        """
+        channel = await self.client._get_channel()
+        return etcdrpc.LeaseStub(channel)
+
+    @_async_handle_errors
+    async def lease(self, ttl, lease_id=None):
+        """Create a new lease."""
+        lease_grant_request = etcdrpc.LeaseGrantRequest(TTL=ttl, ID=lease_id)
+        lease_stub = await self._get_lease_stub()
+        lease_grant_response = await lease_stub.LeaseGrant(
+            lease_grant_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+        
+        # Set trace attributes for lease
+        if is_tracing_available():
+            span = get_current_span()
+            if span:
+                attrs = build_lease_attributes(
+                    lease_id=lease_grant_response.ID,
+                    ttl=lease_grant_response.TTL
+                )
+                endpoint = str(self.client.endpoint_in_use) if self.client.endpoint_in_use else None
+                if endpoint:
+                    attrs[TraceAttributes.ETCD_ENDPOINT] = endpoint
+                for key, value in attrs.items():
+                    span.set_attribute(key, value)
+        
+        return AsyncLease(
+            lease_id=lease_grant_response.ID,
+            ttl=lease_grant_response.TTL,
+            etcd_client=self.client,
+        )
+
+    @_async_handle_errors
+    async def revoke_lease(self, lease_id):
+        """Revoke a lease."""
+        lease_revoke_request = etcdrpc.LeaseRevokeRequest(ID=lease_id)
+        lease_stub = await self._get_lease_stub()
+        await lease_stub.LeaseRevoke(
+            lease_revoke_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+        
+        # Set trace attributes for lease revocation
+        if is_tracing_available():
+            span = get_current_span()
+            if span:
+                attrs = build_lease_attributes(lease_id=lease_id)
+                endpoint = str(self.client.endpoint_in_use) if self.client.endpoint_in_use else None
+                if endpoint:
+                    attrs[TraceAttributes.ETCD_ENDPOINT] = endpoint
+                for key, value in attrs.items():
+                    span.set_attribute(key, value)
+
+    @_async_handle_keepalive_errors
+    async def refresh_lease(self, lease_id):
+        """Refresh a lease continuously with automatic reconnection.
+
+        This method creates a long-lived gRPC stream that sends keep-alive requests
+        periodically to maintain the lease alive. The caller should iterate over
+        this generator as long as they want the lease to be kept alive.
+
+        The stream will yield a response whenever the lease TTL is refreshed by
+        the server (typically when TTL drops below a threshold).
+
+        This method automatically retries on connection errors, allowing recovery
+        from temporary ETCD server restarts (as long as the lease hasn't expired).
+
+        Raises:
+            grpc.RpcError: When connection errors exceed max retry attempts or
+                          lease is permanently not found
+        """
+        consecutive_errors = 0
+        max_consecutive_errors = 10  # Max consecutive errors before giving up
+        base_retry_delay = 1.0  # Base retry delay in seconds
+        max_retry_delay = 30.0  # Max retry delay in seconds
+
+        while True:
+            try:
+                lease_stub = await self._get_lease_stub()
+                keep_alive_request = etcdrpc.LeaseKeepAliveRequest(ID=lease_id)
+                # Note: We use iter([request]) for each iteration. The gRPC stream
+                # remains open, and etcd server will send keep-alive responses
+                # periodically when the TTL approaches expiration.
+                async for response in lease_stub.LeaseKeepAlive(
+                    iter([keep_alive_request]),
+                    timeout=self.client.timeout,
+                    credentials=self.client.call_credentials,
+                    metadata=self.client.metadata,
+                ):
+                    # Reset error count on successful response
+                    consecutive_errors = 0
+                    yield response
+            except grpc.RpcError as exc:
+                code = exc.code()
+
+                # Lease not found (expired or revoked) - permanent failure, stop retrying
+                if code == grpc.StatusCode.NOT_FOUND:
+                    logger.debug(f"Lease {lease_id} not found, stopping keepalive")
+                    break
+
+                # Connection errors - retry with exponential backoff
+                if code in (
+                    grpc.StatusCode.UNAVAILABLE,
+                    grpc.StatusCode.DEADLINE_EXCEEDED,
+                    grpc.StatusCode.INTERNAL,
+                ):
+                    consecutive_errors += 1
+
+                    if consecutive_errors >= max_consecutive_errors:
+                        logger.error(
+                            f"Lease {lease_id} keepalive failed after "
+                            f"{max_consecutive_errors} consecutive errors"
+                        )
+                        raise  # Exceeded max retries, propagate the error
+
+                    # Exponential backoff with jitter
+                    retry_delay = min(
+                        base_retry_delay * (2 ** (consecutive_errors - 1)),
+                        max_retry_delay
+                    )
+                    jitter = random.uniform(0, retry_delay * 0.1)
+                    actual_delay = retry_delay + jitter
+
+                    logger.warning(
+                        f"Lease {lease_id} keepalive connection error ({code}), "
+                        f"retrying in {actual_delay:.1f}s "
+                        f"(attempt {consecutive_errors}/{max_consecutive_errors})"
+                    )
+
+                    await asyncio.sleep(actual_delay)
+                    continue  # Retry the loop
+
+                # Other errors - propagate to let the decorator handle them
+                raise
+            except GeneratorExit:
+                # Generator was closed externally (e.g., garbage collected or iteration stopped)
+                break
+
+    @_async_handle_errors
+    async def get_lease_info(self, lease_id):
+        """Get information about a lease."""
+        lease_timing_request = etcdrpc.LeaseTimeToLiveRequest(ID=lease_id, keys=True)
+        lease_stub = await self._get_lease_stub()
+        lease_timing_response = await lease_stub.LeaseTimeToLive(
+            lease_timing_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+        
+        # Set trace attributes for lease info
+        if is_tracing_available():
+            span = get_current_span()
+            if span:
+                attrs = build_lease_attributes(
+                    lease_id=lease_id,
+                    ttl=lease_timing_response.TTL,
+                    remaining_ttl=getattr(lease_timing_response, 'TTL', None)
+                )
+                endpoint = str(self.client.endpoint_in_use) if self.client.endpoint_in_use else None
+                if endpoint:
+                    attrs[TraceAttributes.ETCD_ENDPOINT] = endpoint
+                for key, value in attrs.items():
+                    if value is not None:
+                        span.set_attribute(key, value)
+        
+        return lease_timing_response
+
+
+class AsyncLockOperations:
+    """Handles lock operations for AsyncMultiEndpointEtcd3Client."""
+
+    def __init__(self, client):
+        self.client = client
+
+    async def lock(self, name, ttl=60):
+        """Create a new async lock."""
+        return AsyncLock(name, ttl=ttl, etcd_client=self.client)
+
+
+class AsyncClusterOperations:
+    """Handles cluster operations for AsyncMultiEndpointEtcd3Client."""
+
+    def __init__(self, client):
+        self.client = client
+
+    async def add_member(self, urls):
+        """Add a member into the cluster."""
+        return await self.client.cluster_manager.add_member(urls)
+
+    async def remove_member(self, member_id):
+        """Remove an existing member from the cluster."""
+        return await self.client.cluster_manager.remove_member(member_id)
+
+    async def update_member(self, member_id, peer_urls):
+        """Update the configuration of an existing member in the cluster."""
+        return await self.client.cluster_manager.update_member(member_id, peer_urls)
+
+    async def members(self):
+        """List of all members associated with the cluster."""
+        async for member in self.client.cluster_manager.members():
+            yield member
+
+
+class AsyncMaintenanceOperations:
+    """Handles maintenance operations for AsyncMultiEndpointEtcd3Client."""
+
+    def __init__(self, client):
+        self.client = client
+
+    async def _get_maintenance_stub(self):
+        """Get the maintenance stub using the current channel with auto-reconnection.
+
+        Creates a new stub each time to ensure it uses the latest channel,
+        which supports automatic reconnection after etcd server restart.
+        """
+        channel = await self.client._get_channel()
+        return etcdrpc.MaintenanceStub(channel)
+
+    async def status(self):
+        """Get the status of the responding member."""
+        return await self.client.maintenance_manager.status()
+
+    @_async_handle_errors
+    async def compact(self, revision, physical=False):
+        """Compact the event history in etcd."""
+        compact_request = etcdrpc.CompactionRequest(
+            revision=revision, physical=physical
+        )
+        maintenance_stub = await self._get_maintenance_stub()
+        await maintenance_stub.Compact(
+            compact_request,
+            timeout=self.client.timeout,
+            credentials=self.client.call_credentials,
+            metadata=self.client.metadata,
+        )
+
+    async def defragment(self):
+        """Defragment a member's backend database to recover storage space."""
+        return await self.client.maintenance_manager.defragment()
+
+    async def hash(self):
+        """Return the hash of the local KV state."""
+        return await self.client.maintenance_manager.hash()
+
+    @_async_handle_errors
+    async def create_alarm(self, member_id=0):
+        """Create an alarm.
+
+        If no member id is given, the alarm is activated for all the
+        members of the cluster. Only the `no space` alarm can be raised.
+
+        :param member_id: The cluster member id to create an alarm to.
+                          If 0, the alarm is created for all the members
+                          of the cluster.
+        :returns: list of :class:`.Alarm`
+        """
+        return await self.client.maintenance_manager.create_alarm(member_id)
+
+    @_async_handle_generator_errors
+    async def list_alarms(self, member_id=0, alarm_type="none"):
+        """List the activated alarms.
+
+        :param member_id: The cluster member id to list alarms from.
+                          If 0, alarms are listed for all the members
+                          of the cluster.
+        :param alarm_type: The type of alarm to filter by.
+                          Can be "none" or "no space".
+        :returns: sequence of :class:`.Alarm`
+        """
+        async for alarm in self.client.maintenance_manager.list_alarms(
+            member_id, alarm_type
+        ):
+            yield alarm
+
+    @_async_handle_errors
+    async def disarm_alarm(self, member_id=0):
+        """Cancel an alarm.
+
+        :param member_id: The cluster member id to cancel an alarm.
+                          If 0, the alarm is canceled for all the members
+                          of the cluster.
+        :returns: List of :class:`.Alarm`
+        """
+        return await self.client.maintenance_manager.disarm_alarm(member_id)
+
+    @_async_handle_errors
+    async def snapshot(self, file_obj):
+        """Take a snapshot of the database.
+
+        :param file_obj: A file-like object to write the database contents in.
+        """
+        return await self.client.maintenance_manager.snapshot(file_obj)
+
+
+class AsyncWatchOperations:
+    """Handles watch operations for AsyncMultiEndpointEtcd3Client."""
+
+    def __init__(self, client):
+        self.client = client
+
+    async def add_watch_callback(self, *args, **kwargs):
+        """Watch a key or range of keys and call a callback on every response."""
+        return await self.client.watch_manager.add_watch_callback(*args, **kwargs)
+
+    async def add_watch_prefix_callback(self, *args, **kwargs):
+        """Watch a prefix and call a callback on every response."""
+        return await self.client.watch_manager.add_watch_prefix_callback(
+            *args, **kwargs
+        )
+
+    async def watch_response(self, *args, **kwargs):
+        """Watch a key."""
+        return await self.client.watch_manager.watch_response(*args, **kwargs)
+
+    async def watch(self, *args, **kwargs):
+        """Watch a key."""
+        return await self.client.watch_manager.watch(*args, **kwargs)
+
+    async def watch_prefix_response(self, *args, **kwargs):
+        """Watch a range of keys with a prefix."""
+        return await self.client.watch_manager.watch_prefix_response(*args, **kwargs)
+
+    async def watch_prefix(self, *args, **kwargs):
+        """Watch a range of keys with a prefix."""
+        return await self.client.watch_manager.watch_prefix(*args, **kwargs)
+
+    async def watch_once_response(self, *args, **kwargs):
+        """Watch a key and stop after the first response."""
+        return await self.client.watch_manager.watch_once_response(*args, **kwargs)
+
+    async def watch_once(self, *args, **kwargs):
+        """Watch a key and stop after the first event."""
+        return await self.client.watch_manager.watch_once(*args, **kwargs)
+
+    async def watch_prefix_once_response(self, *args, **kwargs):
+        """Watch a range of keys with a prefix and stop after the first response."""
+        return await self.client.watch_manager.watch_prefix_once_response(
+            *args, **kwargs
+        )
+
+    async def watch_prefix_once(self, *args, **kwargs):
+        """Watch a range of keys with a prefix and stop after the first event."""
+        return await self.client.watch_manager.watch_prefix_once(*args, **kwargs)
+
+    async def cancel_watch(self, watch_id):
+        """Cancel a watch."""
+        return await self.client.watch_manager.cancel_watch(watch_id)