diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py b/sdks/python/apache_beam/io/gcp/pubsub.py index 55856f504787..16c758604eb4 100644 --- a/sdks/python/apache_beam/io/gcp/pubsub.py +++ b/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//topics/". 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) diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py b/sdks/python/apache_beam/io/gcp/pubsub_test.py index 5a8f2bf617b8..a4564c293521 100644 --- a/sdks/python/apache_beam/io/gcp/pubsub_test.py +++ b/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' diff --git a/sdks/python/apache_beam/yaml/yaml_io.py b/sdks/python/apache_beam/yaml/yaml_io.py index 3fa15044b60d..71f33098182f 100644 --- a/sdks/python/apache_beam/yaml/yaml_io.py +++ b/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 diff --git a/sdks/python/apache_beam/yaml/yaml_io_test.py b/sdks/python/apache_beam/yaml/yaml_io_test.py index dd1ce3ecaf85..d53d890d224a 100644 --- a/sdks/python/apache_beam/yaml/yaml_io_test.py +++ b/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: