Skip to content

Commit 5a0cad1

Browse files
fateleifatelei
committed
chore: compatiable flask 3.1+
1 parent bd3c1f2 commit 5a0cad1

File tree

2 files changed

+536
-7
lines changed

2 files changed

+536
-7
lines changed

‎instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py‎

Lines changed: 163 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,7 @@ def response_hook(span: Span, status: str, response_headers: List):
254254
---
255255
"""
256256

257+
import sys
257258
import weakref
258259
from logging import getLogger
259260
from time import time_ns
@@ -333,6 +334,71 @@ def get_default_span_name():
333334
return span_name
334335

335336

337+
def _ensure_streaming_context_cleanup(environ):
338+
"""
339+
Ensure proper context cleanup for streaming responses in Flask 3.1+.
340+
341+
This function checks if the response is a streaming response and ensures
342+
that context tokens are properly cleaned up to prevent token reuse issues.
343+
Only applies to Flask 3.1+ and Python 3.10+ for compatibility reasons.
344+
"""
345+
# Double-check Flask version - this should only run in Flask 3.1+
346+
# Use the same check method as other places for consistency
347+
if not (
348+
hasattr(flask, "__version__")
349+
and package_version.parse(flask.__version__)
350+
>= package_version.parse("3.1.0")
351+
):
352+
return
353+
354+
# Only enable streaming context cleanup for Python 3.10+ to avoid compatibility issues
355+
# with older Python versions that have different context management behavior
356+
if sys.version_info < (3, 10):
357+
return
358+
359+
activation = environ.get(_ENVIRON_ACTIVATION_KEY)
360+
token = environ.get(_ENVIRON_TOKEN)
361+
362+
if not activation or not token:
363+
return
364+
365+
# Additional safety check - only proceed if we haven't already cleaned up
366+
if (
367+
environ.get(_ENVIRON_ACTIVATION_KEY) is None
368+
or environ.get(_ENVIRON_TOKEN) is None
369+
):
370+
return
371+
372+
try:
373+
# Clean up the context token safely
374+
if token:
375+
try:
376+
context.detach(token)
377+
except RuntimeError as exc:
378+
# Token has already been used - this can happen in Flask 3.1+
379+
# with streaming responses, so we just log and continue
380+
_logger.debug("Token already detached, continuing: %s", exc)
381+
# If detach failed, don't proceed with activation cleanup
382+
return
383+
384+
# Clean up the activation
385+
if hasattr(activation, "__exit__"):
386+
try:
387+
activation.__exit__(None, None, None)
388+
except (RuntimeError, AttributeError) as exc:
389+
_logger.debug("Error during activation cleanup: %s", exc)
390+
391+
# Mark that we've handled the cleanup to prevent double cleanup in teardown
392+
environ[_ENVIRON_ACTIVATION_KEY] = None
393+
environ[_ENVIRON_TOKEN] = None
394+
395+
except (RuntimeError, ValueError, TypeError) as exc:
396+
# Log the error but don't raise it to avoid breaking the response
397+
_logger.debug(
398+
"Error during streaming context cleanup: %s", exc, exc_info=True
399+
)
400+
401+
336402
def _rewrapped_app(
337403
wsgi_app,
338404
active_requests_counter,
@@ -408,6 +474,46 @@ def _start_response(status, response_headers, *args, **kwargs):
408474
return start_response(status, response_headers, *args, **kwargs)
409475

410476
result = wsgi_app(wrapped_app_environ, _start_response)
477+
478+
# For Flask 3.1+, check if we need to handle streaming response context cleanup
479+
# Only run this logic in Flask 3.1+ and Python 3.10+ to avoid any interference with older versions
480+
# Use very conservative checks to ensure we never interfere with Flask < 3.1
481+
if (
482+
should_trace
483+
and hasattr(
484+
flask, "__version__"
485+
) # Ensure Flask has version attribute
486+
and package_version.parse(flask.__version__)
487+
>= package_version.parse("3.1.0") # Only Flask 3.1+
488+
and sys.version_info >= (3, 10) # Only Python 3.10+
489+
):
490+
# Only call streaming context cleanup for actual streaming responses
491+
# Add additional safety checks to ensure we're really in Flask 3.1+
492+
try:
493+
# Additional safety check: verify we're in a Flask request context
494+
if hasattr(flask, "request") and hasattr(
495+
flask.request, "response"
496+
):
497+
is_streaming = (
498+
hasattr(flask.request, "response")
499+
and flask.request.response
500+
and hasattr(flask.request.response, "stream")
501+
and flask.request.response.stream
502+
)
503+
if is_streaming:
504+
_ensure_streaming_context_cleanup(wrapped_app_environ)
505+
except (
506+
RuntimeError,
507+
ValueError,
508+
TypeError,
509+
AttributeError,
510+
) as exc:
511+
# Ensure our Flask 3.1+ logic never interferes with normal request processing
512+
_logger.debug(
513+
"Flask 3.1+ streaming context cleanup failed, continuing: %s",
514+
exc,
515+
)
516+
411517
if should_trace:
412518
duration_s = default_timer() - start
413519
if duration_histogram_old:
@@ -433,6 +539,7 @@ def _start_response(status, response_headers, *args, **kwargs):
433539
duration_histogram_new.record(
434540
max(duration_s, 0), duration_attrs_new
435541
)
542+
436543
active_requests_counter.add(-1, active_requests_count_attrs)
437544
return result
438545

@@ -537,6 +644,7 @@ def _teardown_request(exc):
537644
return
538645

539646
activation = flask.request.environ.get(_ENVIRON_ACTIVATION_KEY)
647+
token = flask.request.environ.get(_ENVIRON_TOKEN)
540648

541649
original_reqctx_ref = flask.request.environ.get(
542650
_ENVIRON_REQCTX_REF_KEY
@@ -554,15 +662,63 @@ def _teardown_request(exc):
554662
# like any decorated with `flask.copy_current_request_context`.
555663

556664
return
557-
if exc is None:
558-
activation.__exit__(None, None, None)
559-
else:
560-
activation.__exit__(
561-
type(exc), exc, getattr(exc, "__traceback__", None)
665+
666+
try:
667+
# For Flask 3.1+, check if this is a streaming response that might
668+
# have already been cleaned up to prevent double cleanup
669+
# Only check for streaming in Flask 3.1+ and Python 3.10+ to avoid interference with older versions
670+
# Use very conservative checks to ensure we never interfere with Flask < 3.1
671+
is_flask_31_plus = (
672+
hasattr(
673+
flask, "__version__"
674+
) # Ensure Flask has version attribute
675+
and package_version.parse(flask.__version__)
676+
>= package_version.parse("3.1.0")
677+
and sys.version_info >= (3, 10)
562678
)
563679

564-
if flask.request.environ.get(_ENVIRON_TOKEN, None):
565-
context.detach(flask.request.environ.get(_ENVIRON_TOKEN))
680+
is_streaming = False
681+
if is_flask_31_plus:
682+
try:
683+
# Additional safety check: verify we're in a Flask request context
684+
if hasattr(flask, "request") and hasattr(
685+
flask.request, "response"
686+
):
687+
is_streaming = (
688+
hasattr(flask.request, "response")
689+
and flask.request.response
690+
and hasattr(flask.request.response, "stream")
691+
and flask.request.response.stream
692+
)
693+
except (RuntimeError, AttributeError):
694+
# Not in a proper Flask request context, don't check for streaming
695+
is_streaming = False
696+
697+
if is_flask_31_plus and is_streaming:
698+
# For streaming responses in Flask 3.1+, the context might have been
699+
# cleaned up already in _ensure_streaming_context_cleanup
700+
# Mark the activation and token as None to prevent double cleanup
701+
flask.request.environ[_ENVIRON_ACTIVATION_KEY] = None
702+
flask.request.environ[_ENVIRON_TOKEN] = None
703+
return
704+
705+
if exc is None:
706+
activation.__exit__(None, None, None)
707+
else:
708+
activation.__exit__(
709+
type(exc), exc, getattr(exc, "__traceback__", None)
710+
)
711+
712+
if token:
713+
context.detach(token)
714+
715+
except (RuntimeError, AttributeError, ValueError) as teardown_exc:
716+
# Log the error but don't raise it to avoid breaking the request handling
717+
_logger.debug(
718+
"Error during request teardown: %s",
719+
teardown_exc,
720+
exc_info=True,
721+
)
566722

567723
return _teardown_request
568724

0 commit comments

Comments
 (0)