From 14e02fa97fd9aac9dcd94388edde1d7375410503 Mon Sep 17 00:00:00 2001 From: Robbe Sneyders Date: Sun, 19 Nov 2023 23:46:21 +0100 Subject: [PATCH] Add interface to add WSGI middleware (#1814) As discussed in #1807. Allowing the injection of WSGI middleware can enable easier migration from Connexion 2 to Connexion 3. The use cases are limited though, as this will only work for middleware that can work at the end of the middleware stack. --- connexion/apps/flask.py | 18 +++++++++++++++++- connexion/types.py | 32 ++++++++++++++++++++++++++++++-- docs/middleware.rst | 16 +++++++++++++++- tests/test_middleware.py | 28 ++++++++++++++++++++++++++++ 4 files changed, 90 insertions(+), 4 deletions(-) diff --git a/connexion/apps/flask.py b/connexion/apps/flask.py index c93ab1f76..057ce8baa 100644 --- a/connexion/apps/flask.py +++ b/connexion/apps/flask.py @@ -23,7 +23,7 @@ from connexion.operations import AbstractOperation from connexion.options import SwaggerUIOptions from connexion.resolver import Resolver -from connexion.types import MaybeAwaitable +from connexion.types import MaybeAwaitable, WSGIApp from connexion.uri_parsing import AbstractURIParser @@ -259,3 +259,19 @@ def add_error_handler( ], ) -> None: self.middleware.add_error_handler(code_or_exception, function) + + def add_wsgi_middleware( + self, middleware: t.Type[WSGIApp], **options: t.Any + ) -> None: + """Wrap the underlying Flask application with a WSGI middleware. Note that it will only be + called at the end of the middleware stack. Middleware that needs to act sooner, needs to + be added as ASGI middleware instead. + + Adding multiple middleware using this method wraps each middleware around the previous one. + + :param middleware: Middleware class to add + :param options: Options to pass to the middleware_class on initialization + """ + self._middleware_app.asgi_app.app = middleware( + self._middleware_app.asgi_app.app, **options # type: ignore + ) diff --git a/connexion/types.py b/connexion/types.py index ee20ecfa3..d6eca36e8 100644 --- a/connexion/types.py +++ b/connexion/types.py @@ -1,4 +1,32 @@ +import types import typing as t -ReturnType = t.TypeVar("ReturnType") -MaybeAwaitable = t.Union[t.Awaitable[ReturnType], ReturnType] +# Maybe Awaitable +_ReturnType = t.TypeVar("_ReturnType") +MaybeAwaitable = t.Union[t.Awaitable[_ReturnType], _ReturnType] + +# WSGIApp +Environ = t.Mapping[str, object] + +_WriteCallable = t.Callable[[bytes], t.Any] +_ExcInfo = t.Tuple[type, BaseException, types.TracebackType] + +_StartResponseCallable = t.Callable[ + [ + str, # status + t.Sequence[t.Tuple[str, str]], # response headers + ], + _WriteCallable, # write() callable +] +_StartResponseCallableWithExcInfo = t.Callable[ + [ + str, # status + t.Sequence[t.Tuple[str, str]], # response headers + t.Optional[_ExcInfo], # exc_info + ], + _WriteCallable, # write() callable +] +StartResponse = t.Union[_StartResponseCallable, _StartResponseCallableWithExcInfo] +ResponseStream = t.Iterable[bytes] + +WSGIApp = t.Callable[[Environ, StartResponse], ResponseStream] diff --git a/docs/middleware.rst b/docs/middleware.rst index cee83463a..ea23c935f 100644 --- a/docs/middleware.rst +++ b/docs/middleware.rst @@ -64,6 +64,20 @@ You can easily add additional ASGI middleware to the middleware stack with the .. automethod:: connexion.FlaskApp.add_middleware :noindex: + You can also add WSGI middleware to a ``FlaskApp``. Note that it will only be called at the + end of the middleware stack. If you need your middleware to act sooner, you will have to + use an ASGI middleware instead. + + .. code-block:: python + + app.add_wsgi_middleware(MiddlewareClass, **options) + + .. dropdown:: View a detailed reference of the :code:`add_middleware` method + :icon: eye + + .. automethod:: connexion.FlaskApp.add_wsgi_middleware + :noindex: + .. tab-item:: ConnexionMiddleware :sync: ConnexionMiddleware @@ -77,7 +91,7 @@ You can easily add additional ASGI middleware to the middleware stack with the app.add_middleware(MiddlewareClass, **options) - .. dropdown:: View a detailed reference of the :code:`add_middleware` method + .. dropdown:: View a detailed reference of the :code:`add_wsgi_middleware` method :icon: eye .. automethod:: connexion.ConnexionMiddleware.add_middleware diff --git a/tests/test_middleware.py b/tests/test_middleware.py index f71ed6844..29c5627e2 100644 --- a/tests/test_middleware.py +++ b/tests/test_middleware.py @@ -1,6 +1,11 @@ +import typing as t +from unittest.mock import Mock + import pytest +from connexion import FlaskApp from connexion.middleware import ConnexionMiddleware, MiddlewarePosition from connexion.middleware.swagger_ui import SwaggerUIMiddleware +from connexion.types import Environ, ResponseStream, StartResponse, WSGIApp from starlette.datastructures import MutableHeaders from conftest import build_app_from_fixture @@ -81,3 +86,26 @@ def test_position(spec, app_class): == f"Could not insert middleware at position BEFORE_SWAGGER. " f"Please make sure you have a {SwaggerUIMiddleware} in your stack." ) + + +def test_add_wsgi_middleware(spec): + app: FlaskApp = build_app_from_fixture("simple", app_class=FlaskApp, spec_file=spec) + + class WSGIMiddleware: + def __init__(self, app_: WSGIApp, mock_counter): + self.next_app = app_ + self.mock_counter = mock_counter + + def __call__( + self, environ: Environ, start_response: StartResponse + ) -> ResponseStream: + self.mock_counter() + return self.next_app(environ, start_response) + + mock = Mock() + app.add_wsgi_middleware(WSGIMiddleware, mock_counter=mock) + + app_client = app.test_client() + app_client.post("/v1.0/greeting/robbe") + + mock.assert_called_once()