pubsub_dataflow_simple

from __future__ import absolute_import

import logging

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions, GoogleCloudOptions, SetupOptions
from apache_beam.options.pipeline_options import StandardOptions

'''
python -m apache_beam.testing.load_tests.dataflow_test
'''

def run():
  pipeline_options = PipelineOptions()
  pipeline_options.view_as(StandardOptions).streaming = True
  pipeline_options.view_as(SetupOptions).sdk_location = 'dist/apache-beam-2.21.0.dev0.tar.gz'
  pipeline_options.view_as(GoogleCloudOptions).temp_location = 'gs://beam-test-pszuberski/tmp'
  pipeline_options.view_as(GoogleCloudOptions).staging_location = 'gs://beam-test-pszuberski/staging'
  pipeline_options.view_as(GoogleCloudOptions).project = 'apache-beam-io-testing'
  pipeline_options.view_as(StandardOptions).runner = 'DataflowRunner'

  topic_name = 'projects/apache-beam-io-testing/topics/abc'
  sub_name = 'projects/apache-beam-io-testing/subscriptions/abc_sub'

  with beam.Pipeline(options=pipeline_options) as p:
    _ = p | beam.Create([b'simplified']).with_output_types(bytes) | beam.io.WriteToPubSub(topic_name)


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  run()