fix: support gzip decompression on metadata requests

testbench/common.py

@@ -16,6 +16,8 @@
 
 import base64
 from functools import wraps
+import gzip
+import io
 import json
 import random
 import re
@@ -801,6 +803,19 @@ def wrapper(*args, **kwargs):
     return retry_test
 
 
+def handle_gzip_request(request):
+    """
+    Handle gzip compressed JSON payloads when Content-Encoding: gzip is present on metadata requests.
+    No decompressions for media uploads when object's metadata includes Content-Encoding: gzip.
+    """
+    if (
+        request.headers.get("Content-Encoding", None) == "gzip"
+        and request.args.get("contentEncoding", None) != "gzip"
+    ):
+        request.data = gzip.decompress(request.data)
+        request.environ["wsgi.input"] = io.BytesIO(request.data)
+
+
 def rest_crc32c_to_proto(crc32c):
     """Convert from the REST representation of crc32c checksums to the proto representation.
 

testbench/rest_server.py

@@ -182,6 +182,11 @@ def start_grpc():
 gcs.register_error_handler(Exception, testbench.error.RestException.handler)
 
 
+@gcs.before_request
+def handle_gzip_compressed_request():
+    return testbench.common.handle_gzip_request(flask.request)
+
+
 # === BUCKET === #
 
 
@@ -888,6 +893,11 @@ def download_object_get(bucket_name, object_name):
 upload.register_error_handler(Exception, testbench.error.RestException.handler)
 
 
+@upload.before_request
+def handle_gzip_compressed_request():
+    return testbench.common.handle_gzip_request(flask.request)
+
+
 @upload.route("/b/<bucket_name>/o", methods=["POST"])
 @retry_test(method="storage.objects.insert")
 def object_insert(bucket_name):

testbench/servers/iam_rest_server.py

@@ -28,6 +28,11 @@
 _iam.register_error_handler(Exception, testbench.error.RestException.handler)
 
 
+@_iam.before_request
+def handle_gzip_compressed_request():
+    return testbench.common.handle_gzip_request(flask.request)
+
+
 @_iam.route("/projects/-/serviceAccounts/<service_account>:signBlob", methods=["POST"])
 def sign_blob(service_account):
     """Implement the `projects.serviceAccounts.signBlob` API."""

testbench/servers/projects_rest_server.py

@@ -30,6 +30,10 @@ def get_projects_app(db):
     projects.debug = False
     projects.register_error_handler(Exception, testbench.error.RestException.handler)
 
+    @projects.before_request
+    def handle_gzip_compressed_request():
+        return testbench.common.handle_gzip_request(flask.request)
+
     @projects.route("/<project_id>/serviceAccount")
     @retry_test("storage.serviceaccount.get")
     def projects_get(project_id):

tests/test_common.py

@@ -17,6 +17,7 @@
 """Unit test for utils"""
 
 import base64
+import gzip
 import hashlib
 import json
 import types
@@ -1205,6 +1206,46 @@ def test_bucket_to_from_proto(self):
             ),
         )
 
+    def test_handle_gzip_request(self):
+        # Test gzip decompresses request payload when Content-Encoding: gzip is present.
+        payload = b'{"name": "bucket-name"}'
+        compressed = gzip.compress(payload)
+        request = testbench.common.FakeRequest(
+            headers={"Content-Encoding": "gzip"},
+            args={},
+            data=compressed,
+            environ={},
+        )
+        testbench.common.handle_gzip_request(request)
+        self.assertEqual(request.data, payload)
+
+        # Test no decompression when object's metadata includes Content-Encoding: gzip,
+        # so that data can be stored in compressed form if applicable (PR#322).
+        media = "How vexingly quick daft zebras jump!"
+        compressed = gzip.compress(media.encode("utf-8"))
+        request = testbench.common.FakeRequest(
+            headers={"Content-Encoding": "gzip"},
+            args={
+                "name": "test_blob",
+                "uploadType": "media",
+                "contentEncoding": "gzip",
+            },
+            data=compressed,
+            environ={},
+        )
+        testbench.common.handle_gzip_request(request)
+        self.assertEqual(request.data, compressed)
+
+        # Test no decompression when Content-Encoding: gzip is not present.
+        request = testbench.common.FakeRequest(
+            headers={},
+            args={},
+            data=payload,
+            environ={},
+        )
+        testbench.common.handle_gzip_request(request)
+        self.assertEqual(request.data, payload)
+
 
 if __name__ == "__main__":
     unittest.main()