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

etcd3/async_watch.py

@@ -0,0 +1,462 @@
+"""Async watch implementation for etcd3 client.
+
+This module provides asynchronous watch functionality for the etcd3 async client.
+"""
+
+import logging
+import asyncio
+
+import grpc
+import etcdrpc
+
+from etcd3 import events
+from etcd3 import exceptions
+from etcd3 import utils
+
+
+class WatchRequestConfig:
+    """Configuration for watch request.
+
+    :param range_end: End of the key range to watch
+    :type range_end: str or bytes
+    :param start_revision: Revision to start watching from
+    :type start_revision: int
+    :param progress_notify: Whether to receive progress notifications
+    :type progress_notify: bool
+    :param filters: Event filters to apply
+    :type filters: list
+    :param prev_kv: Whether to include previous key values
+    :type prev_kv: bool
+    """
+
+    def __init__(
+        self,
+        range_end=None,
+        start_revision=None,
+        progress_notify=False,
+        filters=None,
+        prev_kv=False,
+    ):
+        self.range_end = range_end
+        self.start_revision = start_revision
+        self.progress_notify = progress_notify
+        self.filters = filters
+        self.prev_kv = prev_kv
+
+
+_log = logging.getLogger(__name__)
+
+
+class AsyncWatch:
+    """Represents an active watch in etcd.
+
+    :param watch_id: The watch ID returned by etcd
+    :type watch_id: int
+    :param event_iterator: An optional iterator for watch events
+    :type event_iterator: async iterator
+    :param etcd_client: The etcd client instance
+    :type etcd_client: AsyncMultiEndpointEtcd3Client
+    """
+
+    def __init__(self, watch_id, event_iterator=None, etcd_client=None):
+        self.watch_id = watch_id
+        self.etcd_client = etcd_client
+        self._event_iterator = event_iterator
+
+    async def cancel(self):
+        """Cancel the watch.
+
+        This method cancels the active watch by calling the etcd client's cancel_watch method.
+        """
+        await self.etcd_client.cancel_watch(self.watch_id)
+
+    def get_event_iterator(self):
+        """Get the iterator for watch events.
+
+        :returns: The iterator for watch events
+        :raises ValueError: If the iterator is undefined
+        """
+        if self._event_iterator is not None:
+            return self._event_iterator
+
+        raise ValueError("Undefined event iterator")
+
+    @property
+    def iterator(self):
+        """Get the iterator for watch events (property version).
+
+        :returns: The iterator for watch events
+        :raises ValueError: If the iterator is undefined
+        """
+        return self.get_event_iterator()
+
+
+class AsyncWatcher:
+    """Manages watch operations for etcd.
+
+    This class handles the creation and management of watch streams, callbacks,
+    and watch requests to etcd.
+
+    :param watchstub: The gRPC watch stub
+    :type watchstub: etcdrpc.WatchStub
+    :param timeout: The timeout for watch operations
+    :type timeout: int or float
+    :param call_credentials: The gRPC call credentials
+    :type call_credentials: grpc.CallCredentials
+    :param metadata: The gRPC metadata
+    :type metadata: tuple
+    """
+
+    def __init__(self, watchstub, timeout=None, call_credentials=None, metadata=None):
+        self.timeout = timeout
+        self._watch_stub = watchstub
+        self._credentials = call_credentials
+        self._metadata = metadata
+
+        self._lock = asyncio.Lock()
+        self._request_queue = asyncio.Queue(maxsize=10)
+        self._callbacks = {}
+        self._task = None
+        self._new_watch_event = asyncio.Event()
+        self._new_watch = None
+        self._stopping = False
+
+    def _create_watch_request(self, key, config=None, **kwargs):
+        """Create a watch request.
+
+        :param key: The key to watch
+        :type key: str or bytes
+        :param config: Watch request configuration
+        :type config: WatchRequestConfig
+        :param kwargs: Additional keyword arguments for backward compatibility
+        :return: WatchRequest instance
+        """
+        create_watch = etcdrpc.WatchCreateRequest()
+        create_watch.key = utils.to_bytes(key)
+
+        # Use config object if provided, otherwise use kwargs or defaults
+        if config is None:
+            # Create config from kwargs for backward compatibility
+            config = WatchRequestConfig(
+                range_end=kwargs.get("range_end"),
+                start_revision=kwargs.get("start_revision"),
+                progress_notify=kwargs.get("progress_notify", False),
+                filters=kwargs.get("filters"),
+                prev_kv=kwargs.get("prev_kv", False),
+            )
+
+        if config.range_end is not None:
+            create_watch.range_end = utils.to_bytes(config.range_end)
+        if config.start_revision is not None:
+            create_watch.start_revision = config.start_revision
+        if config.progress_notify:
+            create_watch.progress_notify = config.progress_notify
+        if config.filters is not None:
+            create_watch.filters.extend(config.filters)
+        if config.prev_kv:
+            create_watch.prev_kv = config.prev_kv
+        return etcdrpc.WatchRequest(create_request=create_watch)
+
+    async def add_callback(self, key, callback, config=None, **kwargs):
+        """Add a callback for watch events.
+
+        :param key: The key to watch
+        :type key: str or bytes
+        :param callback: The callback function to call on watch events
+        :type callback: callable
+        :param config: Watch request configuration
+        :type config: WatchRequestConfig
+        :param kwargs: Additional keyword arguments for backward compatibility
+        :returns: The watch ID
+        :raises WatchTimedOut: If the watch creation times out
+        """
+        rq = self._create_watch_request(key, config=config, **kwargs)
+
+        # Start the task if it is not yet running.
+        async with self._lock:
+            # Wait for exiting task to close
+            if self._stopping:
+                if self._task:
+                    await self._task
+                    self._task = None
+                self._stopping = False
+
+            # Start the task if it is not yet running.
+            if not self._task:
+                self._task = asyncio.create_task(self._run())
+
+            # Only one create watch request can be pending at a time
+            while self._new_watch is not None:
+                await self._new_watch_event.wait()
+                self._new_watch_event.clear()
+
+            # Submit a create watch request.
+            new_watch = _NewWatch(callback)
+            await self._request_queue.put(rq)
+            self._new_watch = new_watch
+
+        try:
+            # Wait for the request to be completed, or timeout.
+            try:
+                await asyncio.wait_for(
+                    self._new_watch_event.wait(), timeout=self.timeout
+                )
+                self._new_watch_event.clear()
+            except asyncio.TimeoutError as exc:
+                # If the request not completed yet, then raise a timeout
+                # exception.
+                if new_watch.id is None and new_watch.err is None:
+                    raise exceptions.WatchTimedOut() from exc
+
+            # Raise an exception if the watch request failed.
+            if new_watch.err:
+                raise new_watch.err
+        finally:
+            # Wake up tasks stuck on add_callback call if any.
+            async with self._lock:
+                self._new_watch = None
+                self._new_watch_event.set()
+
+        return new_watch.id
+
+    async def cancel(self, watch_id):
+        """Cancel a watch.
+
+        :param watch_id: The ID of the watch to cancel
+        :type watch_id: int
+        """
+        async with self._lock:
+            callback = self._callbacks.pop(watch_id, None)
+            if not callback:
+                return
+
+            await self._cancel_no_lock(watch_id)
+
+    async def _run(self):
+        callback_err = None
+        try:
+            response_iter = self._watch_stub.Watch(
+                self._new_request_iter(),
+                credentials=self._credentials,
+                metadata=self._metadata,
+            )
+            async for rs in response_iter:
+                await self._handle_response(rs)
+
+        except grpc.RpcError as err:
+            callback_err = err
+
+        finally:
+            async with self._lock:
+                self._stopping = True
+                if self._new_watch:
+                    self._new_watch.err = callback_err
+                    self._new_watch_event.set()
+
+                callbacks = self._callbacks
+                self._callbacks = {}
+
+                # Rotate request queue
+                await self._request_queue.put(None)
+                self._request_queue = asyncio.Queue(maxsize=10)
+
+            for callback in callbacks.values():
+                await _safe_callback(callback, callback_err)
+
+    async def _handle_response(self, rs):
+        async with self._lock:
+            if rs.created:
+                # If the new watch request has already expired then cancel the
+                # created watch right away.
+                if not self._new_watch:
+                    await self._cancel_no_lock(rs.watch_id)
+                    return
+
+                if rs.compact_revision != 0:
+                    self._new_watch.err = exceptions.RevisionCompactedError(
+                        rs.compact_revision
+                    )
+                    self._new_watch_event.set()
+                    return
+
+                self._callbacks[rs.watch_id] = self._new_watch.callback
+                self._new_watch.id = rs.watch_id
+                self._new_watch_event.set()
+
+            callback = self._callbacks.get(rs.watch_id)
+
+        # Ignore leftovers from canceled watches.
+        if not callback:
+            return
+
+        # The watcher can be safely reused, but adding a new event
+        # to indicate that the revision is already compacted
+        # requires api change which would break all users of this
+        # module. So, raising an exception if a watcher is still
+        # alive.
+        if rs.compact_revision != 0:
+            err = exceptions.RevisionCompactedError(rs.compact_revision)
+            await _safe_callback(callback, err)
+            await self.cancel(rs.watch_id)
+            return
+
+        # Call the callback even when there are no events in the watch
+        # response so as not to ignore progress notify responses.
+        if rs.events or not (rs.created or rs.canceled):
+            new_events = [events.new_event(event) for event in rs.events]
+            response = WatchResponse(rs.header, new_events)
+            await _safe_callback(callback, response)
+
+    async def _cancel_no_lock(self, watch_id):
+        cancel_watch = etcdrpc.WatchCancelRequest()
+        cancel_watch.watch_id = watch_id
+        rq = etcdrpc.WatchRequest(cancel_request=cancel_watch)
+        await self._request_queue.put(rq)
+
+    async def close(self):
+        """Close the watcher and clean up resources.
+
+        This method stops the watcher task and cleans up any resources used by the watcher.
+        """
+        task_to_wait = None
+        async with self._lock:
+            if self._task and not self._stopping:
+                # Signal the task to stop
+                await self._request_queue.put(None)
+                self._stopping = True
+                task_to_wait = self._task
+                self._task = None
+            else:
+                self._stopping = True
+
+        # Wait for the task to complete if we have one
+        if task_to_wait is not None:
+            try:
+                await task_to_wait
+            except asyncio.CancelledError:
+                pass
+            except Exception:
+                _log.exception("Error waiting for watcher task to complete")
+
+    async def watch(self, key, **kwargs):
+        """Watch a key and return an event iterator.
+
+        :param key: The key to watch
+        :type key: str or bytes
+        :param config: Watch request configuration
+        :type config: WatchRequestConfig
+        :param kwargs: Additional keyword arguments for backward compatibility
+        :returns: An async iterator of watch responses
+        """
+        rq = self._create_watch_request(key, config=None, **kwargs)
+
+        response_iter = self._watch_stub.Watch(
+            self._watch_request_iter(rq),
+            credentials=self._credentials,
+            metadata=self._metadata,
+        )
+
+        async for rs in response_iter:
+            if rs.created:
+                continue
+            if rs.compact_revision != 0:
+                raise exceptions.RevisionCompactedError(rs.compact_revision)
+            if rs.canceled:
+                break
+            if rs.events:
+                new_events = [events.new_event(event) for event in rs.events]
+                yield WatchResponse(rs.header, new_events)
+
+    async def watch_once(self, key, timeout=None, **kwargs):
+        """Watch a key and return the first response.
+
+        :param key: The key to watch
+        :type key: str or bytes
+        :param timeout: Timeout in seconds
+        :type timeout: float or None
+        :param config: Watch request configuration
+        :type config: WatchRequestConfig
+        :param kwargs: Additional keyword arguments for backward compatibility
+        :returns: The first WatchResponse
+        :raises WatchTimedOut: If timeout is reached
+        """
+        rq = self._create_watch_request(key, config=None, **kwargs)
+
+        response_iter = self._watch_stub.Watch(
+            self._watch_request_iter(rq),
+            credentials=self._credentials,
+            metadata=self._metadata,
+        )
+
+        async def iter_with_timeout():
+            async for rs in response_iter:
+                yield rs
+
+        try:
+            async with asyncio.timeout(timeout):
+                async for rs in iter_with_timeout():
+                    if rs.created:
+                        continue
+                    if rs.compact_revision != 0:
+                        raise exceptions.RevisionCompactedError(rs.compact_revision)
+                    if rs.canceled:
+                        raise exceptions.WatchTimedOut()
+                    if rs.events:
+                        new_events = [events.new_event(event) for event in rs.events]
+                        return WatchResponse(rs.header, new_events)
+        except asyncio.TimeoutError:
+            raise exceptions.WatchTimedOut()
+
+    async def _watch_request_iter(self, initial_request):
+        """Create an iterator that yields watch requests.
+
+        :param initial_request: The initial watch request to send
+        :type initial_request: etcdrpc.WatchRequest
+        :yields: WatchRequest instances
+        """
+        yield initial_request
+        # Keep the stream open but don't yield more requests
+        # This is a simple implementation - a more robust one would handle
+        # additional requests coming in while watching
+        try:
+            async for _ in self._new_request_iter():
+                pass
+        except asyncio.CancelledError:
+            pass
+
+    async def _new_request_iter(self):
+        while True:
+            rq = await self._request_queue.get()
+            if rq is None:
+                return
+
+            yield rq
+
+
+class WatchResponse:
+    """Represents a watch response from etcd.
+
+    :param header: The response header
+    :type header: etcdrpc.ResponseHeader
+    :param events_list: The list of events in the response
+    :type events_list: list
+    """
+
+    def __init__(self, header, events_list):
+        self.header = header
+        self.events = events_list
+
+
+class _NewWatch:
+    def __init__(self, callback):
+        self.callback = callback
+        self.id = None
+        self.err = None
+
+
+async def _safe_callback(callback, response_or_err):
+    try:
+        await callback(response_or_err)
+
+    except Exception as exc:
+        _log.exception("Watch callback failed: %s", exc)