support code attributes for custom span

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

@@ -22,6 +22,7 @@
 from .utils import add_code_attributes_to_span_from_frame
 
 
+# pylint: disable=no-self-use
 class CodeAttributesSpanProcessor(SpanProcessor):
     """
     A SpanProcessor that captures and attaches code attributes to spans.
@@ -68,26 +69,47 @@ def on_start(
         # Capture code attributes from stack trace
         self._capture_code_attributes(span)
 
-    @staticmethod
-    def _should_process_span(span: Span) -> bool:
+    def _should_process_span(self, span: Span) -> bool:
         """
         Determine if span should be processed for code attributes.
 
         Returns False if:
+        - Span is library instrumentation SERVER or INTERNAL span
         - Span already has code attributes
-        - Span is SERVER or INTERNAL span
+
+        Note: Library instrumentation CLIENT/PRODUCER/CONSUMER spans are still processed
+        as they provide valuable context for tracing call chains.
         """
-        # Skip if span already has code attributes
+
+        if span.kind in (SpanKind.SERVER, SpanKind.INTERNAL) and self._is_library_instrumentation_span(span):
+            return False
+
         if span.attributes is not None and CODE_FUNCTION_NAME in span.attributes:
             return False
 
-        # Process spans except SERVER and INTERNAL spans
-        return span.kind not in (SpanKind.SERVER, SpanKind.INTERNAL)
+        return True
+
+    def _is_library_instrumentation_span(self, span: Span) -> bool:
+        """
+        Check if span is created by library instrumentation.
+
+        Args:
+            span: The span to check
+
+        Returns:
+            True if span is from library instrumentation, False otherwise
+        """
+        scope = span.instrumentation_scope
+
+        if scope is None or scope.name is None:
+            return False  # No scope info, assume user-created
+
+        return scope.name.startswith("opentelemetry.instrumentation")
 
     def _capture_code_attributes(self, span: Span) -> None:
         """Capture and attach code attributes from current stack trace."""
         try:
-            current_frame = sys._getframe(1)
+            current_frame = sys._getframe(1)  # pylint: disable=protected-access
 
             for frame_index, frame in enumerate(self._iter_stack_frames(current_frame)):
                 if frame_index >= self.MAX_STACK_FRAMES:
@@ -112,6 +134,6 @@ def shutdown(self) -> None:
         """Called when the processor is shutdown. No cleanup needed."""
         # No cleanup needed for code attributes processor
 
-    def force_flush(self, timeout_millis: int = 30000) -> bool:  # pylint: disable=no-self-use,unused-argument
+    def force_flush(self, timeout_millis: int = 30000) -> bool:  # pylint: disable=unused-argument
         """Force flush any pending spans. Always returns True as no pending work."""
         return True

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_attributes_span_processor.py

@@ -8,6 +8,7 @@
 from amazon.opentelemetry.distro.code_correlation.code_attributes_span_processor import CodeAttributesSpanProcessor
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import ReadableSpan, Span
+from opentelemetry.sdk.util.instrumentation import InstrumentationScope
 from opentelemetry.semconv.attributes.code_attributes import CODE_FUNCTION_NAME
 from opentelemetry.trace import SpanKind
 
@@ -82,57 +83,85 @@ def test_max_stack_frames_constant(self):
         """Test that MAX_STACK_FRAMES is set to expected value."""
         self.assertEqual(CodeAttributesSpanProcessor.MAX_STACK_FRAMES, 50)
 
-    def create_mock_span(self, span_kind=SpanKind.CLIENT, attributes=None):
+    def create_mock_span(self, span_kind=SpanKind.CLIENT, attributes=None, instrumentation_scope_name=None):
         """Helper to create a mock span with specified properties."""
         mock_span = MagicMock(spec=Span)
         mock_span.kind = span_kind
         mock_span.attributes = attributes
+
+        # Set up instrumentation scope
+        if instrumentation_scope_name is not None:
+            mock_scope = MagicMock(spec=InstrumentationScope)
+            mock_scope.name = instrumentation_scope_name
+            mock_span.instrumentation_scope = mock_scope
+        else:
+            mock_span.instrumentation_scope = None
+
         return mock_span
 
-    def test_should_process_span_client_span_without_attributes(self):
-        """Test that CLIENT spans without code attributes should be processed."""
-        span = self.create_mock_span(SpanKind.CLIENT, attributes=None)
-        result = CodeAttributesSpanProcessor._should_process_span(span)
+    def test_should_process_span_user_client_span_without_attributes(self):
+        """Test that user CLIENT spans without code attributes should be processed."""
+        span = self.create_mock_span(SpanKind.CLIENT, attributes=None, instrumentation_scope_name="my-app")
+        result = self.processor._should_process_span(span)
         self.assertTrue(result)
 
-    def test_should_process_span_client_span_with_empty_attributes(self):
-        """Test that CLIENT spans with empty attributes should be processed."""
-        span = self.create_mock_span(SpanKind.CLIENT, attributes={})
-        result = CodeAttributesSpanProcessor._should_process_span(span)
+    def test_should_process_span_user_client_span_with_empty_attributes(self):
+        """Test that user CLIENT spans with empty attributes should be processed."""
+        span = self.create_mock_span(SpanKind.CLIENT, attributes={}, instrumentation_scope_name="my-app")
+        result = self.processor._should_process_span(span)
         self.assertTrue(result)
 
     def test_should_process_span_client_span_with_existing_code_attributes(self):
-        """Test that CLIENT spans with existing code attributes should not be processed."""
+        """Test that spans with existing code attributes should not be processed."""
         attributes = {CODE_FUNCTION_NAME: "existing.function"}
-        span = self.create_mock_span(SpanKind.CLIENT, attributes=attributes)
-        result = CodeAttributesSpanProcessor._should_process_span(span)
+        span = self.create_mock_span(SpanKind.CLIENT, attributes=attributes, instrumentation_scope_name="my-app")
+        result = self.processor._should_process_span(span)
         self.assertFalse(result)
 
-    def test_should_process_span_server_and_internal_spans(self):
-        """Test that SERVER and INTERNAL spans should not be processed."""
-        test_cases = [
-            SpanKind.SERVER,
-            SpanKind.INTERNAL,
-        ]
+    def test_should_process_span_user_server_spans(self):
+        """Test that user SERVER spans should be processed (new logic)."""
+        span = self.create_mock_span(SpanKind.SERVER, attributes=None, instrumentation_scope_name="my-app")
+        result = self.processor._should_process_span(span)
+        self.assertTrue(result)
 
-        for span_kind in test_cases:
-            with self.subTest(span_kind=span_kind):
-                span = self.create_mock_span(span_kind, attributes=None)
-                result = CodeAttributesSpanProcessor._should_process_span(span)
-                self.assertFalse(result)
+    def test_should_process_span_library_server_spans_not_processed(self):
+        """Test that library instrumentation SERVER spans should not be processed."""
+        span = self.create_mock_span(
+            SpanKind.SERVER, attributes=None, instrumentation_scope_name="opentelemetry.instrumentation.flask"
+        )
+        result = self.processor._should_process_span(span)
+        self.assertFalse(result)
 
-    def test_should_process_span_client_producer_consumer_spans(self):
-        """Test that CLIENT, PRODUCER, and CONSUMER spans should be processed."""
+    def test_should_process_span_library_internal_spans_not_processed(self):
+        """Test that library instrumentation INTERNAL spans should not be processed."""
+        span = self.create_mock_span(
+            SpanKind.INTERNAL, attributes=None, instrumentation_scope_name="opentelemetry.instrumentation.botocore"
+        )
+        result = self.processor._should_process_span(span)
+        self.assertFalse(result)
+
+    def test_should_process_span_library_client_spans_processed(self):
+        """Test that library instrumentation CLIENT spans should be processed."""
+        span = self.create_mock_span(
+            SpanKind.CLIENT, attributes=None, instrumentation_scope_name="opentelemetry.instrumentation.requests"
+        )
+        result = self.processor._should_process_span(span)
+        self.assertTrue(result)
+
+    def test_should_process_span_user_spans_all_kinds(self):
+        """Test that user spans of all kinds should be processed."""
         test_cases = [
             SpanKind.CLIENT,
+            SpanKind.SERVER,
             SpanKind.PRODUCER,
             SpanKind.CONSUMER,
+            SpanKind.INTERNAL,
         ]
 
         for span_kind in test_cases:
             with self.subTest(span_kind=span_kind):
-                span = self.create_mock_span(span_kind, attributes=None)
-                result = CodeAttributesSpanProcessor._should_process_span(span)
+                span = self.create_mock_span(span_kind, attributes=None, instrumentation_scope_name="my-app")
+                result = self.processor._should_process_span(span)
                 self.assertTrue(result)
 
     @patch("amazon.opentelemetry.distro.code_correlation.code_attributes_span_processor.sys._getframe")
@@ -300,7 +329,8 @@ def is_user_code_side_effect(filename):
 
     def test_on_start_should_not_process_span(self):
         """Test on_start when span should not be processed."""
-        span = self.create_mock_span(SpanKind.SERVER)  # Non-client span
+        # Library instrumentation SERVER span should not be processed
+        span = self.create_mock_span(SpanKind.SERVER, instrumentation_scope_name="opentelemetry.instrumentation.flask")
 
         with patch.object(self.processor, "_capture_code_attributes") as mock_capture:
             self.processor.on_start(span)

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_aio_pika_patches.py

@@ -46,7 +46,7 @@ def test_patch_callback_decorator_decorate(self):
         # Verify we got a function back (the enhanced decorated callback)
         self.assertTrue(callable(result))
 
-    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches." "get_code_correlation_enabled_status")
+    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches.get_code_correlation_enabled_status")
     def test_apply_aio_pika_instrumentation_patches_disabled(self, mock_get_status):
         """Test patches are not applied when code correlation is disabled."""
         mock_get_status.return_value = False
@@ -61,7 +61,7 @@ def test_apply_aio_pika_instrumentation_patches_disabled(self, mock_get_status):
             # Should not raise exception when code correlation is disabled
             _apply_aio_pika_instrumentation_patches()
 
-    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches." "get_code_correlation_enabled_status")
+    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches.get_code_correlation_enabled_status")
     def test_apply_aio_pika_instrumentation_patches_enabled(self, mock_get_status):
         """Test patches are applied when code correlation is enabled."""
         mock_get_status.return_value = True
@@ -83,7 +83,7 @@ def test_apply_aio_pika_instrumentation_patches_enabled(self, mock_get_status):
             # Verify the decorate method was patched
             self.assertTrue(hasattr(mock_callback_decorator, "decorate"))
 
-    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches." "get_code_correlation_enabled_status")
+    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches.get_code_correlation_enabled_status")
     def test_apply_aio_pika_instrumentation_patches_import_error(self, mock_get_status):
         """Test patches handle import errors gracefully."""
         mock_get_status.return_value = True
@@ -99,7 +99,7 @@ def test_apply_aio_pika_instrumentation_patches_import_error(self, mock_get_stat
             _apply_aio_pika_instrumentation_patches()
 
     @patch("amazon.opentelemetry.distro.patches._aio_pika_patches.logger")
-    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches." "get_code_correlation_enabled_status")
+    @patch("amazon.opentelemetry.distro.patches._aio_pika_patches.get_code_correlation_enabled_status")
     def test_apply_aio_pika_instrumentation_patches_exception_handling(self, mock_get_status, mock_logger):
         """Test patches handle general exceptions gracefully."""
         mock_get_status.side_effect = Exception("Test exception")

contract-tests/tests/test/amazon/misc/code_attributes_test.py

@@ -0,0 +1,81 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Dict, List
+
+from mock_collector_client import ResourceScopeSpan
+from requests import Response
+from typing_extensions import override
+
+from amazon.base.contract_test_base import ContractTestBase
+from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
+from opentelemetry.proto.trace.v1.trace_pb2 import Span
+from opentelemetry.semconv.attributes.code_attributes import CODE_LINE_NUMBER
+
+
+class CodeAttributesTest(ContractTestBase):
+    @override
+    @staticmethod
+    def get_application_image_name() -> str:
+        return "aws-application-signals-tests-requests-app"
+
+    @override
+    def get_application_network_aliases(self) -> List[str]:
+        """
+        This will be the target hostname of the clients making http requests in the application image, so that they
+        don't use localhost.
+        """
+        return ["backend"]
+
+    @override
+    def get_application_extra_environment_variables(self) -> Dict[str, str]:
+        """
+        This does not appear to do anything, as it does not seem that OTEL supports peer service for Python. Keeping
+        for consistency with Java contract tests at this time.
+        """
+        return {
+            "OTEL_INSTRUMENTATION_COMMON_PEER_SERVICE_MAPPING": "backend=backend:8080",
+            "OTEL_AWS_CODE_CORRELATION_ENABLED": "true",
+        }
+
+    def test_success(self) -> None:
+        self.do_test_requests("success", "GET", 200, 0, 0, request_method="GET")
+
+    def test_error(self) -> None:
+        self.do_test_requests("error", "GET", 400, 1, 0, request_method="GET")
+
+    def test_fault(self) -> None:
+        self.do_test_requests("fault", "GET", 500, 0, 1, request_method="GET")
+
+    def test_success_post(self) -> None:
+        self.do_test_requests("success/postmethod", "POST", 200, 0, 0, request_method="POST")
+
+    def test_error_post(self) -> None:
+        self.do_test_requests("error/postmethod", "POST", 400, 1, 0, request_method="POST")
+
+    def test_fault_post(self) -> None:
+        self.do_test_requests("fault/postmethod", "POST", 500, 0, 1, request_method="POST")
+
+    def do_test_requests(
+        self, path: str, method: str, status_code: int, expected_error: int, expected_fault: int, **kwargs
+    ) -> None:
+        response: Response = self.send_request(method, path)
+        self.assertEqual(status_code, response.status_code)
+
+        resource_scope_spans: List[ResourceScopeSpan] = self.mock_collector_client.get_traces()
+        self._assert_span_code_attributes(resource_scope_spans, path, **kwargs)
+
+    @override
+    def _assert_span_code_attributes(self, resource_scope_spans: List[ResourceScopeSpan], path: str, **kwargs) -> None:
+        target_spans: List[Span] = []
+        for resource_scope_span in resource_scope_spans:
+            # pylint: disable=no-member
+            if resource_scope_span.span.kind == Span.SPAN_KIND_CLIENT:
+                target_spans.append(resource_scope_span.span)
+
+        self.assertEqual(len(target_spans), 1)
+        self._assert_code_attribute(target_spans[0].attributes)
+
+    def _assert_code_attribute(self, attributes_list: List[KeyValue]):
+        attributes_dict: Dict[str, AnyValue] = self._get_attributes_dict(attributes_list)
+        # The line number of calling requests.request in requests_server.py
+        self._assert_int_attribute(attributes_dict, CODE_LINE_NUMBER, 41)