🐛 Custom tags for Middlewares (#185)

Author: perdy

flama/authentication/middleware.py

@@ -18,8 +18,9 @@
 
 
 class AuthenticationMiddleware:
-    def __init__(self, app: "types.App", *, ignored: list[str] = []):
+    def __init__(self, app: "types.App", *, tag: str = "permissions", ignored: list[str] = []):
         self.app: Flama = t.cast("Flama", app)
+        self._tag = tag
         self._ignored = [re.compile(x) for x in ignored]
 
     async def __call__(self, scope: "types.Scope", receive: "types.Receive", send: "types.Send") -> None:
@@ -34,7 +35,7 @@ async def __call__(self, scope: "types.Scope", receive: "types.Receive", send: "
     def _get_permissions(self, app: "Flama", scope: "types.Scope") -> set[str]:
         try:
             route, _ = app.router.resolve_route(scope)
-            permissions = set(route.tags.get("permissions", []))
+            permissions = set(route.tags.get(self._tag, []))
         except (exceptions.MethodNotAllowedException, exceptions.NotFoundException):
             permissions = []
 

flama/telemetry/middleware.py

@@ -3,7 +3,7 @@
 import re
 import typing as t
 
-from flama import Flama, concurrency, types
+from flama import Flama, concurrency, exceptions, types
 from flama.telemetry.data_structures import Error, Response, TelemetryData
 
 logger = logging.getLogger(__name__)
@@ -126,12 +126,14 @@ def __init__(
         log_level: int = logging.NOTSET,
         before: t.Optional[HookFunction] = None,
         after: t.Optional[HookFunction] = None,
+        tag: str = "telemetry",
         ignored: list[str] = [],
     ) -> None:
         self.app: Flama = t.cast(Flama, app)
         self._log_level = log_level
         self._before = before
         self._after = after
+        self._tag = tag
         self._ignored = [re.compile(x) for x in ignored]
 
     async def before(self, data: TelemetryData):
@@ -142,8 +144,20 @@ async def after(self, data: TelemetryData):
         if self._after:
             await concurrency.run(self._after, data)
 
+    def _get_tag(self, scope: "types.Scope") -> bool:
+        try:
+            app: Flama = scope["app"]
+            route, _ = app.router.resolve_route(scope)
+            return route.tags.get(self._tag, True)
+        except (exceptions.MethodNotAllowedException, exceptions.NotFoundException):
+            return False
+
     async def __call__(self, scope: types.Scope, receive: types.Receive, send: types.Send) -> None:
-        if scope["type"] not in ("http", "websocket") or any(pattern.match(scope["path"]) for pattern in self._ignored):
+        if (
+            scope["type"] not in ("http", "websocket")
+            or any(pattern.match(scope["path"]) for pattern in self._ignored)
+            or not self._get_tag(scope)
+        ):
             await self.app(scope, receive, send)
             return
 

tests/telemetry/test_middleware.py

@@ -47,12 +47,17 @@ def error():
         def ignored():
             return "ignored"
 
+        @app.post("/explicit-off/", name="explicit-off", tags={"telemetry": False})
+        def explicit_off():
+            return "explicit_off"
+
     @pytest.mark.parametrize(
         [
             "path",
             "request_params",
             "request_body",
             "request_cookies",
+            "status_code",
             "response",
             "exception",
             "before",
@@ -65,6 +70,7 @@ def ignored():
                 {"y": 1},
                 b"body",
                 {"access_token": TOKEN},
+                http.HTTPStatus.OK,
                 {"x": 1, "y": 1, "body": "body"},
                 None,
                 None,
@@ -77,6 +83,7 @@ def ignored():
                 {"y": 1},
                 b"body",
                 {"access_token": TOKEN},
+                http.HTTPStatus.OK,
                 {"x": 1, "y": 1, "body": "body"},
                 None,
                 MagicMock(),
@@ -128,6 +135,7 @@ def ignored():
                 {"y": 1},
                 b"body",
                 {"access_token": TOKEN},
+                http.HTTPStatus.OK,
                 {"x": 1, "y": 1, "body": "body"},
                 None,
                 AsyncMock(),
@@ -179,6 +187,7 @@ def ignored():
                 {},
                 None,
                 {},
+                http.HTTPStatus.OK,
                 None,
                 ValueError("foo"),
                 None,
@@ -191,6 +200,7 @@ def ignored():
                 {},
                 None,
                 {},
+                http.HTTPStatus.OK,
                 None,
                 ValueError("foo"),
                 MagicMock(),
@@ -222,6 +232,7 @@ def ignored():
                 {},
                 None,
                 {},
+                http.HTTPStatus.OK,
                 None,
                 ValueError("foo"),
                 AsyncMock(),
@@ -248,12 +259,62 @@ def ignored():
                 ),
                 id="error_async_hooks",
             ),
-            pytest.param("/ignored/", {}, None, {}, "ignored", None, AsyncMock(), AsyncMock(), None, id="ignored"),
+            pytest.param(
+                "/ignored/",
+                {},
+                None,
+                {},
+                http.HTTPStatus.OK,
+                "ignored",
+                None,
+                AsyncMock(),
+                AsyncMock(),
+                None,
+                id="ignored",
+            ),
+            pytest.param(
+                "/explicit-off/",
+                {},
+                None,
+                {},
+                http.HTTPStatus.OK,
+                "explicit_off",
+                None,
+                AsyncMock(),
+                AsyncMock(),
+                None,
+                id="explicit_off",
+            ),
+            pytest.param(
+                "/not-found/",
+                {},
+                None,
+                {},
+                http.HTTPStatus.NOT_FOUND,
+                {"status_code": 404, "detail": "Not Found", "error": "HTTPException"},
+                None,
+                None,
+                None,
+                None,
+                id="not_found",
+            ),
         ],
         indirect=["exception"],
     )
     async def test_request(
-        self, app, client, path, request_params, request_body, request_cookies, response, exception, before, after, data
+        self,
+        app,
+        client,
+        path,
+        request_params,
+        request_body,
+        request_cookies,
+        status_code,
+        response,
+        exception,
+        before,
+        after,
+        data,
     ):
         app.add_middleware(Middleware(TelemetryMiddleware, before=before, after=after, ignored=[r"/ignored.*"]))
 
@@ -271,7 +332,7 @@ async def test_request(
         with exception, patch("datetime.datetime", MagicMock(now=MagicMock(return_value=now))):
             r = await client.post(path, params=request_params, content=request_body)
 
-            assert r.status_code == http.HTTPStatus.OK
+            assert r.status_code == status_code
             assert r.json() == response
 
         if before: