Skip to content

Commit c51cef4

Browse files
committed
feat: Add werkzeug instrumentation and unittests
Signed-off-by: Cagri Yonca <cagri@ibm.com>
1 parent c60f7c9 commit c51cef4

5 files changed

Lines changed: 748 additions & 0 deletions

File tree

src/instana/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,7 @@ def boot_agent() -> None:
184184
sqlalchemy, # noqa: F401
185185
starlette, # noqa: F401
186186
urllib3, # noqa: F401
187+
werkzeug, # noqa: F401
187188
gevent, # noqa: F401
188189
)
189190
from instana.instrumentation.aiohttp import (
Lines changed: 354 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,354 @@
1+
# (c) Copyright IBM Corp. 2026
2+
3+
"""
4+
Instana Werkzeug Instrumentation
5+
6+
This module provides automatic instrumentation for Werkzeug-based applications.
7+
Werkzeug is a comprehensive WSGI web application library used by Flask and other frameworks.
8+
"""
9+
10+
try:
11+
from contextlib import contextmanager
12+
from typing import (
13+
TYPE_CHECKING,
14+
Any,
15+
Callable,
16+
Dict,
17+
Generator,
18+
Iterable,
19+
List,
20+
Optional,
21+
Tuple,
22+
)
23+
24+
from opentelemetry import context, trace
25+
from opentelemetry.semconv.trace import SpanAttributes
26+
27+
from werkzeug.wrappers import Request
28+
29+
from instana.log import logger
30+
from instana.propagators.format import Format
31+
from instana.singletons import agent, get_tracer
32+
from instana.util.secrets import strip_secrets_from_query
33+
from instana.util.traceutils import extract_custom_headers
34+
35+
if TYPE_CHECKING:
36+
from instana.span.span import InstanaSpan
37+
38+
def _create_span_with_context(environ: Dict[str, Any]) -> Tuple["InstanaSpan", Any]:
39+
"""
40+
Create and configure a span with context for the request.
41+
42+
Args:
43+
environ: WSGI environment dictionary
44+
45+
Returns:
46+
Tuple of (span, context_token)
47+
"""
48+
tracer = get_tracer()
49+
parent_context = tracer.extract(Format.HTTP_HEADERS, environ)
50+
span = tracer.start_span("wsgi", context=parent_context)
51+
52+
ctx = trace.set_span_in_context(span)
53+
token = context.attach(ctx)
54+
55+
extract_custom_headers(span, environ, format=True)
56+
_set_request_attributes(span, environ)
57+
58+
return span, token
59+
60+
def _build_start_response(
61+
span: "InstanaSpan",
62+
start_response: Callable,
63+
) -> Callable:
64+
"""Create an instrumented start_response callable."""
65+
66+
def new_start_response(
67+
status: str,
68+
headers: List[Tuple[str, str]],
69+
exc_info: Optional[Tuple[Any, Any, Any]] = None,
70+
) -> Callable:
71+
"""Modified start_response with trace context injection."""
72+
try:
73+
extract_custom_headers(span, headers)
74+
tracer = get_tracer()
75+
tracer.inject(
76+
span.context,
77+
Format.HTTP_HEADERS,
78+
headers,
79+
)
80+
81+
status_code = _parse_status_code(status)
82+
if status_code is not None:
83+
if status_code >= 500:
84+
span.mark_as_errored()
85+
span.set_attribute(
86+
SpanAttributes.HTTP_STATUS_CODE,
87+
status_code,
88+
)
89+
90+
return start_response(
91+
status,
92+
_normalize_headers(headers),
93+
exc_info,
94+
)
95+
except Exception:
96+
logger.debug("Error in Werkzeug start_response wrapper", exc_info=True)
97+
return start_response(status, headers, exc_info)
98+
99+
return new_start_response
100+
101+
def _normalize_headers(
102+
headers: List[Tuple[str, Any]],
103+
) -> List[Tuple[str, str]]:
104+
"""
105+
Ensure all header values are strings for WSGI compliance.
106+
107+
Args:
108+
headers: List of (name, value) tuples
109+
110+
Returns:
111+
List of (name, str_value) tuples
112+
"""
113+
return [
114+
(name, value if isinstance(value, str) else str(value))
115+
for name, value in headers
116+
]
117+
118+
def _parse_status_code(status: str) -> Optional[int]:
119+
"""Safely parse the HTTP status code from a WSGI status string."""
120+
try:
121+
return int(status.split()[0])
122+
except (AttributeError, IndexError, TypeError, ValueError):
123+
return None
124+
125+
def _end_span_after_iterating(
126+
iterable: Iterable[bytes],
127+
span: "InstanaSpan",
128+
token: Any,
129+
) -> Iterable[bytes]:
130+
"""
131+
Generator that yields from the iterable and ensures span cleanup.
132+
133+
Args:
134+
iterable: The response iterable from the application
135+
span: The active span
136+
token: The context token
137+
138+
Yields:
139+
Response chunks from the iterable
140+
"""
141+
try:
142+
yield from iterable
143+
finally:
144+
# Ensure iterable cleanup (important for generators)
145+
if hasattr(iterable, "close"):
146+
try:
147+
iterable.close() # type: ignore
148+
except Exception:
149+
logger.debug("Error closing iterable", exc_info=True)
150+
151+
# End span and detach token after iteration completes
152+
if span and span.is_recording():
153+
span.end()
154+
if token:
155+
context.detach(token) # type: ignore
156+
157+
def _scrub_query_params(query_string: str) -> Optional[str]:
158+
"""
159+
Scrub secrets from query string parameters.
160+
161+
Args:
162+
query_string: The query string to scrub
163+
164+
Returns:
165+
Scrubbed query string or None if agent not available
166+
"""
167+
if agent is not None:
168+
return strip_secrets_from_query(
169+
query_string,
170+
agent.options.secrets_matcher, # type: ignore
171+
agent.options.secrets_list, # type: ignore
172+
)
173+
return None
174+
175+
def _set_request_attributes(span: "InstanaSpan", environ: Dict[str, Any]) -> None:
176+
"""
177+
Extract and set HTTP attributes from the WSGI environ.
178+
179+
Args:
180+
span: The active span
181+
environ: WSGI environment dictionary
182+
"""
183+
try:
184+
# Set HTTP method
185+
if "REQUEST_METHOD" in environ:
186+
span.set_attribute(
187+
SpanAttributes.HTTP_METHOD, environ["REQUEST_METHOD"]
188+
)
189+
190+
# Set HTTP path
191+
if "PATH_INFO" in environ:
192+
span.set_attribute("http.path", environ["PATH_INFO"])
193+
194+
# Set HTTP query parameters (with secrets scrubbed)
195+
if environ.get("QUERY_STRING"):
196+
scrubbed_params = _scrub_query_params(environ["QUERY_STRING"])
197+
if scrubbed_params is not None:
198+
span.set_attribute("http.params", scrubbed_params)
199+
200+
# Set HTTP host
201+
if "HTTP_HOST" in environ:
202+
span.set_attribute(
203+
SpanAttributes.HTTP_HOST,
204+
environ["HTTP_HOST"],
205+
)
206+
207+
# Set HTTP URL (without query string to avoid exposing secrets)
208+
if "wsgi.url_scheme" in environ:
209+
scheme = environ["wsgi.url_scheme"]
210+
host = environ.get("HTTP_HOST", "")
211+
script_name = environ.get("SCRIPT_NAME", "")
212+
path = environ.get("PATH_INFO", "")
213+
214+
url = f"{scheme}://{host}{script_name}{path}"
215+
span.set_attribute(SpanAttributes.HTTP_URL, url)
216+
217+
except Exception:
218+
logger.debug("Error setting request attributes", exc_info=True)
219+
220+
class InstanaWerkzeugMiddleware:
221+
"""
222+
Instana WSGI middleware for Werkzeug applications.
223+
224+
This middleware automatically traces Werkzeug-based applications by:
225+
- Extracting trace context from incoming requests
226+
- Creating spans for each request
227+
- Capturing HTTP attributes (method, path, query, headers)
228+
- Injecting trace context into responses
229+
- Handling errors and exceptions
230+
231+
Usage:
232+
from werkzeug.wrappers import Request, Response
233+
from instana.instrumentation.werkzeug import InstanaWerkzeugMiddleware
234+
235+
@Request.application
236+
def application(request):
237+
return Response('Hello World!')
238+
239+
# Wrap with Instana middleware
240+
app = InstanaWerkzeugMiddleware(application)
241+
"""
242+
243+
def __init__(self, app: Callable) -> None:
244+
"""
245+
Initialize the Instana Werkzeug middleware.
246+
247+
Args:
248+
app: The WSGI application to wrap
249+
"""
250+
self.app = app
251+
252+
def __call__(
253+
self, env: Dict[str, Any], start_response: Callable
254+
) -> Iterable[bytes]:
255+
"""
256+
WSGI application interface.
257+
258+
Args:
259+
env: WSGI environment dictionary
260+
start_response: WSGI start_response callable
261+
262+
Returns:
263+
Iterable response body
264+
"""
265+
try:
266+
span, token = _create_span_with_context(env)
267+
wrapped_start_response = _build_start_response(span, start_response)
268+
except Exception as exc:
269+
# Instrumentation failed - fail silently, run original app
270+
logger.debug("werkzeug instrumentation setup failed", exc_info=exc)
271+
return self.app(env, start_response)
272+
273+
try:
274+
iterable = self.app(env, wrapped_start_response)
275+
return _end_span_after_iterating(iterable, span, token)
276+
except Exception as exc:
277+
# App exception - record in span and propagate
278+
try:
279+
if span and span.is_recording():
280+
span.record_exception(exc)
281+
span.end()
282+
if token:
283+
context.detach(token)
284+
except Exception:
285+
logger.debug("werkzeug span cleanup failed", exc_info=True)
286+
raise
287+
288+
# Context manager for manual instrumentation
289+
@contextmanager
290+
def trace_werkzeug_request(
291+
request: Request,
292+
) -> Generator["InstanaSpan", None, None]:
293+
"""
294+
Context manager for manually tracing Werkzeug requests.
295+
296+
Usage:
297+
@Request.application
298+
def application(request):
299+
with trace_werkzeug_request(request) as span:
300+
# Your application logic
301+
span.set_attribute("custom.attribute", "value")
302+
return Response('Hello World!')
303+
304+
Args:
305+
request: Werkzeug Request object
306+
307+
Yields:
308+
The active span for the request
309+
"""
310+
tracer = get_tracer()
311+
span = None
312+
token = None
313+
314+
try:
315+
# Extract context from request environ (preserves multi-value headers)
316+
parent_context = tracer.extract(
317+
Format.HTTP_HEADERS,
318+
request.environ,
319+
)
320+
321+
# Start span
322+
span = tracer.start_span("wsgi", context=parent_context)
323+
ctx = trace.set_span_in_context(span)
324+
token = context.attach(ctx)
325+
326+
# Set request attributes
327+
span.set_attribute(SpanAttributes.HTTP_METHOD, request.method)
328+
span.set_attribute(SpanAttributes.HTTP_URL, request.url)
329+
span.set_attribute(SpanAttributes.HTTP_HOST, request.host)
330+
331+
if request.query_string:
332+
scrubbed_params = _scrub_query_params(
333+
request.query_string.decode("utf-8")
334+
)
335+
if scrubbed_params is not None:
336+
span.set_attribute("http.params", scrubbed_params)
337+
338+
yield span
339+
except Exception as exc:
340+
if span and span.is_recording():
341+
span.record_exception(exc)
342+
raise
343+
finally:
344+
if span and span.is_recording():
345+
span.end()
346+
if token:
347+
context.detach(token)
348+
349+
logger.debug("Instrumenting werkzeug")
350+
351+
except ImportError:
352+
pass
353+
354+
# Made with Bob
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
# (c) Copyright IBM Corp. 2026

0 commit comments

Comments
 (0)