Skip to content

Commit 5b1529d

Browse files
florentinlflorentinl
committed
feat(proxy): update tags for proxy inferred spans
1 parent 01850f2 commit 5b1529d

File tree

21 files changed

+320
-84
lines changed

21 files changed

+320
-84
lines changed

‎ddtrace/_trace/_inferred_proxy.py‎

Lines changed: 118 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
1+
from dataclasses import dataclass
12
import logging
3+
from typing import Callable
24
from typing import Dict
3-
from typing import Union
5+
from typing import Optional
46

57
from ddtrace import config
68
from ddtrace._trace.span import Span
9+
from ddtrace.ext import SpanKind
710
from ddtrace.ext import SpanTypes
811
from ddtrace.ext import http
912
from ddtrace.internal.constants import COMPONENT
@@ -13,18 +16,61 @@
1316

1417
log = logging.getLogger(__name__)
1518

19+
20+
@dataclass
21+
class ProxyHeaderContext:
22+
system_name: str
23+
request_time: str
24+
method: Optional[str]
25+
path: Optional[str]
26+
resource_path: Optional[str]
27+
domain_name: Optional[str]
28+
stage: Optional[str]
29+
account_id: Optional[str]
30+
api_id: Optional[str]
31+
region: Optional[str]
32+
user: Optional[str]
33+
34+
35+
@dataclass
36+
class ProxyInfo:
37+
span_name: str
38+
component: str
39+
resource_arn_builder: Optional[Callable[[ProxyHeaderContext], Optional[str]]] = None
40+
41+
42+
def _api_gateway_rest_api_arn(proxy_context: ProxyHeaderContext) -> Optional[str]:
43+
if proxy_context.region and proxy_context.api_id:
44+
return f"arn:aws:apigateway:{proxy_context.region}::/restapis/{proxy_context.api_id}"
45+
return None
46+
47+
48+
def _api_gateway_http_api_arn(proxy_context: ProxyHeaderContext) -> Optional[str]:
49+
if proxy_context.region and proxy_context.api_id:
50+
return f"arn:aws:apigateway:{proxy_context.region}::/apis/{proxy_context.api_id}"
51+
return None
52+
53+
54+
supported_proxies: Dict[str, ProxyInfo] = {
55+
"aws-apigateway": ProxyInfo("aws.apigateway", "aws-apigateway", _api_gateway_rest_api_arn),
56+
"aws-httpapi": ProxyInfo("aws.httpapi", "aws-httpapi", _api_gateway_http_api_arn),
57+
}
58+
59+
SUPPORTED_PROXY_SPAN_NAMES = {info.span_name for info in supported_proxies.values()}
60+
1661
# Checking lower case and upper case versions per WSGI spec following ddtrace/propagation/http.py's
1762
# logic to extract http headers
1863
POSSIBLE_PROXY_HEADER_SYSTEM = _possible_header("x-dd-proxy")
1964
POSSIBLE_PROXY_HEADER_START_TIME_MS = _possible_header("x-dd-proxy-request-time-ms")
2065
POSSIBLE_PROXY_HEADER_PATH = _possible_header("x-dd-proxy-path")
66+
POSSIBLE_PROXY_HEADER_RESOURCE_PATH = _possible_header("x-dd-proxy-resource-path")
2167
POSSIBLE_PROXY_HEADER_HTTPMETHOD = _possible_header("x-dd-proxy-httpmethod")
2268
POSSIBLE_PROXY_HEADER_DOMAIN = _possible_header("x-dd-proxy-domain-name")
2369
POSSIBLE_PROXY_HEADER_STAGE = _possible_header("x-dd-proxy-stage")
24-
25-
supported_proxies: Dict[str, Dict[str, str]] = {
26-
"aws-apigateway": {"span_name": "aws.apigateway", "component": "aws-apigateway"}
27-
}
70+
POSSIBLE_PROXY_HEADER_ACCOUNT_ID = _possible_header("x-dd-proxy-account-id")
71+
POSSIBLE_PROXY_HEADER_API_ID = _possible_header("x-dd-proxy-api-id")
72+
POSSIBLE_PROXY_HEADER_REGION = _possible_header("x-dd-proxy-region")
73+
POSSIBLE_PROXY_HEADER_USER = _possible_header("x-dd-proxy-user")
2874

2975

3076
def create_inferred_proxy_span_if_headers_exist(ctx, headers, child_of, tracer) -> None:
@@ -38,19 +84,23 @@ def create_inferred_proxy_span_if_headers_exist(ctx, headers, child_of, tracer)
3884
if not proxy_context:
3985
return None
4086

41-
proxy_span_info = supported_proxies[proxy_context["proxy_system_name"]]
87+
proxy_info = supported_proxies[proxy_context.system_name]
88+
89+
method = proxy_context.method
90+
route_or_path = proxy_context.resource_path or proxy_context.path
91+
resource = f"{method or ''} {route_or_path or ''}"
4292

4393
span = tracer.start_span(
44-
proxy_span_info["span_name"],
45-
service=proxy_context.get("domain_name", config._get_service()),
46-
resource=proxy_context["method"] + " " + proxy_context["path"],
94+
proxy_info.span_name,
95+
service=proxy_context.domain_name or config._get_service(),
96+
resource=resource,
4797
span_type=SpanTypes.WEB,
4898
activate=True,
4999
child_of=child_of,
50100
)
51-
span.start_ns = int(proxy_context["request_time"]) * 1000000
101+
span.start_ns = int(proxy_context.request_time) * 1000000
52102

53-
set_inferred_proxy_span_tags(span, proxy_context)
103+
set_inferred_proxy_span_tags(span, proxy_context, proxy_info)
54104

55105
# we need a callback to finish the api gateway span, this callback will be added to the child spans finish callbacks
56106
def finish_callback(_):
@@ -62,24 +112,56 @@ def finish_callback(_):
62112
ctx.set_item("headers", headers)
63113

64114

65-
def set_inferred_proxy_span_tags(span, proxy_context) -> Span:
66-
span._set_tag_str(COMPONENT, supported_proxies[proxy_context["proxy_system_name"]]["component"])
115+
def set_inferred_proxy_span_tags(span: Span, proxy_context: ProxyHeaderContext, proxy_info: ProxyInfo) -> Span:
116+
span._set_tag_str(COMPONENT, proxy_info.component)
117+
span._set_tag_str("span.kind", SpanKind.SERVER)
118+
119+
span._set_tag_str(http.URL, f"https://{proxy_context.domain_name or ''}{proxy_context.path or ''}")
120+
121+
if proxy_context.method:
122+
span._set_tag_str(http.METHOD, proxy_context.method)
123+
124+
if proxy_context.resource_path:
125+
span._set_tag_str(http.ROUTE, proxy_context.resource_path)
126+
127+
if proxy_context.stage:
128+
span._set_tag_str("stage", proxy_context.stage)
129+
130+
if proxy_context.account_id:
131+
span._set_tag_str("account_id", proxy_context.account_id)
67132

68-
span._set_tag_str(http.METHOD, proxy_context["method"])
69-
span._set_tag_str(http.URL, f"{proxy_context['domain_name']}{proxy_context['path']}")
70-
span._set_tag_str("stage", proxy_context["stage"])
133+
if proxy_context.api_id:
134+
span._set_tag_str("apiid", proxy_context.api_id)
135+
136+
if proxy_context.region:
137+
span._set_tag_str("region", proxy_context.region)
138+
139+
if proxy_context.user:
140+
span._set_tag_str("user", proxy_context.user)
141+
142+
if proxy_info.resource_arn_builder:
143+
resource_arn = proxy_info.resource_arn_builder(proxy_context)
144+
if resource_arn:
145+
span._set_tag_str("dd_resource_key", resource_arn)
71146

72147
span.set_metric("_dd.inferred_span", 1)
73148
return span
74149

75150

76-
def extract_inferred_proxy_context(headers) -> Union[None, Dict[str, str]]:
77-
proxy_header_system = str(_extract_header_value(POSSIBLE_PROXY_HEADER_SYSTEM, headers))
78-
proxy_header_start_time_ms = str(_extract_header_value(POSSIBLE_PROXY_HEADER_START_TIME_MS, headers))
79-
proxy_header_path = str(_extract_header_value(POSSIBLE_PROXY_HEADER_PATH, headers))
80-
proxy_header_httpmethod = str(_extract_header_value(POSSIBLE_PROXY_HEADER_HTTPMETHOD, headers))
81-
proxy_header_domain = str(_extract_header_value(POSSIBLE_PROXY_HEADER_DOMAIN, headers))
82-
proxy_header_stage = str(_extract_header_value(POSSIBLE_PROXY_HEADER_STAGE, headers))
151+
def extract_inferred_proxy_context(headers) -> Optional[ProxyHeaderContext]:
152+
proxy_header_system = _extract_header_value(POSSIBLE_PROXY_HEADER_SYSTEM, headers)
153+
proxy_header_start_time_ms = _extract_header_value(POSSIBLE_PROXY_HEADER_START_TIME_MS, headers)
154+
proxy_header_path = _extract_header_value(POSSIBLE_PROXY_HEADER_PATH, headers)
155+
proxy_header_resource_path = _extract_header_value(POSSIBLE_PROXY_HEADER_RESOURCE_PATH, headers)
156+
157+
proxy_header_httpmethod = _extract_header_value(POSSIBLE_PROXY_HEADER_HTTPMETHOD, headers)
158+
proxy_header_domain = _extract_header_value(POSSIBLE_PROXY_HEADER_DOMAIN, headers)
159+
proxy_header_stage = _extract_header_value(POSSIBLE_PROXY_HEADER_STAGE, headers)
160+
161+
proxy_header_account_id = _extract_header_value(POSSIBLE_PROXY_HEADER_ACCOUNT_ID, headers)
162+
proxy_header_api_id = _extract_header_value(POSSIBLE_PROXY_HEADER_API_ID, headers)
163+
proxy_header_region = _extract_header_value(POSSIBLE_PROXY_HEADER_REGION, headers)
164+
proxy_header_user = _extract_header_value(POSSIBLE_PROXY_HEADER_USER, headers)
83165

84166
# Exit if start time header is not present
85167
if proxy_header_start_time_ms is None:
@@ -92,14 +174,19 @@ def extract_inferred_proxy_context(headers) -> Union[None, Dict[str, str]]:
92174
)
93175
return None
94176

95-
return {
96-
"request_time": proxy_header_start_time_ms,
97-
"method": proxy_header_httpmethod,
98-
"path": proxy_header_path,
99-
"stage": proxy_header_stage,
100-
"domain_name": proxy_header_domain,
101-
"proxy_system_name": proxy_header_system,
102-
}
177+
return ProxyHeaderContext(
178+
proxy_header_system,
179+
proxy_header_start_time_ms,
180+
proxy_header_httpmethod,
181+
proxy_header_path,
182+
proxy_header_resource_path,
183+
proxy_header_domain,
184+
proxy_header_stage,
185+
proxy_header_account_id,
186+
proxy_header_api_id,
187+
proxy_header_region,
188+
proxy_header_user,
189+
)
103190

104191

105192
def normalize_headers(headers) -> Dict[str, str]:

‎ddtrace/_trace/trace_handlers.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import ddtrace
1616
from ddtrace import config
17+
from ddtrace._trace._inferred_proxy import SUPPORTED_PROXY_SPAN_NAMES
1718
from ddtrace._trace._inferred_proxy import create_inferred_proxy_span_if_headers_exist
1819
from ddtrace._trace._span_link import SpanLinkKind as _SpanLinkKind
1920
from ddtrace._trace._span_pointer import _SpanPointerDescription
@@ -244,7 +245,7 @@ def _on_web_framework_finish_request(
244245

245246

246247
def _set_inferred_proxy_tags(span, status_code):
247-
if span._parent and span._parent.name == "aws.apigateway":
248+
if span._parent and span._parent.name in SUPPORTED_PROXY_SPAN_NAMES:
248249
inferred_span = span._parent
249250
status_code = status_code if status_code else span.get_tag("http.status_code")
250251
if status_code:

‎releasenotes/notes/feat-inferred-proxy-httpapi-headers-33a0d4d4c1d6c5a1.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
feature:
3+
- |
4+
tracing: Ensures inferred proxy spans for AWS API Gateway HTTP APIs are created when the
5+
``x-dd-proxy`` header reports ``aws-httpapi``.
6+
AAP: Update inferred proxy span tags to ensure that inferred services are discovered by the App and API Protection API Catalog.

‎tests/contrib/aiohttp/test_middleware.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -558,6 +558,7 @@ async def test_parenting_200_dd(app_tracer, aiohttp_client):
558558
"x-dd-proxy": "aws-apigateway",
559559
"x-dd-proxy-request-time-ms": "1736973768000",
560560
"x-dd-proxy-path": "/",
561+
"x-dd-proxy-resource-path": "/",
561562
"x-dd-proxy-httpmethod": "GET",
562563
"x-dd-proxy-domain-name": "local",
563564
"x-dd-proxy-stage": "stage",
@@ -569,6 +570,7 @@ async def test_parenting_200_dd(app_tracer, aiohttp_client):
569570
"x-dd-proxy": "aws-apigateway",
570571
"x-dd-proxy-request-time-ms": "1736973768000",
571572
"x-dd-proxy-path": "/",
573+
"x-dd-proxy-resource-path": "/",
572574
"x-dd-proxy-httpmethod": "GET",
573575
"x-dd-proxy-domain-name": "local",
574576
"x-dd-proxy-stage": "stage",
@@ -616,10 +618,11 @@ async def test_inferred_spans_api_gateway(app_tracer, aiohttp_client, test_app,
616618
web_span_service_name="aiohttp-web",
617619
web_span_resource=test_app["http_method"] + " " + test_app["path"],
618620
api_gateway_service_name="local",
619-
api_gateway_resource="GET /",
620621
method="GET",
621622
status_code=str(test_app["status_code"]),
622-
url="local/",
623+
url="https://local/",
624+
path="/",
625+
stage="stage",
623626
start=1736973768,
624627
is_distributed=test_headers["type"] == "distributed",
625628
distributed_trace_id=1,

‎tests/contrib/asgi/test_asgi.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -747,6 +747,7 @@ async def test_inferred_spans_api_gateway_default(scope, tracer, test_spans, app
747747
("x-dd-proxy", "aws-apigateway"),
748748
("x-dd-proxy-request-time-ms", "1736973768000"),
749749
("x-dd-proxy-path", "/"),
750+
("x-dd-proxy-resource-path", "/"),
750751
("x-dd-proxy-httpmethod", "GET"),
751752
("x-dd-proxy-domain-name", "local"),
752753
("x-dd-proxy-stage", "stage"),
@@ -756,6 +757,7 @@ async def test_inferred_spans_api_gateway_default(scope, tracer, test_spans, app
756757
("x-dd-proxy", "aws-apigateway"),
757758
("x-dd-proxy-request-time-ms", "1736973768000"),
758759
("x-dd-proxy-path", "/"),
760+
("x-dd-proxy-resource-path", "/"),
759761
("x-dd-proxy-httpmethod", "GET"),
760762
("x-dd-proxy-domain-name", "local"),
761763
("x-dd-proxy-stage", "stage"),
@@ -805,10 +807,12 @@ async def test_inferred_spans_api_gateway_default(scope, tracer, test_spans, app
805807
web_span_service_name="tests.contrib.asgi",
806808
web_span_resource="GET /",
807809
api_gateway_service_name="local",
808-
api_gateway_resource="GET /",
809810
method="GET",
810811
status_code=app_type["status_code"],
811-
url="local/",
812+
url="https://local/",
813+
route="/",
814+
path="/",
815+
stage="stage",
812816
start=1736973768,
813817
is_distributed=headers == distributed_headers,
814818
distributed_trace_id=1,

‎tests/contrib/bottle/test.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -524,6 +524,7 @@ def handled_error_endpoint():
524524
"x-dd-proxy": "aws-apigateway",
525525
"x-dd-proxy-request-time-ms": "1736973768000",
526526
"x-dd-proxy-path": "/",
527+
"x-dd-proxy-resource-path": "/",
527528
"x-dd-proxy-httpmethod": "GET",
528529
"x-dd-proxy-domain-name": "local",
529530
"x-dd-proxy-stage": "stage",
@@ -555,10 +556,11 @@ def handled_error_endpoint():
555556
web_span_service_name=SERVICE,
556557
web_span_resource="GET " + test_endpoint["endpoint"],
557558
api_gateway_service_name="local",
558-
api_gateway_resource="GET /",
559559
method="GET",
560560
status_code=str(test_endpoint["status"]),
561-
url="local/",
561+
url="https://local/",
562+
path="/",
563+
stage="stage",
562564
start=1736973768,
563565
is_distributed=False,
564566
distributed_trace_id=1,

‎tests/contrib/bottle/test_distributed.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ def default_endpoint():
142142
"x-dd-proxy": "aws-apigateway",
143143
"x-dd-proxy-request-time-ms": "1736973768000",
144144
"x-dd-proxy-path": "/",
145+
"x-dd-proxy-resource-path": "/",
145146
"x-dd-proxy-httpmethod": "GET",
146147
"x-dd-proxy-domain-name": "local",
147148
"x-dd-proxy-stage": "stage",
@@ -166,10 +167,11 @@ def default_endpoint():
166167
web_span_service_name=SERVICE,
167168
web_span_resource="GET " + "/",
168169
api_gateway_service_name="local",
169-
api_gateway_resource="GET /",
170170
method="GET",
171171
status_code=200,
172-
url="local/",
172+
url="https://local/",
173+
path="/",
174+
stage="stage",
173175
start=1736973768,
174176
is_distributed=True,
175177
distributed_trace_id=1,
@@ -190,6 +192,7 @@ def default_endpoint():
190192
"x-dd-proxy": "aws-apigateway",
191193
"x-dd-proxy-request-time-ms": "1736973768000",
192194
"x-dd-proxy-path": "/",
195+
"x-dd-proxy-resource-path": "/",
193196
"x-dd-proxy-httpmethod": "GET",
194197
"x-dd-proxy-domain-name": "local",
195198
"x-dd-proxy-stage": "stage",

‎tests/contrib/cherrypy/test_middleware.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def test_and_emit_get_version(self):
6161
from ddtrace.contrib.internal.cherrypy.patch import get_version
6262

6363
version = get_version()
64-
assert type(version) == str
64+
assert isinstance(version, str)
6565
assert version != ""
6666

6767
emit_integration_and_version_to_test_agent("cherrypy", version)
@@ -487,6 +487,7 @@ def test_inferred_spans_api_gateway_default(self):
487487
("x-dd-proxy", "aws-apigateway"),
488488
("x-dd-proxy-request-time-ms", "1736973768000"),
489489
("x-dd-proxy-path", "/"),
490+
("x-dd-proxy-resource-path", "/"),
490491
("x-dd-proxy-httpmethod", "GET"),
491492
("x-dd-proxy-domain-name", "local"),
492493
("x-dd-proxy-stage", "stage"),
@@ -496,6 +497,7 @@ def test_inferred_spans_api_gateway_default(self):
496497
("x-dd-proxy", "aws-apigateway"),
497498
("x-dd-proxy-request-time-ms", "1736973768000"),
498499
("x-dd-proxy-path", "/"),
500+
("x-dd-proxy-resource-path", "/"),
499501
("x-dd-proxy-httpmethod", "GET"),
500502
("x-dd-proxy-domain-name", "local"),
501503
("x-dd-proxy-stage", "stage"),
@@ -539,10 +541,11 @@ def test_inferred_spans_api_gateway_default(self):
539541
web_span_service_name="test.cherrypy.service",
540542
web_span_resource="GET " + test_endpoint["endpoint"],
541543
api_gateway_service_name="local",
542-
api_gateway_resource="GET /",
543544
method="GET",
544545
status_code=test_endpoint["status"],
545-
url="local/",
546+
url="https://local/",
547+
path="/",
548+
stage="stage",
546549
start=1736973768,
547550
is_distributed=test_headers == distributed_headers,
548551
distributed_trace_id=1,

0 commit comments

Comments
 (0)