Merge pull request #17794 from [#21252] Enforce pubsub message publishing limits in the python SDK

* [#21252] Enforce pubsub message publishing limits in the python SDK

This improves visiblity as otherwise messages can become stuck publishing
inside the dataflow runner. This also allows users to handle errors by
catching them.

Fixes issue #21252 for python sdk.

* fixtest

Co-authored-by: Pablo <pabloem@users.noreply.github.com>

Author: scwhittle

sdks/python/apache_beam/io/gcp/pubsub.py

@@ -135,6 +135,9 @@ def _from_proto_str(proto_msg):
   def _to_proto_str(self, for_publish=False):
     """Get serialized form of ``PubsubMessage``.
 
+    The serialized message is validated against pubsub message limits specified
+    at https://cloud.google.com/pubsub/quotas#resource_limits
+
     Args:
       proto_msg: str containing a serialized protobuf.
       for_publish: bool, if True strip out message fields which cannot be
@@ -147,15 +150,39 @@ def _to_proto_str(self, for_publish=False):
       containing the payload of this object.
     """
     msg = pubsub.types.PubsubMessage()
+    if len(self.data) > (10 << 20):
+      raise ValueError('A pubsub message data field must not exceed 10MB')
     msg.data = self.data
-    for key, value in self.attributes.items():
-      msg.attributes[key] = value
-    if self.message_id and not for_publish:
-      msg.message_id = self.message_id
-    if self.publish_time and not for_publish:
-      msg.publish_time = self.publish_time
+
+    if self.attributes:
+      if len(self.attributes) > 100:
+        raise ValueError(
+            'A pubsub message must not have more than 100 attributes.')
+      for key, value in self.attributes.items():
+        if len(key) > 256:
+          raise ValueError(
+              'A pubsub message attribute key must not exceed 256 bytes.')
+        if len(value) > 1024:
+          raise ValueError(
+              'A pubsub message attribute value must not exceed 1024 bytes')
+        msg.attributes[key] = value
+
+    if not for_publish:
+      if self.message_id:
+        msg.message_id = self.message_id
+        if self.publish_time:
+          msg.publish_time = self.publish_time
+
+    if len(self.ordering_key) > 1024:
+      raise ValueError(
+          'A pubsub message ordering key must not exceed 1024 bytes.')
     msg.ordering_key = self.ordering_key
-    return pubsub.types.PubsubMessage.serialize(msg)
+
+    serialized = pubsub.types.PubsubMessage.serialize(msg)
+    if len(serialized) > (10 << 20):
+      raise ValueError(
+          'Serialized pubsub message exceeds the publish request limit of 10MB')
+    return serialized
 
   @staticmethod
   def _from_message(msg):
@@ -340,9 +367,8 @@ def message_to_proto_str(element):
   @staticmethod
   def bytes_to_proto_str(element):
     # type: (bytes) -> bytes
-    msg = pubsub.types.PubsubMessage()
-    msg.data = element
-    return pubsub.types.PubsubMessage.serialize(msg)
+    msg = PubsubMessage(element, {})
+    return msg._to_proto_str(for_publish=True)
 
   def expand(self, pcoll):
     if self.with_attributes:

sdks/python/apache_beam/io/gcp/pubsub_test.py

@@ -87,6 +87,27 @@ def test_proto_conversion(self):
     self.assertEqual(m_converted.data, data)
     self.assertEqual(m_converted.attributes, attributes)
 
+  @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
+  def test_payload_publish_invalid(self):
+    with self.assertRaisesRegex(ValueError, r'data field.*10MB'):
+      msg = PubsubMessage(b'0' * 1024 * 1024 * 11, None)
+      msg._to_proto_str(for_publish=True)
+    with self.assertRaisesRegex(ValueError, 'attribute key'):
+      msg = PubsubMessage(b'0', {'0' * 257: '0'})
+      msg._to_proto_str(for_publish=True)
+    with self.assertRaisesRegex(ValueError, 'attribute value'):
+      msg = PubsubMessage(b'0', {'0' * 100: '0' * 1025})
+      msg._to_proto_str(for_publish=True)
+    with self.assertRaisesRegex(ValueError, '100 attributes'):
+      attributes = {}
+      for i in range(0, 101):
+        attributes[str(i)] = str(i)
+      msg = PubsubMessage(b'0', attributes)
+      msg._to_proto_str(for_publish=True)
+    with self.assertRaisesRegex(ValueError, 'ordering key'):
+      msg = PubsubMessage(b'0', None, ordering_key='0' * 1301)
+      msg._to_proto_str(for_publish=True)
+
   def test_eq(self):
     a = PubsubMessage(b'abc', {1: 2, 3: 4})
     b = PubsubMessage(b'abc', {1: 2, 3: 4})