support code attributes for starlette, fastapi, flask and aws lambda

.pylintrc

@@ -7,7 +7,7 @@ extension-pkg-whitelist=cassandra
 
 # Add list of files or directories to be excluded. They should be base names, not
 # paths.
-ignore=CVS,gen,Dockerfile,docker-compose.yml,README.md,requirements.txt,mock_collector_service_pb2.py,mock_collector_service_pb2.pyi,mock_collector_service_pb2_grpc.py
+ignore=CVS,gen,Dockerfile,docker-compose.yml,README.md,requirements.txt,mock_collector_service_pb2.py,mock_collector_service_pb2.pyi,mock_collector_service_pb2_grpc.py,pyproject.toml,db.sqlite3
 
 # Add files or directories matching the regex patterns to be excluded. The
 # regex matches against base names, not paths.

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py

@@ -616,7 +616,7 @@ def _is_application_signals_runtime_enabled():
     )
 
 
-def _get_code_correlation_enabled_status() -> Optional[bool]:
+def get_code_correlation_enabled_status() -> Optional[bool]:
     """
     Get the code correlation enabled status from environment variable.
 

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py

@@ -21,7 +21,7 @@
 CODE_LINE_NUMBER = "code.line.number"
 
 
-def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None:
+def add_code_attributes_to_span(span, func: Callable[..., Any]) -> None:
     """
     Add code-related attributes to a span based on a Python function.
 
@@ -68,7 +68,7 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None:
         pass
 
 
-def add_code_attributes_to_span(func: Callable[..., Any]) -> Callable[..., Any]:
+def record_code_attributes(func: Callable[..., Any]) -> Callable[..., Any]:
     """
     Decorator to automatically add code attributes to the current OpenTelemetry span.
 
@@ -81,12 +81,12 @@ def add_code_attributes_to_span(func: Callable[..., Any]) -> Callable[..., Any]:
     This decorator supports both synchronous and asynchronous functions.
 
     Usage:
-        @add_code_attributes_to_span
+        @record_code_attributes
         def my_sync_function():
             # Sync function implementation
             pass
 
-        @add_code_attributes_to_span
+        @record_code_attributes
         async def my_async_function():
             # Async function implementation
             pass
@@ -109,7 +109,7 @@ async def async_wrapper(*args, **kwargs):
             try:
                 current_span = trace.get_current_span()
                 if current_span:
-                    _add_code_attributes_to_span(current_span, func)
+                    add_code_attributes_to_span(current_span, func)
             except Exception:  # pylint: disable=broad-exception-caught
                 # Silently handle any unexpected errors
                 pass
@@ -126,7 +126,7 @@ def sync_wrapper(*args, **kwargs):
         try:
             current_span = trace.get_current_span()
             if current_span:
-                _add_code_attributes_to_span(current_span, func)
+                add_code_attributes_to_span(current_span, func)
         except Exception:  # pylint: disable=broad-exception-caught
             # Silently handle any unexpected errors
             pass

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_fastapi_patches.py

@@ -0,0 +1,119 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+# Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
+
+from logging import getLogger
+
+from amazon.opentelemetry.distro.aws_opentelemetry_configurator import get_code_correlation_enabled_status
+
+_logger = getLogger(__name__)
+
+
+def _apply_fastapi_instrumentation_patches() -> None:
+    """FastAPI instrumentation patches
+
+    Applies patches to provide code attributes support for FastAPI instrumentation.
+    This patches the FastAPI instrumentation to automatically add code attributes
+    to spans by decorating view functions with record_code_attributes.
+    """
+    if get_code_correlation_enabled_status() is True:
+        _apply_fastapi_code_attributes_patch()
+
+
+def _apply_fastapi_code_attributes_patch() -> None:
+    """FastAPI instrumentation patch for code attributes
+
+    This patch modifies the FastAPI instrumentation to automatically apply
+    the current_span_code_attributes decorator to all endpoint functions when
+    the FastAPI app is instrumented.
+
+    The patch:
+    1. Imports current_span_code_attributes decorator from AWS distro utils
+    2. Hooks FastAPI's APIRouter.add_api_route method during instrumentation
+    3. Automatically decorates endpoint functions as they are registered
+    4. Adds code.function.name, code.file.path, and code.line.number to spans
+    5. Provides cleanup during uninstrumentation
+    """
+    try:
+        # Import FastAPI instrumentation classes and AWS decorator
+        from fastapi import routing  # pylint: disable=import-outside-toplevel
+
+        from amazon.opentelemetry.distro.code_correlation import (  # pylint: disable=import-outside-toplevel
+            record_code_attributes,
+        )
+        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor  # pylint: disable=import-outside-toplevel
+
+        # Store the original _instrument and _uninstrument methods
+        original_instrument = FastAPIInstrumentor._instrument
+        original_uninstrument = FastAPIInstrumentor._uninstrument
+
+        def _wrapped_add_api_route(original_add_api_route_method):
+            """Wrapper for APIRouter.add_api_route method."""
+
+            def wrapper(self, *args, **kwargs):
+                # Apply current_span_code_attributes decorator to endpoint function
+                try:
+                    # Get endpoint function from args or kwargs
+                    endpoint = None
+                    if len(args) >= 2:
+                        endpoint = args[1]
+                    else:
+                        endpoint = kwargs.get("endpoint")
+
+                    if endpoint and callable(endpoint):
+                        # Check if function is already decorated (avoid double decoration)
+                        if not hasattr(endpoint, "_current_span_code_attributes_decorated"):
+                            # Apply decorator
+                            decorated_endpoint = record_code_attributes(endpoint)
+                            # Mark as decorated to avoid double decoration
+                            decorated_endpoint._current_span_code_attributes_decorated = True
+                            decorated_endpoint._original_endpoint = endpoint
+
+                            # Replace endpoint in args or kwargs
+                            if len(args) >= 2:
+                                args = list(args)
+                                args[1] = decorated_endpoint
+                                args = tuple(args)
+                            elif "endpoint" in kwargs:
+                                kwargs["endpoint"] = decorated_endpoint
+
+                except Exception as exc:  # pylint: disable=broad-exception-caught
+                    _logger.warning("Failed to apply code attributes decorator to endpoint: %s", exc)
+
+                return original_add_api_route_method(self, *args, **kwargs)
+
+            return wrapper
+
+        def patched_instrument(self, **kwargs):
+            """Patched _instrument method with APIRouter.add_api_route wrapping"""
+            # Store original add_api_route method if not already stored
+            if not hasattr(self, "_original_apirouter"):
+                self._original_apirouter = routing.APIRouter.add_api_route
+
+            # Wrap APIRouter.add_api_route with code attributes decoration
+            routing.APIRouter.add_api_route = _wrapped_add_api_route(self._original_apirouter)
+
+            # Call the original _instrument method
+            original_instrument(self, **kwargs)
+
+        def patched_uninstrument(self, **kwargs):
+            """Patched _uninstrument method with APIRouter.add_api_route restoration"""
+            # Call the original _uninstrument method first
+            original_uninstrument(self, **kwargs)
+
+            # Restore original APIRouter.add_api_route method if it exists
+            if hasattr(self, "_original_apirouter"):
+                try:
+                    routing.APIRouter.add_api_route = self._original_apirouter
+                    delattr(self, "_original_apirouter")
+                except Exception as exc:  # pylint: disable=broad-exception-caught
+                    _logger.warning("Failed to restore original APIRouter.add_api_route method: %s", exc)
+
+        # Apply the patches to FastAPIInstrumentor
+        FastAPIInstrumentor._instrument = patched_instrument
+        FastAPIInstrumentor._uninstrument = patched_uninstrument
+
+        _logger.debug("FastAPI instrumentation code attributes patch applied successfully")
+
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        _logger.warning("Failed to apply FastAPI code attributes patch: %s", exc)

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_flask_patches.py

@@ -0,0 +1,145 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+# Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
+
+from logging import getLogger
+
+from amazon.opentelemetry.distro.aws_opentelemetry_configurator import get_code_correlation_enabled_status
+
+_logger = getLogger(__name__)
+
+
+def _apply_flask_instrumentation_patches() -> None:
+    """Flask instrumentation patches
+
+    Applies patches to provide code attributes support for Flask instrumentation.
+    This patches the Flask instrumentation to automatically add code attributes
+    to spans by decorating view functions with record_code_attributes.
+    """
+    if get_code_correlation_enabled_status() is True:
+        _apply_flask_code_attributes_patch()
+
+
+def _apply_flask_code_attributes_patch() -> None:  # pylint: disable=too-many-statements
+    """Flask instrumentation patch for code attributes
+
+    This patch modifies the Flask instrumentation to automatically apply
+    the current_span_code_attributes decorator to all view functions when
+    the Flask app is instrumented.
+
+    The patch:
+    1. Imports current_span_code_attributes decorator from AWS distro utils
+    2. Hooks Flask's add_url_rule method during _instrument by patching Flask class
+    3. Hooks Flask's dispatch_request method to handle deferred view function binding
+    4. Automatically decorates view functions as they are registered or at request time
+    5. Adds code.function.name, code.file.path, and code.line.number to spans
+    6. Provides cleanup during _uninstrument
+    """
+    try:
+        # Import Flask instrumentation classes and AWS decorator
+        import flask  # pylint: disable=import-outside-toplevel
+
+        from amazon.opentelemetry.distro.code_correlation import (  # pylint: disable=import-outside-toplevel
+            record_code_attributes,
+        )
+        from opentelemetry.instrumentation.flask import FlaskInstrumentor  # pylint: disable=import-outside-toplevel
+
+        # Store the original _instrument and _uninstrument methods
+        original_instrument = FlaskInstrumentor._instrument
+        original_uninstrument = FlaskInstrumentor._uninstrument
+
+        # Store reference to original Flask methods
+        original_flask_add_url_rule = flask.Flask.add_url_rule
+        original_flask_dispatch_request = flask.Flask.dispatch_request
+
+        def _decorate_view_func(view_func, endpoint=None):
+            """Helper function to decorate a view function with code attributes."""
+            try:
+                if view_func and callable(view_func):
+                    # Check if function is already decorated (avoid double decoration)
+                    if not hasattr(view_func, "_current_span_code_attributes_decorated"):
+                        # Apply decorator
+                        decorated_view_func = record_code_attributes(view_func)
+                        # Mark as decorated to avoid double decoration
+                        decorated_view_func._current_span_code_attributes_decorated = True
+                        decorated_view_func._original_view_func = view_func
+                        return decorated_view_func
+                return view_func
+            except Exception as exc:  # pylint: disable=broad-exception-caught
+                _logger.warning("Failed to apply code attributes decorator to view function %s: %s", endpoint, exc)
+                return view_func
+
+        def _wrapped_add_url_rule(self, rule, endpoint=None, view_func=None, **options):
+            """Wrapped Flask.add_url_rule method with code attributes decoration."""
+            # Apply decorator to view function if available
+            if view_func:
+                view_func = _decorate_view_func(view_func, endpoint)
+
+            return original_flask_add_url_rule(self, rule, endpoint, view_func, **options)
+
+        def _wrapped_dispatch_request(self):
+            """Wrapped Flask.dispatch_request method to handle deferred view function binding."""
+            try:
+                # Get the current request context
+                from flask import request  # pylint: disable=import-outside-toplevel
+
+                # Check if there's an endpoint for this request
+                endpoint = request.endpoint
+                if endpoint and endpoint in self.view_functions:
+                    view_func = self.view_functions[endpoint]
+
+                    # Check if the view function needs decoration
+                    if view_func and callable(view_func):
+                        if not hasattr(view_func, "_current_span_code_attributes_decorated"):
+                            # Decorate the view function and replace it in view_functions
+                            decorated_view_func = _decorate_view_func(view_func, endpoint)
+                            if decorated_view_func != view_func:
+                                self.view_functions[endpoint] = decorated_view_func
+                                _logger.debug(
+                                    "Applied code attributes decorator to deferred view function for endpoint: %s",
+                                    endpoint,
+                                )
+
+            except Exception as exc:  # pylint: disable=broad-exception-caught
+                _logger.warning("Failed to process deferred view function decoration: %s", exc)
+
+            # Call the original dispatch_request method
+            return original_flask_dispatch_request(self)
+
+        def patched_instrument(self, **kwargs):
+            """Patched _instrument method with Flask method wrapping"""
+            # Store original methods if not already stored
+            if not hasattr(self, "_original_flask_add_url_rule"):
+                self._original_flask_add_url_rule = flask.Flask.add_url_rule
+                self._original_flask_dispatch_request = flask.Flask.dispatch_request
+
+                # Wrap Flask methods with code attributes decoration
+                flask.Flask.add_url_rule = _wrapped_add_url_rule
+                flask.Flask.dispatch_request = _wrapped_dispatch_request
+
+            # Call the original _instrument method
+            original_instrument(self, **kwargs)
+
+        def patched_uninstrument(self, **kwargs):
+            """Patched _uninstrument method with Flask method restoration"""
+            # Call the original _uninstrument method first
+            original_uninstrument(self, **kwargs)
+
+            # Restore original Flask methods if they exist
+            if hasattr(self, "_original_flask_add_url_rule"):
+                try:
+                    flask.Flask.add_url_rule = self._original_flask_add_url_rule
+                    flask.Flask.dispatch_request = self._original_flask_dispatch_request
+                    delattr(self, "_original_flask_add_url_rule")
+                    delattr(self, "_original_flask_dispatch_request")
+                except Exception as exc:  # pylint: disable=broad-exception-caught
+                    _logger.warning("Failed to restore original Flask methods: %s", exc)
+
+        # Apply the patches to FlaskInstrumentor
+        FlaskInstrumentor._instrument = patched_instrument
+        FlaskInstrumentor._uninstrument = patched_uninstrument
+
+        _logger.debug("Flask instrumentation code attributes patch applied successfully")
+
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        _logger.warning("Failed to apply Flask code attributes patch: %s", exc)

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_instrumentation_patch.py

@@ -71,6 +71,20 @@ def apply_instrumentation_patches() -> None:
         # TODO: Remove patch after syncing with upstream v1.34.0 or later
         _apply_starlette_instrumentation_patches()
 
+    if is_installed("flask"):
+        # pylint: disable=import-outside-toplevel
+        # Delay import to only occur if patches is safe to apply (e.g. the instrumented library is installed).
+        from amazon.opentelemetry.distro.patches._flask_patches import _apply_flask_instrumentation_patches
+
+        _apply_flask_instrumentation_patches()
+
+    if is_installed("fastapi"):
+        # pylint: disable=import-outside-toplevel
+        # Delay import to only occur if patches is safe to apply (e.g. the instrumented library is installed).
+        from amazon.opentelemetry.distro.patches._fastapi_patches import _apply_fastapi_instrumentation_patches
+
+        _apply_fastapi_instrumentation_patches()
+
     # No need to check if library is installed as this patches opentelemetry.sdk,
     # which must be installed for the distro to work at all.
     _apply_resource_detector_patches()