Feature: Add grpc compression on publish

google/cloud/pubsub_v1/publisher/_batch/thread.py

@@ -26,6 +26,7 @@
 from google.cloud.pubsub_v1.publisher import futures
 from google.cloud.pubsub_v1.publisher._batch import base
 from google.pubsub_v1 import types as gapic_types
+import grpc
 
 if typing.TYPE_CHECKING:  # pragma: NO COVER
     from google.cloud import pubsub_v1
@@ -118,6 +119,12 @@ def __init__(
 
         self._commit_retry = commit_retry
         self._commit_timeout = commit_timeout
+        self._enable_grpc_compression = (
+            self.client.publisher_options.enable_grpc_compression
+        )
+        self._compression_bytes_threshold = (
+            self.client.publisher_options.compression_bytes_threshold
+        )
 
     @staticmethod
     def make_lock() -> threading.Lock:
@@ -269,13 +276,25 @@ def _commit(self) -> None:
         start = time.time()
 
         batch_transport_succeeded = True
+
+        # Set compression if enabled.
+        compression = None
+
+        if (
+            self._enable_grpc_compression
+            and gapic_types.PublishRequest(messages=self._messages)._pb.ByteSize()
+            >= self._compression_bytes_threshold
+        ):
+            compression = grpc.Compression.Gzip
+
         try:
             # Performs retries for errors defined by the retry configuration.
             response = self._client._gapic_publish(
                 topic=self._topic,
                 messages=self._messages,
                 retry=self._commit_retry,
                 timeout=self._commit_timeout,
+                compression=compression,
             )
         except google.api_core.exceptions.GoogleAPIError as exc:
             # We failed to publish, even after retries, so set the exception on

google/cloud/pubsub_v1/publisher/client.py

@@ -153,6 +153,11 @@ def __init__(
         # The object controlling the message publishing flow
         self._flow_controller = FlowController(self.publisher_options.flow_control)
 
+        self._enable_grpc_compression = self.publisher_options.enable_grpc_compression
+        self._compression_bytes_threshold = (
+            self.publisher_options.compression_bytes_threshold
+        )
+
     @classmethod
     def from_service_account_file(  # type: ignore[override]
         cls,

google/cloud/pubsub_v1/types.py

@@ -173,6 +173,9 @@ class PublisherOptions(NamedTuple):
         "Timeout settings for message publishing by the client. It should be "
         "compatible with :class:`~.pubsub_v1.types.TimeoutType`."
     )
+    enable_grpc_compression: bool = False
+
+    compression_bytes_threshold: int = 240
 
 
 # Define the type class and default values for flow control settings.

google/pubsub_v1/services/publisher/client.py

@@ -737,6 +737,7 @@ def publish(
         retry: OptionalRetry = gapic_v1.method.DEFAULT,
         timeout: TimeoutType = gapic_v1.method.DEFAULT,
         metadata: Sequence[Tuple[str, str]] = (),
+        compression: grpc.Compression = None,
     ) -> pubsub.PublishResponse:
         r"""Adds one or more messages to the topic. Returns ``NOT_FOUND`` if
         the topic does not exist.
@@ -833,6 +834,7 @@ def sample_publish():
             retry=retry,
             timeout=timeout,
             metadata=metadata,
+            compression=compression,
         )
 
         # Done; return the response.

samples/snippets/main.py

@@ -0,0 +1,5 @@
+from google.cloud import pubsub_v1
+pub_sub_client = pubsub_v1.PublisherClient()
+
+if __name__ == '__main__':
+    result = pub_sub_client.publish("projects/annaco-python-lib-test/topics/annaco-python-lib-test-topic", bytes("Some message here!", "utf8")).result()

samples/snippets/publisher.py

@@ -216,6 +216,76 @@ def callback(future: pubsub_v1.publisher.futures.Future) -> None:
     # [END pubsub_publisher_batch_settings]
 
 
+def publish_messages_with_default_compression_threshold(project_id: str, topic_id: str) -> None:
+    """Publishes messages to a Pub/Sub topic with grpc compression enabled."""
+    # [START pubsub_publisher_compression_settings]
+    from concurrent import futures
+    from google.cloud import pubsub_v1
+
+    # TODO(developer)
+    # project_id = "your-project-id"
+    # topic_id = "your-topic-id"
+
+    # Configure publisher with compression
+    publisher = pubsub_v1.PublisherClient(publisher_options=pubsub_v1.types.PublisherOptions(enable_grpc_compression=True))
+    topic_path = publisher.topic_path(project_id, topic_id)
+    publish_futures = []
+
+    # Resolve the publish future in a separate thread.
+    def callback(future: pubsub_v1.publisher.futures.Future) -> None:
+        message_id = future.result()
+        print(message_id)
+
+    for n in range(1, 10):
+        data_str = f"Message number {n}"
+        # Data must be a bytestring
+        data = data_str.encode("utf-8")
+        publish_future = publisher.publish(topic_path, data)
+        # Non-blocking. Allow the publisher client to batch multiple messages.
+        publish_future.add_done_callback(callback)
+        publish_futures.append(publish_future)
+
+    futures.wait(publish_futures, return_when=futures.ALL_COMPLETED)
+
+    print(f"Published messages with compression settings to {topic_path}.")
+    # [END pubsub_publisher_compression_settings]
+
+def publish_messages_with_low_compression_threshold(project_id: str, topic_id: str) -> None:
+    """Publishes messages to a Pub/Sub topic with grpc compression enabled."""
+    # [START pubsub_publisher_compression_settings]
+    from concurrent import futures
+    from google.cloud import pubsub_v1
+
+    # TODO(developer)
+    # project_id = "your-project-id"
+    # topic_id = "your-topic-id"
+
+    # Configure publisher with compression
+    publisher = pubsub_v1.PublisherClient(publisher_options=pubsub_v1.types.PublisherOptions(enable_grpc_compression=True, compression_bytes_threshold=0))
+    topic_path = publisher.topic_path(project_id, topic_id)
+    publish_futures = []
+
+    # Resolve the publish future in a separate thread.
+    def callback(future: pubsub_v1.publisher.futures.Future) -> None:
+        message_id = future.result()
+        print(message_id)
+
+    for n in range(1, 10):
+        data_str = f"Message number {n}"
+        # Data must be a bytestring
+        data = data_str.encode("utf-8")
+        publish_future = publisher.publish(topic_path, data)
+        # Non-blocking. Allow the publisher client to batch multiple messages.
+        publish_future.add_done_callback(callback)
+        publish_futures.append(publish_future)
+
+    futures.wait(publish_futures, return_when=futures.ALL_COMPLETED)
+
+    print(f"Published messages with compression settings to {topic_path}.")
+    # [END pubsub_publisher_compression_settings]
+
+
+
 def publish_messages_with_flow_control_settings(project_id: str, topic_id: str) -> None:
     """Publishes messages to a Pub/Sub topic with flow control settings."""
     # [START pubsub_publisher_flow_control]

samples/snippets/publisher_test.py

@@ -192,6 +192,24 @@ def test_publish_with_ordering_keys(
     assert f"Published messages with ordering keys to {topic_path}." in out
 
 
+def test_publish_with_default_compression(
+    topic_path: str, capsys: CaptureFixture[str]
+) -> None:
+    publisher.publish_messages_with_default_compression_threshold(PROJECT_ID, TOPIC_ID)
+
+    out, _ = capsys.readouterr()
+    assert f"Published messages with compression settings to {topic_path}." in out
+
+
+def test_publish_with_low_compression(
+    topic_path: str, capsys: CaptureFixture[str]
+) -> None:
+    publisher.publish_messages_with_default_compression_threshold(PROJECT_ID, TOPIC_ID)
+
+    out, _ = capsys.readouterr()
+    assert f"Published messages with compression settings to {topic_path}." in out
+
+
 def test_resume_publish_with_error_handler(
     topic_path: str, capsys: CaptureFixture[str]
 ) -> None:

samples/snippets/sponge_log.xml -v

@@ -0,0 +1 @@
+<?xml version="1.0" encoding="utf-8"?><testsuites><testsuite name="pytest" errors="0" failures="0" skipped="0" tests="24" time="23.343" timestamp="2022-04-22T15:12:12.032784" hostname="cs-872183299652-default"><testcase classname="samples.snippets.iam_test" name="test_get_topic_policy" time="3.244" /><testcase classname="samples.snippets.iam_test" name="test_get_subscription_policy" time="1.090" /><testcase classname="samples.snippets.iam_test" name="test_set_topic_policy" time="0.401" /><testcase classname="samples.snippets.iam_test" name="test_set_subscription_policy" time="0.400" /><testcase classname="samples.snippets.iam_test" name="test_check_topic_permissions" time="0.053" /><testcase classname="samples.snippets.iam_test" name="test_check_subscription_permissions" time="0.497" /><testcase classname="samples.snippets.publisher_test" name="test_create" time="1.904" /><testcase classname="samples.snippets.publisher_test" name="test_list" time="0.124" /><testcase classname="samples.snippets.publisher_test" name="test_publish" time="1.317" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_custom_attributes" time="0.858" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_batch_settings" time="1.261" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_flow_control_settings" time="0.836" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_retry_settings" time="0.572" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_error_handler" time="0.064" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_ordering_keys" time="1.423" /><testcase classname="samples.snippets.publisher_test" name="test_resume_publish_with_error_handler" time="0.562" /><testcase classname="samples.snippets.publisher_test" name="test_detach_subscription" time="1.520" /><testcase classname="samples.snippets.publisher_test" name="test_delete" time="0.607" /><testcase classname="samples.snippets.schema_test" name="test_create_avro_schema" time="0.522" /><testcase classname="samples.snippets.schema_test" name="test_create_proto_schema" time="0.225" /><testcase classname="samples.snippets.schema_test" name="test_get_schema" time="0.097" /><testcase classname="samples.snippets.schema_test" name="test_list_schemas" time="0.076" /><testcase classname="samples.snippets.schema_test" name="test_create_topic_with_schema" time="2.364" /><testcase classname="samples.snippets.schema_test" name="test_publish_avro_records" time="0.499" /><testcase time="0.925" /></testsuite></testsuites>

samples/snippets/sponge_log.xml -v -s

@@ -0,0 +1 @@
+<?xml version="1.0" encoding="utf-8"?><testsuites><testsuite name="pytest" errors="0" failures="0" skipped="0" tests="51" time="244.740" timestamp="2022-04-22T15:13:22.043387" hostname="cs-872183299652-default"><testcase classname="samples.snippets.iam_test" name="test_get_topic_policy" time="2.263" /><testcase classname="samples.snippets.iam_test" name="test_get_subscription_policy" time="1.700" /><testcase classname="samples.snippets.iam_test" name="test_set_topic_policy" time="0.394" /><testcase classname="samples.snippets.iam_test" name="test_set_subscription_policy" time="0.498" /><testcase classname="samples.snippets.iam_test" name="test_check_topic_permissions" time="0.067" /><testcase classname="samples.snippets.iam_test" name="test_check_subscription_permissions" time="0.688" /><testcase classname="samples.snippets.publisher_test" name="test_create" time="2.598" /><testcase classname="samples.snippets.publisher_test" name="test_list" time="0.118" /><testcase classname="samples.snippets.publisher_test" name="test_publish" time="1.426" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_custom_attributes" time="1.010" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_batch_settings" time="1.071" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_flow_control_settings" time="0.671" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_retry_settings" time="0.435" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_error_handler" time="0.072" /><testcase classname="samples.snippets.publisher_test" name="test_publish_with_ordering_keys" time="1.472" /><testcase classname="samples.snippets.publisher_test" name="test_resume_publish_with_error_handler" time="0.587" /><testcase classname="samples.snippets.publisher_test" name="test_detach_subscription" time="1.885" /><testcase classname="samples.snippets.publisher_test" name="test_delete" time="0.749" /><testcase classname="samples.snippets.schema_test" name="test_create_avro_schema" time="0.395" /><testcase classname="samples.snippets.schema_test" name="test_create_proto_schema" time="0.330" /><testcase classname="samples.snippets.schema_test" name="test_get_schema" time="0.092" /><testcase classname="samples.snippets.schema_test" name="test_list_schemas" time="0.087" /><testcase classname="samples.snippets.schema_test" name="test_create_topic_with_schema" time="2.238" /><testcase classname="samples.snippets.schema_test" name="test_publish_avro_records" time="0.398" /><testcase classname="samples.snippets.schema_test" name="test_subscribe_with_avro_schema" time="10.028" /><testcase classname="samples.snippets.schema_test" name="test_publish_proto_records" time="2.963" /><testcase classname="samples.snippets.schema_test" name="test_subscribe_with_proto_schema" time="10.451" /><testcase classname="samples.snippets.schema_test" name="test_delete_schema" time="2.091" /><testcase classname="samples.snippets.subscriber_test" name="test_list_in_topic" time="3.186" /><testcase classname="samples.snippets.subscriber_test" name="test_list_in_project" time="0.079" /><testcase classname="samples.snippets.subscriber_test" name="test_create_subscription" time="1.097" /><testcase classname="samples.snippets.subscriber_test" name="test_create_subscription_with_dead_letter_policy" time="4.799" /><testcase classname="samples.snippets.subscriber_test" name="test_receive_with_delivery_attempts" time="90.934" /><testcase classname="samples.snippets.subscriber_test" name="test_update_dead_letter_policy" time="0.234" /><testcase classname="samples.snippets.subscriber_test" name="test_remove_dead_letter_policy" time="0.285" /><testcase classname="samples.snippets.subscriber_test" name="test_create_subscription_with_ordering" time="1.163" /><testcase classname="samples.snippets.subscriber_test" name="test_create_subscription_with_filtering" time="1.167" /><testcase classname="samples.snippets.subscriber_test" name="test_create_subscription_with_exactly_once_delivery" time="6.503" /><testcase classname="samples.snippets.subscriber_test" name="test_create_push_subscription" time="2.292" /><testcase classname="samples.snippets.subscriber_test" name="test_update_push_suscription" time="0.380" /><testcase classname="samples.snippets.subscriber_test" name="test_delete_subscription" time="0.303" /><testcase classname="samples.snippets.subscriber_test" name="test_receive" time="6.855" /><testcase classname="samples.snippets.subscriber_test" name="test_receive_with_custom_attributes" time="5.221" /><testcase classname="samples.snippets.subscriber_test" name="test_receive_with_flow_control" time="5.257" /><testcase classname="samples.snippets.subscriber_test" name="test_receive_with_blocking_shutdown" time="8.700" /><testcase classname="samples.snippets.subscriber_test" name="test_receive_messages_with_exactly_once_delivery_enabled" time="12.036" /><testcase classname="samples.snippets.subscriber_test" name="test_listen_for_errors" time="5.240" /><testcase classname="samples.snippets.subscriber_test" name="test_receive_synchronously" time="3.541" /><testcase classname="samples.snippets.subscriber_test" name="test_receive_synchronously_with_lease" time="23.804" /><testcase classname="samples.snippets.quickstart.quickstart_test" name="test_pub" time="2.808" /><testcase classname="samples.snippets.quickstart.quickstart_test" name="test_sub" time="11.505" /></testsuite></testsuites>

samples/snippets/subscriber_test.py

@@ -61,6 +61,13 @@
 def publisher_client() -> Generator[pubsub_v1.PublisherClient, None, None]:
     yield pubsub_v1.PublisherClient()
 
+@pytest.fixture(scope="module")
+def publisher_client_with_default_compression() -> Generator[pubsub_v1.PublisherClient, None, None]:
+    yield pubsub_v1.PublisherClient(publisher_options=pubsub_v1.types.PublisherOptions(enable_grpc_compression=True)) 
+
+@pytest.fixture(scope="module")
+def publisher_client_with_low_compression() -> Generator[pubsub_v1.PublisherClient, None, None]:
+    yield pubsub_v1.PublisherClient(publisher_options=pubsub_v1.types.PublisherOptions(enable_grpc_compression=True, compression_bytes_threshold=0))
 
 @pytest.fixture(scope="module")
 def regional_publisher_client() -> Generator[pubsub_v1.PublisherClient, None, None]:
@@ -974,6 +981,36 @@ def test_listen_for_errors(
     subscriber_client.delete_subscription(request={"subscription": subscription_path})
 
 
+def test_listen_for_errors_default_compression(
+    publisher_client_with_default_compression: pubsub_v1.PublisherClient,
+    topic: str,
+    subscription_async: str,
+    capsys: CaptureFixture[str],
+) -> None:
+    _ = _publish_messages(publisher_client_with_default_compression, topic)
+
+    subscriber.listen_for_errors(PROJECT_ID, SUBSCRIPTION_ASYNC, 5)
+
+    out, _ = capsys.readouterr()
+    assert subscription_async in out
+    assert "threw an exception" in out
+
+
+def test_listen_for_errors_low_compression(
+    publisher_client_with_low_compression: pubsub_v1.PublisherClient,
+    topic: str,
+    subscription_async: str,
+    capsys: CaptureFixture[str],
+) -> None:
+    _ = _publish_messages(publisher_client_with_low_compression, topic)
+
+    subscriber.listen_for_errors(PROJECT_ID, SUBSCRIPTION_ASYNC, 5)
+
+    out, _ = capsys.readouterr()
+    assert subscription_async in out
+    assert "threw an exception" in out
+
+
 def test_receive_synchronously(
     subscriber_client: pubsub_v1.SubscriberClient,
     publisher_client: pubsub_v1.PublisherClient,