adding zstd compression to fastapi

python/api/middleware.py

@@ -0,0 +1,49 @@
+"""Middleware for the FastAPI application."""
+
+from compression import zstd
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+from starlette.responses import Response
+
+MINIMUM_RESPONSE_SIZE = 500
+
+
+class ZstdMiddleware(BaseHTTPMiddleware):
+    """Middleware that compresses responses with zstd when the client supports it."""
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        """Compress the response with zstd if the client accepts it."""
+        accepted_encodings = request.headers.get("accept-encoding", "")
+        if "zstd" not in accepted_encodings:
+            return await call_next(request)
+
+        response = await call_next(request)
+
+        if response.headers.get("content-encoding") or "text/event-stream" in response.headers.get("content-type", ""):
+            return response
+
+        body = b""
+        async for chunk in response.body_iterator:
+            body += chunk if isinstance(chunk, bytes) else chunk.encode()
+
+        if len(body) < MINIMUM_RESPONSE_SIZE:
+            return Response(
+                content=body,
+                status_code=response.status_code,
+                headers=dict(response.headers),
+                media_type=response.media_type,
+            )
+
+        compressed = zstd.compress(body)
+
+        headers = dict(response.headers)
+        headers["content-encoding"] = "zstd"
+        headers["content-length"] = str(len(compressed))
+        headers.pop("transfer-encoding", None)
+
+        return Response(
+            content=compressed,
+            status_code=response.status_code,
+            headers=headers,
+            media_type=response.media_type,
+        )