feat: Add werkzeug instrumentation and unittests

Signed-off-by: Cagri Yonca <cagri@ibm.com>

Author: CagriYonca

src/instana/__init__.py

@@ -184,6 +184,7 @@ def boot_agent() -> None:
             sqlalchemy,  # noqa: F401
             starlette,  # noqa: F401
             urllib3,  # noqa: F401
+            werkzeug,  # noqa: F401
             gevent,  # noqa: F401
         )
         from instana.instrumentation.aiohttp import (

src/instana/instrumentation/werkzeug.py

@@ -0,0 +1,138 @@
+# (C) Copyright IBM Corp. 2026
+
+"""
+Instana Werkzeug Instrumentation
+
+This module provides automatic instrumentation for Werkzeug-based applications.
+Werkzeug is a comprehensive WSGI web application library used by Flask and other frameworks.
+
+This module automatically patches Werkzeug applications when imported via autowrapt.
+"""
+
+try:
+    from typing import Any, Callable
+
+    import wrapt
+    from opentelemetry import context
+
+    from instana.log import logger
+    from instana.util.wsgi_utils import (
+        build_start_response,
+        create_span_with_context,
+        end_span_after_iterating,
+    )
+
+    # Autowrapt patching for automatic instrumentation
+    class _TracedWSGIApp:
+        """Wrapper that traces WSGI applications."""
+
+        def __init__(self, app: Callable) -> None:
+            self.app = app
+
+        def __call__(self, environ: dict[str, Any], start_response: Callable) -> Any:
+            try:
+                span, token = create_span_with_context(environ)
+                wrapped_start_response = build_start_response(span, start_response)
+            except Exception:
+                logger.debug("werkzeug setup failed", exc_info=True)
+                return self.app(environ, start_response)
+
+            try:
+                iterable = self.app(environ, wrapped_start_response)
+                return end_span_after_iterating(iterable, span, token)
+            except Exception as exc:
+                try:
+                    if span and span.is_recording():
+                        span.record_exception(exc)
+                        span.end()
+                    if token:
+                        context.detach(token)
+                except Exception:
+                    logger.debug("werkzeug cleanup failed", exc_info=True)
+                raise
+
+    def _is_flask_app(app: Any) -> bool:
+        """
+        Check if the application is a Flask app.
+
+        Flask apps have their own instrumentation, so we skip wrapping them
+        to avoid double instrumentation (2 spans per request).
+
+        Args:
+            app: The WSGI application to check
+
+        Returns:
+            True if app is a Flask application, False otherwise
+        """
+        try:
+            # Check if it's a Flask app by class name
+            if hasattr(app, "__class__"):
+                class_name = app.__class__.__name__
+                module_name = getattr(app.__class__, "__module__", "")
+
+                # Direct Flask app check
+                if class_name == "Flask" and "flask" in module_name:
+                    return True
+
+                # Check for Flask app wrapped in middleware
+                if hasattr(app, "wsgi_app"):
+                    return _is_flask_app(app.wsgi_app)
+
+            return False
+        except Exception:
+            logger.debug("Error checking if app is Flask", exc_info=True)
+            return False
+
+    @wrapt.patch_function_wrapper("werkzeug.serving", "run_simple")
+    def run_simple_with_instana(
+        wrapped: Callable,
+        instance: Any,
+        args: tuple,
+        kwargs: dict[str, Any],
+    ) -> Any:
+        """
+        Patch werkzeug.serving.run_simple to wrap WSGI applications.
+
+        Skips Flask applications as they have their own instrumentation.
+        """
+        try:
+            # run_simple(hostname, port, application, ...)
+            if len(args) >= 3:
+                hostname, port, application = args[0], args[1], args[2]
+
+                # Skip Flask apps (they have their own instrumentation)
+                if _is_flask_app(application):
+                    logger.debug(
+                        f"Skipping Werkzeug instrumentation for Flask app at {hostname}:{port}"
+                    )
+                    return wrapped(*args, **kwargs)
+
+                # Wrap non-Flask WSGI apps
+                instrumented_app = _TracedWSGIApp(application)
+                logger.debug(f"Werkzeug app wrapped: {hostname}:{port}")
+                args = (hostname, port, instrumented_app) + args[3:]
+            elif "application" in kwargs:
+                application = kwargs["application"]
+
+                # Skip Flask apps (they have their own instrumentation)
+                if _is_flask_app(application):
+                    logger.debug(
+                        "Skipping Werkzeug instrumentation for Flask app (kwargs)"
+                    )
+                    return wrapped(*args, **kwargs)
+
+                # Wrap non-Flask WSGI apps
+                instrumented_app = _TracedWSGIApp(application)
+                kwargs["application"] = instrumented_app
+                logger.debug("Werkzeug app wrapped (kwargs)")
+        except Exception:
+            logger.debug("Failed to wrap Werkzeug app", exc_info=True)
+
+        return wrapped(*args, **kwargs)
+
+    logger.debug("Instrumenting werkzeug")
+
+except ImportError:
+    pass
+
+# Made with Bob

src/instana/instrumentation/wsgi.py

@@ -5,114 +5,39 @@
 Instana WSGI Middleware
 """
 
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple
+from typing import Any, Callable
 
-from opentelemetry import context, trace
-from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry import context
 
-from instana.propagators.format import Format
-from instana.singletons import agent, get_tracer
-from instana.util.secrets import strip_secrets_from_query
-from instana.util.traceutils import extract_custom_headers
-
-if TYPE_CHECKING:
-    from instana.span.span import InstanaSpan
+from instana.util.wsgi_utils import (
+    build_start_response,
+    create_span_with_context,
+    end_span_after_iterating,
+)
 
 
 class InstanaWSGIMiddleware(object):
     """Instana WSGI middleware"""
 
-    def __init__(self, app: object) -> None:
+    def __init__(self, app: Callable) -> None:
         self.app = app
 
-    def __call__(self, environ: Dict[str, Any], start_response: Callable) -> object:
-        env = environ
-        tracer = get_tracer()
-
-        # Extract context and start span
-        parent_context = tracer.extract(Format.HTTP_HEADERS, env)
-        span = tracer.start_span("wsgi", context=parent_context)
-
-        # Attach context - this makes the span current
-        ctx = trace.set_span_in_context(span)
-        token = context.attach(ctx)
-
-        # Extract custom headers from request
-        extract_custom_headers(span, env, format=True)
-
-        # Set request attributes
-        _set_request_attributes(span, env)
-
-        def new_start_response(
-            status: str,
-            headers: List[Tuple[object, ...]],
-            exc_info: Optional[Exception] = None,
-        ) -> object:
-            """Modified start response with additional headers."""
-            extract_custom_headers(span, headers)
-
-            tracer.inject(span.context, Format.HTTP_HEADERS, headers)
-
-            headers_str = [
-                (header[0], str(header[1]))
-                if not isinstance(header[1], str)
-                else header
-                for header in headers
-            ]
-
-            # Set status code attribute
-            sc = status.split(" ")[0]
-            if int(sc) >= 500:
-                span.mark_as_errored()
-
-            span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, sc)
-
-            return start_response(status, headers_str, exc_info)
-
+    def __call__(self, environ: dict[str, Any], start_response: Callable) -> object:
         try:
-            iterable = self.app(environ, new_start_response)
-
-            # Wrap the iterable to ensure span ends after iteration completes
-            return _end_span_after_iterating(iterable, span, token)
+            span, token = create_span_with_context(environ)
+            wrapped_start_response = build_start_response(
+                span, start_response, status_as_string=True
+            )
+        except Exception:
+            return self.app(environ, start_response)
 
+        try:
+            iterable = self.app(environ, wrapped_start_response)
+            return end_span_after_iterating(iterable, span, token)
         except Exception as exc:
-            # If exception occurs before iteration completes, end span and detach token
             if span and span.is_recording():
                 span.record_exception(exc)
                 span.end()
             if token:
                 context.detach(token)
             raise exc
-
-
-def _end_span_after_iterating(
-    iterable: Iterable[object], span: "InstanaSpan", token: object
-) -> Iterable[object]:
-    try:
-        yield from iterable
-    finally:
-        # Ensure iterable cleanup (important for generators)
-        if hasattr(iterable, "close"):
-            iterable.close()
-
-        # End span and detach token after iteration completes
-        if span and span.is_recording():
-            span.end()
-        if token:
-            context.detach(token)
-
-
-def _set_request_attributes(span: "InstanaSpan", env: Dict[str, Any]) -> None:
-    if "PATH_INFO" in env:
-        span.set_attribute("http.path", env["PATH_INFO"])
-    if "QUERY_STRING" in env and len(env["QUERY_STRING"]):
-        scrubbed_params = strip_secrets_from_query(
-            env["QUERY_STRING"],
-            agent.options.secrets_matcher,
-            agent.options.secrets_list,
-        )
-        span.set_attribute("http.params", scrubbed_params)
-    if "REQUEST_METHOD" in env:
-        span.set_attribute(SpanAttributes.HTTP_METHOD, env["REQUEST_METHOD"])
-    if "HTTP_HOST" in env:
-        span.set_attribute(SpanAttributes.HTTP_HOST, env["HTTP_HOST"])