Skip to content

Commit 65d1818

Browse files
committed
add check if the app is starlette app
1 parent b5e08ae commit 65d1818

File tree

2 files changed

+22
-10
lines changed

2 files changed

+22
-10
lines changed

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
from logging import Logger, getLogger
55
from typing import Collection
66

7+
from starlette import applications
8+
79
from amazon.opentelemetry.distro._utils import is_agent_observability_enabled
810

911
_logger: Logger = getLogger(__name__)
@@ -31,17 +33,20 @@ def patched_instrumentation_dependencies(self) -> Collection[str]:
3133
# Apply the patch
3234
StarletteInstrumentor.instrumentation_dependencies = patched_instrumentation_dependencies
3335

34-
# Patch to exclude http receive/send ASGI event spans from Bedrock AgentCore,
36+
# Patch to exclude http receive/send ASGI event spans from Bedrock AgentCore Starlette applications,
3537
# this Middleware instrumentation is injected internally by Starlette Instrumentor, see:
3638
# https://github.com/open-telemetry/opentelemetry-python-contrib/blob/51da0a766e5d3cbc746189e10c9573163198cfcd/instrumentation/opentelemetry-instrumentation-asgi/src/opentelemetry/instrumentation/asgi/__init__.py#L573
3739
if is_agent_observability_enabled():
3840
original_init = OpenTelemetryMiddleware.__init__
3941

4042
def patched_init(self, app, **kwargs):
4143
original_init(self, app, **kwargs)
42-
if hasattr(self, "exclude_receive_span"):
44+
# See: https://github.com/encode/starlette/blob/cb8f84f5284dc301ca7a31eb732b9e140769dd48/starlette/applications.py#L77-L78
45+
is_starlette_app = hasattr(app, "user_middleware") and hasattr(app, "middleware_stack")
46+
47+
# Only apply exclude_spans for Starlette apps
48+
if is_starlette_app and hasattr(self, "exclude_receive_span") and hasattr(self, "exclude_send_span"):
4349
self.exclude_receive_span = True
44-
if hasattr(self, "exclude_send_span"):
4550
self.exclude_send_span = True
4651

4752
OpenTelemetryMiddleware.__init__ = patched_init

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

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,22 @@
11
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
# SPDX-License-Identifier: Apache-2.0
3-
import unittest
3+
from unittest import TestCase
44
from unittest.mock import MagicMock, patch
55

66
from amazon.opentelemetry.distro.patches._starlette_patches import _apply_starlette_instrumentation_patches
77

88

9-
class TestStarlettePatch(unittest.TestCase):
9+
class TestStarlettePatch(TestCase):
1010
"""Test the Starlette instrumentation patches."""
1111

1212
@patch("amazon.opentelemetry.distro.patches._starlette_patches._logger")
1313
def test_starlette_patch_applied_successfully(self, mock_logger):
1414
"""Test that the Starlette instrumentation patch is applied successfully."""
15-
for agent_enabled in [True, False]:
16-
with self.subTest(agent_enabled=agent_enabled):
17-
with patch.dict("os.environ", {"AGENT_OBSERVABILITY_ENABLED": "true" if agent_enabled else "false"}):
15+
for patched_starlette_enabled in [True, False]:
16+
with self.subTest(agent_enabled=patched_starlette_enabled):
17+
with patch.dict(
18+
"os.environ", {"AGENT_OBSERVABILITY_ENABLED": "true" if patched_starlette_enabled else "false"}
19+
):
1820
# Create a mock StarletteInstrumentor class
1921
mock_instrumentor_class = MagicMock()
2022
mock_instrumentor_class.__name__ = "StarletteInstrumentor"
@@ -52,13 +54,18 @@ def __init__(self, app, **kwargs):
5254
result = mock_instrumentor_class.instrumentation_dependencies(mock_instance)
5355
self.assertEqual(result, ("starlette >= 0.13",))
5456

57+
mock_app = MagicMock()
58+
if patched_starlette_enabled:
59+
mock_app.user_middleware = []
60+
mock_app.middleware_stack = None
61+
5562
mock_middleware_instance = MagicMock()
5663
mock_middleware_instance.exclude_receive_span = False
5764
mock_middleware_instance.exclude_send_span = False
58-
mock_middleware_class.__init__(mock_middleware_instance, "app")
65+
mock_middleware_class.__init__(mock_middleware_instance, mock_app)
5966

6067
# Test middleware patching sets exclude flags
61-
if agent_enabled:
68+
if patched_starlette_enabled:
6269
self.assertTrue(mock_middleware_instance.exclude_receive_span)
6370
self.assertTrue(mock_middleware_instance.exclude_send_span)
6471
else:

0 commit comments

Comments
 (0)