feat: allow gzip compression on request bodies

Author: rmkeezer

ibm_cloud_sdk_core/base_service.py

@@ -20,6 +20,7 @@
 from os.path import basename
 import platform
 import sys
+import zlib
 from typing import Dict, List, Optional, Tuple, Union
 
 import requests
@@ -41,6 +42,8 @@
 # import http.client as http_client
 # http_client.HTTPConnection.debuglevel = 1
 
+#pylint: disable=too-many-instance-attributes
+#pylint: disable=too-many-locals
 class BaseService:
     """Common functionality shared by generated service classes.
 
@@ -52,6 +55,7 @@ class BaseService:
         authenticator: Adds authentication data to service requests. Defaults to None.
         disable_ssl_verification: A flag that indicates whether verification of the server's SSL
             certificate should be disabled or not. Defaults to False.
+        enable_gzip_compression: A flag that indicates whether to enable gzip compression on request bodies
 
     Attributes:
         service_url (str): Url to the service endpoint.
@@ -61,6 +65,7 @@ class BaseService:
         default_headers (dict): A dictionary of headers to be sent with every HTTP request to the service endpoint.
         jar (http.cookiejar.CookieJar): Stores cookies received from the service.
         http_config (dict): A dictionary containing values that control the timeout, proxies, and etc of HTTP requests.
+        enable_gzip_compression (bool): A flag that indicates whether to enable gzip compression on request bodies
     Raises:
         ValueError: If Authenticator is not provided or invalid type.
     """
@@ -74,13 +79,15 @@ def __init__(self,
                  *,
                  service_url: str = None,
                  authenticator: Authenticator = None,
-                 disable_ssl_verification: bool = False) -> None:
+                 disable_ssl_verification: bool = False,
+                 enable_gzip_compression: bool = False) -> None:
         self.set_service_url(service_url)
         self.http_config = {}
         self.jar = CookieJar()
         self.authenticator = authenticator
         self.disable_ssl_verification = disable_ssl_verification
         self.default_headers = None
+        self.enable_gzip_compression = enable_gzip_compression
         self._set_user_agent_header(self._build_user_agent())
         if not self.authenticator:
             raise ValueError('authenticator must be provided')
@@ -124,6 +131,8 @@ def configure_service(self, service_name: str) -> None:
             self.set_disable_ssl_verification(
                 bool(config.get('DISABLE_SSL'))
             )
+        if config.get('ENABLE_GZIP') is not None:
+            self.set_enable_gzip_compression(config.get('ENABLE_GZIP') == 'True')
 
     def _set_user_agent_header(self, user_agent_string: str) -> None:
         self.user_agent_header = {'User-Agent': user_agent_string}
@@ -245,6 +254,14 @@ def send(self, request: requests.Request, **kwargs) -> DetailedResponse:
             logging.exception('Error in service call')
             raise
 
+    def set_enable_gzip_compression(self, should_enable_compression: bool = False) -> None:
+        """Set value to enable gzip compression on request bodies"""
+        self.enable_gzip_compression = should_enable_compression
+
+    def get_enable_gzip_compression(self) -> bool:
+        """Get value for enabling gzip compression on request bodies"""
+        return self.enable_gzip_compression
+
     def prepare_request(self,
                         method: str,
                         url: str,
@@ -309,6 +326,16 @@ def prepare_request(self,
 
         self.authenticator.authenticate(request)
 
+        # Compress the request body if applicable
+        if (self.get_enable_gzip_compression() and
+                'content-encoding' not in headers and
+                request['data'] is not None):
+            headers['content-encoding'] = 'gzip'
+            uncompressed_data = request['data']
+            request_body = zlib.compress(uncompressed_data)
+            request['data'] = request_body
+            request['headers'] = headers
+
         # Next, we need to process the 'files' argument to try to fill in
         # any missing filenames where possible.
         # 'files' can be a dictionary (i.e { '<part-name>': (<tuple>)} )

resources/ibm-credentials-gzip.env

@@ -0,0 +1,4 @@
+INCLUDE_EXTERNAL_CONFIG_APIKEY=mockkey
+INCLUDE_EXTERNAL_CONFIG_AUTH_TYPE=iam
+INCLUDE_EXTERNAL_CONFIG_URL=https://mockurl
+INCLUDE_EXTERNAL_CONFIG_ENABLE_GZIP=True

test/test_base_service.py

@@ -5,6 +5,7 @@
 import os
 from shutil import copyfile
 from typing import Optional
+import zlib
 import pytest
 import responses
 import requests
@@ -475,6 +476,46 @@ def test_get_authenticator():
     service = AnyServiceV1('2018-11-20', authenticator=auth)
     assert service.get_authenticator() is not None
 
+def test_gzip_compression():
+    # Should return uncompressed data when gzip is off
+    service = AnyServiceV1('2018-11-20', authenticator=NoAuthAuthenticator())
+    assert not service.get_enable_gzip_compression()
+    prepped = service.prepare_request('GET', url='', data=json.dumps({"foo": "bar"}))
+    assert prepped['data'] == b'{"foo": "bar"}'
+    assert prepped['headers'].get('content-encoding') != 'gzip'
+
+    # Should return compressed data when gzip is on
+    service.set_enable_gzip_compression(True)
+    assert service.get_enable_gzip_compression()
+    prepped = service.prepare_request('GET', url='', data=json.dumps({"foo": "bar"}))
+    assert prepped['data'] == zlib.compress(b'{"foo": "bar"}')
+    assert prepped['headers'].get('content-encoding') == 'gzip'
+
+    # Should return compressed data when gzip is on for non-json data
+    assert service.get_enable_gzip_compression()
+    prepped = service.prepare_request('GET', url='', data=b'rawdata')
+    assert prepped['data'] == zlib.compress(b'rawdata')
+    assert prepped['headers'].get('content-encoding') == 'gzip'
+
+    # Should return uncompressed data when content-encoding is set
+    assert service.get_enable_gzip_compression()
+    prepped = service.prepare_request('GET', url='', headers={"content-encoding": "gzip"},
+                                      data=json.dumps({"foo": "bar"}))
+    assert prepped['data'] == b'{"foo": "bar"}'
+    assert prepped['headers'].get('content-encoding') == 'gzip'
+
+def test_gzip_compression_external():
+    # Should set gzip compression from external config
+    file_path = os.path.join(
+        os.path.dirname(__file__), '../resources/ibm-credentials-gzip.env')
+    os.environ['IBM_CREDENTIALS_FILE'] = file_path
+    service = IncludeExternalConfigService('v1', authenticator=NoAuthAuthenticator())
+    assert service.service_url == 'https://mockurl'
+    assert service.get_enable_gzip_compression() is True
+    prepped = service.prepare_request('GET', url='', data=json.dumps({"foo": "bar"}))
+    assert prepped['data'] == zlib.compress(b'{"foo": "bar"}')
+    assert prepped['headers'].get('content-encoding') == 'gzip'
+
 @responses.activate
 def test_user_agent_header():
     service = AnyServiceV1('2018-11-20', authenticator=NoAuthAuthenticator())