Apply yapf formatting

Author: Nikita Grover

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/ExternalWrite.java

@@ -53,6 +53,7 @@ public static class Configuration {
     private String topic;
     private @Nullable String idAttribute;
     private @Nullable String timestampAttribute;
+    private boolean publishWithOrderingKey = false;
 
     public void setTopic(String topic) {
       this.topic = topic;
@@ -65,6 +66,10 @@ public void setIdLabel(@Nullable String idAttribute) {
     public void setTimestampAttribute(@Nullable String timestampAttribute) {
       this.timestampAttribute = timestampAttribute;
     }
+
+    public void setPublishWithOrderingKey(Boolean publishWithOrderingKey) {
+      this.publishWithOrderingKey = publishWithOrderingKey != null && publishWithOrderingKey;
+    }
   }
 
   public static class WriteBuilder
@@ -85,6 +90,9 @@ public PTransform<PCollection<byte[]>, PDone> buildExternal(Configuration config
       if (config.timestampAttribute != null) {
         writeBuilder.setTimestampAttribute(config.timestampAttribute);
       }
+      if (config.publishWithOrderingKey) {
+        writeBuilder.setPublishWithOrderingKey(true);
+      }
       writeBuilder.setDynamicDestinations(false);
       return writeBuilder.build();
     }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java

@@ -1727,7 +1727,10 @@ public void startBundle(StartBundleContext c) throws IOException {
         this.pubsubClient =
             getPubsubClientFactory()
                 .newClient(
-                    getTimestampAttribute(), null, c.getPipelineOptions().as(PubsubOptions.class));
+                    getTimestampAttribute(),
+                    null,
+                    c.getPipelineOptions().as(PubsubOptions.class),
+                    Write.this.getPubsubRootUrl());
       }
 
       @ProcessElement

sdks/python/apache_beam/io/external/gcp/pubsub.py

@@ -117,6 +117,7 @@ def expand(self, pbegin):
         # this is not implemented yet on the Java side:
         # ('with_attributes', bool),
         ('timestamp_attribute', typing.Optional[str]),
+        ('publish_with_ordering_key', bool),
     ])
 
 
@@ -135,6 +136,7 @@ def __init__(
       with_attributes=False,
       id_label=None,
       timestamp_attribute=None,
+      publish_with_ordering_key=False,
       expansion_service=None):
     """Initializes ``WriteToPubSub``.
 
@@ -150,18 +152,24 @@ def __init__(
         in a ReadFromPubSub PTransform to deduplicate messages.
       timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub
         message with the given name and the message's publish time as the value.
+      publish_with_ordering_key: If True, enables ordering key support when
+        publishing messages. The ordering key must be set on each
+        PubsubMessage via the ``ordering_key`` attribute. Requires
+        messages to be routed to the same region.
     """
     self.params = WriteToPubsubSchema(
         topic=topic,
         id_label=id_label,
         # with_attributes=with_attributes,
-        timestamp_attribute=timestamp_attribute)
+        timestamp_attribute=timestamp_attribute,
+        publish_with_ordering_key=publish_with_ordering_key)
     self.expansion_service = expansion_service
     self.with_attributes = with_attributes
 
   def expand(self, pvalue):
     if self.with_attributes:
-      pcoll = pvalue | 'ToProto' >> Map(pubsub.WriteToPubSub.to_proto_str)
+      pcoll = pvalue | 'ToProto' >> Map(
+          lambda m: m._to_proto_str(for_publish=True))
     else:
       pcoll = pvalue | 'ToProto' >> Map(
           lambda x: pubsub.PubsubMessage(x, {})._to_proto_str())