python-etcd3-async 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

etcd3/__init__.py

@@ -8,6 +8,7 @@ from etcd3.shared import Transactions
 from etcd3.async_single_client import AsyncEtcd3Client
 from etcd3.async_multi_client import AsyncMultiEndpointEtcd3Client
 from etcd3.async_endpoint import AsyncEndpoint
+from etcd3.async_endpoint import AsyncEndpointConfig
 from etcd3.async_client_extensions import get_request_context, set_request_context
 from etcd3.exceptions import Etcd3Exception
 from etcd3.observability import ObservabilityConfig
@@ -50,7 +51,7 @@ from etcd3.async_single_client import async_client
 
 __author__ = "owasd"
 __email__ = ""
-__version__ = "0.1.6"
+__version__ = "0.2.0"
 
 __all__ = (
     "etcdrpc",
@@ -71,6 +72,7 @@ __all__ = (
     "Permission",
     "MultiEndpointEtcd3Client",
     "AsyncEndpoint",
+    "AsyncEndpointConfig",
     "AsyncEtcd3Client",
     "AsyncMultiEndpointEtcd3Client",
     "async_client",

etcd3/async_client.py

@@ -8,6 +8,9 @@ For new code, prefer importing from:
     - etcd3.async_endpoint: AsyncEndpoint
 """
 
+import sys
+import types
+
 # Re-export config classes
 from etcd3.async_config import AsyncClientConfig, AsyncMultiEndpointConfig
 from etcd3.config import RangeRequestConfig
@@ -48,3 +51,21 @@ __all__ = [
     "_async_handle_keepalive_errors",
     "async_client",
 ]
+
+
+class _CallableModule(types.ModuleType):
+    """Module that can also be awaited like the factory function it re-exports.
+
+    ``import etcd3.async_client`` makes Python rebind the
+    ``etcd3.async_client`` attribute from the factory function to this
+    module. Without this, the documented ``await etcd3.async_client(...)``
+    call would break for anyone who imported the compatibility shim.
+    Calling the module delegates to the factory function, so both
+    spellings keep working.
+    """
+
+    def __call__(self, *args, **kwargs):
+        return async_client(*args, **kwargs)
+
+
+sys.modules[__name__].__class__ = _CallableModule

etcd3/async_client_extensions.py

@@ -117,7 +117,55 @@ class AsyncKVOperations:
 
     def __init__(self, client, max_concurrent=100):
         self.client = client
-        self._semaphore = asyncio.Semaphore(max_concurrent)
+        # Client-wide ceiling for concurrent key-value operations. A batch
+        # call may ask for less concurrency, never for more.
+        self._max_concurrent = max_concurrent
+
+    def _limit(self, concurrency):
+        """Resolve the concurrency limit for one batch call."""
+        if concurrency is None:
+            concurrency = self._max_concurrent
+        return max(1, min(int(concurrency), self._max_concurrent))
+
+    async def _run_bounded(self, items, operation, concurrency=None):
+        """Run ``operation`` over ``items`` with bounded concurrency.
+
+        Shared by ``batch_put``, ``batch_delete`` and ``put_many`` so that
+        all three have one concurrency policy and one failure policy:
+        per-item errors are captured and never abort the whole batch.
+
+        :param items: Items to process
+        :type items: list
+        :param operation: Async callable applied to a single item
+        :type operation: callable
+        :param concurrency: Maximum number of items processed at the same time
+        :type concurrency: int, optional
+        :returns: List of ``(item, value, error)`` tuples in input order
+        :rtype: list[tuple]
+        """
+        limit = self._limit(concurrency)
+        # The semaphore is created per call so it always belongs to the
+        # running event loop (asyncio.Semaphore binds itself on first use).
+        semaphore = asyncio.Semaphore(limit)
+        results = [(item, None, None) for item in items]
+
+        async def run(index, item):
+            async with semaphore:
+                try:
+                    value = await operation(item)
+                except Exception as exc:  # noqa: BLE001 - per-item isolation
+                    results[index] = (item, None, exc)
+                else:
+                    results[index] = (item, value, None)
+
+        # asyncio.TaskGroup keeps every task inside this scope: no task can
+        # outlive the call unobserved, and one failing item cannot cancel
+        # the others (errors are captured inside ``run``).
+        async with asyncio.TaskGroup() as tg:
+            for index, item in enumerate(items):
+                tg.create_task(run(index, item))
+
+        return results
 
     async def _get_kv_stub(self):
         """Get the KV stub using the current channel with auto-reconnection.
@@ -139,13 +187,19 @@ class AsyncKVOperations:
         """
         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,
-        )
+        async def _range():
+            # The stub is created per attempt so that a retry picks up a
+            # freshly built channel.
+            kv_stub = await self._get_kv_stub()
+            return await kv_stub.Range(
+                range_request,
+                timeout=self.client.timeout,
+                credentials=self.client.call_credentials,
+                metadata=self.client.metadata,
+            )
+
+        # Reads are idempotent, so a transient failure can be retried.
+        response = await self.client._retry_with_backoff(_range)
 
         # Set trace attributes from response
         endpoint = (
@@ -329,8 +383,18 @@ class AsyncKVOperations:
         return delete_response.deleted
 
     @_async_handle_errors
-    async def delete_prefix(self, prefix):
-        """Delete a range of keys with a prefix."""
+    async def delete_prefix(self, prefix, return_response=False):
+        """
+        Delete a range of keys with a prefix.
+
+        :param prefix: prefix of the keys to delete
+        :type prefix: str or bytes
+        :param return_response: return the full response
+        :type return_response: bool
+        :returns: the number of deleted keys when ``return_response`` is
+            False, and the full response when ``return_response`` is True
+        :rtype: int or :class:`.rpc_pb2.DeleteRangeResponse`
+        """
         delete_request = RequestBuilder.build_delete_request(
             key=prefix,
             range_end=utils.prefix_range_end(utils.to_bytes(prefix)),
@@ -345,8 +409,24 @@ class AsyncKVOperations:
             metadata=self.client.metadata,
         )
 
+        if return_response:
+            return delete_response
         return delete_response.deleted
 
+    async def transaction(self, compare, success=None, failure=None):
+        """Perform a transaction.
+
+        :param compare: A list of comparisons to make
+        :param success: A list of operations to perform if all the comparisons
+            are true
+        :param failure: A list of operations to perform if any of the
+            comparisons are false
+        :return: A tuple of (operation status, responses)
+        """
+        return await self.client.transaction(
+            compare=compare, success=success, failure=failure
+        )
+
     async def batch_put(self, items, concurrency=10, atomic=True):
         """Batch put multiple key-value pairs.
 
@@ -354,8 +434,8 @@ class AsyncKVOperations:
         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 items: Iterable of (key, value) tuples or a mapping to put
+        :type items: Iterable[tuple[str | bytes, str | bytes]] or dict
         :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.
@@ -364,6 +444,9 @@ class AsyncKVOperations:
         :rtype: int
         :raises BatchOperationError: If atomic=True and some operations fail
         """
+        # Accept a mapping as well as an iterable of pairs
+        if isinstance(items, dict):
+            items = list(items.items())
         # Convert to list to allow multiple iterations
         items_list = list(items)
         if not items_list:
@@ -397,23 +480,14 @@ class AsyncKVOperations:
                 ) 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 def put_item(item):
+                key, value = item
+                return await self.put(key, value)
 
-            async with asyncio.TaskGroup() as tg:
-                for k, v in items_list:
-                    tg.create_task(put_with_semaphore(k, v))
+            results = await self._run_bounded(items_list, put_item, concurrency)
 
             # Count successful operations
-            return sum(1 for _, success in results if success == 1)
+            return sum(1 for _, _, error in results if error is None)
 
     async def batch_delete(self, keys, concurrency=10, atomic=True):
         """Batch delete multiple keys.
@@ -464,54 +538,54 @@ class AsyncKVOperations:
                 ) 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))
+            async def delete_item(key):
+                # delete returns the number of deleted keys
+                return await self.delete(key)
+
+            results = await self._run_bounded(keys_list, delete_item, concurrency)
 
             # Count successful operations
-            return sum(success for _, success in results)
+            return sum(count for _, count, error in results if error is None)
 
-    async def put_many(self, items, timeout=None):
+    async def put_many(self, items, timeout=None, concurrency=10):
         """Put multiple key-value pairs with controlled concurrency.
 
         This is a simpler interface for batch puts that handles all
-        the asyncio complexity internally.
+        the asyncio complexity internally. It shares the concurrency and
+        failure policy of :meth:`batch_put`: items are processed in
+        parallel up to ``concurrency`` at a time, and a failing item never
+        aborts the rest of the batch.
 
         :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
+        :param concurrency: Maximum number of concurrent put operations
+        :type concurrency: int
         :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))
+        items_list = list(items)
+        if not items_list:
+            return 0, []
+
+        async def put_item(item):
+            key, value = item
+            if timeout is None:
+                return await self.put(key, value)
+            # asyncio.timeout applies the deadline to the current task, so
+            # it composes with an enclosing timeout.
+            async with asyncio.timeout(timeout):
+                return await self.put(key, value)
+
+        results = await self._run_bounded(items_list, put_item, concurrency)
+
+        success_count = sum(1 for _, _, error in results if error is None)
+        failed_keys = [
+            (item[0], error) for item, _, error in results if error is not None
+        ]
 
         return success_count, failed_keys
 

etcd3/async_endpoint.py

@@ -41,6 +41,12 @@ class AsyncEndpointConfig:
 class AsyncEndpoint(BaseEndpoint):
     """Represents an etcd cluster endpoint for async operations.
 
+    The signature mirrors the synchronous ``Endpoint`` class: ``host`` and
+    ``port`` come first, so both ``AsyncEndpoint('localhost', 2379)`` and
+    ``AsyncEndpoint(host='localhost', port=2379)`` work. An
+    ``AsyncEndpointConfig`` can be passed as the keyword-only ``config``
+    argument instead.
+
     :param config: AsyncEndpointConfig object with endpoint configuration
     :param str host: Endpoint host (ignored if config is provided)
     :param int port: Endpoint port (ignored if config is provided)
@@ -57,13 +63,14 @@ class AsyncEndpoint(BaseEndpoint):
 
     def __init__(
         self,
-        config=None,
         host=None,
         port=None,
         secure=True,
         creds=None,
         time_retry=300.0,
         opts=None,
+        *,
+        config=None,
     ):
         # Use config object if provided, otherwise use individual parameters
         if config is not None:
@@ -96,17 +103,24 @@ class AsyncEndpoint(BaseEndpoint):
         except Exception:
             return False
 
-    async def health_check_async(self):
+    async def health_check_async(self, timeout=2.0):
         """Perform an asynchronous health check using gRPC status check.
 
         This method uses the channel's channel_ready() for a more robust
         connectivity check.
 
+        :param timeout: Maximum time to wait for the channel to become ready
+        :type timeout: float, optional
         :returns: True if the endpoint is healthy
         :rtype: bool
         """
         try:
-            await asyncio.wait_for(self.channel.channel_ready(), timeout=2.0)
+            # asyncio.timeout (3.11+) is used instead of asyncio.wait_for:
+            # it applies a deadline to the current task, so it composes
+            # correctly with an enclosing timeout instead of cancelling a
+            # nested task.
+            async with asyncio.timeout(timeout):
+                await self.channel.channel_ready()
             return True
         except asyncio.TimeoutError:
             return False

etcd3/async_managers.py

@@ -39,6 +39,12 @@ class AsyncWatchManager:
             )
         return self._watcher
 
+    async def close(self):
+        """Close the watcher and release its background task."""
+        if self._watcher is not None:
+            await self._watcher.close()
+            self._watcher = None
+
     async def add_watch_callback(self, *args, **kwargs):
         """
         Watch a key or range of keys and call a callback on every response.
@@ -90,23 +96,42 @@ class AsyncWatchManager:
                   Use ``cancel`` to cancel the watch request.
         """
         watcher = await self._get_watcher()
-        event_iter = watcher.watch(key, **kwargs)
+
+        # The watch runs on the shared watcher (one background task and one
+        # gRPC stream for the whole client), so the real watch id is known
+        # and the watch can actually be cancelled server-side.
+        response_queue = asyncio.Queue()
+
+        async def callback(response):
+            await response_queue.put(response)
+
+        watch_id = await watcher.add_callback(key, callback, **kwargs)
 
         cancel_event = asyncio.Event()
 
         async def cancel():
+            if cancel_event.is_set():
+                return
             cancel_event.set()
-            await watcher.cancel(0)
+            # Wake up a pending iteration so it can finish.
+            await response_queue.put(None)
+            await watcher.cancel(watch_id)
 
         async def event_iterator():
             try:
-                async for event in event_iter:
-                    if cancel_event.is_set():
+                while not cancel_event.is_set():
+                    response = await response_queue.get()
+                    if response is None:
                         break
-                    yield event
+                    if isinstance(response, Exception):
+                        cancel_event.set()
+                        raise response
+                    yield response
             except grpc.RpcError as exc:
                 if not cancel_event.is_set():
                     self.etcd_client._manage_grpc_errors(exc)
+            finally:
+                await cancel()
 
         return event_iterator(), cancel