function decorator for adding code level attributes to span

wangzlei · wangzlei · commit 83edee58a388 · 2025-10-01T15:44:06.000-07:00
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
@@ -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
@@ -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
@@ -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
