Stream messages from Pub/Sub by using Dataflow and Cloud Storage

importjava.io.IOException;
importorg.apache.beam.examples.common.WriteOneFilePerWindow;
importorg.apache.beam.sdk.Pipeline;
importorg.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
importorg.apache.beam.sdk.options.Default;
importorg.apache.beam.sdk.options.Description;
importorg.apache.beam.sdk.options.PipelineOptionsFactory;
importorg.apache.beam.sdk.options.StreamingOptions;
importorg.apache.beam.sdk.options.Validation.Required;
importorg.apache.beam.sdk.transforms.windowing.FixedWindows;
importorg.apache.beam.sdk.transforms.windowing.Window;
importorg.joda.time.Duration;
publicclass PubSubToGcs{
/*
 * Define your own configuration options. Add your own arguments to be processed
 * by the command-line parser, and specify default values for them.
 */
publicinterface PubSubToGcsOptionsextendsStreamingOptions{
@Description("The Cloud Pub/Sub topic to read from.")
@Required
StringgetInputTopic();
voidsetInputTopic(Stringvalue);
@Description("Output file's window size in number of minutes.")
@Default.Integer(1)
IntegergetWindowSize();
voidsetWindowSize(Integervalue);
@Description("Path of the output file including its filename prefix.")
@Required
StringgetOutput();
voidsetOutput(Stringvalue);
}
publicstaticvoidmain(String[]args)throwsIOException{
// The maximum number of shards when writing output.
intnumShards=1;
PubSubToGcsOptionsoptions=
PipelineOptionsFactory.fromArgs(args).withValidation().as(PubSubToGcsOptions.class);
options.setStreaming(true);
Pipelinepipeline=Pipeline.create(options);
pipeline
// 1) Read string messages from a Pub/Sub topic.
.apply("Read PubSub Messages",PubsubIO.readStrings().fromTopic(options.getInputTopic()))
// 2) Group the messages into fixed-sized minute intervals.
.apply(Window.into(FixedWindows.of(Duration.standardMinutes(options.getWindowSize()))))
// 3) Write one file to GCS for every window of messages.
.apply("Write Files to GCS",newWriteOneFilePerWindow(options.getOutput(),numShards));
// Execute the pipeline and wait until it finishes running.
pipeline.run().waitUntilFinish();
}
}

importargparse
fromdatetimeimport datetime
importlogging
importrandom
fromapache_beamimport (
 DoFn,
 GroupByKey,
 io,
 ParDo,
 Pipeline,
 PTransform,
 WindowInto,
 WithKeys,
)
fromapache_beam.options.pipeline_optionsimport PipelineOptions
fromapache_beam.transforms.windowimport FixedWindows
classGroupMessagesByFixedWindows(PTransform):
"""A composite transform that groups Pub/Sub messages based on publish time
 and outputs a list of tuples, each containing a message and its publish time.
 """
 def__init__(self, window_size, num_shards=5):
 # Set window size to 60 seconds.
 self.window_size = int(window_size * 60)
 self.num_shards = num_shards
 defexpand(self, pcoll):
 return (
 pcoll
 # Bind window info to each element using element timestamp (or publish time).
 | "Window into fixed intervals"
 >> WindowInto(FixedWindows(self.window_size))
 | "Add timestamp to windowed elements" >> ParDo(AddTimestamp())
 # Assign a random key to each windowed element based on the number of shards.
 | "Add key" >> WithKeys(lambda _: random.randint(0, self.num_shards - 1))
 # Group windowed elements by key. All the elements in the same window must fit
 # memory for this. If not, you need to use `beam.util.BatchElements`.
 | "Group by key" >> GroupByKey()
 )
classAddTimestamp(DoFn):
 defprocess(self, element, publish_time=DoFn.TimestampParam):
"""Processes each windowed element by extracting the message body and its
 publish time into a tuple.
 """
 yield (
 element.decode("utf-8"),
 datetime.utcfromtimestamp(float(publish_time)).strftime(
 "%Y-%m-%d %H:%M:%S.%f"
 ),
 )
classWriteToGCS(DoFn):
 def__init__(self, output_path):
 self.output_path = output_path
 defprocess(self, key_value, window=DoFn.WindowParam):
"""Write messages in a batch to Google Cloud Storage."""
 ts_format = "%H:%M"
 window_start = window.start.to_utc_datetime().strftime(ts_format)
 window_end = window.end.to_utc_datetime().strftime(ts_format)
 shard_id, batch = key_value
 filename = "-".join([self.output_path, window_start, window_end, str(shard_id)])
 with io.gcsio.GcsIO().open(filename=filename, mode="w") as f:
 for message_body, publish_time in batch:
 f.write(f"{message_body},{publish_time}\n".encode())
defrun(input_topic, output_path, window_size=1.0, num_shards=5, pipeline_args=None):
 # Set `save_main_session` to True so DoFns can access globally imported modules.
 pipeline_options = PipelineOptions(
 pipeline_args, streaming=True, save_main_session=True
 )
 with Pipeline(options=pipeline_options) as pipeline:
 (
 pipeline
 # Because `timestamp_attribute` is unspecified in `ReadFromPubSub`, Beam
 # binds the publish time returned by the Pub/Sub server for each message
 # to the element's timestamp parameter, accessible via `DoFn.TimestampParam`.
 # https://beam.apache.org/releases/pydoc/current/apache_beam.io.gcp.pubsub.html#apache_beam.io.gcp.pubsub.ReadFromPubSub
 | "Read from Pub/Sub" >> io.ReadFromPubSub(topic=input_topic)
 | "Window into" >> GroupMessagesByFixedWindows(window_size, num_shards)
 | "Write to GCS" >> ParDo(WriteToGCS(output_path))
 )
if __name__ == "__main__":
 logging.getLogger().setLevel(logging.INFO)
 parser = argparse.ArgumentParser()
 parser.add_argument(
 "--input_topic",
 help="The Cloud Pub/Sub topic to read from."
 '"projects/<PROJECT_ID>/topics/<TOPIC_ID>".',
 )
 parser.add_argument(
 "--window_size",
 type=float,
 default=1.0,
 help="Output file's window size in minutes.",
 )
 parser.add_argument(
 "--output_path",
 help="Path of the output GCS file including the prefix.",
 )
 parser.add_argument(
 "--num_shards",
 type=int,
 default=5,
 help="Number of shards to use when writing windowed elements to GCS.",
 )
 known_args, pipeline_args = parser.parse_known_args()
 run(
 known_args.input_topic,
 known_args.output_path,
 known_args.window_size,
 known_args.num_shards,
 pipeline_args,
 )