Add support for gzip Content-Encoding

Signed-off-by: ivan-valkov <iv.v.valkov@gmail.com>

prometheus_client/asgi.py

@@ -14,20 +14,25 @@ async def prometheus_app(scope, receive, send):
         params = parse_qs(scope.get('query_string', b''))
         accept_header = "Accept: " + ",".join([
             value.decode("utf8") for (name, value) in scope.get('headers')
-            if name.decode("utf8") == 'accept'
+            if name.decode("utf8").lower() == 'accept'
+        ])
+        accept_encoding_header = ",".join([
+            value.decode("utf8") for (name, value) in scope.get('headers')
+            if name.decode("utf8").lower() == 'accept-encoding'
         ])
         # Bake output
-        status, header, output = _bake_output(registry, accept_header, params)
+        status, headers, output = _bake_output(registry, accept_header, accept_encoding_header, params)
+        formatted_headers = []
+        for header in headers:
+            formatted_headers.append(tuple(x.encode('utf8') for x in header))
         # Return output
         payload = await receive()
         if payload.get("type") == "http.request":
             await send(
                 {
                     "type": "http.response.start",
                     "status": int(status.split(' ')[0]),
-                    "headers": [
-                        tuple(x.encode('utf8') for x in header)
-                    ]
+                    "headers": formatted_headers,
                 }
             )
             await send({"type": "http.response.body", "body": output})

prometheus_client/exposition.py

@@ -1,4 +1,5 @@
 import base64
+import gzip
 from contextlib import closing
 from http.server import BaseHTTPRequestHandler
 import os
@@ -93,13 +94,19 @@ def redirect_request(self, req, fp, code, msg, headers, newurl):
         return new_request
 
 
-def _bake_output(registry, accept_header, params):
+def _bake_output(registry, accept_header, accept_encoding_header, params):
     """Bake output for metrics output."""
-    encoder, content_type = choose_encoder(accept_header)
+    # Choose the correct plain text format of the output.
+    formatter, content_type = choose_formatter(accept_header)
     if 'name[]' in params:
         registry = registry.restricted_registry(params['name[]'])
-    output = encoder(registry)
-    return '200 OK', ('Content-Type', content_type), output
+    output = formatter(registry)
+    headers = [('Content-Type', content_type)]
+    # If gzip encoding required, gzip the output.
+    if gzip_accepted(accept_encoding_header):
+        output = gzip.compress(output)
+        headers.append(('Content-Encoding', 'gzip'))
+    return '200 OK', headers, output
 
 
 def make_wsgi_app(registry: CollectorRegistry = REGISTRY) -> Callable:
@@ -108,17 +115,18 @@ def make_wsgi_app(registry: CollectorRegistry = REGISTRY) -> Callable:
     def prometheus_app(environ, start_response):
         # Prepare parameters
         accept_header = environ.get('HTTP_ACCEPT')
+        accept_encoding_header = environ.get('HTTP_ACCEPT_ENCODING')
         params = parse_qs(environ.get('QUERY_STRING', ''))
         if environ['PATH_INFO'] == '/favicon.ico':
             # Serve empty response for browsers
             status = '200 OK'
-            header = ('', '')
+            headers = [('', '')]
             output = b''
         else:
             # Bake output
-            status, header, output = _bake_output(registry, accept_header, params)
+            status, headers, output = _bake_output(registry, accept_header, accept_encoding_header, params)
         # Return output
-        start_response(status, [header])
+        start_response(status, headers)
         return [output]
 
     return prometheus_app
@@ -227,15 +235,23 @@ def sample_line(line):
     return ''.join(output).encode('utf-8')
 
 
-def choose_encoder(accept_header: str) -> Tuple[Callable[[CollectorRegistry], bytes], str]:
+def choose_formatter(accept_header: str) -> Tuple[Callable[[CollectorRegistry], bytes], str]:
     accept_header = accept_header or ''
     for accepted in accept_header.split(','):
-        if accepted.split(';')[0].strip() == 'application/openmetrics-text':
+        if accepted.split(';')[0].strip().lower() == 'application/openmetrics-text':
             return (openmetrics.generate_latest,
                     openmetrics.CONTENT_TYPE_LATEST)
     return generate_latest, CONTENT_TYPE_LATEST
 
 
+def gzip_accepted(accept_encoding_header: str) -> bool:
+    accept_encoding_header = accept_encoding_header or ''
+    for accepted in accept_encoding_header.split(','):
+        if accepted.split(';')[0].strip().lower() == 'gzip':
+            return True
+    return False
+
+
 class MetricsHandler(BaseHTTPRequestHandler):
     """HTTP handler that gives metrics from ``REGISTRY``."""
     registry: CollectorRegistry = REGISTRY
@@ -244,12 +260,14 @@ def do_GET(self) -> None:
         # Prepare parameters
         registry = self.registry
         accept_header = self.headers.get('Accept')
+        accept_encoding_header = self.headers.get('Accept-Encoding')
         params = parse_qs(urlparse(self.path).query)
         # Bake output
-        status, header, output = _bake_output(registry, accept_header, params)
+        status, headers, output = _bake_output(registry, accept_header, accept_encoding_header, params)
         # Return output
         self.send_response(int(status.split(' ')[0]))
-        self.send_header(*header)
+        for header in headers:
+            self.send_header(*header)
         self.end_headers()
         self.wfile.write(output)
 

tests/test_asgi.py

@@ -1,3 +1,4 @@
+import gzip
 from unittest import skipUnless, TestCase
 
 from prometheus_client import CollectorRegistry, Counter
@@ -115,3 +116,36 @@ def test_report_metrics_3(self):
 
     def test_report_metrics_4(self):
         self.validate_metrics("failed_requests", "Number of failed requests", 7)
+
+    def test_gzip(self):
+        # Increment a metric.
+        metric_name = "counter"
+        help_text = "A counter"
+        increments = 2
+        c = Counter(metric_name, help_text, registry=self.registry)
+        for _ in range(increments):
+            c.inc()
+        app = make_asgi_app(self.registry)
+        self.seed_app(app)
+        # Send input with gzip header.
+        self.scope["headers"] = [(b"accept-encoding", b"gzip")]
+        self.send_input({"type": "http.request", "body": b""})
+        # Assert outputs
+        outputs = self.get_all_output()
+        # Assert outputs
+        self.assertEqual(len(outputs), 2)
+        response_start = outputs[0]
+        self.assertEqual(response_start['type'], 'http.response.start')
+        response_body = outputs[1]
+        self.assertEqual(response_body['type'], 'http.response.body')
+        # Status code
+        self.assertEqual(response_start['status'], 200)
+        # Headers
+        self.assertEqual(len(response_start['headers']), 2)
+        self.assertIn((b"Content-Type", CONTENT_TYPE_LATEST.encode('utf8')), response_start['headers'])
+        self.assertIn((b"Content-Encoding", b"gzip"), response_start['headers'])
+        # Body
+        output = gzip.decompress(response_body['body']).decode('utf8')
+        self.assertIn("# HELP " + metric_name + "_total " + help_text + "\n", output)
+        self.assertIn("# TYPE " + metric_name + "_total counter\n", output)
+        self.assertIn(metric_name + "_total " + str(increments) + ".0\n", output)

tests/test_wsgi.py

@@ -1,3 +1,4 @@
+import gzip
 from unittest import TestCase
 from wsgiref.util import setup_testing_defaults
 
@@ -70,3 +71,30 @@ def test_favicon_path(self):
             # Try accessing normal paths
             app(self.environ, self.capture)
             self.assertEqual(mock.call_count, 1)
+
+    def test_gzip(self):
+        # Increment a metric
+        metric_name = "counter"
+        help_text = "A counter"
+        increments = 2
+        c = Counter(metric_name, help_text, registry=self.registry)
+        for _ in range(increments):
+            c.inc()
+        app = make_wsgi_app(self.registry)
+        # Try accessing metrics using the gzip Accept-Content header.
+        gzip_environ = dict(self.environ)
+        gzip_environ['HTTP_ACCEPT_ENCODING'] = 'gzip'
+        outputs = app(gzip_environ, self.capture)
+        # Assert outputs
+        self.assertEqual(len(outputs), 1)
+        output = gzip.decompress(outputs[0]).decode(encoding="utf-8")
+        # Status code
+        self.assertEqual(self.captured_status, "200 OK")
+        # Headers
+        self.assertEqual(len(self.captured_headers), 2)
+        self.assertIn(("Content-Type", CONTENT_TYPE_LATEST), self.captured_headers)
+        self.assertIn(("Content-Encoding", "gzip"), self.captured_headers)
+        # Body
+        self.assertIn("# HELP " + metric_name + "_total " + help_text + "\n", output)
+        self.assertIn("# TYPE " + metric_name + "_total counter\n", output)
+        self.assertIn(metric_name + "_total " + str(increments) + ".0\n", output)