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

etcd3/shared.py

@@ -0,0 +1,486 @@
+"""Shared utilities for etcd3 client.
+
+This module contains shared classes and functions used by both
+synchronous and asynchronous etcd3 client implementations.
+"""
+
+import functools
+
+import grpc
+import etcdrpc
+
+import etcd3.exceptions as exceptions
+import etcd3.transactions as transactions
+import etcd3.utils as utils
+from etcd3.trace import create_span, get_tracer
+
+
+_EXCEPTIONS_BY_CODE = {
+    grpc.StatusCode.INTERNAL: exceptions.InternalServerError,
+    grpc.StatusCode.UNAVAILABLE: exceptions.ConnectionFailedError,
+    grpc.StatusCode.DEADLINE_EXCEEDED: exceptions.ConnectionTimeoutError,
+    grpc.StatusCode.FAILED_PRECONDITION: exceptions.PreconditionFailedError,
+    grpc.StatusCode.RESOURCE_EXHAUSTED: exceptions.ConnectionFailedError,
+    grpc.StatusCode.ABORTED: exceptions.PreconditionFailedError,
+    grpc.StatusCode.OUT_OF_RANGE: exceptions.RevisionCompactedError,
+}
+
+_FAILED_EP_CODES = [
+    grpc.StatusCode.UNAVAILABLE,
+    grpc.StatusCode.DEADLINE_EXCEEDED,
+    grpc.StatusCode.INTERNAL,
+]
+
+
+class Transactions(object):
+    """Transaction operations for etcd3."""
+
+    def __init__(self):
+        self.value = transactions.Value
+        self.version = transactions.Version
+        self.create = transactions.Create
+        self.mod = transactions.Mod
+        self.lease = transactions.Lease
+
+        self.put = transactions.Put
+        self.get = transactions.Get
+        self.delete = transactions.Delete
+        self.txn = transactions.Txn
+
+
+class KVMetadata(object):
+    """Metadata for key-value pairs."""
+
+    __slots__ = (
+        "key",
+        "value",
+        "create_revision",
+        "mod_revision",
+        "version",
+        "lease_id",
+        "response_header",
+    )
+
+    def __init__(self, keyvalue, header):
+        self.key = keyvalue.key
+        self.value = keyvalue.value
+        self.create_revision = keyvalue.create_revision
+        self.mod_revision = keyvalue.mod_revision
+        self.version = keyvalue.version
+        self.lease_id = keyvalue.lease
+        self.response_header = header
+
+
+class Status(object):
+    """Etcd cluster status information."""
+
+    __slots__ = ("version", "db_size", "leader", "raft_index", "raft_term")
+
+    def __init__(self, version, db_size, leader, raft_index, raft_term):
+        self.version = version
+        self.db_size = db_size
+        self.leader = leader
+        self.raft_index = raft_index
+        self.raft_term = raft_term
+
+
+class Alarm(object):
+    """Etcd alarm information."""
+
+    __slots__ = ("alarm_type", "member_id")
+
+    def __init__(self, alarm_type, member_id):
+        self.alarm_type = alarm_type
+        self.member_id = member_id
+
+
+class EtcdTokenCallCredentials(grpc.AuthMetadataPlugin):
+    """Metadata wrapper for raw access token credentials."""
+
+    def __init__(self, access_token):
+        self._access_token = access_token
+
+    def __call__(self, context, callback):
+        metadata = (("token", self._access_token),)
+        callback(metadata, None)
+
+
+def _handle_errors(payload):
+    """Decorator to handle gRPC errors for synchronous clients."""
+
+    @functools.wraps(payload)
+    def handler(self, *args, **kwargs):
+        try:
+            return payload(self, *args, **kwargs)
+        except grpc.RpcError as exc:
+            self._manage_grpc_errors(exc)
+
+    return handler
+
+
+def _handle_generator_errors(payload):
+    """Decorator to handle gRPC errors for generators in synchronous clients."""
+
+    @functools.wraps(payload)
+    def handler(self, *args, **kwargs):
+        try:
+            for item in payload(self, *args, **kwargs):
+                yield item
+        except grpc.RpcError as exc:
+            self._manage_grpc_errors(exc)
+
+    return handler
+
+
+def get_secure_creds(ca_cert, cert_key=None, cert_cert=None):
+    """Get gRPC secure channel credentials.
+
+    :param ca_cert: Path to CA certificate file
+    :type ca_cert: str
+    :param cert_key: Path to client certificate key file
+    :type cert_key: str, optional
+    :param cert_cert: Path to client certificate file
+    :type cert_cert: str, optional
+    :return: gRPC channel credentials
+    :rtype: grpc.ChannelCredentials
+    """
+    cert_key_file = None
+    cert_cert_file = None
+
+    with open(ca_cert, "rb") as f:
+        ca_cert_file = f.read()
+
+    if cert_key is not None:
+        with open(cert_key, "rb") as f:
+            cert_key_file = f.read()
+
+    if cert_cert is not None:
+        with open(cert_cert, "rb") as f:
+            cert_cert_file = f.read()
+
+    return grpc.ssl_channel_credentials(ca_cert_file, cert_key_file, cert_cert_file)
+
+
+EXCEPTIONS_BY_CODE = _EXCEPTIONS_BY_CODE
+FAILED_EP_CODES = _FAILED_EP_CODES
+
+
+class RequestBuilder:
+    """Shared request builder for gRPC requests."""
+
+    @staticmethod
+    def build_get_range_request(
+        key,
+        range_end=None,
+        limit=None,
+        revision=None,
+        sort_order=None,
+        sort_target="key",
+        serializable=False,
+        keys_only=False,
+        count_only=False,
+        min_mod_revision=None,
+        max_mod_revision=None,
+        min_create_revision=None,
+        max_create_revision=None,
+    ):
+        """Build a RangeRequest.
+
+        :param key: The key to query
+        :type key: str or bytes
+        :param range_end: End of the key range
+        :type range_end: str or bytes, optional
+        :param limit: Maximum number of keys to return
+        :type limit: int, optional
+        :param revision: Revision to get
+        :type revision: int, optional
+        :param sort_order: Sort order ('ascend', 'descend', or None)
+        :type sort_order: str, optional
+        :param sort_target: Sort target ('key', 'version', 'create', 'mod')
+        :type sort_target: str
+        :param serializable: Whether to use serializable read
+        :type serializable: bool
+        :param keys_only: Whether to return only keys
+        :type keys_only: bool
+        :param count_only: Whether to return only count
+        :type count_only: bool
+        :param min_mod_revision: Minimum modification revision
+        :type min_mod_revision: int, optional
+        :param max_mod_revision: Maximum modification revision
+        :type max_mod_revision: int, optional
+        :param min_create_revision: Minimum creation revision
+        :type min_create_revision: int, optional
+        :param max_create_revision: Maximum creation revision
+        :type max_create_revision: int, optional
+        :return: RangeRequest instance
+        :rtype: etcdrpc.RangeRequest
+        """
+        range_request = etcdrpc.RangeRequest()
+        range_request.key = utils.to_bytes(key)
+        range_request.keys_only = keys_only
+        range_request.count_only = count_only
+        range_request.serializable = serializable
+
+        if range_end is not None:
+            range_request.range_end = utils.to_bytes(range_end)
+        if limit is not None:
+            range_request.limit = limit
+        if revision is not None:
+            range_request.revision = revision
+        if min_mod_revision is not None:
+            range_request.min_mod_revision = min_mod_revision
+        if max_mod_revision is not None:
+            range_request.max_mod_revision = max_mod_revision
+        if min_create_revision is not None:
+            range_request.min_create_revision = min_create_revision
+        if max_create_revision is not None:
+            range_request.max_create_revision = max_create_revision
+
+        sort_orders = {
+            None: etcdrpc.RangeRequest.NONE,
+            "ascend": etcdrpc.RangeRequest.ASCEND,
+            "descend": etcdrpc.RangeRequest.DESCEND,
+        }
+        request_sort_order = sort_orders.get(sort_order)
+        if request_sort_order is None:
+            raise ValueError(f'unknown sort order: "{sort_order}"')
+        range_request.sort_order = request_sort_order
+
+        sort_targets = {
+            None: etcdrpc.RangeRequest.KEY,
+            "key": etcdrpc.RangeRequest.KEY,
+            "version": etcdrpc.RangeRequest.VERSION,
+            "create": etcdrpc.RangeRequest.CREATE,
+            "mod": etcdrpc.RangeRequest.MOD,
+        }
+        request_sort_target = sort_targets.get(sort_target)
+        if request_sort_target is None:
+            raise ValueError(
+                'sort_target must be one of "key", "version", "create" or "mod"'
+            )
+        range_request.sort_target = request_sort_target
+
+        return range_request
+
+    @staticmethod
+    def build_put_request(
+        key,
+        value=None,
+        lease=None,
+        prev_kv=False,
+        ignore_value=False,
+        ignore_lease=False,
+    ):
+        """Build a PutRequest.
+
+        :param key: The key to set
+        :type key: str or bytes
+        :param value: The value to set
+        :type value: str or bytes, optional
+        :param lease: Lease ID to associate with the key
+        :type lease: int, optional
+        :param prev_kv: Whether to return the previous key-value pair
+        :type prev_kv: bool
+        :param ignore_value: Whether to ignore the value field
+        :type ignore_value: bool
+        :param ignore_lease: Whether to ignore the lease field
+        :type ignore_lease: bool
+        :return: PutRequest instance
+        :rtype: etcdrpc.PutRequest
+        """
+        put_request = etcdrpc.PutRequest()
+        put_request.key = utils.to_bytes(key)
+        if value is not None:
+            put_request.value = utils.to_bytes(value)
+        put_request.lease = utils.lease_to_id(lease)
+        put_request.prev_kv = prev_kv
+        put_request.ignore_value = ignore_value
+        put_request.ignore_lease = ignore_lease
+        return put_request
+
+    @staticmethod
+    def build_delete_request(key, range_end=None, prev_kv=False):
+        """Build a DeleteRangeRequest.
+
+        :param key: The key to delete
+        :type key: str or bytes
+        :param range_end: End of the key range (for prefix deletes)
+        :type range_end: str or bytes, optional
+        :param prev_kv: Whether to return the deleted key-value pairs
+        :type prev_kv: bool
+        :return: DeleteRangeRequest instance
+        :rtype: etcdrpc.DeleteRangeRequest
+        """
+        delete_request = etcdrpc.DeleteRangeRequest()
+        delete_request.key = utils.to_bytes(key)
+        delete_request.prev_kv = prev_kv
+
+        if range_end is not None:
+            delete_request.range_end = utils.to_bytes(range_end)
+
+        return delete_request
+
+
+def _make_async_handle_errors_with_tracing(error_handler_func, operation_name: str):
+    """Factory function to create async error handlers with tracing support.
+
+    Args:
+        error_handler_func: A callable that takes (client, exc) to handle errors
+        operation_name: Name of the operation for tracing
+
+    Returns:
+        A decorator for async methods that handles gRPC errors and creates spans
+    """
+
+    def decorator(payload):
+        @functools.wraps(payload)
+        async def handler(self, *args, **kwargs):
+            tracer = get_tracer()
+            if tracer is None:
+                # No tracing available, use basic error handling
+                try:
+                    return await payload(self, *args, **kwargs)
+                except grpc.RpcError as exc:
+                    error_handler_func(self, exc)
+
+            # Create span for the operation
+            with create_span(operation_name, {"operation": operation_name}):
+                try:
+                    return await payload(self, *args, **kwargs)
+                except grpc.RpcError as exc:
+                    error_handler_func(self, exc)
+
+        return handler
+
+    return decorator
+
+
+def _make_async_handle_generator_errors_with_tracing(
+    error_handler_func,
+    operation_name: str,
+    propagate_errors: bool = False,
+):
+    """Factory function to create async generator error handlers with tracing.
+
+    Args:
+        error_handler_func: A callable that takes (client, exc) to handle errors
+        operation_name: Name of the operation for tracing
+        propagate_errors: If True, re-raise gRPC errors after calling error_handler_func
+
+    Returns:
+        A decorator for async generators that handles gRPC errors and creates spans
+    """
+
+    def decorator(payload):
+        @functools.wraps(payload)
+        async def handler(self, *args, **kwargs):
+            tracer = get_tracer()
+            try:
+                result = payload(self, *args, **kwargs)
+                if hasattr(result, "__aiter__"):
+                    async for item in result:
+                        if tracer is not None:
+                            with create_span(
+                                operation_name, {"operation": operation_name}
+                            ):
+                                yield item
+                        else:
+                            yield item
+                else:
+                    async for item in await result:
+                        if tracer is not None:
+                            with create_span(
+                                operation_name, {"operation": operation_name}
+                            ):
+                                yield item
+                        else:
+                            yield item
+            except grpc.RpcError as exc:
+                error_handler_func(self, exc)
+                if propagate_errors:
+                    raise
+            except StopIteration:
+                # Normal iteration completion, just exit
+                return
+
+        return handler
+
+    return decorator
+
+
+def _make_sync_handle_errors_with_tracing(error_handler_func, operation_name: str):
+    """Factory function to create sync error handlers with tracing support.
+
+    Args:
+        error_handler_func: A callable that takes (client, exc) to handle errors
+        operation_name: Name of the operation for tracing
+
+    Returns:
+        A decorator for sync methods that handles gRPC errors and creates spans
+    """
+
+    def decorator(payload):
+        @functools.wraps(payload)
+        def handler(self, *args, **kwargs):
+            tracer = get_tracer()
+            if tracer is None:
+                # No tracing available, use basic error handling
+                try:
+                    return payload(self, *args, **kwargs)
+                except grpc.RpcError as exc:
+                    error_handler_func(self, exc)
+
+            # Create span for the operation
+            with create_span(operation_name, {"operation": operation_name}):
+                try:
+                    return payload(self, *args, **kwargs)
+                except grpc.RpcError as exc:
+                    error_handler_func(self, exc)
+
+        return handler
+
+    return decorator
+
+
+def _make_sync_handle_generator_errors_with_tracing(
+    error_handler_func, operation_name: str, propagate_errors: bool = False
+):
+    """Factory function to create sync generator error handlers with tracing.
+
+    Args:
+        error_handler_func: A callable that takes (client, exc) to handle errors
+        operation_name: Name of the operation for tracing
+        propagate_errors: If True, re-raise gRPC errors after calling error_handler_func
+
+    Returns:
+        A decorator for sync generators that handles gRPC errors and creates spans
+    """
+
+    def decorator(payload):
+        @functools.wraps(payload)
+        def handler(self, *args, **kwargs):
+            tracer = get_tracer()
+            if tracer is None:
+                # No tracing available, use basic error handling
+                try:
+                    for item in payload(self, *args, **kwargs):
+                        yield item
+                except grpc.RpcError as exc:
+                    error_handler_func(self, exc)
+                    if propagate_errors:
+                        raise
+                return
+
+            # Create span for the operation
+            with create_span(operation_name, {"operation": operation_name}):
+                try:
+                    for item in payload(self, *args, **kwargs):
+                        yield item
+                except grpc.RpcError as exc:
+                    error_handler_func(self, exc)
+                    if propagate_errors:
+                        raise
+
+        return handler
+
+    return decorator

etcd3/trace.py

@@ -0,0 +1,231 @@
+"""OpenTelemetry tracing utilities for etcd3 client.
+
+This module provides OpenTelemetry tracing support for etcd3 operations.
+It automatically gets the tracer from the global tracer provider.
+"""
+
+import os
+from contextlib import contextmanager
+from typing import Optional, Dict, Any
+
+OTEL_AVAILABLE = False
+try:
+    from opentelemetry import trace
+    from opentelemetry.trace import SpanKind, Status, StatusCode
+    from opentelemetry.trace.propagation.tracecontext import (
+        TraceContextTextMapPropagator,
+    )
+
+    OTEL_AVAILABLE = True
+except ImportError:
+    pass
+
+_tracer = None
+_propagator = None
+_tracing_initialized = False
+
+ENV_PREFIX = "ETCD_OBS_"
+
+
+def init_tracing(
+    service_name: str = "etcd3-client",
+    endpoint: Optional[str] = None,
+    insecure: bool = True,
+    sample_rate: float = 1.0,
+):
+    """Initialize tracing with optional OTLP exporter.
+
+    Args:
+        service_name: Name of the service for tracing
+        endpoint: OTLP exporter endpoint (e.g., "localhost:4317")
+        insecure: Whether to use insecure connection for OTLP
+        sample_rate: Trace sampling rate (0.0-1.0)
+
+    Returns:
+        True if tracing was initialized successfully, False otherwise
+    """
+    global _tracer, _propagator, _tracing_initialized
+
+    if not OTEL_AVAILABLE:
+        return False
+
+    try:
+        # Check if user already set a global provider
+        try:
+            provider = trace.get_tracer_provider()
+            # If user configured their own provider, use it
+            if provider and hasattr(provider, "get_tracer"):
+                _tracer = provider.get_tracer(service_name)
+                _tracing_initialized = True
+                return True
+        except Exception:
+            pass
+
+        # Try to auto-configure OTLP exporter if endpoint is provided
+        if endpoint is None:
+            endpoint = os.environ.get(f"{ENV_PREFIX}OTLP_ENDPOINT")
+
+        if endpoint:
+            try:
+                from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+                    OTLPSpanExporter,
+                )
+                from opentelemetry.sdk.trace import TracerProvider
+                from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+                exporter = OTLPSpanExporter(
+                    endpoint=endpoint,
+                    insecure=insecure,
+                )
+                provider = TracerProvider()
+
+                # Apply sampling if rate < 1.0
+                if sample_rate < 1.0:
+                    from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
+
+                    sampler = TraceIdRatioBased(sample_rate)
+                    provider = TracerProvider(sampler=sampler)
+
+                provider.add_span_processor(BatchSpanProcessor(exporter))
+                trace.set_tracer_provider(provider)
+            except ImportError:
+                pass  # OTLP exporter not installed
+        elif sample_rate < 1.0:
+            # Configure sampler even without OTLP endpoint
+            try:
+                from opentelemetry.sdk.trace import TracerProvider
+                from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
+
+                sampler = TraceIdRatioBased(sample_rate)
+                provider = TracerProvider(sampler=sampler)
+                trace.set_tracer_provider(provider)
+            except ImportError:
+                pass
+
+        # Get tracer from provider
+        provider = trace.get_tracer_provider()
+        _tracer = provider.get_tracer(service_name)
+        _propagator = TraceContextTextMapPropagator()
+        _tracing_initialized = True
+        return True
+    except Exception:  # pylint: disable=broad-except
+        _tracer = None
+        _propagator = None
+        _tracing_initialized = False
+        return False
+
+
+def get_tracer():
+    """Get the tracer instance.
+
+    Returns:
+        Tracer instance or None if not available
+    """
+    global _tracer, _tracing_initialized
+    if _tracer is None and OTEL_AVAILABLE and not _tracing_initialized:
+        init_tracing()
+    return _tracer
+
+
+def is_tracing_available() -> bool:
+    """Check if OpenTelemetry tracing is available.
+
+    Returns:
+        True if OpenTelemetry is installed and initialized, False otherwise
+    """
+    return OTEL_AVAILABLE and get_tracer() is not None
+
+
+@contextmanager
+def create_span(
+    operation_name: str, attributes: Optional[Dict[str, Any]] = None, kind=None
+):
+    """Create a span context manager for tracing an operation.
+
+    Args:
+        operation_name: Name of the operation/span
+        attributes: Optional attributes to add to the span
+        kind: Optional span kind (default: CLIENT)
+
+    Yields:
+        Span object if tracing is available, None otherwise
+    """
+    tracer = get_tracer()
+    if tracer is None:
+        yield None
+        return
+
+    if kind is None:
+        kind = SpanKind.CLIENT
+
+    with tracer.start_as_current_span(operation_name, kind=kind) as span:
+        if attributes:
+            for key, value in attributes.items():
+                if value is not None:
+                    span.set_attribute(key, value)
+        try:
+            yield span
+            span.set_status(Status(StatusCode.OK))
+        except Exception as e:
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.record_exception(e)
+            raise
+
+
+def inject_context(carrier: dict):
+    """Inject trace context into a carrier dictionary.
+
+    Args:
+        carrier: Dictionary to inject context into
+    """
+    if _propagator:
+        _propagator.inject(carrier)
+
+
+def extract_context(carrier: dict):
+    """Extract trace context from a carrier dictionary.
+
+    Args:
+        carrier: Dictionary containing trace context
+
+    Returns:
+        Extracted context or None
+    """
+    if _propagator:
+        return _propagator.extract(carrier)
+    return None
+
+
+def get_current_span():
+    """Get the current active span.
+
+    Returns:
+        Current span or None
+    """
+    if not is_tracing_available():
+        return None
+    return trace.get_current_span()
+
+
+def add_span_attribute(key: str, value: Any):
+    """Add an attribute to the current span.
+
+    Args:
+        key: Attribute key
+        value: Attribute value
+    """
+    span = get_current_span()
+    if span:
+        span.set_attribute(key, value)
+
+
+def add_span_event(name: str, attributes: Optional[Dict[str, Any]] = None):
+    """Add an event to the current span.
+
+    Args:
+        name: Event name
+        attributes: Optional event attributes
+    """
+    span = get_current_span()
+    if span:
+        span.add_event(name, attributes=attributes or {})