Merge pull request #537 from alphagov/ris-dm-gzip-middleware

add DMGzipMiddleware, a thin wrapper around Flask_gzip adding header-controllability & instrumentation

dmutils/__init__.py

@@ -2,4 +2,4 @@
 from .flask_init import init_app, init_manager
 
 
-__version__ = '48.5.0'
+__version__ = '48.6.0'

dmutils/flask.py

@@ -1,6 +1,7 @@
 from functools import partial
 
 from flask import render_template, render_template_string
+from flask_gzip import Gzip
 
 from dmutils.timing import logged_duration
 
@@ -32,3 +33,24 @@
 timed_render_template_string.__doc__ = """
     See ``timed_render_template``, only for ``render_template_string``.
 """
+
+
+class DMGzipMiddleware(Gzip):
+    compress_by_default: bool
+
+    def __init__(self, *args, compress_by_default: bool = False, **kwargs):
+        kwargs.setdefault("minimum_size", 8192)
+        self.compress_by_default = compress_by_default
+        super().__init__(*args, **kwargs)
+
+    def after_request(self, response):
+        x_compression_safe = response.headers.pop("X-Compression-Safe", None)
+        compress = {"0": False, "1": True}.get(x_compression_safe, self.compress_by_default)
+
+        # flask_gzip makes the minimum_size comparison itself, but we want to avoid outputting a misleading
+        # logged_duration message if it's going to be prevented in the superclass.
+        if compress and len(response.get_data()) >= self.minimum_size:
+            with logged_duration(message="Spent {duration_real}s compressing response"):
+                return super().after_request(response)
+        else:
+            return response

setup.py

@@ -26,6 +26,7 @@
          'Flask-Script>=2.0.6',
          'Flask-WTF>=0.14.2',
          'Flask<1.1,>=1.0.2',
+         'Flask-gzip>=0.2',
          'Flask-Login>=0.2.11',
          'boto3<2,>=1.7.83',
          'contextlib2>=0.4.0',

tests/test_flask.py

@@ -0,0 +1,55 @@
+from dmutils.flask import DMGzipMiddleware
+
+
+class TestDMGzipMiddleware:
+    def test_with_compression_safe_header(self, app):
+        DMGzipMiddleware(app, compress_by_default=False)
+
+        @app.route('/')
+        def some_route():
+            return "!" * 9000, 200, {"X-Compression-Safe": "1"}
+
+        response = app.test_client().get('/', headers={"Accept-Encoding": "gzip"})
+        assert response.status_code == 200
+        assert response.headers["Content-Encoding"] == "gzip"
+        assert len(response.get_data()) < 9000
+        assert "X-Compression-Safe" not in response.headers
+
+    def test_with_compression_unsafe_header(self, app):
+        DMGzipMiddleware(app, compress_by_default=False)
+
+        @app.route('/')
+        def some_route():
+            return "!" * 9000, 200, {"X-Compression-Safe": "0"}
+
+        response = app.test_client().get('/', headers={"Accept-Encoding": "gzip"})
+        assert response.status_code == 200
+        assert response.headers.get("Content-Encoding") != "gzip"
+        assert len(response.get_data()) == 9000
+        assert "X-Compression-Safe" not in response.headers
+
+    def test_with_compression_default_false(self, app):
+        DMGzipMiddleware(app)
+
+        @app.route('/')
+        def some_route():
+            return "!" * 9000, 200
+
+        response = app.test_client().get('/', headers={"Accept-Encoding": "gzip"})
+        assert response.status_code == 200
+        assert response.headers.get("Content-Encoding") != "gzip"
+        assert len(response.get_data()) == 9000
+        assert "X-Compression-Safe" not in response.headers
+
+    def test_with_compression_default_true(self, app):
+        DMGzipMiddleware(app, compress_by_default=True)
+
+        @app.route('/')
+        def some_route():
+            return "!" * 9000, 200
+
+        response = app.test_client().get('/', headers={"Accept-Encoding": "gzip"})
+        assert response.status_code == 200
+        assert response.headers["Content-Encoding"] == "gzip"
+        assert len(response.get_data()) < 9000
+        assert "X-Compression-Safe" not in response.headers