[DE-716] Compression of requests and responses (#328)

* Compression of requests works

* Accepting encoded requests

* Refactoring request compression as a plugin

* Setting default compression values

* Updating docs

* Adding content_encoding

* Adding content_encoding

* Small fix

* Skip interface from coverage

Author: apetenchea

arango/client.py

@@ -13,7 +13,12 @@
 )
 from arango.database import StandardDatabase
 from arango.exceptions import ServerConnectionError
-from arango.http import DEFAULT_REQUEST_TIMEOUT, DefaultHTTPClient, HTTPClient
+from arango.http import (
+    DEFAULT_REQUEST_TIMEOUT,
+    DefaultHTTPClient,
+    HTTPClient,
+    RequestCompression,
+)
 from arango.resolver import (
     FallbackHostResolver,
     HostResolver,
@@ -33,7 +38,7 @@ def default_serializer(x: Any) -> str:
     :return: The object serialized as a JSON string
     :rtype: str
     """
-    return dumps(x)
+    return dumps(x, separators=(",", ":"))
 
 
 def default_deserializer(x: str) -> Any:
@@ -85,6 +90,12 @@ class ArangoClient:
        None: No timeout.
        int: Timeout value in seconds.
     :type request_timeout: int | float
+    :param request_compression: Will compress requests to the server according to
+        the given algorithm. No compression happens by default.
+    :type request_compression: arango.http.RequestCompression | None
+    :param response_compression: Tells the server what compression algorithm is
+        acceptable for the response. No compression happens by default.
+    :type response_compression: str | None
     """
 
     def __init__(
@@ -97,6 +108,8 @@ def __init__(
         deserializer: Callable[[str], Any] = default_deserializer,
         verify_override: Union[bool, str, None] = None,
         request_timeout: Union[int, float, None] = DEFAULT_REQUEST_TIMEOUT,
+        request_compression: Optional[RequestCompression] = None,
+        response_compression: Optional[str] = None,
     ) -> None:
         if isinstance(hosts, str):
             self._hosts = [host.strip("/") for host in hosts.split(",")]
@@ -133,6 +146,9 @@ def __init__(
             for session in self._sessions:
                 session.verify = verify_override
 
+        self._request_compression = request_compression
+        self._response_compression = response_compression
+
     def __repr__(self) -> str:
         return f"<ArangoClient {','.join(self._hosts)}>"
 
@@ -231,6 +247,8 @@ def db(
                 serializer=self._serializer,
                 deserializer=self._deserializer,
                 superuser_token=superuser_token,
+                request_compression=self._request_compression,
+                response_compression=self._response_compression,
             )
         elif user_token is not None:
             connection = JwtConnection(
@@ -242,6 +260,8 @@ def db(
                 serializer=self._serializer,
                 deserializer=self._deserializer,
                 user_token=user_token,
+                request_compression=self._request_compression,
+                response_compression=self._response_compression,
             )
         elif auth_method.lower() == "basic":
             connection = BasicConnection(
@@ -254,6 +274,8 @@ def db(
                 http_client=self._http,
                 serializer=self._serializer,
                 deserializer=self._deserializer,
+                request_compression=self._request_compression,
+                response_compression=self._response_compression,
             )
         elif auth_method.lower() == "jwt":
             connection = JwtConnection(
@@ -266,6 +288,8 @@ def db(
                 http_client=self._http,
                 serializer=self._serializer,
                 deserializer=self._deserializer,
+                request_compression=self._request_compression,
+                response_compression=self._response_compression,
             )
         else:
             raise ValueError(f"invalid auth_method: {auth_method}")

arango/collection.py

@@ -1813,7 +1813,7 @@ def insert_many(
             index caches if document insertions affect the edge index or
             cache-enabled persistent indexes.
         :type refill_index_caches: bool | None
-        param version_attribute: support for simple external versioning to
+        :param version_attribute: support for simple external versioning to
             document operations.
         :type version_attribute: str
         :return: List of document metadata (e.g. document keys, revisions) and
@@ -1939,7 +1939,7 @@ def update_many(
             as opposed to returning the error as an object in the result list.
             Defaults to False.
         :type raise_on_document_error: bool
-        param version_attribute: support for simple external versioning to
+        :param version_attribute: support for simple external versioning to
             document operations.
         :type version_attribute: str
         :return: List of document metadata (e.g. document keys, revisions) and
@@ -2138,7 +2138,7 @@ def replace_many(
             index caches if document operations affect the edge index or
             cache-enabled persistent indexes.
         :type refill_index_caches: bool | None
-        param version_attribute: support for simple external versioning to
+        :param version_attribute: support for simple external versioning to
             document operations.
         :type version_attribute: str
         :return: List of document metadata (e.g. document keys, revisions) and
@@ -2670,7 +2670,7 @@ def insert(
             index caches if document insertions affect the edge index or
             cache-enabled persistent indexes.
         :type refill_index_caches: bool | None
-        param version_attribute: support for simple external versioning to
+        :param version_attribute: support for simple external versioning to
             document operations.
         :type version_attribute: str
         :return: Document metadata (e.g. document key, revision) or True if
@@ -2765,7 +2765,7 @@ def update(
             index caches if document insertions affect the edge index or
             cache-enabled persistent indexes.
         :type refill_index_caches: bool | None
-        param version_attribute: support for simple external versioning
+        :param version_attribute: support for simple external versioning
             to document operations.
         :type version_attribute: str
         :return: Document metadata (e.g. document key, revision) or True if
@@ -2850,7 +2850,7 @@ def replace(
             index caches if document insertions affect the edge index or
             cache-enabled persistent indexes.
         :type refill_index_caches: bool | None
-        param version_attribute: support for simple external versioning to
+        :param version_attribute: support for simple external versioning to
             document operations.
         :type version_attribute: str
         :return: Document metadata (e.g. document key, revision) or True if

arango/connection.py

@@ -23,7 +23,7 @@
     JWTRefreshError,
     ServerConnectionError,
 )
-from arango.http import HTTPClient
+from arango.http import HTTPClient, RequestCompression
 from arango.request import Request
 from arango.resolver import HostResolver
 from arango.response import Response
@@ -44,6 +44,8 @@ def __init__(
         http_client: HTTPClient,
         serializer: Callable[..., str],
         deserializer: Callable[[str], Any],
+        request_compression: Optional[RequestCompression] = None,
+        response_compression: Optional[str] = None,
     ) -> None:
         self._url_prefixes = [f"{host}/_db/{db_name}" for host in hosts]
         self._host_resolver = host_resolver
@@ -53,6 +55,8 @@ def __init__(
         self._serializer = serializer
         self._deserializer = deserializer
         self._username: Optional[str] = None
+        self._request_compression = request_compression
+        self._response_compression = response_compression
 
     @property
     def db_name(self) -> str:
@@ -133,14 +137,27 @@ def process_request(
         """
         tries = 0
         indexes_to_filter: Set[int] = set()
+
+        data = self.normalize_data(request.data)
+        if (
+            self._request_compression is not None
+            and isinstance(data, str)
+            and self._request_compression.needs_compression(data)
+        ):
+            request.headers["content-encoding"] = self._request_compression.encoding()
+            data = self._request_compression.compress(data)
+
+        if self._response_compression is not None:
+            request.headers["accept-encoding"] = self._response_compression
+
         while tries < self._host_resolver.max_tries:
             try:
                 resp = self._http.send_request(
                     session=self._sessions[host_index],
                     method=request.method,
                     url=self._url_prefixes[host_index] + request.endpoint,
                     params=request.params,
-                    data=self.normalize_data(request.data),
+                    data=data,
                     headers=request.headers,
                     auth=auth,
                 )
@@ -243,6 +260,10 @@ class BasicConnection(BaseConnection):
     :type password: str
     :param http_client: User-defined HTTP client.
     :type http_client: arango.http.HTTPClient
+    :param: request_compression: The request compression algorithm.
+    :type request_compression: arango.http.RequestCompression | None
+    :param: response_compression: The response compression algorithm.
+    :type response_compression: str | None
     """
 
     def __init__(
@@ -256,6 +277,8 @@ def __init__(
         http_client: HTTPClient,
         serializer: Callable[..., str],
         deserializer: Callable[[str], Any],
+        request_compression: Optional[RequestCompression] = None,
+        response_compression: Optional[str] = None,
     ) -> None:
         super().__init__(
             hosts,
@@ -265,6 +288,8 @@ def __init__(
             http_client,
             serializer,
             deserializer,
+            request_compression,
+            response_compression,
         )
         self._username = username
         self._auth = (username, password)
@@ -298,6 +323,10 @@ class JwtConnection(BaseConnection):
     :type password: str
     :param http_client: User-defined HTTP client.
     :type http_client: arango.http.HTTPClient
+    :param request_compression: The request compression algorithm.
+    :type request_compression: arango.http.RequestCompression | None
+    :param response_compression: The response compression algorithm.
+    :type response_compression: str | None
     """
 
     def __init__(
@@ -312,6 +341,8 @@ def __init__(
         username: Optional[str] = None,
         password: Optional[str] = None,
         user_token: Optional[str] = None,
+        request_compression: Optional[RequestCompression] = None,
+        response_compression: Optional[str] = None,
     ) -> None:
         super().__init__(
             hosts,
@@ -321,6 +352,8 @@ def __init__(
             http_client,
             serializer,
             deserializer,
+            request_compression,
+            response_compression,
         )
         self._username = username
         self._password = password
@@ -439,6 +472,10 @@ class JwtSuperuserConnection(BaseConnection):
     :type http_client: arango.http.HTTPClient
     :param superuser_token: User generated token for superuser access.
     :type superuser_token: str
+    :param request_compression: The request compression algorithm.
+    :type request_compression: arango.http.RequestCompression | None
+    :param response_compression: The response compression algorithm.
+    :type response_compression: str | None
     """
 
     def __init__(
@@ -451,6 +488,8 @@ def __init__(
         serializer: Callable[..., str],
         deserializer: Callable[[str], Any],
         superuser_token: str,
+        request_compression: Optional[RequestCompression] = None,
+        response_compression: Optional[str] = None,
     ) -> None:
         super().__init__(
             hosts,
@@ -460,6 +499,8 @@ def __init__(
             http_client,
             serializer,
             deserializer,
+            request_compression,
+            response_compression,
         )
         self._auth_header = f"bearer {superuser_token}"