From c823c51721434747fb92e7390c047efeb0de95c3 Mon Sep 17 00:00:00 2001
From: ducklet <ducklet@noreply.code.dumpr.org>
Date: Wed, 23 Jun 2021 22:57:09 +0200
Subject: [PATCH] add elapsed processing time to response

---
 unwind/middleware/responsetime.py | 26 ++++++++++++++++++++++++++
 unwind/web.py                     |  6 +++++-
 2 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 unwind/middleware/responsetime.py

diff --git a/unwind/middleware/responsetime.py b/unwind/middleware/responsetime.py
new file mode 100644
index 0000000..1fcad62
--- /dev/null
+++ b/unwind/middleware/responsetime.py
@@ -0,0 +1,26 @@
+from time import perf_counter
+
+from starlette.datastructures import MutableHeaders
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+
+class ResponseTimeMiddleware:
+    def __init__(self, app: ASGIApp, header_name: str = "Elapsed") -> None:
+        self.app = app
+        self.header_name = header_name
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        start = perf_counter()
+
+        if scope["type"] not in ("http", "websocket"):
+            await self.app(scope, receive, send)
+            return
+
+        async def send_wrapper(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                headers = MutableHeaders(scope=message)
+                elapsed = perf_counter() - start
+                headers.append(self.header_name, str(elapsed))
+            await send(message)
+
+        await self.app(scope, receive, send_wrapper)
diff --git a/unwind/web.py b/unwind/web.py
index 411c519..3f1e61a 100644
--- a/unwind/web.py
+++ b/unwind/web.py
@@ -19,6 +19,7 @@ from starlette.routing import Mount, Route
 
 from . import config, db
 from .db import close_connection_pool, find_ratings, open_connection_pool
+from .middleware.responsetime import ResponseTimeMiddleware
 from .models import Movie, asplain
 
 log = logging.getLogger(__name__)
@@ -190,5 +191,8 @@ def create_app():
                 ],
             ),
         ],
-        middleware=[Middleware(AuthenticationMiddleware, backend=BasicAuthBackend())],
+        middleware=[
+            Middleware(ResponseTimeMiddleware, header_name="Unwind-Elapsed"),
+            Middleware(AuthenticationMiddleware, backend=BasicAuthBackend()),
+        ],
     )