From 83edee58a388eebf0b2bd71cf952d196aa5c42fd Mon Sep 17 00:00:00 2001 From: wangzlei Date: Wed, 1 Oct 2025 15:44:06 -0700 Subject: [PATCH 1/3] function decorator for adding code level attributes to span --- .../distro/aws_opentelemetry_configurator.py | 27 ++ .../distro/code_correlation/__init__.py | 142 +++++++ .../test_aws_opentelementry_configurator.py | 67 +++ .../distro/test_code_correlation.py | 388 ++++++++++++++++++ 4 files changed, 624 insertions(+) create mode 100644 aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py create mode 100644 aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py index 3f16e5dca..ce409ae7d 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py @@ -95,6 +95,7 @@ OTEL_EXPORTER_OTLP_TRACES_ENDPOINT = "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT" OTEL_EXPORTER_OTLP_LOGS_ENDPOINT = "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" OTEL_EXPORTER_OTLP_LOGS_HEADERS = "OTEL_EXPORTER_OTLP_LOGS_HEADERS" +CODE_CORRELATION_ENABLED_CONFIG = "OTEL_AWS_CODE_CORRELATION_ENABLED" XRAY_SERVICE = "xray" LOGS_SERIVCE = "logs" @@ -615,6 +616,32 @@ def _is_application_signals_runtime_enabled(): ) +def _get_code_correlation_enabled_status() -> Optional[bool]: + """ + Get the code correlation enabled status from environment variable. + + Returns: + True if OTEL_AWS_CODE_CORRELATION_ENABLED is set to 'true' + False if OTEL_AWS_CODE_CORRELATION_ENABLED is set to 'false' + None if OTEL_AWS_CODE_CORRELATION_ENABLED is not set (default state) + """ + env_value = os.environ.get(CODE_CORRELATION_ENABLED_CONFIG) + + if env_value is None: + return None # Default state - environment variable not set + + env_value_lower = env_value.strip().lower() + if env_value_lower == "true": + return True + elif env_value_lower == "false": + return False + else: + # Invalid value, treat as default and log warning + _logger.warning("Invalid value for %s: %s. Expected 'true' or 'false'. Using default.", + CODE_CORRELATION_ENABLED_CONFIG, env_value) + return None + + def _is_lambda_environment(): # detect if running in AWS Lambda environment return AWS_LAMBDA_FUNCTION_NAME_CONFIG in os.environ diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py new file mode 100644 index 000000000..cf3c880da --- /dev/null +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py @@ -0,0 +1,142 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Code correlation module for AWS OpenTelemetry Python Instrumentation. + +This module provides functionality for correlating code execution with telemetry data. +""" + +__version__ = "1.0.0" + + +""" +Utility functions for adding code information to OpenTelemetry spans. +""" + +from typing import Any, Callable +from functools import wraps +from opentelemetry import trace + + +# Code correlation attribute constants +CODE_FUNCTION_NAME = "code.function.name" +CODE_FILE_PATH = "code.file.path" +CODE_LINE_NUMBER = "code.line.number" + + +def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None: + """ + Add code-related attributes to a span based on a Python function. + + This utility method extracts function metadata and adds the following + span attributes: + - CODE_FUNCTION_NAME: The name of the function + - CODE_FILE_PATH: The file path where the function is defined + - CODE_LINE_NUMBER: The line number where the function is defined + + Args: + span: The OpenTelemetry span to add attributes to + func: The Python function to extract metadata from + """ + if not span.is_recording(): + return + + try: + # Get function name + function_name = getattr(func, '__name__', str(func)) + span.set_attribute(CODE_FUNCTION_NAME, function_name) + + # Get function source file from code object + try: + if hasattr(func, '__code__'): + source_file = func.__code__.co_filename + span.set_attribute(CODE_FILE_PATH, source_file) + except (AttributeError, TypeError): + # Handle cases where code object is not available + # (e.g., built-in functions, C extensions) + pass + + # Get function line number from code object + try: + if hasattr(func, '__code__'): + line_number = func.__code__.co_firstlineno + span.set_attribute(CODE_LINE_NUMBER, line_number) + except (AttributeError, TypeError): + # Handle cases where code object is not available + pass + + except Exception: + # Silently handle any unexpected errors to avoid breaking + # the instrumentation flow + pass + + +def add_code_attributes_to_span(func: Callable[..., Any]) -> Callable[..., Any]: + """ + Decorator to automatically add code attributes to the current OpenTelemetry span. + + This decorator extracts metadata from the decorated function and adds it as + attributes to the current active span. The attributes added are: + - code.function.name: The name of the function + - code.file.path: The file path where the function is defined + - code.line.number: The line number where the function is defined + + This decorator supports both synchronous and asynchronous functions. + + Usage: + @add_code_attributes_to_span + def my_sync_function(): + # Sync function implementation + pass + + @add_code_attributes_to_span + async def my_async_function(): + # Async function implementation + pass + + Args: + func: The function to be decorated + + Returns: + The wrapped function with current span code attributes tracing + """ + # Detect async functions: check function code object flags or special attributes + # CO_ITERABLE_COROUTINE = 0x80, async functions will have this flag set + is_async = (hasattr(func, '__code__') and + func.__code__.co_flags & 0x80) or hasattr(func, '_is_coroutine') + + if is_async: + # Async function wrapper + @wraps(func) + async def async_wrapper(*args, **kwargs): + # Add code attributes to current span + try: + current_span = trace.get_current_span() + if current_span: + _add_code_attributes_to_span(current_span, func) + except Exception: + # Silently handle any unexpected errors + pass + + # Call and await the original async function + return await func(*args, **kwargs) + + return async_wrapper + else: + # Sync function wrapper + @wraps(func) + def sync_wrapper(*args, **kwargs): + # Add code attributes to current span + try: + current_span = trace.get_current_span() + if current_span: + _add_code_attributes_to_span(current_span, func) + except Exception: + # Silently handle any unexpected errors + pass + + # Call the original sync function + return func(*args, **kwargs) + + return sync_wrapper diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py index efcc5a317..025e47f61 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py @@ -39,6 +39,7 @@ _export_unsampled_span_for_agent_observability, _export_unsampled_span_for_lambda, _fetch_logs_header, + _get_code_correlation_enabled_status, _init_logging, _is_application_signals_enabled, _is_application_signals_runtime_enabled, @@ -1425,6 +1426,72 @@ def test_create_emf_exporter_cloudwatch_exporter_import_error( self.assertIsNone(result) mock_logger.error.assert_called_once() + def test_get_code_correlation_enabled_status(self): + """Test _get_code_correlation_enabled_status function with various environment variable values""" + # Import the constant we need + from amazon.opentelemetry.distro.aws_opentelemetry_configurator import CODE_CORRELATION_ENABLED_CONFIG + + # Test when environment variable is not set (default state) + os.environ.pop(CODE_CORRELATION_ENABLED_CONFIG, None) + result = _get_code_correlation_enabled_status() + self.assertIsNone(result) + + # Test when environment variable is set to 'true' (case insensitive) + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "true" + result = _get_code_correlation_enabled_status() + self.assertTrue(result) + + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "TRUE" + result = _get_code_correlation_enabled_status() + self.assertTrue(result) + + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "True" + result = _get_code_correlation_enabled_status() + self.assertTrue(result) + + # Test when environment variable is set to 'false' (case insensitive) + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "false" + result = _get_code_correlation_enabled_status() + self.assertFalse(result) + + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "FALSE" + result = _get_code_correlation_enabled_status() + self.assertFalse(result) + + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "False" + result = _get_code_correlation_enabled_status() + self.assertFalse(result) + + # Test with leading/trailing whitespace + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = " true " + result = _get_code_correlation_enabled_status() + self.assertTrue(result) + + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = " false " + result = _get_code_correlation_enabled_status() + self.assertFalse(result) + + # Test invalid values (should return None and log warning) + # We'll use caplog to capture log messages instead of mocking + import logging + + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "invalid" + result = _get_code_correlation_enabled_status() + self.assertIsNone(result) + + # Test another invalid value + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "yes" + result = _get_code_correlation_enabled_status() + self.assertIsNone(result) + + # Test empty string (invalid) + os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "" + result = _get_code_correlation_enabled_status() + self.assertIsNone(result) + + # Clean up + os.environ.pop(CODE_CORRELATION_ENABLED_CONFIG, None) + def validate_distro_environ(): tc: TestCase = TestCase() diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py new file mode 100644 index 000000000..671be8d05 --- /dev/null +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py @@ -0,0 +1,388 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from unittest import TestCase +from unittest.mock import MagicMock, patch, PropertyMock + +from amazon.opentelemetry.distro.code_correlation import ( + CODE_FUNCTION_NAME, + CODE_FILE_PATH, + CODE_LINE_NUMBER, + _add_code_attributes_to_span, + add_code_attributes_to_span, +) +from opentelemetry.trace import Span + + +class TestCodeCorrelationConstants(TestCase): + """Test code correlation attribute constants.""" + + def test_constants_values(self): + """Test that constants have the expected values.""" + self.assertEqual(CODE_FUNCTION_NAME, "code.function.name") + self.assertEqual(CODE_FILE_PATH, "code.file.path") + self.assertEqual(CODE_LINE_NUMBER, "code.line.number") + + +class TestAddCodeAttributesToSpan(TestCase): + """Test the _add_code_attributes_to_span function.""" + + def setUp(self): + """Set up test fixtures.""" + self.mock_span = MagicMock(spec=Span) + self.mock_span.is_recording.return_value = True + + def test_add_code_attributes_to_recording_span(self): + """Test adding code attributes to a recording span.""" + def test_function(): + pass + + _add_code_attributes_to_span(self.mock_span, test_function) + + # Verify function name attribute is set + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_function") + + # Verify file path attribute is set + expected_file_path = test_function.__code__.co_filename + self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, expected_file_path) + + # Verify line number attribute is set + expected_line_number = test_function.__code__.co_firstlineno + self.mock_span.set_attribute.assert_any_call(CODE_LINE_NUMBER, expected_line_number) + + def test_add_code_attributes_to_non_recording_span(self): + """Test that no attributes are added to a non-recording span.""" + self.mock_span.is_recording.return_value = False + + def test_function(): + pass + + _add_code_attributes_to_span(self.mock_span, test_function) + + # Verify no attributes are set + self.mock_span.set_attribute.assert_not_called() + + def test_add_code_attributes_function_without_code(self): + """Test handling of functions without __code__ attribute.""" + # Create a mock function without __code__ attribute + mock_func = MagicMock() + mock_func.__name__ = "mock_function" + delattr(mock_func, '__code__') + + _add_code_attributes_to_span(self.mock_span, mock_func) + + # Verify only function name attribute is set + self.mock_span.set_attribute.assert_called_once_with(CODE_FUNCTION_NAME, "mock_function") + + def test_add_code_attributes_builtin_function(self): + """Test handling of built-in functions.""" + # Use a built-in function like len + _add_code_attributes_to_span(self.mock_span, len) + + # Verify only function name attribute is set + self.mock_span.set_attribute.assert_called_once_with(CODE_FUNCTION_NAME, "len") + + def test_add_code_attributes_function_without_name(self): + """Test handling of functions without __name__ attribute.""" + # Create an object without __name__ attribute + mock_func = MagicMock() + delattr(mock_func, '__name__') + mock_func.__code__ = MagicMock() + mock_func.__code__.co_filename = "/test/file.py" + mock_func.__code__.co_firstlineno = 10 + + _add_code_attributes_to_span(self.mock_span, mock_func) + + # Verify function name uses str() representation + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, str(mock_func)) + self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, "/test/file.py") + self.mock_span.set_attribute.assert_any_call(CODE_LINE_NUMBER, 10) + + def test_add_code_attributes_exception_handling(self): + """Test that exceptions are handled gracefully.""" + # Create a function that will cause an exception when accessing attributes + mock_func = MagicMock() + mock_func.__name__ = "test_func" + mock_func.__code__ = MagicMock() + mock_func.__code__.co_filename = "/test/file.py" + mock_func.__code__.co_firstlineno = MagicMock(side_effect=Exception("Test exception")) + + # This should not raise an exception + _add_code_attributes_to_span(self.mock_span, mock_func) + + # Verify function name and file path are still set + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_func") + self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, "/test/file.py") + + @patch('amazon.opentelemetry.distro.code_correlation.getattr') + def test_add_code_attributes_getattr_exception(self, mock_getattr): + """Test exception handling when getattr fails.""" + mock_getattr.side_effect = Exception("Test exception") + + def test_function(): + pass + + # This should not raise an exception + _add_code_attributes_to_span(self.mock_span, test_function) + + # Verify no attributes are set due to exception + self.mock_span.set_attribute.assert_not_called() + + def test_add_code_attributes_co_filename_exception(self): + """Test exception handling when accessing co_filename raises exception.""" + # Create a mock function with __code__ that raises exception on co_filename access + mock_func = MagicMock() + mock_func.__name__ = "test_func" + mock_code = MagicMock() + mock_code.co_firstlineno = 10 + + # Make co_filename raise AttributeError + type(mock_code).co_filename = PropertyMock(side_effect=AttributeError("Test exception")) + mock_func.__code__ = mock_code + + # This should not raise an exception + _add_code_attributes_to_span(self.mock_span, mock_func) + + # Verify function name and line number are still set, but not file path + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_func") + self.mock_span.set_attribute.assert_any_call(CODE_LINE_NUMBER, 10) + # File path should not be called due to exception + with self.assertRaises(AssertionError): + self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, MagicMock()) + + def test_add_code_attributes_co_firstlineno_exception(self): + """Test exception handling when accessing co_firstlineno raises exception.""" + # Create a mock function with __code__ that raises exception on co_firstlineno access + mock_func = MagicMock() + mock_func.__name__ = "test_func" + mock_code = MagicMock() + mock_code.co_filename = "/test/file.py" + + # Make co_firstlineno raise TypeError + type(mock_code).co_firstlineno = PropertyMock(side_effect=TypeError("Test exception")) + mock_func.__code__ = mock_code + + # This should not raise an exception + _add_code_attributes_to_span(self.mock_span, mock_func) + + # Verify function name and file path are still set, but not line number + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_func") + self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, "/test/file.py") + # Line number should not be called due to exception + with self.assertRaises(AssertionError): + self.mock_span.set_attribute.assert_any_call(CODE_LINE_NUMBER, MagicMock()) + + def test_add_code_attributes_co_filename_type_error(self): + """Test exception handling when accessing co_filename raises TypeError.""" + # Create a mock function with __code__ that raises TypeError on co_filename access + mock_func = MagicMock() + mock_func.__name__ = "test_func" + mock_code = MagicMock() + mock_code.co_firstlineno = 10 + + # Make co_filename raise TypeError + type(mock_code).co_filename = PropertyMock(side_effect=TypeError("Test exception")) + mock_func.__code__ = mock_code + + # This should not raise an exception + _add_code_attributes_to_span(self.mock_span, mock_func) + + # Verify function name and line number are still set, but not file path + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_func") + self.mock_span.set_attribute.assert_any_call(CODE_LINE_NUMBER, 10) + # File path should not be called due to TypeError + with self.assertRaises(AssertionError): + self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, MagicMock()) + + +class TestAddCodeAttributesToSpanDecorator(TestCase): + """Test the add_code_attributes_to_span decorator.""" + + def setUp(self): + """Set up test fixtures.""" + self.mock_span = MagicMock(spec=Span) + self.mock_span.is_recording.return_value = True + + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_sync_function(self, mock_get_current_span): + """Test decorator with synchronous function.""" + mock_get_current_span.return_value = self.mock_span + + @add_code_attributes_to_span + def test_sync_function(arg1, arg2=None): + return f"sync result: {arg1}, {arg2}" + + # Call the decorated function + result = test_sync_function("test_arg", arg2="test_kwarg") + + # Verify the function still works correctly + self.assertEqual(result, "sync result: test_arg, test_kwarg") + + # Verify span attributes were set + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_sync_function") + + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_async_function(self, mock_get_current_span): + """Test decorator with asynchronous function.""" + mock_get_current_span.return_value = self.mock_span + + @add_code_attributes_to_span + async def test_async_function(arg1, arg2=None): + return f"async result: {arg1}, {arg2}" + + # Call the decorated async function + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + result = loop.run_until_complete(test_async_function("test_arg", arg2="test_kwarg")) + finally: + loop.close() + + # Verify the function still works correctly + self.assertEqual(result, "async result: test_arg, test_kwarg") + + # Verify span attributes were set + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_async_function") + + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_no_current_span(self, mock_get_current_span): + """Test decorator when there's no current span.""" + mock_get_current_span.return_value = None + + @add_code_attributes_to_span + def test_function(): + return "test result" + + # Call the decorated function + result = test_function() + + # Verify the function still works correctly + self.assertEqual(result, "test result") + + # Verify no span attributes were set + self.mock_span.set_attribute.assert_not_called() + + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_exception_handling(self, mock_get_current_span): + """Test decorator handles exceptions gracefully.""" + mock_get_current_span.side_effect = Exception("Test exception") + + @add_code_attributes_to_span + def test_function(): + return "test result" + + # Call the decorated function - should not raise exception + result = test_function() + + # Verify the function still works correctly + self.assertEqual(result, "test result") + + def test_decorator_preserves_function_metadata(self): + """Test that decorator preserves original function metadata.""" + @add_code_attributes_to_span + def test_function(): + """Test function docstring.""" + return "test result" + + # Verify function metadata is preserved + self.assertEqual(test_function.__name__, "test_function") + self.assertEqual(test_function.__doc__, "Test function docstring.") + + def test_async_function_detection(self): + """Test that async functions are properly detected.""" + # Create a regular function + def sync_func(): + pass + + # Create an async function + async def async_func(): + pass + + # Apply decorator to both + decorated_sync = add_code_attributes_to_span(sync_func) + decorated_async = add_code_attributes_to_span(async_func) + + # Check that sync function returns a regular function + self.assertFalse(asyncio.iscoroutinefunction(decorated_sync)) + + # Check that async function returns a coroutine function + self.assertTrue(asyncio.iscoroutinefunction(decorated_async)) + + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_with_function_that_raises_exception(self, mock_get_current_span): + """Test decorator with function that raises exception.""" + mock_get_current_span.return_value = self.mock_span + + @add_code_attributes_to_span + def test_function(): + raise ValueError("Test function exception") + + # Verify exception is still raised + with self.assertRaises(ValueError): + test_function() + + # Verify span attributes were still set before exception + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_function") + + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_with_async_function_that_raises_exception(self, mock_get_current_span): + """Test decorator with async function that raises exception.""" + mock_get_current_span.return_value = self.mock_span + + @add_code_attributes_to_span + async def test_async_function(): + raise ValueError("Test async function exception") + + # Verify exception is still raised + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + with self.assertRaises(ValueError): + loop.run_until_complete(test_async_function()) + finally: + loop.close() + + # Verify span attributes were still set before exception + self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_async_function") + + @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span') + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_internal_exception_handling_sync(self, mock_get_current_span, mock_add_attributes): + """Test that decorator handles internal exceptions gracefully in sync function.""" + mock_get_current_span.return_value = self.mock_span + # Make _add_code_attributes_to_span raise an exception + mock_add_attributes.side_effect = Exception("Internal exception") + + @add_code_attributes_to_span + def test_function(): + return "test result" + + # Call the decorated function - should not raise exception + result = test_function() + + # Verify the function still works correctly despite internal exception + self.assertEqual(result, "test result") + + @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span') + @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + def test_decorator_internal_exception_handling_async(self, mock_get_current_span, mock_add_attributes): + """Test that decorator handles internal exceptions gracefully in async function.""" + mock_get_current_span.return_value = self.mock_span + # Make _add_code_attributes_to_span raise an exception + mock_add_attributes.side_effect = Exception("Internal exception") + + @add_code_attributes_to_span + async def test_async_function(): + return "async test result" + + # Call the decorated async function - should not raise exception + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + result = loop.run_until_complete(test_async_function()) + finally: + loop.close() + + # Verify the function still works correctly despite internal exception + self.assertEqual(result, "async test result") From 3f9d2b7cb512bd7d90c236be70c3ae7ed749ae7a Mon Sep 17 00:00:00 2001 From: wangzlei Date: Wed, 1 Oct 2025 16:10:17 -0700 Subject: [PATCH 2/3] function decorator for adding code level attributes to span --- .../distro/code_correlation/__init__.py | 48 +++++++++---------- .../distro/code_correlation/__init__.py | 2 + .../test_code_correlation.py | 12 ++--- 3 files changed, 30 insertions(+), 32 deletions(-) create mode 100644 aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/__init__.py rename aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/{ => code_correlation}/test_code_correlation.py (99%) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py index cf3c880da..a3706a558 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py @@ -7,17 +7,13 @@ This module provides functionality for correlating code execution with telemetry data. """ -__version__ = "1.0.0" - - -""" -Utility functions for adding code information to OpenTelemetry spans. -""" - -from typing import Any, Callable from functools import wraps +from typing import Any, Callable + from opentelemetry import trace +__version__ = "1.0.0" + # Code correlation attribute constants CODE_FUNCTION_NAME = "code.function.name" @@ -28,25 +24,25 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None: """ Add code-related attributes to a span based on a Python function. - + This utility method extracts function metadata and adds the following span attributes: - CODE_FUNCTION_NAME: The name of the function - CODE_FILE_PATH: The file path where the function is defined - CODE_LINE_NUMBER: The line number where the function is defined - + Args: span: The OpenTelemetry span to add attributes to func: The Python function to extract metadata from """ if not span.is_recording(): return - + try: # Get function name function_name = getattr(func, '__name__', str(func)) span.set_attribute(CODE_FUNCTION_NAME, function_name) - + # Get function source file from code object try: if hasattr(func, '__code__'): @@ -56,7 +52,7 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None: # Handle cases where code object is not available # (e.g., built-in functions, C extensions) pass - + # Get function line number from code object try: if hasattr(func, '__code__'): @@ -65,7 +61,7 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None: except (AttributeError, TypeError): # Handle cases where code object is not available pass - + except Exception: # Silently handle any unexpected errors to avoid breaking # the instrumentation flow @@ -75,37 +71,37 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None: def add_code_attributes_to_span(func: Callable[..., Any]) -> Callable[..., Any]: """ Decorator to automatically add code attributes to the current OpenTelemetry span. - + This decorator extracts metadata from the decorated function and adds it as attributes to the current active span. The attributes added are: - code.function.name: The name of the function - code.file.path: The file path where the function is defined - code.line.number: The line number where the function is defined - + This decorator supports both synchronous and asynchronous functions. - + Usage: @add_code_attributes_to_span def my_sync_function(): # Sync function implementation pass - + @add_code_attributes_to_span async def my_async_function(): # Async function implementation pass - + Args: func: The function to be decorated - + Returns: The wrapped function with current span code attributes tracing """ # Detect async functions: check function code object flags or special attributes # CO_ITERABLE_COROUTINE = 0x80, async functions will have this flag set - is_async = (hasattr(func, '__code__') and + is_async = (hasattr(func, '__code__') and func.__code__.co_flags & 0x80) or hasattr(func, '_is_coroutine') - + if is_async: # Async function wrapper @wraps(func) @@ -118,10 +114,10 @@ async def async_wrapper(*args, **kwargs): except Exception: # Silently handle any unexpected errors pass - + # Call and await the original async function return await func(*args, **kwargs) - + return async_wrapper else: # Sync function wrapper @@ -135,8 +131,8 @@ def sync_wrapper(*args, **kwargs): except Exception: # Silently handle any unexpected errors pass - + # Call the original sync function return func(*args, **kwargs) - + return sync_wrapper diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/__init__.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/__init__.py new file mode 100644 index 000000000..04f8b7b76 --- /dev/null +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/__init__.py @@ -0,0 +1,2 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py similarity index 99% rename from aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py rename to aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py index 671be8d05..8ff1b78d5 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py @@ -42,11 +42,11 @@ def test_function(): # Verify function name attribute is set self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_function") - + # Verify file path attribute is set expected_file_path = test_function.__code__.co_filename self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, expected_file_path) - + # Verify line number attribute is set expected_line_number = test_function.__code__.co_firstlineno self.mock_span.set_attribute.assert_any_call(CODE_LINE_NUMBER, expected_line_number) @@ -136,7 +136,7 @@ def test_add_code_attributes_co_filename_exception(self): mock_func.__name__ = "test_func" mock_code = MagicMock() mock_code.co_firstlineno = 10 - + # Make co_filename raise AttributeError type(mock_code).co_filename = PropertyMock(side_effect=AttributeError("Test exception")) mock_func.__code__ = mock_code @@ -158,7 +158,7 @@ def test_add_code_attributes_co_firstlineno_exception(self): mock_func.__name__ = "test_func" mock_code = MagicMock() mock_code.co_filename = "/test/file.py" - + # Make co_firstlineno raise TypeError type(mock_code).co_firstlineno = PropertyMock(side_effect=TypeError("Test exception")) mock_func.__code__ = mock_code @@ -180,7 +180,7 @@ def test_add_code_attributes_co_filename_type_error(self): mock_func.__name__ = "test_func" mock_code = MagicMock() mock_code.co_firstlineno = 10 - + # Make co_filename raise TypeError type(mock_code).co_filename = PropertyMock(side_effect=TypeError("Test exception")) mock_func.__code__ = mock_code @@ -364,7 +364,7 @@ def test_function(): # Verify the function still works correctly despite internal exception self.assertEqual(result, "test result") - @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span') + @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span') @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') def test_decorator_internal_exception_handling_async(self, mock_get_current_span, mock_add_attributes): """Test that decorator handles internal exceptions gracefully in async function.""" From aa3d0a345a18f35cfee43832fcdbd70605e36409 Mon Sep 17 00:00:00 2001 From: wangzlei Date: Wed, 1 Oct 2025 16:35:18 -0700 Subject: [PATCH 3/3] fix pylint --- .../distro/aws_opentelemetry_configurator.py | 20 +++++---- .../distro/code_correlation/__init__.py | 43 +++++++++---------- .../code_correlation/test_code_correlation.py | 33 +++++++------- .../test_aws_opentelementry_configurator.py | 7 +-- 4 files changed, 51 insertions(+), 52 deletions(-) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py index ce409ae7d..5cce1ced0 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py @@ -619,27 +619,29 @@ def _is_application_signals_runtime_enabled(): def _get_code_correlation_enabled_status() -> Optional[bool]: """ Get the code correlation enabled status from environment variable. - + Returns: True if OTEL_AWS_CODE_CORRELATION_ENABLED is set to 'true' False if OTEL_AWS_CODE_CORRELATION_ENABLED is set to 'false' None if OTEL_AWS_CODE_CORRELATION_ENABLED is not set (default state) """ env_value = os.environ.get(CODE_CORRELATION_ENABLED_CONFIG) - + if env_value is None: return None # Default state - environment variable not set - + env_value_lower = env_value.strip().lower() if env_value_lower == "true": return True - elif env_value_lower == "false": + if env_value_lower == "false": return False - else: - # Invalid value, treat as default and log warning - _logger.warning("Invalid value for %s: %s. Expected 'true' or 'false'. Using default.", - CODE_CORRELATION_ENABLED_CONFIG, env_value) - return None + # Invalid value, treat as default and log warning + _logger.warning( + "Invalid value for %s: %s. Expected 'true' or 'false'. Using default.", + CODE_CORRELATION_ENABLED_CONFIG, + env_value, + ) + return None def _is_lambda_environment(): diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py index a3706a558..420145419 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py @@ -40,12 +40,12 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None: try: # Get function name - function_name = getattr(func, '__name__', str(func)) + function_name = getattr(func, "__name__", str(func)) span.set_attribute(CODE_FUNCTION_NAME, function_name) # Get function source file from code object try: - if hasattr(func, '__code__'): + if hasattr(func, "__code__"): source_file = func.__code__.co_filename span.set_attribute(CODE_FILE_PATH, source_file) except (AttributeError, TypeError): @@ -55,14 +55,14 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None: # Get function line number from code object try: - if hasattr(func, '__code__'): + if hasattr(func, "__code__"): line_number = func.__code__.co_firstlineno span.set_attribute(CODE_LINE_NUMBER, line_number) except (AttributeError, TypeError): # Handle cases where code object is not available pass - except Exception: + except Exception: # pylint: disable=broad-exception-caught # Silently handle any unexpected errors to avoid breaking # the instrumentation flow pass @@ -99,8 +99,7 @@ async def my_async_function(): """ # Detect async functions: check function code object flags or special attributes # CO_ITERABLE_COROUTINE = 0x80, async functions will have this flag set - is_async = (hasattr(func, '__code__') and - func.__code__.co_flags & 0x80) or hasattr(func, '_is_coroutine') + is_async = (hasattr(func, "__code__") and func.__code__.co_flags & 0x80) or hasattr(func, "_is_coroutine") if is_async: # Async function wrapper @@ -111,7 +110,7 @@ async def async_wrapper(*args, **kwargs): current_span = trace.get_current_span() if current_span: _add_code_attributes_to_span(current_span, func) - except Exception: + except Exception: # pylint: disable=broad-exception-caught # Silently handle any unexpected errors pass @@ -119,20 +118,20 @@ async def async_wrapper(*args, **kwargs): return await func(*args, **kwargs) return async_wrapper - else: - # Sync function wrapper - @wraps(func) - def sync_wrapper(*args, **kwargs): - # Add code attributes to current span - try: - current_span = trace.get_current_span() - if current_span: - _add_code_attributes_to_span(current_span, func) - except Exception: - # Silently handle any unexpected errors - pass - # Call the original sync function - return func(*args, **kwargs) + # Sync function wrapper + @wraps(func) + def sync_wrapper(*args, **kwargs): + # Add code attributes to current span + try: + current_span = trace.get_current_span() + if current_span: + _add_code_attributes_to_span(current_span, func) + except Exception: # pylint: disable=broad-exception-caught + # Silently handle any unexpected errors + pass + + # Call the original sync function + return func(*args, **kwargs) - return sync_wrapper + return sync_wrapper diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py index 8ff1b78d5..f75d7e0f3 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py @@ -3,11 +3,11 @@ import asyncio from unittest import TestCase -from unittest.mock import MagicMock, patch, PropertyMock +from unittest.mock import MagicMock, PropertyMock, patch from amazon.opentelemetry.distro.code_correlation import ( - CODE_FUNCTION_NAME, CODE_FILE_PATH, + CODE_FUNCTION_NAME, CODE_LINE_NUMBER, _add_code_attributes_to_span, add_code_attributes_to_span, @@ -35,6 +35,7 @@ def setUp(self): def test_add_code_attributes_to_recording_span(self): """Test adding code attributes to a recording span.""" + def test_function(): pass @@ -68,7 +69,7 @@ def test_add_code_attributes_function_without_code(self): # Create a mock function without __code__ attribute mock_func = MagicMock() mock_func.__name__ = "mock_function" - delattr(mock_func, '__code__') + delattr(mock_func, "__code__") _add_code_attributes_to_span(self.mock_span, mock_func) @@ -87,7 +88,7 @@ def test_add_code_attributes_function_without_name(self): """Test handling of functions without __name__ attribute.""" # Create an object without __name__ attribute mock_func = MagicMock() - delattr(mock_func, '__name__') + delattr(mock_func, "__name__") mock_func.__code__ = MagicMock() mock_func.__code__.co_filename = "/test/file.py" mock_func.__code__.co_firstlineno = 10 @@ -115,7 +116,7 @@ def test_add_code_attributes_exception_handling(self): self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_func") self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, "/test/file.py") - @patch('amazon.opentelemetry.distro.code_correlation.getattr') + @patch("amazon.opentelemetry.distro.code_correlation.getattr") def test_add_code_attributes_getattr_exception(self, mock_getattr): """Test exception handling when getattr fails.""" mock_getattr.side_effect = Exception("Test exception") @@ -204,7 +205,7 @@ def setUp(self): self.mock_span = MagicMock(spec=Span) self.mock_span.is_recording.return_value = True - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_sync_function(self, mock_get_current_span): """Test decorator with synchronous function.""" mock_get_current_span.return_value = self.mock_span @@ -222,7 +223,7 @@ def test_sync_function(arg1, arg2=None): # Verify span attributes were set self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_sync_function") - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_async_function(self, mock_get_current_span): """Test decorator with asynchronous function.""" mock_get_current_span.return_value = self.mock_span @@ -245,7 +246,7 @@ async def test_async_function(arg1, arg2=None): # Verify span attributes were set self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_async_function") - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_no_current_span(self, mock_get_current_span): """Test decorator when there's no current span.""" mock_get_current_span.return_value = None @@ -263,7 +264,7 @@ def test_function(): # Verify no span attributes were set self.mock_span.set_attribute.assert_not_called() - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_exception_handling(self, mock_get_current_span): """Test decorator handles exceptions gracefully.""" mock_get_current_span.side_effect = Exception("Test exception") @@ -280,6 +281,7 @@ def test_function(): def test_decorator_preserves_function_metadata(self): """Test that decorator preserves original function metadata.""" + @add_code_attributes_to_span def test_function(): """Test function docstring.""" @@ -291,6 +293,7 @@ def test_function(): def test_async_function_detection(self): """Test that async functions are properly detected.""" + # Create a regular function def sync_func(): pass @@ -309,7 +312,7 @@ async def async_func(): # Check that async function returns a coroutine function self.assertTrue(asyncio.iscoroutinefunction(decorated_async)) - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_with_function_that_raises_exception(self, mock_get_current_span): """Test decorator with function that raises exception.""" mock_get_current_span.return_value = self.mock_span @@ -325,7 +328,7 @@ def test_function(): # Verify span attributes were still set before exception self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_function") - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_with_async_function_that_raises_exception(self, mock_get_current_span): """Test decorator with async function that raises exception.""" mock_get_current_span.return_value = self.mock_span @@ -346,8 +349,8 @@ async def test_async_function(): # Verify span attributes were still set before exception self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_async_function") - @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span') - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span") + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_internal_exception_handling_sync(self, mock_get_current_span, mock_add_attributes): """Test that decorator handles internal exceptions gracefully in sync function.""" mock_get_current_span.return_value = self.mock_span @@ -364,8 +367,8 @@ def test_function(): # Verify the function still works correctly despite internal exception self.assertEqual(result, "test result") - @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span') - @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span') + @patch("amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span") + @patch("amazon.opentelemetry.distro.code_correlation.trace.get_current_span") def test_decorator_internal_exception_handling_async(self, mock_get_current_span, mock_add_attributes): """Test that decorator handles internal exceptions gracefully in async function.""" mock_get_current_span.return_value = self.mock_span diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py index 025e47f61..7ab15a626 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py @@ -17,6 +17,7 @@ from amazon.opentelemetry.distro.aws_lambda_span_processor import AwsLambdaSpanProcessor from amazon.opentelemetry.distro.aws_metric_attributes_span_exporter import AwsMetricAttributesSpanExporter from amazon.opentelemetry.distro.aws_opentelemetry_configurator import ( + CODE_CORRELATION_ENABLED_CONFIG, LAMBDA_SPAN_EXPORT_BATCH_SIZE, OTEL_EXPORTER_OTLP_LOGS_ENDPOINT, OTEL_EXPORTER_OTLP_LOGS_HEADERS, @@ -1428,9 +1429,6 @@ def test_create_emf_exporter_cloudwatch_exporter_import_error( def test_get_code_correlation_enabled_status(self): """Test _get_code_correlation_enabled_status function with various environment variable values""" - # Import the constant we need - from amazon.opentelemetry.distro.aws_opentelemetry_configurator import CODE_CORRELATION_ENABLED_CONFIG - # Test when environment variable is not set (default state) os.environ.pop(CODE_CORRELATION_ENABLED_CONFIG, None) result = _get_code_correlation_enabled_status() @@ -1472,9 +1470,6 @@ def test_get_code_correlation_enabled_status(self): self.assertFalse(result) # Test invalid values (should return None and log warning) - # We'll use caplog to capture log messages instead of mocking - import logging - os.environ[CODE_CORRELATION_ENABLED_CONFIG] = "invalid" result = _get_code_correlation_enabled_status() self.assertIsNone(result)