Move all key access to a single spot and deduplicate context managers

Author: adriangb

instrumentation/opentelemetry-instrumentation-aiohttp-client/src/opentelemetry/instrumentation/aiohttp_client/__init__.py

@@ -94,8 +94,8 @@ def response_hook(span: Span, params: typing.Union[
 from opentelemetry.instrumentation.aiohttp_client.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import (
-    _SUPPRESS_INSTRUMENTATION_KEY,
     http_status_to_status_code,
+    is_instrumentation_enabled,
     unwrap,
 )
 from opentelemetry.propagate import inject
@@ -179,7 +179,7 @@ async def on_request_start(
         trace_config_ctx: types.SimpleNamespace,
         params: aiohttp.TraceRequestStartParams,
     ):
-        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        if not is_instrumentation_enabled():
             trace_config_ctx.span = None
             return
 
@@ -282,7 +282,7 @@ def _instrument(
 
     # pylint:disable=unused-argument
     def instrumented_init(wrapped, instance, args, kwargs):
-        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        if not is_instrumentation_enabled():
             return wrapped(*args, **kwargs)
 
         client_trace_configs = list(kwargs.get("trace_configs") or [])

instrumentation/opentelemetry-instrumentation-aiohttp-client/tests/test_aiohttp_client_integration.py

@@ -32,7 +32,7 @@
 from opentelemetry.instrumentation.aiohttp_client import (
     AioHttpClientInstrumentor,
 )
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.instrumentation.utils import suppress_instrumentation
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.trace import Span, StatusCode
@@ -506,25 +506,17 @@ async def uninstrument_request(server: aiohttp.test_utils.TestServer):
         self.assert_spans(1)
 
     def test_suppress_instrumentation(self):
-        token = context.attach(
-            context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)
-        )
-        try:
+        with suppress_instrumentation():
             run_with_test_server(
                 self.get_default_request(), self.URL, self.default_handler
             )
-        finally:
-            context.detach(token)
         self.assert_spans(0)
 
     @staticmethod
     async def suppressed_request(server: aiohttp.test_utils.TestServer):
         async with aiohttp.test_utils.TestClient(server) as client:
-            token = context.attach(
-                context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)
-            )
-            await client.get(TestAioHttpClientInstrumentor.URL)
-            context.detach(token)
+            with suppress_instrumentation():
+                await client.get(TestAioHttpClientInstrumentor.URL)
 
     def test_suppress_instrumentation_after_creation(self):
         run_with_test_server(

instrumentation/opentelemetry-instrumentation-boto3sqs/src/opentelemetry/instrumentation/boto3sqs/__init__.py

@@ -38,7 +38,7 @@
 from opentelemetry import context, propagate, trace
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import (
-    _SUPPRESS_INSTRUMENTATION_KEY,
+    is_instrumentation_enabled,
     unwrap,
 )
 from opentelemetry.propagators.textmap import CarrierT, Getter, Setter
@@ -218,7 +218,7 @@ def _create_processing_span(
 
     def _wrap_send_message(self, sqs_class: type) -> None:
         def send_wrapper(wrapped, instance, args, kwargs):
-            if context.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+            if not is_instrumentation_enabled():
                 return wrapped(*args, **kwargs)
             queue_url = kwargs.get("QueueUrl")
             # The method expect QueueUrl and Entries params, so if they are None, we call wrapped to receive the
@@ -252,7 +252,7 @@ def send_batch_wrapper(wrapped, instance, args, kwargs):
             # The method expect QueueUrl and Entries params, so if they are None, we call wrapped to receive the
             # original exception
             if (
-                context.get_value(_SUPPRESS_INSTRUMENTATION_KEY)
+                not is_instrumentation_enabled()
                 or not queue_url
                 or not entries
             ):

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py

@@ -87,9 +87,6 @@ def response_hook(span, service_name, operation_name, result):
 from wrapt import wrap_function_wrapper
 
 from opentelemetry import context as context_api
-
-# FIXME: fix the importing of this private attribute when the location of the _SUPPRESS_HTTP_INSTRUMENTATION_KEY is defined.
-from opentelemetry.context import _SUPPRESS_HTTP_INSTRUMENTATION_KEY
 from opentelemetry.instrumentation.botocore.extensions import _find_extension
 from opentelemetry.instrumentation.botocore.extensions.types import (
     _AwsSdkCallContext,
@@ -98,7 +95,8 @@ def response_hook(span, service_name, operation_name, result):
 from opentelemetry.instrumentation.botocore.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import (
-    _SUPPRESS_INSTRUMENTATION_KEY,
+    is_instrumentation_enabled,
+    suppress_http_instrumentation,
     unwrap,
 )
 from opentelemetry.propagators.aws.aws_xray_propagator import AwsXRayPropagator
@@ -171,7 +169,7 @@ def _patched_endpoint_prepare_request(
 
     # pylint: disable=too-many-branches
     def _patched_api_call(self, original_func, instance, args, kwargs):
-        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        if not is_instrumentation_enabled():
             return original_func(*args, **kwargs)
 
         call_context = _determine_call_context(instance, args)
@@ -200,25 +198,22 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
             _safe_invoke(extension.before_service_call, span)
             self._call_request_hook(span, call_context)
 
-            token = context_api.attach(
-                context_api.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True)
-            )
-
-            result = None
             try:
-                result = original_func(*args, **kwargs)
-            except ClientError as error:
-                result = getattr(error, "response", None)
-                _apply_response_attributes(span, result)
-                _safe_invoke(extension.on_error, span, error)
-                raise
-            else:
-                _apply_response_attributes(span, result)
-                _safe_invoke(extension.on_success, span, result)
+                with suppress_http_instrumentation():
+                    result = None
+                    try:
+                        result = original_func(*args, **kwargs)
+                    except ClientError as error:
+                        result = getattr(error, "response", None)
+                        _apply_response_attributes(span, result)
+                        _safe_invoke(extension.on_error, span, error)
+                        raise
+                    else:
+                        _apply_response_attributes(span, result)
+                        _safe_invoke(extension.on_success, span, result)
+                    finally:
+                        _safe_invoke(extension.after_service_call)
             finally:
-                context_api.detach(token)
-                _safe_invoke(extension.after_service_call)
-
                 self._call_response_hook(span, call_context, result)
 
             return result

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py

@@ -27,14 +27,8 @@
 )
 
 from opentelemetry import trace as trace_api
-from opentelemetry.context import (
-    _SUPPRESS_HTTP_INSTRUMENTATION_KEY,
-    attach,
-    detach,
-    set_value,
-)
 from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.instrumentation.utils import suppress_instrumentation
 from opentelemetry.propagate import get_global_textmap, set_global_textmap
 from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_HEADER_KEY
 from opentelemetry.semconv.trace import SpanAttributes
@@ -341,23 +335,16 @@ def check_headers(**kwargs):
     @mock_xray
     def test_suppress_instrumentation_xray_client(self):
         xray_client = self._make_client("xray")
-        token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
-        try:
+        with suppress_instrumentation():
             xray_client.put_trace_segments(TraceSegmentDocuments=["str1"])
             xray_client.put_trace_segments(TraceSegmentDocuments=["str2"])
-        finally:
-            detach(token)
         self.assertEqual(0, len(self.get_finished_spans()))
 
     @mock_xray
     def test_suppress_http_instrumentation_xray_client(self):
-        xray_client = self._make_client("xray")
-        token = attach(set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True))
-        try:
+        with suppress_instrumentation():
             xray_client.put_trace_segments(TraceSegmentDocuments=["str1"])
             xray_client.put_trace_segments(TraceSegmentDocuments=["str2"])
-        finally:
-            detach(token)
         self.assertEqual(2, len(self.get_finished_spans()))
 
     @mock_s3

instrumentation/opentelemetry-instrumentation-cassandra/src/opentelemetry/instrumentation/cassandra/__init__.py

@@ -43,9 +43,9 @@
 from wrapt import wrap_function_wrapper
 
 from opentelemetry import trace
-from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.cassandra.package import _instruments
 from opentelemetry.instrumentation.cassandra.version import __version__
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.semconv.trace import SpanAttributes
 

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py

@@ -63,6 +63,7 @@ def add(x, y):
 from timeit import default_timer
 from typing import Collection, Iterable
 
+from billiard import VERSION
 from billiard.einfo import ExceptionInfo
 from celery import signals  # pylint: disable=no-name-in-module
 
@@ -76,8 +77,6 @@ def add(x, y):
 from opentelemetry.propagators.textmap import Getter
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
-from billiard import VERSION
-
 
 if VERSION >= (4, 0, 1):
     from billiard.einfo import ExceptionWithTraceback

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py

@@ -14,8 +14,8 @@
 
 import logging
 
-from celery import registry  # pylint: disable=no-name-in-module
 from billiard import VERSION
+from celery import registry  # pylint: disable=no-name-in-module
 
 from opentelemetry.semconv.trace import SpanAttributes
 

instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/__init__.py

@@ -112,8 +112,8 @@ def instrument_consumer(consumer: Consumer, tracer_provider=None)
 from .package import _instruments
 from .utils import (
     KafkaPropertiesExtractor,
-    _end_current_consume_span,
     _create_new_consume_span,
+    _end_current_consume_span,
     _enrich_span,
     _get_span_name,
     _kafka_getter,

instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/utils.py

@@ -2,13 +2,13 @@
 from typing import List, Optional
 
 from opentelemetry import context, propagate
-from opentelemetry.trace import SpanKind, Link
 from opentelemetry.propagators import textmap
 from opentelemetry.semconv.trace import (
     MessagingDestinationKindValues,
     MessagingOperationValues,
     SpanAttributes,
 )
+from opentelemetry.trace import Link, SpanKind
 
 _LOG = getLogger(__name__)
 

instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py

@@ -14,13 +14,6 @@
 
 # pylint: disable=no-name-in-module
 
-from opentelemetry.semconv.trace import (
-    SpanAttributes,
-    MessagingDestinationKindValues,
-)
-from opentelemetry.test.test_base import TestBase
-from .utils import MockConsumer, MockedMessage
-
 from confluent_kafka import Consumer, Producer
 
 from opentelemetry.instrumentation.confluent_kafka import (
@@ -32,6 +25,13 @@
     KafkaContextGetter,
     KafkaContextSetter,
 )
+from opentelemetry.semconv.trace import (
+    MessagingDestinationKindValues,
+    SpanAttributes,
+)
+from opentelemetry.test.test_base import TestBase
+
+from .utils import MockConsumer, MockedMessage
 
 
 class TestConfluentKafka(TestBase):

instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py

@@ -40,7 +40,6 @@
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS,
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
-    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
     OTEL_PYTHON_INSTRUMENTATION_HTTP_CAPTURE_ALL_METHODS,
     get_excluded_urls,
 )

instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_aio_client.py

@@ -19,12 +19,11 @@
 import grpc
 from grpc.aio import ClientCallDetails
 
-from opentelemetry import context
 from opentelemetry.instrumentation.grpc._client import (
     OpenTelemetryClientInterceptor,
     _carrier_setter,
 )
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.instrumentation.utils import is_instrumentation_enabled
 from opentelemetry.propagate import inject
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
@@ -139,9 +138,10 @@ async def _wrap_stream_response(self, span, call):
             span.end()
 
     def tracing_skipped(self, client_call_details):
-        return context.get_value(
-            _SUPPRESS_INSTRUMENTATION_KEY
-        ) or not self.rpc_matches_filters(client_call_details)
+        return (
+            not is_instrumentation_enabled()
+            or not self.rpc_matches_filters(client_call_details)
+        )
 
     def rpc_matches_filters(self, client_call_details):
         return self._filter is None or self._filter(client_call_details)

instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_client.py

@@ -28,7 +28,7 @@
 from opentelemetry import context, trace
 from opentelemetry.instrumentation.grpc import grpcext
 from opentelemetry.instrumentation.grpc._utilities import RpcInfo
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.instrumentation.utils import is_instrumentation_enabled
 from opentelemetry.propagate import inject
 from opentelemetry.propagators.textmap import Setter
 from opentelemetry.semconv.trace import SpanAttributes
@@ -123,7 +123,7 @@ def _trace_result(self, span, rpc_info, result):
         return result
 
     def _intercept(self, request, metadata, client_info, invoker):
-        if context.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        if not is_instrumentation_enabled():
             return invoker(request, metadata)
 
         if not metadata:
@@ -219,7 +219,7 @@ def _intercept_server_stream(
     def intercept_stream(
         self, request_or_iterator, metadata, client_info, invoker
     ):
-        if context.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        if not is_instrumentation_enabled():
             return invoker(request_or_iterator, metadata)
 
         if self._filter is not None and not self._filter(client_info):