Add a STRING format to PubSub reading that interpretes the payload as utf-8 encoded. (#34301)

Author: robertwb

sdks/python/apache_beam/yaml/yaml_io.py

@@ -224,6 +224,12 @@ def _validate_schema():
     return (
         schema_pb2.Schema(fields=[schemas.schema_field('payload', bytes)]),
         lambda payload: beam.Row(payload=payload))
+  if format == 'STRING':
+    if schema:
+      raise ValueError('STRING format does not take a schema')
+    return (
+        schema_pb2.Schema(fields=[schemas.schema_field('payload', str)]),
+        lambda payload: beam.Row(payload=payload.decode('utf-8')))
   elif format == 'JSON':
     _validate_schema()
     beam_schema = json_utils.json_schema_to_beam_schema(schema)
@@ -313,6 +319,7 @@ def read_from_pubsub(
 
         - RAW: Produces records with a single `payload` field whose contents
             are the raw bytes of the pubsub message.
+        - STRING: Like RAW, but the bytes are decoded as a UTF-8 string.
         - AVRO: Parses records with a given Avro schema.
         - JSON: Parses records with a given JSON schema.
 

sdks/python/apache_beam/yaml/yaml_io_test.py

@@ -101,6 +101,28 @@ def test_simple_read(self):
             result,
             equal_to([beam.Row(payload=b'msg1'), beam.Row(payload=b'msg2')]))
 
+  def test_simple_read_string(self):
+    with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+        pickle_library='cloudpickle')) as p:
+      with mock.patch('apache_beam.io.ReadFromPubSub',
+                      FakeReadFromPubSub(
+                          topic='my_topic',
+                          messages=[PubsubMessage('äter'.encode('utf-8'),
+                                                  {'attr': 'value1'}),
+                                    PubsubMessage('köttbullar'.encode('utf-8'),
+                                                  {'attr': 'value2'})])):
+        result = p | YamlTransform(
+            '''
+            type: ReadFromPubSub
+            config:
+              topic: my_topic
+              format: STRING
+            ''')
+        assert_that(
+            result,
+            equal_to([beam.Row(payload='äter'),
+                      beam.Row(payload='köttbullar')]))
+
   def test_read_with_attribute(self):
     with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
         pickle_library='cloudpickle')) as p: