feat: Add grpc Compression argument to channels and methods (#451)

* (feat): Add grpc Compression argument

* Add compression arg to channel creation

* fix linter errors

* fix linter errors

* refactor with new lib

* reformat

* fix tests

* add compression after refactor:

* fix lint

* fix unit tests

* fix unit tests

* fix operation

* remove unused import

* remove compression for grpc_gcp.secure_channel call

* fix method.py comment

* Update grpc_helpers.py

Remove experimental disclaimer

* Update grpc_helpers_async.py

Remove experimental disclaimer

* Update google/api_core/operations_v1/operations_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

* Update google/api_core/operations_v1/operations_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

* Update google/api_core/operations_v1/operations_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

* Update google/api_core/operations_v1/operations_async_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

* Update google/api_core/operations_v1/operations_async_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

* Update google/api_core/operations_v1/operations_async_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

* Update google/api_core/operations_v1/operations_async_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

* Update google/api_core/operations_v1/operations_client.py

Co-authored-by: Anthonios Partheniou <[email protected]>

---------

Co-authored-by: Anthonios Partheniou <[email protected]>

Author: acocuzzo

google/api_core/gapic_v1/method.py

@@ -15,7 +15,7 @@
 """Helpers for wrapping low-level gRPC methods with common functionality.
 
 This is used by gapic clients to provide common error mapping, retry, timeout,
-pagination, and long-running operations to gRPC methods.
+compression, pagination, and long-running operations to gRPC methods.
 """
 
 import enum
@@ -38,7 +38,7 @@ class _MethodDefault(enum.Enum):
 
 
 DEFAULT = _MethodDefault._DEFAULT_VALUE
-"""Sentinel value indicating that a retry or timeout argument was unspecified,
+"""Sentinel value indicating that a retry, timeout, or compression argument was unspecified,
 so the default should be used."""
 
 
@@ -72,27 +72,43 @@ class _GapicCallable(object):
             after its start, not to be confused with deadline). If ``None``,
             this callable will not specify a timeout argument to the low-level
             RPC method.
+        compression (grpc.Compression): The default compression for the callable.
+            If ``None``, this callable will not specify a compression argument
+            to the low-level RPC method.
         metadata (Sequence[Tuple[str, str]]): Additional metadata that is
             provided to the RPC method on every invocation. This is merged with
             any metadata specified during invocation. If ``None``, no
             additional metadata will be passed to the RPC method.
     """
 
-    def __init__(self, target, retry, timeout, metadata=None):
+    def __init__(
+        self,
+        target,
+        retry,
+        timeout,
+        compression,
+        metadata=None,
+    ):
         self._target = target
         self._retry = retry
         self._timeout = timeout
+        self._compression = compression
         self._metadata = metadata
 
-    def __call__(self, *args, timeout=DEFAULT, retry=DEFAULT, **kwargs):
-        """Invoke the low-level RPC with retry, timeout, and metadata."""
+    def __call__(
+        self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
+    ):
+        """Invoke the low-level RPC with retry, timeout, compression, and metadata."""
 
         if retry is DEFAULT:
             retry = self._retry
 
         if timeout is DEFAULT:
             timeout = self._timeout
 
+        if compression is DEFAULT:
+            compression = self._compression
+
         if isinstance(timeout, (int, float)):
             timeout = TimeToDeadlineTimeout(timeout=timeout)
 
@@ -109,6 +125,8 @@ def __call__(self, *args, timeout=DEFAULT, retry=DEFAULT, **kwargs):
             metadata = list(metadata)
             metadata.extend(self._metadata)
             kwargs["metadata"] = metadata
+        if self._compression is not None:
+            kwargs["compression"] = compression
 
         return wrapped_func(*args, **kwargs)
 
@@ -117,19 +135,21 @@ def wrap_method(
     func,
     default_retry=None,
     default_timeout=None,
+    default_compression=None,
     client_info=client_info.DEFAULT_CLIENT_INFO,
 ):
     """Wrap an RPC method with common behavior.
 
-    This applies common error wrapping, retry, and timeout behavior a function.
-    The wrapped function will take optional ``retry`` and ``timeout``
+    This applies common error wrapping, retry, timeout, and compression behavior to a function.
+    The wrapped function will take optional ``retry``, ``timeout``, and ``compression``
     arguments.
 
     For example::
 
         import google.api_core.gapic_v1.method
         from google.api_core import retry
         from google.api_core import timeout
+        from grpc import Compression
 
         # The original RPC method.
         def get_topic(name, timeout=None):
@@ -138,6 +158,7 @@ def get_topic(name, timeout=None):
 
         default_retry = retry.Retry(deadline=60)
         default_timeout = timeout.Timeout(deadline=60)
+        default_compression = Compression.NoCompression
         wrapped_get_topic = google.api_core.gapic_v1.method.wrap_method(
             get_topic, default_retry)
 
@@ -186,6 +207,9 @@ def get_topic(name, timeout=None):
         default_timeout (Optional[google.api_core.Timeout]): The default
             timeout strategy. Can also be specified as an int or float. If
             ``None``, the method will not have timeout specified by default.
+        default_compression (Optional[grpc.Compression]): The default
+            grpc.Compression. If ``None``, the method will not have
+            compression specified by default.
         client_info
             (Optional[google.api_core.gapic_v1.client_info.ClientInfo]):
                 Client information used to create a user-agent string that's
@@ -194,19 +218,23 @@ def get_topic(name, timeout=None):
                 metadata will be provided to the RPC method.
 
     Returns:
-        Callable: A new callable that takes optional ``retry`` and ``timeout``
-            arguments and applies the common error mapping, retry, timeout,
+        Callable: A new callable that takes optional ``retry``, ``timeout``,
+            and ``compression``
+            arguments and applies the common error mapping, retry, timeout, compression,
             and metadata behavior to the low-level RPC method.
     """
     func = grpc_helpers.wrap_errors(func)
-
     if client_info is not None:
         user_agent_metadata = [client_info.to_grpc_metadata()]
     else:
         user_agent_metadata = None
 
     return functools.wraps(func)(
         _GapicCallable(
-            func, default_retry, default_timeout, metadata=user_agent_metadata
+            func,
+            default_retry,
+            default_timeout,
+            default_compression,
+            metadata=user_agent_metadata,
         )
     )

google/api_core/gapic_v1/method_async.py

@@ -14,7 +14,7 @@
 """AsyncIO helpers for wrapping gRPC methods with common functionality.
 
 This is used by gapic clients to provide common error mapping, retry, timeout,
-pagination, and long-running operations to gRPC methods.
+compression, pagination, and long-running operations to gRPC methods.
 """
 
 import functools
@@ -30,19 +30,26 @@ def wrap_method(
     func,
     default_retry=None,
     default_timeout=None,
+    default_compression=None,
     client_info=client_info.DEFAULT_CLIENT_INFO,
 ):
     """Wrap an async RPC method with common behavior.
 
     Returns:
-        Callable: A new callable that takes optional ``retry`` and ``timeout``
-            arguments and applies the common error mapping, retry, timeout,
-            and metadata behavior to the low-level RPC method.
+        Callable: A new callable that takes optional ``retry``, ``timeout``,
+            and ``compression`` arguments and applies the common error mapping,
+            retry, timeout, metadata, and compression behavior to the low-level RPC method.
     """
     func = grpc_helpers_async.wrap_errors(func)
 
     metadata = [client_info.to_grpc_metadata()] if client_info is not None else None
 
     return functools.wraps(func)(
-        _GapicCallable(func, default_retry, default_timeout, metadata=metadata)
+        _GapicCallable(
+            func,
+            default_retry,
+            default_timeout,
+            default_compression,
+            metadata=metadata,
+        )
     )

google/api_core/grpc_helpers.py

@@ -16,6 +16,7 @@
 
 import collections
 import functools
+import logging
 import warnings
 
 import grpc
@@ -51,6 +52,8 @@
 # The list of gRPC Callable interfaces that return iterators.
 _STREAM_WRAP_CLASSES = (grpc.UnaryStreamMultiCallable, grpc.StreamStreamMultiCallable)
 
+_LOGGER = logging.getLogger(__name__)
+
 
 def _patch_callable_name(callable_):
     """Fix-up gRPC callable attributes.
@@ -276,7 +279,8 @@ def create_channel(
     quota_project_id=None,
     default_scopes=None,
     default_host=None,
-    **kwargs
+    compression=None,
+    **kwargs,
 ):
     """Create a secure channel with credentials.
 
@@ -297,6 +301,8 @@ def create_channel(
         default_scopes (Sequence[str]): Default scopes passed by a Google client
             library. Use 'scopes' for user-defined scopes.
         default_host (str): The default endpoint. e.g., "pubsub.googleapis.com".
+        compression (grpc.Compression): An optional value indicating the
+            compression method to be used over the lifetime of the channel.
         kwargs: Additional key-word args passed to
             :func:`grpc_gcp.secure_channel` or :func:`grpc.secure_channel`.
             Note: `grpc_gcp` is only supported in environments with protobuf < 4.0.0.
@@ -319,12 +325,18 @@ def create_channel(
     )
 
     if HAS_GRPC_GCP:  # pragma: NO COVER
+        if compression is not None and compression != grpc.Compression.NoCompression:
+            _LOGGER.debug(
+                "Compression argument is being ignored for grpc_gcp.secure_channel creation."
+            )
         return grpc_gcp.secure_channel(target, composite_credentials, **kwargs)
-    return grpc.secure_channel(target, composite_credentials, **kwargs)
+    return grpc.secure_channel(
+        target, composite_credentials, compression=compression, **kwargs
+    )
 
 
 _MethodCall = collections.namedtuple(
-    "_MethodCall", ("request", "timeout", "metadata", "credentials")
+    "_MethodCall", ("request", "timeout", "metadata", "credentials", "compression")
 )
 
 _ChannelRequest = collections.namedtuple("_ChannelRequest", ("method", "request"))
@@ -351,11 +363,15 @@ def __init__(self, method, channel):
         """List[protobuf.Message]: All requests sent to this callable."""
         self.calls = []
         """List[Tuple]: All invocations of this callable. Each tuple is the
-        request, timeout, metadata, and credentials."""
+        request, timeout, metadata, compression, and credentials."""
 
-    def __call__(self, request, timeout=None, metadata=None, credentials=None):
+    def __call__(
+        self, request, timeout=None, metadata=None, credentials=None, compression=None
+    ):
         self._channel.requests.append(_ChannelRequest(self._method, request))
-        self.calls.append(_MethodCall(request, timeout, metadata, credentials))
+        self.calls.append(
+            _MethodCall(request, timeout, metadata, credentials, compression)
+        )
         self.requests.append(request)
 
         response = self.response

google/api_core/grpc_helpers_async.py

@@ -212,6 +212,7 @@ def create_channel(
     quota_project_id=None,
     default_scopes=None,
     default_host=None,
+    compression=None,
     **kwargs
 ):
     """Create an AsyncIO secure channel with credentials.
@@ -233,6 +234,8 @@ def create_channel(
         default_scopes (Sequence[str]): Default scopes passed by a Google client
             library. Use 'scopes' for user-defined scopes.
         default_host (str): The default endpoint. e.g., "pubsub.googleapis.com".
+        compression (grpc.Compression): An optional value indicating the
+            compression method to be used over the lifetime of the channel.
         kwargs: Additional key-word args passed to :func:`aio.secure_channel`.
 
     Returns:
@@ -252,7 +255,9 @@ def create_channel(
         default_host=default_host,
     )
 
-    return aio.secure_channel(target, composite_credentials, **kwargs)
+    return aio.secure_channel(
+        target, composite_credentials, compression=compression, **kwargs
+    )
 
 
 class FakeUnaryUnaryCall(_WrappedUnaryUnaryCall):

google/api_core/operation.py

@@ -315,10 +315,16 @@ def from_grpc(operation, operations_stub, result_type, grpc_metadata=None, **kwa
             operation.
     """
     refresh = functools.partial(
-        _refresh_grpc, operations_stub, operation.name, metadata=grpc_metadata
+        _refresh_grpc,
+        operations_stub,
+        operation.name,
+        metadata=grpc_metadata,
     )
     cancel = functools.partial(
-        _cancel_grpc, operations_stub, operation.name, metadata=grpc_metadata
+        _cancel_grpc,
+        operations_stub,
+        operation.name,
+        metadata=grpc_metadata,
     )
     return Operation(operation, refresh, cancel, result_type, **kwargs)
 
@@ -347,9 +353,13 @@ def from_gapic(operation, operations_client, result_type, grpc_metadata=None, **
             operation.
     """
     refresh = functools.partial(
-        operations_client.get_operation, operation.name, metadata=grpc_metadata
+        operations_client.get_operation,
+        operation.name,
+        metadata=grpc_metadata,
     )
     cancel = functools.partial(
-        operations_client.cancel_operation, operation.name, metadata=grpc_metadata
+        operations_client.cancel_operation,
+        operation.name,
+        metadata=grpc_metadata,
     )
     return Operation(operation, refresh, cancel, result_type, **kwargs)

google/api_core/operation_async.py

@@ -213,9 +213,13 @@ def from_gapic(operation, operations_client, result_type, grpc_metadata=None, **
             operation.
     """
     refresh = functools.partial(
-        operations_client.get_operation, operation.name, metadata=grpc_metadata
+        operations_client.get_operation,
+        operation.name,
+        metadata=grpc_metadata,
     )
     cancel = functools.partial(
-        operations_client.cancel_operation, operation.name, metadata=grpc_metadata
+        operations_client.cancel_operation,
+        operation.name,
+        metadata=grpc_metadata,
     )
     return AsyncOperation(operation, refresh, cancel, result_type, **kwargs)