[Bug]: WriteToPubSub failing on Beam YAML

[ffernandez92]

What happened?

I attempted to write some events to PubSub using Beam YAML, and I encountered the following issue:

2024-02-28 10:20:48.089 CET
  File "/usr/local/lib/python3.11/site-packages/apache_beam/yaml/yaml_transform.py", line 96, in wrapper
2024-02-28 10:20:48.089 CET
    self._cache[key] = func(self, *args)
2024-02-28 10:20:48.089 CET
                       ^^^^^^^^^^^^^^^^^
2024-02-28 10:20:48.089 CET
  File "/usr/local/lib/python3.11/site-packages/apache_beam/yaml/yaml_transform.py", line 283, in compute_outputs
2024-02-28 10:20:48.089 CET
    return expand_transform(self._transforms_by_uuid[transform_id], self)
2024-02-28 10:20:48.089 CET
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2024-02-28 10:20:48.089 CET
  File "/usr/local/lib/python3.11/site-packages/apache_beam/yaml/yaml_transform.py", line 457, in expand_transform
2024-02-28 10:20:48.090 CET
    return expand_leaf_transform(spec, scope)
2024-02-28 10:20:48.090 CET
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2024-02-28 10:20:48.090 CET
  File "/usr/local/lib/python3.11/site-packages/apache_beam/yaml/yaml_transform.py", line 497, in expand_leaf_transform
2024-02-28 10:20:48.090 CET
    raise ValueError(
2024-02-28 10:20:48.090 CET
ValueError: Transform "WriteToPubSub" at line 9 returned an unexpected type <class 'apache_beam.pvalue.PDone'>

I utilized the latest Flex template from Google templates (Dataflow runner). To simplify the setup and isolate any potential external factors affecting the overall behavior, I modified my YAML file to its simplest form. Below is the YAML configuration I ultimately used to reproduce the error:

pipeline:
  type: chain
  transforms:
    - type: ReadFromPubSub
      config:
        topic: <TOPIC_READ>
        format: json
        schema: <EVENT_SCHEMA>
    - type: WriteToPubSub
      config:
        topic: <TOPIC_WRITE>
        format: json

My assumption is that WriteToPubSub returns a PDone and it should be returning a PTransform. I managed to workaround the situation adding this to the YAML config:

pipeline:
  transforms:
    - type: ReadFromPubSub
      config:
        topic: <TOPIC_READ>
        format: json
        schema: <EVENT_SCHEMA>
    - type: WindowInto
      input: ReadFromPubSub
      windowing:
        type: fixed
        size: 60s
    - type: WriteToPubSub
      input: WindowInto
      config:
        topic: <TOPIC_WRITE>
        format: json
        error_handling:
          output: my_error_output
    - type: WriteToJson
      name: WriteErrorsToJson
      input: WriteToPubSub.my_error_output
      config:
        path: /path/to/errors.json

However, I think it's still necessary to fix this issue since error_handling is optional.

Issue Priority

Priority: 3 (minor)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner