Repository: mozilla/bugbug

Issue:
Compress HTTP service responses
The JSON responses could be gzipped or brotlified.


Discussion:

Resolving patch:
diff --git a/http_service/bugbug_http/app.py b/http_service/bugbug_http/app.py
--- a/http_service/bugbug_http/app.py
+++ b/http_service/bugbug_http/app.py
@@ -3,6 +3,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import gzip
 import logging
 import os
 import uuid
@@ -16,7 +17,7 @@
 from apispec.ext.marshmallow import MarshmallowPlugin
 from apispec_webframeworks.flask import FlaskPlugin
 from cerberus import Validator
-from flask import Flask, jsonify, render_template, request
+from flask import Flask, Response, jsonify, render_template, request
 from flask_cors import cross_origin
 from marshmallow import Schema, fields
 from redis import Redis
@@ -284,6 +285,28 @@ def get_result(job):
     return None
 
 
+def compress_response(data: dict, status_code: int):
+    """Compress data using gzip compressor and frame response
+
+    :param data: data
+    :type data: dict
+    :param status_code: response status code
+    :type status_code: int
+    :return: response with gzip compressed data
+    :rtype: Response
+    """
+
+    gzip_buffer = gzip.compress(orjson.dumps(data), compresslevel=9)
+
+    response = Response(status=status_code)
+    response.set_data(gzip_buffer)
+    response.headers["Content-Encoding"] = "gzip"
+    response.headers["Content-Length"] = len(gzip_buffer)
+    response.headers["Content-Type"] = "application/json"
+
+    return response
+
+
 @application.route("/<model_name>/predict/<int:bug_id>")
 @cross_origin()
 def model_prediction(model_name, bug_id):
@@ -348,7 +371,7 @@ def model_prediction(model_name, bug_id):
         status_code = 202
         data = {"ready": False}
 
-    return jsonify(**data), status_code
+    return compress_response(data, status_code)
 
 
 @application.route("/<model_name>/predict/batch", methods=["POST"])
@@ -522,7 +545,7 @@ def batch_prediction(model_name):
         # not like getting 1 million bug at a time
         schedule_bug_classification(model_name, missing_bugs)
 
-    return jsonify({"bugs": data}), status_code
+    return compress_response({"bugs": data}, status_code)
 
 
 @application.route("/push/<path:branch>/<rev>/schedules")
@@ -576,7 +599,7 @@ def push_schedules(branch, rev):
     job = JobInfo(schedule_tests, branch, rev)
     data = get_result(job)
     if data:
-        return jsonify(data), 200
+        return compress_response(data, 200)
 
     if not is_pending(job):
         schedule_job(job)


Test patch:
diff --git a/http_service/tests/test_bug_classification.py b/http_service/tests/test_bug_classification.py
--- a/http_service/tests/test_bug_classification.py
+++ b/http_service/tests/test_bug_classification.py
@@ -3,12 +3,23 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import gzip
 import json
 import time
 
+import orjson
+
 from bugbug_http.app import API_TOKEN
 
 
+def retrieve_compressed_reponse(response):
+    # Response is of type "<class 'flask.wrappers.Response'>" -  Flask Client's  Response
+    # Not applicable for "<class 'requests.models.Response'> "
+    if response.headers["Content-Encoding"] == "gzip":
+        return orjson.loads(gzip.decompress(response.data))
+    return response.json
+
+
 def test_model_predict_id(client, jobs, add_result, responses):
     bug_id = "123456"
     result = {
@@ -28,12 +39,12 @@ def do_request():
 
     rv = do_request()
     assert rv.status_code == 202
-    assert rv.json == {"ready": False}
+    assert retrieve_compressed_reponse(rv) == {"ready": False}
 
     # request still not ready
     rv = do_request()
     assert rv.status_code == 202
-    assert rv.json == {"ready": False}
+    assert retrieve_compressed_reponse(rv) == {"ready": False}
     assert len(jobs) == 1
 
     # now it's ready
@@ -42,7 +53,7 @@ def do_request():
 
     rv = do_request()
     assert rv.status_code == 200
-    assert rv.json == result
+    assert retrieve_compressed_reponse(rv) == result
 
 
 def test_model_predict_batch(client, jobs, add_result, add_change_time, responses):
@@ -73,7 +84,9 @@ def do_request():
 
     rv = do_request()
     assert rv.status_code == 202
-    assert rv.json == {"bugs": {str(bug_id): {"ready": False} for bug_id in bug_ids}}
+    assert retrieve_compressed_reponse(rv) == {
+        "bugs": {str(bug_id): {"ready": False} for bug_id in bug_ids}
+    }
     assert len(jobs) == 1
 
     # one of the bugs is ready
@@ -86,7 +99,7 @@ def do_request():
 
     rv = do_request()
     assert rv.status_code == 202
-    assert rv.json == {
+    assert retrieve_compressed_reponse(rv) == {
         "bugs": {str(bug_ids[0]): result, str(bug_ids[1]): {"ready": False}}
     }
 
@@ -95,7 +108,9 @@ def do_request():
 
     rv = do_request()
     assert rv.status_code == 200
-    assert rv.json == {"bugs": {str(bug_id): result for bug_id in bug_ids}}
+    assert retrieve_compressed_reponse(rv) == {
+        "bugs": {str(bug_id): result for bug_id in bug_ids}
+    }
 
 
 def test_empty_batch(client):
diff --git a/http_service/tests/test_push_schedules.py b/http_service/tests/test_push_schedules.py
--- a/http_service/tests/test_push_schedules.py
+++ b/http_service/tests/test_push_schedules.py
@@ -3,9 +3,21 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import gzip
+
+import orjson
+
 from bugbug_http.app import API_TOKEN
 
 
+def retrieve_compressed_reponse(response):
+    # Response is of type "<class 'flask.wrappers.Response'>" -  Flask Client's  Response
+    # Not applicable for "<class 'requests.models.Response'> "
+    if response.headers["Content-Encoding"] == "gzip":
+        return orjson.loads(gzip.decompress(response.data))
+    return response.json
+
+
 def test_queue_job_valid(client, add_result, jobs):
     # schedule job
     rv = client.get("/push/autoland/abcdef/schedules", headers={API_TOKEN: "test"},)
@@ -29,7 +41,7 @@ def test_queue_job_valid(client, add_result, jobs):
 
     rv = client.get("/push/autoland/abcdef/schedules", headers={API_TOKEN: "test"},)
     assert rv.status_code == 200
-    assert rv.json == result
+    assert retrieve_compressed_reponse(rv) == result
 
 
 def test_no_api_key(client):