diff --git a/instrumentation/opentelemetry-instrumentation-aiohttp-client/src/opentelemetry/instrumentation/aiohttp_client/__init__.py b/instrumentation/opentelemetry-instrumentation-aiohttp-client/src/opentelemetry/instrumentation/aiohttp_client/__init__.py index ef3b667c98..9f842bde79 100644 --- a/instrumentation/opentelemetry-instrumentation-aiohttp-client/src/opentelemetry/instrumentation/aiohttp_client/__init__.py +++ b/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 []) diff --git a/instrumentation/opentelemetry-instrumentation-aiohttp-client/tests/test_aiohttp_client_integration.py b/instrumentation/opentelemetry-instrumentation-aiohttp-client/tests/test_aiohttp_client_integration.py index 9c73071465..9b27519185 100644 --- a/instrumentation/opentelemetry-instrumentation-aiohttp-client/tests/test_aiohttp_client_integration.py +++ b/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( diff --git a/instrumentation/opentelemetry-instrumentation-boto3sqs/src/opentelemetry/instrumentation/boto3sqs/__init__.py b/instrumentation/opentelemetry-instrumentation-boto3sqs/src/opentelemetry/instrumentation/boto3sqs/__init__.py index 137c570ac6..ee7f4a59a6 100644 --- a/instrumentation/opentelemetry-instrumentation-boto3sqs/src/opentelemetry/instrumentation/boto3sqs/__init__.py +++ b/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 ): diff --git a/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py b/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py index 686b040b13..6a4f337ec5 100644 --- a/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py +++ b/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,21 @@ 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: - context_api.detach(token) _safe_invoke(extension.after_service_call) - self._call_response_hook(span, call_context, result) return result diff --git a/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py b/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py index 3d25dcbf2d..993faa97a4 100644 --- a/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py +++ b/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_http_instrumentation, 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_http_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 diff --git a/instrumentation/opentelemetry-instrumentation-cassandra/src/opentelemetry/instrumentation/cassandra/__init__.py b/instrumentation/opentelemetry-instrumentation-cassandra/src/opentelemetry/instrumentation/cassandra/__init__.py index 202bf03712..75d2bcfe38 100644 --- a/instrumentation/opentelemetry-instrumentation-cassandra/src/opentelemetry/instrumentation/cassandra/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-cassandra/src/opentelemetry/instrumentation/cassandra/__init__.py @@ -70,10 +70,7 @@ def _traced_execute_async(func, instance, args, kwargs): if span.is_recording(): span.set_attribute(SpanAttributes.DB_NAME, instance.keyspace) span.set_attribute(SpanAttributes.DB_SYSTEM, "cassandra") - span.set_attribute( - SpanAttributes.NET_PEER_NAME, - instance.cluster.contact_points, - ) + span.set_attribute(SpanAttributes.NET_PEER_NAME, instance.cluster.contact_points) if include_db_statement: query = args[0] @@ -82,9 +79,7 @@ def _traced_execute_async(func, instance, args, kwargs): response = func(*args, **kwargs) return response - wrap_function_wrapper( - "cassandra.cluster", "Session.execute_async", _traced_execute_async - ) + wrap_function_wrapper("cassandra.cluster", "Session.execute_async", _traced_execute_async) class CassandraInstrumentor(BaseInstrumentor): diff --git a/instrumentation/opentelemetry-instrumentation-cassandra/tests/test_cassandra_integration.py b/instrumentation/opentelemetry-instrumentation-cassandra/tests/test_cassandra_integration.py index ed488ab07f..6977e1b2a2 100644 --- a/instrumentation/opentelemetry-instrumentation-cassandra/tests/test_cassandra_integration.py +++ b/instrumentation/opentelemetry-instrumentation-cassandra/tests/test_cassandra_integration.py @@ -46,25 +46,15 @@ def tearDown(self): def test_instrument_uninstrument(self): instrumentation = CassandraInstrumentor() instrumentation.instrument() - self.assertTrue( - isinstance( - cassandra.cluster.Session.execute_async, BoundFunctionWrapper - ) - ) + self.assertTrue(isinstance(cassandra.cluster.Session.execute_async, BoundFunctionWrapper)) instrumentation.uninstrument() - self.assertFalse( - isinstance( - cassandra.cluster.Session.execute_async, BoundFunctionWrapper - ) - ) + self.assertFalse(isinstance(cassandra.cluster.Session.execute_async, BoundFunctionWrapper)) @mock.patch("cassandra.cluster.Cluster.connect") @mock.patch("cassandra.cluster.Session.__init__") @mock.patch("cassandra.cluster.Session._create_response_future") - def test_instrumentor( - self, mock_create_response_future, mock_session_init, mock_connect - ): + def test_instrumentor(self, mock_create_response_future, mock_session_init, mock_connect): mock_create_response_future.return_value = mock.Mock() mock_session_init.return_value = None mock_connect.return_value = cassandra.cluster.Session() @@ -95,9 +85,7 @@ def test_instrumentor( @mock.patch("cassandra.cluster.Cluster.connect") @mock.patch("cassandra.cluster.Session.__init__") @mock.patch("cassandra.cluster.Session._create_response_future") - def test_custom_tracer_provider( - self, mock_create_response_future, mock_session_init, mock_connect - ): + def test_custom_tracer_provider(self, mock_create_response_future, mock_session_init, mock_connect): mock_create_response_future.return_value = mock.Mock() mock_session_init.return_value = None mock_connect.return_value = cassandra.cluster.Session() @@ -119,9 +107,7 @@ def test_custom_tracer_provider( @mock.patch("cassandra.cluster.Cluster.connect") @mock.patch("cassandra.cluster.Session.__init__") @mock.patch("cassandra.cluster.Session._create_response_future") - def test_instrument_connection_no_op_tracer_provider( - self, mock_create_response_future, mock_session_init, mock_connect - ): + def test_instrument_connection_no_op_tracer_provider(self, mock_create_response_future, mock_session_init, mock_connect): mock_create_response_future.return_value = mock.Mock() mock_session_init.return_value = None mock_connect.return_value = cassandra.cluster.Session() diff --git a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py index a86bc3166a..bf641aaed4 100644 --- a/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py +++ b/instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py @@ -328,9 +328,7 @@ def test_flask_metric_values(self): if isinstance(point, NumberDataPoint): self.assertEqual(point.value, 0) - def _assert_basic_metric( - self, expected_duration_attributes, expected_requests_count_attributes - ): + def _assert_basic_metric(self, expected_duration_attributes, expected_requests_count_attributes): metrics_list = self.memory_metrics_reader.get_metrics_data() for resource_metric in metrics_list.resource_metrics: for scope_metrics in resource_metric.scope_metrics: @@ -396,7 +394,7 @@ def test_basic_metric_nonstandard_http_method_success(self): ) @patch.dict( - "os.environ", + "os.environ", { OTEL_PYTHON_INSTRUMENTATION_HTTP_CAPTURE_ALL_METHODS: "1", }, diff --git a/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_aio_client.py b/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_aio_client.py index 5d5a5ccc46..8fc992be73 100644 --- a/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_aio_client.py +++ b/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) diff --git a/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_client.py b/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_client.py index b966fff4db..0a4f8df866 100644 --- a/instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_client.py +++ b/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): diff --git a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_aio_client_interceptor.py b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_aio_client_interceptor.py index 6ca5ce92d5..ad95711aa1 100644 --- a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_aio_client_interceptor.py +++ b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_aio_client_interceptor.py @@ -39,7 +39,7 @@ def run(self, result=None): from opentelemetry.instrumentation.grpc._aio_client import ( UnaryUnaryAioClientInterceptor, ) -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.semconv.trace import SpanAttributes from opentelemetry.test.mock_textmap import MockTextMapPropagator @@ -314,53 +314,33 @@ async def test_client_interceptor_trace_context_propagation(self): set_global_textmap(previous_propagator) async def test_unary_unary_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): response = await simple_method(self._stub) assert response.response_data == "data" spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 0) - finally: - context.detach(token) async def test_unary_stream_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): async for response in server_streaming_method(self._stub): self.assertEqual(response.response_data, "data") spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 0) - finally: - context.detach(token) async def test_stream_unary_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): response = await client_streaming_method(self._stub) assert response.response_data == "data" spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 0) - finally: - context.detach(token) async def test_stream_stream_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): async for response in bidirectional_streaming_method(self._stub): self.assertEqual(response.response_data, "data") spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 0) - finally: - context.detach(token) diff --git a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py index 810ee930dd..ba04277b5e 100644 --- a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py +++ b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py @@ -26,7 +26,7 @@ from opentelemetry.instrumentation.grpc.grpcext._interceptor import ( _UnaryClientInfo, ) -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.semconv.trace import SpanAttributes from opentelemetry.test.mock_textmap import MockTextMapPropagator @@ -306,45 +306,25 @@ def invoker(request, metadata): set_global_textmap(previous_propagator) def test_unary_unary_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): simple_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) def test_unary_stream_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): server_streaming_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) def test_stream_unary_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): client_streaming_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) def test_stream_stream_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): bidirectional_streaming_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) diff --git a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor_filter.py b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor_filter.py index a15268464b..bf6b1bda7f 100644 --- a/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor_filter.py +++ b/instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor_filter.py @@ -29,7 +29,7 @@ from opentelemetry.instrumentation.grpc.grpcext._interceptor import ( _UnaryClientInfo, ) -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.semconv.trace import SpanAttributes from opentelemetry.test.mock_textmap import MockTextMapPropagator @@ -638,45 +638,25 @@ def invoker(request, metadata): set_global_textmap(previous_propagator) def test_unary_unary_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): simple_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) def test_unary_stream_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): server_streaming_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) def test_stream_unary_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): client_streaming_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) def test_stream_stream_with_suppress_key(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): bidirectional_streaming_method(self._stub) spans = self.memory_exporter.get_finished_spans() - finally: - context.detach(token) self.assertEqual(len(spans), 0) diff --git a/instrumentation/opentelemetry-instrumentation-httpx/src/opentelemetry/instrumentation/httpx/__init__.py b/instrumentation/opentelemetry-instrumentation-httpx/src/opentelemetry/instrumentation/httpx/__init__.py index 55209bcd4f..f5d34b3c40 100644 --- a/instrumentation/opentelemetry-instrumentation-httpx/src/opentelemetry/instrumentation/httpx/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-httpx/src/opentelemetry/instrumentation/httpx/__init__.py @@ -170,7 +170,7 @@ def response_hook(span, request, response): from opentelemetry.instrumentation.httpx.package import _instruments from opentelemetry.instrumentation.httpx.version import __version__ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor -from opentelemetry.instrumentation.utils import http_status_to_status_code, is_http_instrumentation_enabled +from opentelemetry.instrumentation.utils import http_status_to_status_code from opentelemetry.propagate import inject from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.trace import SpanKind, TracerProvider, get_tracer @@ -316,7 +316,7 @@ def handle_request( httpx.Response, ]: """Add request info to span.""" - if not is_http_instrumentation_enabled(): + if context.get_value("suppress_instrumentation"): return self._transport.handle_request(*args, **kwargs) method, url, headers, stream, extensions = _extract_parameters( @@ -409,7 +409,7 @@ async def handle_async_request( httpx.Response, ]: """Add request info to span.""" - if not is_http_instrumentation_enabled(): + if context.get_value("suppress_instrumentation"): return await self._transport.handle_async_request(*args, **kwargs) method, url, headers, stream, extensions = _extract_parameters( diff --git a/instrumentation/opentelemetry-instrumentation-pymongo/src/opentelemetry/instrumentation/pymongo/__init__.py b/instrumentation/opentelemetry-instrumentation-pymongo/src/opentelemetry/instrumentation/pymongo/__init__.py index 041ff6b928..1f2c000036 100644 --- a/instrumentation/opentelemetry-instrumentation-pymongo/src/opentelemetry/instrumentation/pymongo/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-pymongo/src/opentelemetry/instrumentation/pymongo/__init__.py @@ -86,7 +86,7 @@ def failed_hook(span, event): COMMAND_TO_ATTRIBUTE_MAPPING, ) from opentelemetry.instrumentation.pymongo.version import __version__ -from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY +from opentelemetry.instrumentation.utils import is_instrumentation_enabled from opentelemetry.semconv.trace import DbSystemValues, SpanAttributes from opentelemetry.trace import SpanKind, get_tracer from opentelemetry.trace.span import Span @@ -122,9 +122,7 @@ def __init__( def started(self, event: monitoring.CommandStartedEvent): """Method to handle a pymongo CommandStartedEvent""" - if not self.is_enabled or context.get_value( - _SUPPRESS_INSTRUMENTATION_KEY - ): + if not self.is_enabled or not is_instrumentation_enabled(): return command_name = event.command_name span_name = f"{event.database_name}.{command_name}" @@ -167,9 +165,7 @@ def started(self, event: monitoring.CommandStartedEvent): def succeeded(self, event: monitoring.CommandSucceededEvent): """Method to handle a pymongo CommandSucceededEvent""" - if not self.is_enabled or context.get_value( - _SUPPRESS_INSTRUMENTATION_KEY - ): + if not self.is_enabled or not is_instrumentation_enabled(): return span = self._pop_span(event) if span is None: @@ -185,9 +181,7 @@ def succeeded(self, event: monitoring.CommandSucceededEvent): def failed(self, event: monitoring.CommandFailedEvent): """Method to handle a pymongo CommandFailedEvent""" - if not self.is_enabled or context.get_value( - _SUPPRESS_INSTRUMENTATION_KEY - ): + if not (self.is_enabled and is_instrumentation_enabled()): return span = self._pop_span(event) if span is None: diff --git a/instrumentation/opentelemetry-instrumentation-pymongo/tests/test_pymongo.py b/instrumentation/opentelemetry-instrumentation-pymongo/tests/test_pymongo.py index 8eab3b701c..ba6597b588 100644 --- a/instrumentation/opentelemetry-instrumentation-pymongo/tests/test_pymongo.py +++ b/instrumentation/opentelemetry-instrumentation-pymongo/tests/test_pymongo.py @@ -20,7 +20,7 @@ CommandTracer, PymongoInstrumentor, ) -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 @@ -112,16 +112,10 @@ def test_suppression_key(self): mock_event.command.get = mock.Mock() mock_event.command.get.return_value = "dummy" - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - - try: + with suppress_instrumentation(): command_tracer = CommandTracer(mock_tracer) command_tracer.started(event=mock_event) command_tracer.succeeded(event=mock_event) - finally: - context.detach(token) # if suppression key is set, CommandTracer methods return immediately, so command.get is not invoked. self.assertFalse( diff --git a/instrumentation/opentelemetry-instrumentation-requests/src/opentelemetry/instrumentation/requests/__init__.py b/instrumentation/opentelemetry-instrumentation-requests/src/opentelemetry/instrumentation/requests/__init__.py index 535b14285f..1c0efc2411 100644 --- a/instrumentation/opentelemetry-instrumentation-requests/src/opentelemetry/instrumentation/requests/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-requests/src/opentelemetry/instrumentation/requests/__init__.py @@ -59,15 +59,13 @@ from requests.structures import CaseInsensitiveDict from opentelemetry import context - -# 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.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.requests.package import _instruments from opentelemetry.instrumentation.requests.version import __version__ from opentelemetry.instrumentation.utils import ( - _SUPPRESS_INSTRUMENTATION_KEY, http_status_to_status_code, + is_http_instrumentation_enabled, + suppress_http_instrumentation, ) from opentelemetry.metrics import Histogram, get_meter from opentelemetry.propagate import inject @@ -126,9 +124,7 @@ def get_or_create_headers(): ) return request.headers - if context.get_value( - _SUPPRESS_INSTRUMENTATION_KEY - ) or context.get_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY): + if not is_http_instrumentation_enabled(): return wrapped_send(self, request, **kwargs) # See @@ -170,22 +166,19 @@ def get_or_create_headers(): headers = get_or_create_headers() inject(headers) - token = context.attach( - context.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True) - ) - - start_time = default_timer() - - try: - result = wrapped_send(self, request, **kwargs) # *** PROCEED - except Exception as exc: # pylint: disable=W0703 - exception = exc - result = getattr(exc, "response", None) - finally: - elapsed_time = max( - round((default_timer() - start_time) * 1000), 0 - ) - context.detach(token) + with suppress_http_instrumentation(): + start_time = default_timer() + try: + result = wrapped_send( + self, request, **kwargs + ) # *** PROCEED + except Exception as exc: # pylint: disable=W0703 + exception = exc + result = getattr(exc, "response", None) + finally: + elapsed_time = max( + round((default_timer() - start_time) * 1000), 0 + ) if isinstance(result, Response): if span.is_recording(): diff --git a/instrumentation/opentelemetry-instrumentation-requests/tests/test_requests_integration.py b/instrumentation/opentelemetry-instrumentation-requests/tests/test_requests_integration.py index 3bd76a6995..16c4035393 100644 --- a/instrumentation/opentelemetry-instrumentation-requests/tests/test_requests_integration.py +++ b/instrumentation/opentelemetry-instrumentation-requests/tests/test_requests_integration.py @@ -22,11 +22,11 @@ import opentelemetry.instrumentation.requests from opentelemetry import context, trace - -# 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.requests import RequestsInstrumentor -from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY +from opentelemetry.instrumentation.utils import ( + suppress_http_instrumentation, + suppress_instrumentation, +) from opentelemetry.propagate import get_global_textmap, set_global_textmap from opentelemetry.sdk import resources from opentelemetry.semconv.trace import SpanAttributes @@ -244,26 +244,16 @@ def test_uninstrument_session(self): self.assert_span() def test_suppress_instrumentation(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): result = self.perform_request(self.URL) self.assertEqual(result.text, "Hello!") - finally: - context.detach(token) self.assert_span(num_spans=0) def test_suppress_http_instrumentation(self): - token = context.attach( - context.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_http_instrumentation(): result = self.perform_request(self.URL) self.assertEqual(result.text, "Hello!") - finally: - context.detach(token) self.assert_span(num_spans=0) diff --git a/instrumentation/opentelemetry-instrumentation-urllib/src/opentelemetry/instrumentation/urllib/__init__.py b/instrumentation/opentelemetry-instrumentation-urllib/src/opentelemetry/instrumentation/urllib/__init__.py index da31bf99fa..05d7425a15 100644 --- a/instrumentation/opentelemetry-instrumentation-urllib/src/opentelemetry/instrumentation/urllib/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-urllib/src/opentelemetry/instrumentation/urllib/__init__.py @@ -85,15 +85,12 @@ def response_hook(span, request_obj, response) Request, ) -from opentelemetry import context - -# 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.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.urllib.package import _instruments from opentelemetry.instrumentation.urllib.version import __version__ from opentelemetry.instrumentation.utils import ( - _SUPPRESS_INSTRUMENTATION_KEY, + is_http_instrumentation_enabled, + suppress_http_instrumentation, http_status_to_status_code, ) from opentelemetry.metrics import Histogram, get_meter @@ -206,9 +203,7 @@ def call_wrapped(): def _instrumented_open_call( _, request, call_wrapped, get_or_create_headers ): # pylint: disable=too-many-locals - if context.get_value( - _SUPPRESS_INSTRUMENTATION_KEY - ) or context.get_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY): + if not is_http_instrumentation_enabled(): return call_wrapped() url = request.full_url @@ -236,18 +231,15 @@ def _instrumented_open_call( headers = get_or_create_headers() inject(headers) - token = context.attach( - context.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_http_instrumentation(): start_time = default_timer() - result = call_wrapped() # *** PROCEED - except Exception as exc: # pylint: disable=W0703 - exception = exc - result = getattr(exc, "file", None) - finally: - elapsed_time = round((default_timer() - start_time) * 1000) - context.detach(token) + try: + result = call_wrapped() # *** PROCEED + except Exception as exc: # pylint: disable=W0703 + exception = exc + result = getattr(exc, "file", None) + finally: + elapsed_time = round((default_timer() - start_time) * 1000) if result is not None: code_ = result.getcode() diff --git a/instrumentation/opentelemetry-instrumentation-urllib/tests/test_urllib_integration.py b/instrumentation/opentelemetry-instrumentation-urllib/tests/test_urllib_integration.py index f27f594a30..36189e12c1 100644 --- a/instrumentation/opentelemetry-instrumentation-urllib/tests/test_urllib_integration.py +++ b/instrumentation/opentelemetry-instrumentation-urllib/tests/test_urllib_integration.py @@ -24,14 +24,14 @@ import httpretty import opentelemetry.instrumentation.urllib # pylint: disable=no-name-in-module,import-error -from opentelemetry import context, trace - -# 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 import trace from opentelemetry.instrumentation.urllib import ( # pylint: disable=no-name-in-module,import-error URLLibInstrumentor, ) -from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY +from opentelemetry.instrumentation.utils import ( + suppress_http_instrumentation, + suppress_instrumentation, +) from opentelemetry.propagate import get_global_textmap, set_global_textmap from opentelemetry.sdk import resources from opentelemetry.semconv.trace import SpanAttributes @@ -255,26 +255,16 @@ def test_uninstrument_session(self): self.assert_span() def test_suppress_instrumentation(self): - token = context.attach( - context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_instrumentation(): result = self.perform_request(self.URL) self.assertEqual(result.read(), b"Hello!") - finally: - context.detach(token) self.assert_span(num_spans=0) def test_suppress_http_instrumentation(self): - token = context.attach( - context.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True) - ) - try: + with suppress_http_instrumentation(): result = self.perform_request(self.URL) self.assertEqual(result.read(), b"Hello!") - finally: - context.detach(token) self.assert_span(num_spans=0) diff --git a/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py b/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py index 45bab7454d..4e1b3dcf0d 100644 --- a/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py @@ -88,16 +88,13 @@ def response_hook(span, request, response): import urllib3.connectionpool import wrapt -from opentelemetry import context - -# 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.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.urllib3.package import _instruments from opentelemetry.instrumentation.urllib3.version import __version__ from opentelemetry.instrumentation.utils import ( - _SUPPRESS_INSTRUMENTATION_KEY, http_status_to_status_code, + is_http_instrumentation_enabled, + suppress_http_instrumentation, unwrap, ) from opentelemetry.metrics import Histogram, get_meter @@ -224,7 +221,7 @@ def _instrument( excluded_urls: ExcludeList = None, ): def instrumented_urlopen(wrapped, instance, args, kwargs): - if _is_instrumentation_suppressed(): + if not is_http_instrumentation_enabled(): return wrapped(*args, **kwargs) url = _get_url(instance, args, kwargs, url_filter) @@ -248,7 +245,7 @@ def instrumented_urlopen(wrapped, instance, args, kwargs): request_hook(span, instance, headers, body) inject(headers) - with _suppress_further_instrumentation(): + with suppress_http_instrumentation(): start_time = default_timer() response = wrapped(*args, **kwargs) elapsed_time = round((default_timer() - start_time) * 1000) @@ -352,13 +349,6 @@ def _apply_response(span: Span, response: urllib3.response.HTTPResponse): span.set_status(Status(http_status_to_status_code(response.status))) -def _is_instrumentation_suppressed() -> bool: - return bool( - context.get_value(_SUPPRESS_INSTRUMENTATION_KEY) - or context.get_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY) - ) - - def _create_metric_attributes( instance: urllib3.connectionpool.HTTPConnectionPool, response: urllib3.response.HTTPResponse, @@ -382,16 +372,5 @@ def _create_metric_attributes( return metric_attributes -@contextlib.contextmanager -def _suppress_further_instrumentation(): - token = context.attach( - context.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True) - ) - try: - yield - finally: - context.detach(token) - - def _uninstrument(): unwrap(urllib3.connectionpool.HTTPConnectionPool, "urlopen") diff --git a/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py b/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py index 27e1b81269..966513d0e7 100644 --- a/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py +++ b/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py @@ -20,11 +20,11 @@ import urllib3.exceptions from opentelemetry import context, trace - -# 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.urllib3 import URLLib3Instrumentor -from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY +from opentelemetry.instrumentation.utils import ( + suppress_http_instrumentation, + suppress_instrumentation, +) from opentelemetry.propagate import get_global_textmap, set_global_textmap from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.test.mock_textmap import MockTextMapPropagator @@ -225,20 +225,17 @@ def test_uninstrument(self): URLLib3Instrumentor().instrument() def test_suppress_instrumentation(self): - suppression_keys = ( - _SUPPRESS_HTTP_INSTRUMENTATION_KEY, - _SUPPRESS_INSTRUMENTATION_KEY, + suppression_cms = ( + suppress_instrumentation, + suppress_http_instrumentation, ) - for key in suppression_keys: + for cm in suppression_cms: self.memory_exporter.clear() with self.subTest(key=key): - token = context.attach(context.set_value(key, True)) - try: + with cm(): response = self.perform_request(self.HTTP_URL) self.assertEqual(b"Hello!", response.data) - finally: - context.detach(token) self.assert_span(num_spans=0) diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py index 87c73cc737..36fe69daec 100644 --- a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py @@ -302,9 +302,7 @@ def collect_request_attributes(environ): """ result = { - SpanAttributes.HTTP_METHOD: sanitize_method( - environ.get("REQUEST_METHOD") - ), + SpanAttributes.HTTP_METHOD: sanitize_method(environ.get("REQUEST_METHOD")), SpanAttributes.HTTP_SERVER_NAME: environ.get("SERVER_NAME"), SpanAttributes.HTTP_SCHEME: environ.get("wsgi.url_scheme"), } diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py b/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py index bc78a787ca..6aef096218 100644 --- a/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py +++ b/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py @@ -286,26 +286,22 @@ def test_wsgi_metrics(self): self.assertTrue(number_data_point_seen and histogram_data_point_seen) def test_nonstandard_http_method(self): - self.environ["REQUEST_METHOD"] = "NONSTANDARD" + self.environ["REQUEST_METHOD"]= "NONSTANDARD" app = otel_wsgi.OpenTelemetryMiddleware(simple_wsgi) response = app(self.environ, self.start_response) - self.validate_response( - response, span_name="UNKNOWN /", http_method="UNKNOWN" - ) + self.validate_response(response, span_name="UNKNOWN /", http_method="UNKNOWN") @mock.patch.dict( - "os.environ", + "os.environ", { OTEL_PYTHON_INSTRUMENTATION_HTTP_CAPTURE_ALL_METHODS: "1", }, ) def test_nonstandard_http_method_allowed(self): - self.environ["REQUEST_METHOD"] = "NONSTANDARD" + self.environ["REQUEST_METHOD"]= "NONSTANDARD" app = otel_wsgi.OpenTelemetryMiddleware(simple_wsgi) response = app(self.environ, self.start_response) - self.validate_response( - response, span_name="NONSTANDARD /", http_method="NONSTANDARD" - ) + self.validate_response(response, span_name="NONSTANDARD /", http_method="NONSTANDARD") def test_default_span_name_missing_path_info(self): """Test that default span_names with missing path info.""" diff --git a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/utils.py b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/utils.py index 7baf124b68..2f5b30675d 100644 --- a/opentelemetry-instrumentation/src/opentelemetry/instrumentation/utils.py +++ b/opentelemetry-instrumentation/src/opentelemetry/instrumentation/utils.py @@ -15,10 +15,17 @@ import os import threading import urllib.parse +from contextlib import contextmanager from enum import Enum from re import escape, sub -from typing import Dict, Sequence +from typing import Dict, Iterable, Sequence +# pylint: disable=E0611 +# FIXME: fix the importing of these private attributes when the location of the _SUPPRESS_HTTP_INSTRUMENTATION_KEY is defined.= +from opentelemetry.context import ( + _SUPPRESS_HTTP_INSTRUMENTATION_KEY, + _SUPPRESS_INSTRUMENTATION_KEY, +) from wrapt import ObjectProxy from opentelemetry import context, trace @@ -213,18 +220,41 @@ def _get_opentelemetry_stability_opt_in( def is_instrumentation_enabled() -> bool: - if ( - context.get_value("suppress_instrumentation") - or - context.get_value(context._SUPPRESS_INSTRUMENTATION_KEY) # type: ignore - ): + if context.get_value(_SUPPRESS_INSTRUMENTATION_KEY): return False return True def is_http_instrumentation_enabled() -> bool: return ( - is_http_instrumentation_enabled() + is_instrumentation_enabled() and not - context.get_value(context._SUPPRESS_HTTP_INSTRUMENTATION_KEY) # type: ignore + context.get_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY) ) + + +@contextmanager +def _suppress_instrumentation(*keys: str) -> Iterable[None]: + """Suppress instrumentation within the context.""" + ctx = context.get_current() + for key in keys: + ctx = context.set_value(key, True, ctx) + token = context.attach(ctx) + try: + yield + finally: + context.detach(token) + + +@contextmanager +def suppress_instrumentation() -> Iterable[None]: + """Suppress instrumentation within the context.""" + with _suppress_instrumentation(_SUPPRESS_INSTRUMENTATION_KEY): + yield + + +@contextmanager +def suppress_http_instrumentation() -> Iterable[None]: + """Suppress instrumentation within the context.""" + with _suppress_instrumentation(_SUPPRESS_HTTP_INSTRUMENTATION_KEY): + yield