[pubsub/yaml] - fix pubsub string issue

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

@@ -150,7 +150,10 @@ def _to_proto_str(self, for_publish=False):
       https://cloud.google.com/pubsub/docs/reference/rpc/google.pubsub.v1#google.pubsub.v1.PubsubMessage
       containing the payload of this object.
     """
-    if len(self.data) > (10_000_000):
+    data = self.data
+    if isinstance(data, str):
+      data = data.encode('utf-8')
+    if data is not None and len(data) > (10_000_000):
       raise ValueError('A pubsub message data field must not exceed 10MB')
 
     if self.attributes:
@@ -178,7 +181,7 @@ def _to_proto_str(self, for_publish=False):
           'A pubsub message ordering key must not exceed 1024 bytes.')
 
     msg = pubsub.types.PubsubMessage(
-        data=self.data,
+        data=data,
         attributes=self.attributes,
         message_id=message_id,
         publish_time=publish_time,
@@ -399,8 +402,8 @@ def __init__(
           "projects/<project>/topics/<topic>".
       with_attributes:
         True - input elements will be :class:`~PubsubMessage` objects.
-        False - input elements will be of type ``bytes`` (message
-        data only).
+        False - input elements will be of type ``bytes`` or ``str`` (UTF-8
+        encoded; message data only).
       id_label: If set, will set an attribute for each Cloud Pub/Sub message
         with the given name and a unique value. This attribute can then be used
         in a ReadFromPubSub PTransform to deduplicate messages.
@@ -430,6 +433,8 @@ def message_to_proto_str(element: PubsubMessage) -> bytes:
 
   @staticmethod
   def bytes_to_proto_str(element: Union[bytes, str]) -> bytes:
+    if isinstance(element, str):
+      element = element.encode('utf-8')
     msg = PubsubMessage(element, {})
     return msg._to_proto_str(for_publish=True)
 

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

@@ -90,6 +90,16 @@ def test_proto_conversion(self):
     self.assertEqual(m_converted.data, data)
     self.assertEqual(m_converted.attributes, attributes)
 
+    m_str = PubsubMessage('hello', attributes)
+    m_str_converted = PubsubMessage._from_proto_str(m_str._to_proto_str())
+    self.assertEqual(m_str_converted.data, b'hello')
+    self.assertEqual(m_str_converted.attributes, attributes)
+
+    m_none = PubsubMessage(None, attributes)
+    m_none_converted = PubsubMessage._from_proto_str(m_none._to_proto_str())
+    self.assertEqual(m_none_converted.data, b'')
+    self.assertEqual(m_none_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'):
@@ -980,6 +990,7 @@ def test_write_messages_success(self, mock_pubsub):
     # Check that the call was made with the topic and some data
     call_args = mock_pubsub.return_value.publish.call_args
     self.assertEqual(len(call_args[0]), 2)  # topic and data
+    self.assertEqual(call_args[0][1], b'data')
 
   def test_write_messages_deprecated(self, mock_pubsub):
     data = 'data'
@@ -997,6 +1008,7 @@ def test_write_messages_deprecated(self, mock_pubsub):
     # Check that the call was made with the topic and some data
     call_args = mock_pubsub.return_value.publish.call_args
     self.assertEqual(len(call_args[0]), 2)  # topic and data
+    self.assertEqual(call_args[0][1], b'data')
 
   def test_write_messages_with_attributes_success(self, mock_pubsub):
     data = b'data'
@@ -1016,6 +1028,7 @@ def test_write_messages_with_attributes_success(self, mock_pubsub):
     # Check that the call was made with the topic and some data
     call_args = mock_pubsub.return_value.publish.call_args
     self.assertEqual(len(call_args[0]), 2)  # topic and data
+    self.assertEqual(call_args[0][1], b'data')
 
   def test_write_messages_batch_mode_success(self, mock_pubsub):
     """Test WriteToPubSub works in batch mode (non-streaming)."""
@@ -1037,6 +1050,44 @@ def test_write_messages_batch_mode_success(self, mock_pubsub):
     # Check that the call was made with the topic and some data
     call_args = mock_pubsub.return_value.publish.call_args
     self.assertEqual(len(call_args[0]), 2)  # topic and data
+    self.assertEqual(call_args[0][1], b'data')
+
+  def test_write_string_messages_utf8_batch_and_streaming(self, mock_pubsub):
+    """Test WriteToPubSub UTF-8 encodes non-base64 and non-ASCII strings."""
+    payloads = ['hello', 'äter']
+    for streaming in (False, True):
+      mock_pubsub.reset_mock()
+      options = PipelineOptions([])
+      options.view_as(StandardOptions).streaming = streaming
+      with TestPipeline(options=options) as p:
+        _ = (
+            p
+            | Create(payloads)
+            | WriteToPubSub(
+                'projects/fakeprj/topics/a_topic', with_attributes=False))
+
+      published_data = [
+          c.args[1] for c in mock_pubsub.return_value.publish.call_args_list
+      ]
+      self.assertCountEqual(published_data, [b'hello', 'äter'.encode('utf-8')])
+
+  def test_write_messages_with_str_data_and_attributes_batch_mode(
+      self, mock_pubsub):
+    """Test PubsubMessage with str data is UTF-8 encoded when published."""
+    payloads = [PubsubMessage('hello', {'key': 'value'})]
+    options = PipelineOptions([])
+    options.view_as(StandardOptions).streaming = False
+    with TestPipeline(options=options) as p:
+      _ = (
+          p
+          | Create(payloads)
+          | WriteToPubSub(
+              'projects/fakeprj/topics/a_topic', with_attributes=True))
+
+    mock_pubsub.return_value.publish.assert_called_once()
+    call_args = mock_pubsub.return_value.publish.call_args
+    self.assertEqual(call_args.args[1], b'hello')
+    self.assertEqual(call_args.kwargs, {'key': 'value'})
 
   def test_write_messages_with_attributes_batch_mode_success(self, mock_pubsub):
     """Test WriteToPubSub with attributes works in batch mode."""
@@ -1059,6 +1110,7 @@ def test_write_messages_with_attributes_batch_mode_success(self, mock_pubsub):
     # Check that the call was made with the topic and some data
     call_args = mock_pubsub.return_value.publish.call_args
     self.assertEqual(len(call_args[0]), 2)  # topic and data
+    self.assertEqual(call_args[0][1], b'data')
 
   def test_write_messages_with_attributes_error(self, mock_pubsub):
     data = 'data'

sdks/python/apache_beam/yaml/yaml_io.py

@@ -299,6 +299,24 @@ def convert_to_bytes(row):
             f"got {repr(output)} of type {type(output)}.")
 
     return convert_to_bytes
+  elif format == 'STRING':
+    if schema:
+      raise ValueError('STRING format does not take a schema')
+    field_names = [field.name for field in beam_schema.fields]
+    if len(field_names) != 1:
+      raise ValueError(f'Expecting exactly one field, found {field_names}')
+
+    def encode_string(row):
+      output = getattr(row, field_names[0])
+      if isinstance(output, str):
+        return output.encode('utf-8')
+      else:
+        raise ValueError(
+            f"Cannot encode payload for WriteToPubSub with STRING format. "
+            f"Expected valid string object, "
+            f"got {repr(output)} of type {type(output)}.")
+
+    return encode_string
   elif format == 'JSON':
     return json_utils.json_formater(beam_schema)
   elif format == 'AVRO':
@@ -459,6 +477,8 @@ def write_to_pubsub(
         - RAW: Expects a message with a single field (excluding
             attribute-related fields) whose contents are used as the raw bytes
             of the pubsub message.
+        - STRING: Like RAW, but expects a single string field and encodes it as
+            UTF-8 bytes.
         - AVRO: Encodes records with a given Avro schema, which may be inferred
             from the input PCollection schema.
         - JSON: Formats records with a given JSON schema, which may be inferred

sdks/python/apache_beam/yaml/yaml_io_test.py

@@ -494,6 +494,60 @@ def test_simple_write(self):
               format: RAW
             '''))
 
+  def test_simple_write_string(self):
+    with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+        pickle_library='cloudpickle')) as p:
+      with mock.patch('apache_beam.io.WriteToPubSub',
+                      FakeWriteToPubSub(
+                          topic='my_topic',
+                          messages=[PubsubMessage('äter'.encode('utf-8'), {}),
+                                    PubsubMessage('köttbullar'.encode('utf-8'),
+                                                  {})])):
+        _ = (
+            p
+            | beam.Create([beam.Row(a='äter'), beam.Row(a='köttbullar')])
+            | YamlTransform(
+                '''
+            type: WriteToPubSub
+            config:
+              topic: my_topic
+              format: STRING
+            '''))
+
+  def test_write_string_rejects_schema(self):
+    with self.assertRaisesRegex(ValueError,
+                                'STRING format does not take a schema'):
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        _ = (
+            p | beam.Create([beam.Row(a='hello')])
+            | YamlTransform(
+                '''
+            type: WriteToPubSub
+            config:
+              topic: my_topic
+              format: STRING
+              schema:
+                type: object
+            '''))
+
+  def test_write_string_rejects_non_string(self):
+    with self.assertRaisesRegex((ValueError, RuntimeError),
+                                'Expected valid string object'):
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        with mock.patch('apache_beam.io.WriteToPubSub',
+                        FakeWriteToPubSub(topic='my_topic', messages=[])):
+          _ = (
+              p | beam.Create([beam.Row(a=b'bytes_not_str')])
+              | YamlTransform(
+                  '''
+              type: WriteToPubSub
+              config:
+                topic: my_topic
+                format: STRING
+              '''))
+
   def test_write_with_attribute(self):
     with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
         pickle_library='cloudpickle')) as p: