aws-advanced-python-wrapper 1.0.0__py3-none-any.whl

aws_advanced_python_wrapper/utils/telemetry/open_telemetry.py

@@ -0,0 +1,189 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Optional, Sequence, Union
+
+if TYPE_CHECKING:
+    from opentelemetry.util.types import AttributeValue
+
+from opentelemetry import context as context_api
+from opentelemetry import trace
+from opentelemetry import trace as trace_api
+from opentelemetry.metrics import (CallbackOptions, Meter, Observation,
+                                   get_meter)
+from opentelemetry.sdk.trace import ReadableSpan, Span, StatusCode, Tracer
+
+from aws_advanced_python_wrapper.utils.log import Logger
+from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.telemetry.telemetry import (
+    TelemetryConst, TelemetryContext, TelemetryCounter, TelemetryFactory,
+    TelemetryGauge, TelemetryTraceLevel)
+
+logger = Logger(__name__)
+
+INSTRUMENTATION_NAME = "aws-advanced-python-wrapper"
+
+
+class OpenTelemetryContext(TelemetryContext):
+    def __init__(self, tracer: Tracer, name: str, trace_level: TelemetryTraceLevel,
+                 start_time: Optional[int] = None, link_span: Optional[Span] = None):
+        self._name = name
+        self._tracer = tracer
+        self._span: Optional[Span]
+        self._meter: Meter
+        self._token: Optional[object] = None
+
+        current_span: Span = trace.get_current_span()  # type: ignore
+        is_root = (current_span is None or current_span == trace.INVALID_SPAN)
+
+        if is_root and trace_level == TelemetryTraceLevel.NESTED:
+            trace_level = TelemetryTraceLevel.TOP_LEVEL
+
+        links: Sequence[trace_api.Link] = ()
+        if trace_level in [TelemetryTraceLevel.FORCE_TOP_LEVEL, TelemetryTraceLevel.TOP_LEVEL]:
+            if link_span is not None:
+                links = [trace.Link(link_span.get_span_context())]
+            else:
+                if not is_root:
+                    links = [trace.Link(current_span.get_span_context())]
+
+            self._span = self._tracer.start_span(self._name, context=context_api.Context(),
+                                                 links=links, start_time=start_time)  # type: ignore
+            if not is_root:
+                self.set_attribute(TelemetryConst.TRACE_NAME_ANNOTATION, self._name)
+
+            ctx = trace.set_span_in_context(self._span)  # type: ignore
+            self._token = context_api.attach(ctx)
+            logger.debug("OpenTelemetryContext.TelemetryTraceID", self._name,
+                         self._span.get_span_context().trace_id)  # type: ignore
+
+        elif trace_level == TelemetryTraceLevel.NESTED:
+            if link_span is not None:
+                links = [trace.Link(link_span.get_span_context())]
+
+            self._span = self._tracer.start_span(self._name, links=links, start_time=start_time)  # type: ignore
+            ctx = trace.set_span_in_context(self._span)  # type: ignore
+            self._token = context_api.attach(ctx)
+            self.set_attribute(TelemetryConst.TRACE_NAME_ANNOTATION, self._name)
+
+        elif trace_level == TelemetryTraceLevel.NO_TRACE:
+            self._span = None
+
+    def set_success(self, success: bool):
+        if self._span is not None:
+            self._span.set_status(StatusCode.OK if success else StatusCode.ERROR)
+
+    def set_attribute(self, key: str, value: AttributeValue):
+        if self._span is not None:
+            self._span.set_attribute(key, value)
+
+    def set_exception(self, exception: Exception):
+        if self._span is not None and exception is not None:
+            self._span.set_attribute(TelemetryConst.EXCEPTION_TYPE_ANNOTATION, exception.__class__.__name__)
+            self._span.set_attribute(TelemetryConst.EXCEPTION_MESSAGE_ANNOTATION, str(exception))
+            self._span.record_exception(exception)
+
+    def get_name(self) -> str:
+        return self._name
+
+    def close_context(self):
+        if self._token is not None:
+            context_api.detach(self._token)
+        if self._span is not None:
+            self._span.end()
+
+    @property
+    def tracer(self) -> Tracer:
+        return self._tracer
+
+    @property
+    def span(self) -> Optional[Span]:
+        return self._span
+
+
+def post_copy(context: OpenTelemetryContext, trace_level: TelemetryTraceLevel):
+    if trace_level == TelemetryTraceLevel.NO_TRACE:
+        return
+
+    _clone_and_close_context(context, trace_level)
+
+
+def _clone_and_close_context(context: OpenTelemetryContext, trace_level: TelemetryTraceLevel) -> OpenTelemetryContext:
+    if not isinstance(context.span, ReadableSpan):
+        raise RuntimeError(Messages.get("OpenTelemetry.InvalidContext"))
+
+    clone = OpenTelemetryContext(
+        context.tracer, TelemetryConst.COPY_TRACE_NAME_PREFIX + context.get_name(),
+        trace_level, context.span.start_time, context.span)
+
+    for key in context.span.attributes:  # type: ignore
+        value = context.span.attributes[key]  # type: ignore
+        clone.set_attribute(key, value)
+
+    clone.span.set_status(context.span.status)  # type: ignore
+    clone.set_attribute(TelemetryConst.SOURCE_TRACE_ANNOTATION, str(context.span.get_span_context().trace_id))
+    clone.span.end(context.span.end_time)  # type: ignore
+
+    return clone
+
+
+class OpenTelemetryCounter(TelemetryCounter):
+    def __init__(self, meter: Meter, name: str):
+        self._meter: Meter = meter
+        self._name: str = name
+        self._counter = meter.create_up_down_counter(self._name, unit="1")
+
+    def add(self, value):
+        self._counter.add(value)
+
+    def inc(self):
+        self._counter.add(1)
+
+    def get_name(self):
+        return self._name
+
+
+class OpenTelemetryGauge(TelemetryGauge):
+    def __init__(self, meter: Meter, name: str, callback: Callable[[], Union[float, int]]):
+        self._meter: Meter = meter
+        self.name: str = name
+        self._counter = meter.create_observable_up_down_counter(name, callbacks=[self._callback_observation], unit="1")
+        self._callback: Callable[[], float] = callback
+
+    def get_name(self):
+        return self.name
+
+    def _callback_observation(self, options: CallbackOptions):
+        value: Union[int, float] = self._callback()
+        observation = Observation(value)
+        yield observation
+
+
+class OpenTelemetryFactory(TelemetryFactory):
+    def open_telemetry_context(self, name: str, trace_level: TelemetryTraceLevel) -> TelemetryContext:
+        return OpenTelemetryContext(trace.get_tracer(INSTRUMENTATION_NAME), name, trace_level)  # type: ignore
+
+    def post_copy(self, context: TelemetryContext, trace_level: TelemetryTraceLevel):
+        if isinstance(context, OpenTelemetryContext):
+            post_copy(context, trace_level)
+        else:
+            raise RuntimeError(Messages.get_formatted("OpenTelemetryFactory.WrongParameterType", type(context)))
+
+    def create_counter(self, name: str) -> TelemetryCounter:
+        return OpenTelemetryCounter(get_meter(INSTRUMENTATION_NAME), name)
+
+    def create_gauge(self, name: str, callback: Callable[[], Union[float, int]]):
+        return OpenTelemetryGauge(get_meter(INSTRUMENTATION_NAME), name, callback)

aws_advanced_python_wrapper/utils/telemetry/telemetry.py

@@ -0,0 +1,85 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from enum import Enum, auto
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from opentelemetry.util.types import AttributeValue
+
+
+class TelemetryConst:
+    TRACE_NAME_ANNOTATION = "trace_name"
+    SOURCE_TRACE_ANNOTATION = "source_trace_id"
+    PARENT_TRACE_ANNOTATION = "parent_trace_id"
+    EXCEPTION_TYPE_ANNOTATION = "exception_type"
+    EXCEPTION_MESSAGE_ANNOTATION = "exception_message"
+    COPY_TRACE_NAME_PREFIX = "copy: "
+
+
+class TelemetryTraceLevel(Enum):
+    FORCE_TOP_LEVEL = auto()
+    TOP_LEVEL = auto()
+    NESTED = auto()
+    NO_TRACE = auto()
+
+
+class TelemetryContext(ABC):
+    def set_success(self, success: bool):
+        pass
+
+    def set_attribute(self, key: str, value: AttributeValue):
+        pass
+
+    def set_exception(self, exception: Exception):
+        pass
+
+    def get_name(self):
+        pass
+
+    def close_context(self):
+        pass
+
+
+class TelemetryCounter(ABC):
+    def add(self, value):
+        pass
+
+    def inc(self):
+        pass
+
+
+class TelemetryGauge(ABC):
+    pass
+
+
+class TelemetryFactory(ABC):
+    @abstractmethod
+    def open_telemetry_context(self, name: str, trace_level: TelemetryTraceLevel) -> TelemetryContext:
+        pass
+
+    @abstractmethod
+    def post_copy(self, context: TelemetryContext, trace_level: TelemetryTraceLevel):
+        pass
+
+    @abstractmethod
+    def create_counter(self, name: str) -> TelemetryCounter:
+        pass
+
+    @abstractmethod
+    def create_gauge(self, name: str, callback):
+        pass

aws_advanced_python_wrapper/utils/telemetry/xray_telemetry.py

@@ -0,0 +1,126 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING
+
+from aws_xray_sdk.core import xray_recorder
+
+if TYPE_CHECKING:
+    from opentelemetry.util.types import AttributeValue
+    from aws_xray_sdk.core.models.segment import Segment
+    from aws_xray_sdk.core.models.subsegment import Subsegment
+
+from aws_advanced_python_wrapper.utils.log import Logger
+from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.telemetry.telemetry import (
+    TelemetryConst, TelemetryContext, TelemetryCounter, TelemetryFactory,
+    TelemetryTraceLevel)
+
+logger = Logger(__name__)
+
+
+class XRayTelemetryContext(TelemetryContext):
+    def __init__(self, name: str, trace_level: TelemetryTraceLevel):
+        self._name: str = name
+        self._trace_entity: Segment | Subsegment
+
+        if trace_level in [TelemetryTraceLevel.FORCE_TOP_LEVEL, TelemetryTraceLevel.TOP_LEVEL]:
+            self._trace_entity = xray_recorder.begin_segment(self._name)
+            self.is_segment = True
+            logger.debug("XRayTelemetryContext.TraceID", self._name, self._trace_entity.trace_id)
+        elif trace_level == TelemetryTraceLevel.NESTED:
+            self._trace_entity = xray_recorder.begin_subsegment(self._name)  # type: ignore
+            self.set_attribute(TelemetryConst.TRACE_NAME_ANNOTATION, self._name)
+            self.is_segment = False
+        elif trace_level == TelemetryTraceLevel.NO_TRACE:
+            pass
+
+    def set_success(self, success: bool):
+        if self._trace_entity is not None:
+            self._trace_entity.error = not success
+
+    def set_attribute(self, key: str, value: AttributeValue):
+        if self._trace_entity is not None:
+            self._trace_entity.put_annotation(key, value)
+
+    def set_exception(self, exception: Exception):
+        if self._trace_entity is not None and exception is not None:
+            self._trace_entity.put_annotation("exception_type", exception.__class__.__name__)
+            self._trace_entity.put_annotation("exception_message", str(exception))
+
+    def get_name(self):
+        return self._name
+
+    def close_context(self, end_time=None):
+        if self._trace_entity is not None:
+            if self.is_segment:
+                xray_recorder.end_segment(end_time)
+            else:
+                xray_recorder.end_subsegment(end_time)
+
+
+def post_copy(context: XRayTelemetryContext, trace_level: TelemetryTraceLevel):
+    if trace_level == TelemetryTraceLevel.NO_TRACE:
+        return
+
+    if trace_level in [TelemetryTraceLevel.FORCE_TOP_LEVEL, TelemetryTraceLevel.TOP_LEVEL]:
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(_clone_and_close_context, context, trace_level)
+            future.result()
+    else:
+        _clone_and_close_context(context, trace_level)
+
+
+def _clone_and_close_context(context: XRayTelemetryContext, trace_level: TelemetryTraceLevel) -> XRayTelemetryContext:
+    clone = XRayTelemetryContext(TelemetryConst.COPY_TRACE_NAME_PREFIX + context.get_name(), trace_level)
+
+    clone._trace_entity.start_time = context._trace_entity.start_time
+
+    for key in context._trace_entity.annotations.items():
+        value = context._trace_entity.annotations[key]
+        if key != TelemetryConst.TRACE_NAME_ANNOTATION and value is not None:
+            clone.set_attribute(key, value)
+
+    if context.is_segment and context._trace_entity.error:
+        clone._trace_entity.add_error_flag()
+
+    clone.set_attribute(TelemetryConst.SOURCE_TRACE_ANNOTATION, str(context._trace_entity.trace_id))
+
+    if context._trace_entity.parent_id is not None:
+        if trace_level == TelemetryTraceLevel.NESTED:
+            clone._trace_entity.parent_id = context._trace_entity.parent_id
+
+    clone.close_context(context._trace_entity.end_time)
+
+    return clone
+
+
+class XRayTelemetryFactory(TelemetryFactory):
+    def open_telemetry_context(self, name: str, trace_level: TelemetryTraceLevel) -> TelemetryContext:
+        return XRayTelemetryContext(name, trace_level)
+
+    def post_copy(self, context: TelemetryContext, trace_level: TelemetryTraceLevel):
+        if isinstance(context, XRayTelemetryContext):
+            post_copy(context, trace_level)
+        else:
+            raise RuntimeError(Messages.get_formatted("XRayTelemetryFactory.WrongParameterType", type(context)))
+
+    def create_counter(self, name: str) -> TelemetryCounter:
+        raise RuntimeError(Messages.get_formatted("XRayTelemetryFactory.MetricsNotSupported"))
+
+    def create_gauge(self, name: str, callback):
+        raise RuntimeError(Messages.get_formatted("XRayTelemetryFactory.MetricsNotSupported"))

aws_advanced_python_wrapper/utils/utils.py

@@ -0,0 +1,89 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+import sys
+from importlib import import_module
+from logging import DEBUG, Formatter, Logger, StreamHandler
+from queue import Empty, Queue
+from typing import TYPE_CHECKING, Optional, Tuple
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.hostinfo import HostInfo
+
+
+class LogUtils:
+    @staticmethod
+    def setup_logger(logger: Logger, level: int = DEBUG, format_string: Optional[str] = None):
+        for handler in logger.handlers:
+            if isinstance(handler, StreamHandler):
+                return
+
+        if format_string is None:
+            format_string = \
+                "%(asctime)s.%(msecs)03d %(name)-12s:%(funcName)s [%(levelname)-8s] - %(threadName)s - %(message)s"
+
+        handler = StreamHandler(stream=sys.stdout)
+        handler.setFormatter(Formatter(format_string))
+        handler.setLevel(level)
+
+        logger.setLevel(level)
+        logger.addHandler(handler)
+
+    @staticmethod
+    def log_topology(hosts: Tuple[HostInfo, ...], message_prefix: Optional[str] = None):
+        """Returns a formatted string of the topology that looks like the following
+        Topology: {
+            HostInfo[host=url, port=123]
+            <null>
+            HostInfo[host=url1, port=234]}
+
+        Args:
+            hosts (list): a list of hosts representing the cluster topology
+            message_prefix (str): messages to prefix to the topology information
+
+        Returns:
+            _type_: a formatted string of the topology
+        """
+        msg = "\n\t".join(["<null>" if not host else str(host) for host in hosts])
+        prefix = "" if not message_prefix else message_prefix + " "
+        return prefix + f": {{\n\t{msg}}}"
+
+
+class QueueUtils:
+    @staticmethod
+    def get(q: Queue):
+        try:
+            return q.get_nowait()
+        except Empty:
+            return None
+
+    @staticmethod
+    def clear(q: Queue):
+        with q.mutex:
+            q.queue.clear()
+            q.all_tasks_done.notify_all()
+            q.unfinished_tasks = 0
+
+
+class Utils:
+    @staticmethod
+    def initialize_class(full_class_name: str, *args):
+        try:
+            parts = full_class_name.split('.')
+            m = import_module(".".join(parts[:-1]))
+            return getattr(m, parts[-1])(*args)
+        except ModuleNotFoundError:
+            return None