Skip to content

Commit ecec084

Browse files
CagriYoncaGSVarsha
authored andcommitted
feat: Add werkzeug instrumentation and unittests
Signed-off-by: Cagri Yonca <cagri@ibm.com>
1 parent e7937e4 commit ecec084

10 files changed

Lines changed: 1718 additions & 93 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: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
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+
This module automatically patches Werkzeug applications when imported via wrapt.
10+
"""
11+
12+
try:
13+
from typing import Any, Callable
14+
15+
import wrapt
16+
17+
from instana.instrumentation.wsgi import InstanaWSGIMiddleware
18+
from instana.log import logger
19+
20+
def _is_flask_app(app: Any) -> bool:
21+
"""
22+
Check if the application is a Flask app.
23+
24+
Flask apps have their own instrumentation, so we skip wrapping them
25+
to avoid double instrumentation (2 spans per request).
26+
27+
Args:
28+
app: The WSGI application to check
29+
30+
Returns:
31+
True if app is a Flask application, False otherwise
32+
"""
33+
try:
34+
# Check if it's a Flask app by class name
35+
if hasattr(app, "__class__"):
36+
class_name = app.__class__.__name__
37+
module_name = getattr(app.__class__, "__module__", "")
38+
39+
# Direct Flask app check
40+
if class_name == "Flask" and "flask" in module_name:
41+
return True
42+
43+
# Check for Flask app wrapped in middleware
44+
if hasattr(app, "wsgi_app"):
45+
return _is_flask_app(app.wsgi_app)
46+
47+
return False
48+
except Exception:
49+
logger.debug("Error checking if app is Flask", exc_info=True)
50+
return False
51+
52+
@wrapt.patch_function_wrapper("werkzeug.serving", "run_simple")
53+
def run_simple_with_instana(
54+
wrapped: Callable,
55+
instance: Any,
56+
args: tuple,
57+
kwargs: dict[str, Any],
58+
) -> Any:
59+
"""
60+
Patch werkzeug.serving.run_simple to wrap WSGI applications.
61+
62+
Skips Flask applications as they have their own instrumentation.
63+
"""
64+
try:
65+
# run_simple(hostname, port, application, ...)
66+
if len(args) >= 3:
67+
hostname, port, application = args[0], args[1], args[2]
68+
69+
# Skip Flask apps (they have their own instrumentation)
70+
if _is_flask_app(application):
71+
logger.debug(
72+
f"Skipping Werkzeug instrumentation for Flask app at {hostname}:{port}"
73+
)
74+
return wrapped(*args, **kwargs)
75+
76+
# Wrap non-Flask WSGI apps
77+
instrumented_app = InstanaWSGIMiddleware(
78+
application, status_as_string=False
79+
)
80+
logger.debug(f"Werkzeug app wrapped: {hostname}:{port}")
81+
args = (hostname, port, instrumented_app) + args[3:]
82+
elif "application" in kwargs:
83+
application = kwargs["application"]
84+
85+
# Skip Flask apps (they have their own instrumentation)
86+
if _is_flask_app(application):
87+
logger.debug(
88+
"Skipping Werkzeug instrumentation for Flask app (kwargs)"
89+
)
90+
return wrapped(*args, **kwargs)
91+
92+
# Wrap non-Flask WSGI apps
93+
instrumented_app = InstanaWSGIMiddleware(
94+
application, status_as_string=False
95+
)
96+
kwargs["application"] = instrumented_app
97+
logger.debug("Werkzeug app wrapped (kwargs)")
98+
except Exception:
99+
logger.debug("Failed to wrap Werkzeug app", exc_info=True)
100+
101+
return wrapped(*args, **kwargs)
102+
103+
logger.debug("Instrumenting werkzeug")
104+
105+
except ImportError:
106+
pass
107+
108+
# Made with Bob

src/instana/instrumentation/wsgi.py

Lines changed: 19 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -5,114 +5,40 @@
55
Instana WSGI Middleware
66
"""
77

8-
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple
8+
from typing import Any, Callable
99

10-
from opentelemetry import context, trace
11-
from opentelemetry.semconv.trace import SpanAttributes
10+
from opentelemetry import context
1211

13-
from instana.propagators.format import Format
14-
from instana.singletons import agent, get_tracer
15-
from instana.util.secrets import strip_secrets_from_query
16-
from instana.util.traceutils import extract_custom_headers
17-
18-
if TYPE_CHECKING:
19-
from instana.span.span import InstanaSpan
12+
from instana.util.wsgi_utils import (
13+
build_start_response,
14+
create_span_with_context,
15+
end_span_after_iterating,
16+
)
2017

2118

2219
class InstanaWSGIMiddleware(object):
2320
"""Instana WSGI middleware"""
2421

25-
def __init__(self, app: object) -> None:
22+
def __init__(self, app: Callable, status_as_string: bool = True) -> None:
2623
self.app = app
24+
self.status_as_string = status_as_string
2725

28-
def __call__(self, environ: Dict[str, Any], start_response: Callable) -> object:
29-
env = environ
30-
tracer = get_tracer()
31-
32-
# Extract context and start span
33-
parent_context = tracer.extract(Format.HTTP_HEADERS, env)
34-
span = tracer.start_span("wsgi", context=parent_context)
35-
36-
# Attach context - this makes the span current
37-
ctx = trace.set_span_in_context(span)
38-
token = context.attach(ctx)
39-
40-
# Extract custom headers from request
41-
extract_custom_headers(span, env, format=True)
42-
43-
# Set request attributes
44-
_set_request_attributes(span, env)
45-
46-
def new_start_response(
47-
status: str,
48-
headers: List[Tuple[object, ...]],
49-
exc_info: Optional[Exception] = None,
50-
) -> object:
51-
"""Modified start response with additional headers."""
52-
extract_custom_headers(span, headers)
53-
54-
tracer.inject(span.context, Format.HTTP_HEADERS, headers)
55-
56-
headers_str = [
57-
(header[0], str(header[1]))
58-
if not isinstance(header[1], str)
59-
else header
60-
for header in headers
61-
]
62-
63-
# Set status code attribute
64-
sc = status.split(" ")[0]
65-
if int(sc) >= 500:
66-
span.mark_as_errored()
67-
68-
span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, sc)
69-
70-
return start_response(status, headers_str, exc_info)
71-
26+
def __call__(self, environ: dict[str, Any], start_response: Callable) -> object:
7227
try:
73-
iterable = self.app(environ, new_start_response)
74-
75-
# Wrap the iterable to ensure span ends after iteration completes
76-
return _end_span_after_iterating(iterable, span, token)
28+
span, token = create_span_with_context(environ)
29+
wrapped_start_response = build_start_response(
30+
span, start_response, status_as_string=self.status_as_string
31+
)
32+
except Exception:
33+
return self.app(environ, start_response)
7734

35+
try:
36+
iterable = self.app(environ, wrapped_start_response)
37+
return end_span_after_iterating(iterable, span, token)
7838
except Exception as exc:
79-
# If exception occurs before iteration completes, end span and detach token
8039
if span and span.is_recording():
8140
span.record_exception(exc)
8241
span.end()
8342
if token:
8443
context.detach(token)
8544
raise exc
86-
87-
88-
def _end_span_after_iterating(
89-
iterable: Iterable[object], span: "InstanaSpan", token: object
90-
) -> Iterable[object]:
91-
try:
92-
yield from iterable
93-
finally:
94-
# Ensure iterable cleanup (important for generators)
95-
if hasattr(iterable, "close"):
96-
iterable.close()
97-
98-
# End span and detach token after iteration completes
99-
if span and span.is_recording():
100-
span.end()
101-
if token:
102-
context.detach(token)
103-
104-
105-
def _set_request_attributes(span: "InstanaSpan", env: Dict[str, Any]) -> None:
106-
if "PATH_INFO" in env:
107-
span.set_attribute("http.path", env["PATH_INFO"])
108-
if "QUERY_STRING" in env and len(env["QUERY_STRING"]):
109-
scrubbed_params = strip_secrets_from_query(
110-
env["QUERY_STRING"],
111-
agent.options.secrets_matcher,
112-
agent.options.secrets_list,
113-
)
114-
span.set_attribute("http.params", scrubbed_params)
115-
if "REQUEST_METHOD" in env:
116-
span.set_attribute(SpanAttributes.HTTP_METHOD, env["REQUEST_METHOD"])
117-
if "HTTP_HOST" in env:
118-
span.set_attribute(SpanAttributes.HTTP_HOST, env["HTTP_HOST"])

0 commit comments

Comments
 (0)