Support writing to Pubsub with ordering key; Add PubsubMessage SchemaCoder

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

@@ -735,6 +735,20 @@ public static Write<PubsubMessage> writeMessages() {
         .build();
   }
 
+  /**
+   * Returns A {@link PTransform} that writes {@link PubsubMessage}s, along with the {@link
+   * PubsubMessage#getMessageId() messageId} and {@link PubsubMessage#getOrderingKey()}, to a Google
+   * Cloud Pub/Sub stream.
+   */
+  public static Write<PubsubMessage> writeMessagesWithOrderingKey() {
+    return Write.newBuilder()
+        .setTopicProvider(null)
+        .setTopicFunction(null)
+        .setDynamicDestinations(false)
+        .setNeedsOrderingKey(true)
+        .build();
+  }
+
   /**
    * Enables dynamic destination topics. The {@link PubsubMessage} elements are each expected to
    * contain a destination topic, which can be set using {@link PubsubMessage#withTopic}. If {@link
@@ -1288,6 +1302,8 @@ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
 
     abstract @Nullable String getPubsubRootUrl();
 
+    abstract boolean getNeedsOrderingKey();
+
     abstract BadRecordRouter getBadRecordRouter();
 
     abstract ErrorHandler<BadRecord, ?> getBadRecordErrorHandler();
@@ -1301,6 +1317,7 @@ static <T> Builder<T> newBuilder(
       builder.setFormatFn(formatFn);
       builder.setBadRecordRouter(BadRecordRouter.THROWING_ROUTER);
       builder.setBadRecordErrorHandler(new DefaultErrorHandler<>());
+      builder.setNeedsOrderingKey(false);
       return builder;
     }
 
@@ -1332,6 +1349,8 @@ abstract Builder<T> setFormatFn(
 
       abstract Builder<T> setPubsubRootUrl(String pubsubRootUrl);
 
+      abstract Builder<T> setNeedsOrderingKey(boolean needsOrderingKey);
+
       abstract Builder<T> setBadRecordRouter(BadRecordRouter badRecordRouter);
 
       abstract Builder<T> setBadRecordErrorHandler(
@@ -1487,8 +1506,12 @@ public PDone expand(PCollection<T> input) {
               pubsubMessageTuple
                   .get(BAD_RECORD_TAG)
                   .setCoder(BadRecord.getCoder(input.getPipeline())));
-      PCollection<PubsubMessage> pubsubMessages =
-          pubsubMessageTuple.get(pubsubMessageTupleTag).setCoder(new PubsubMessageWithTopicCoder());
+      PCollection<PubsubMessage> pubsubMessages = pubsubMessageTuple.get(pubsubMessageTupleTag);
+      if (getNeedsOrderingKey()) {
+        pubsubMessages.setCoder(PubsubMessageSchemaCoder.getSchemaCoder());
+      } else {
+        pubsubMessages.setCoder(new PubsubMessageWithTopicCoder());
+      }
       switch (input.isBounded()) {
         case BOUNDED:
           pubsubMessages.apply(

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

@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsub;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
+
+/**
+ * Provides a {@link SchemaCoder} for {@link PubsubMessage}, including the topic and all fields of a
+ * PubSub message from server.
+ *
+ * <p>{@link SchemaCoder} is used so that fields can be added in the future without breaking update
+ * compatibility.
+ */
+public class PubsubMessageSchemaCoder {
+  private static final Schema PUBSUB_MESSAGE_SCHEMA =
+      Schema.builder()
+          .addByteArrayField("payload")
+          .addNullableStringField("topic")
+          .addNullableMapField("attributes", Schema.FieldType.STRING, Schema.FieldType.STRING)
+          .addNullableStringField("message_id")
+          .addNullableStringField("ordering_key")
+          .build();
+
+  private static final SerializableFunction<PubsubMessage, Row> TO_ROW =
+      (PubsubMessage message) -> {
+        Map<String, Object> fieldValues = new HashMap<>();
+        fieldValues.put("payload", message.getPayload());
+
+        String topic = message.getTopic();
+        if (topic != null) {
+          fieldValues.put("topic", topic);
+        }
+        Map<String, String> attributeMap = message.getAttributeMap();
+        if (attributeMap != null) {
+          fieldValues.put("attributes", attributeMap);
+        }
+        String messageId = message.getMessageId();
+        if (messageId != null) {
+          fieldValues.put("message_id", messageId);
+        }
+        String orderingKey = message.getOrderingKey();
+        if (orderingKey != null) {
+          fieldValues.put("ordering_key", orderingKey);
+        }
+        return Row.withSchema(PUBSUB_MESSAGE_SCHEMA).withFieldValues(fieldValues).build();
+      };
+
+  private static final SerializableFunction<Row, PubsubMessage> FROM_ROW =
+      (Row row) -> {
+        PubsubMessage message =
+            new PubsubMessage(
+                Preconditions.checkNotNull(row.getBytes("payload")),
+                row.getMap("attributes"),
+                row.getString("message_id"),
+                row.getString("ordering_key"));
+
+        String topic = row.getString("topic");
+        if (topic != null) {
+          message = message.withTopic(topic);
+        }
+        return message;
+      };
+
+  public static SchemaCoder<PubsubMessage> getSchemaCoder() {
+    return SchemaCoder.of(
+        PUBSUB_MESSAGE_SCHEMA, TypeDescriptor.of(PubsubMessage.class), TO_ROW, FROM_ROW);
+  }
+}

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

@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsub;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.testing.CoderProperties;
+import org.apache.beam.sdk.util.SerializableUtils;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.junit.Test;
+
+public class PubsubMessageSchemaCoderTest {
+  private static final String DATA = "testData";
+  private static final String TOPIC = "testTopic";
+  private static final String MESSAGE_ID = "testMessageId";
+  private static final Map<String, String> ATTRIBUTES =
+      new ImmutableMap.Builder<String, String>().put("1", "hello").build();
+  private static final String ORDERING_KEY = "key123";
+  private static final Coder<PubsubMessage> TEST_CODER = PubsubMessageSchemaCoder.getSchemaCoder();
+  private static final PubsubMessage TEST_VALUE =
+      new PubsubMessage(DATA.getBytes(StandardCharsets.UTF_8), ATTRIBUTES, MESSAGE_ID, ORDERING_KEY)
+          .withTopic(TOPIC);
+  private static final PubsubMessage TEST_MINIMAL_VALUE =
+      new PubsubMessage(DATA.getBytes(StandardCharsets.UTF_8), null);
+
+  @Test
+  public void testValueEncodable() {
+    SerializableUtils.ensureSerializableByCoder(TEST_CODER, TEST_VALUE, "error");
+    SerializableUtils.ensureSerializableByCoder(TEST_CODER, TEST_MINIMAL_VALUE, "error");
+  }
+
+  @Test
+  public void testCoderDecodeEncodeEqual() throws Exception {
+    CoderProperties.structuralValueDecodeEncodeEqual(TEST_CODER, TEST_VALUE);
+    CoderProperties.structuralValueDecodeEncodeEqual(TEST_CODER, TEST_MINIMAL_VALUE);
+  }
+
+  @Test
+  public void testEncodedTypeDescriptor() {
+    TypeDescriptor<PubsubMessage> typeDescriptor = new TypeDescriptor<PubsubMessage>() {};
+    assertThat(TEST_CODER.getEncodedTypeDescriptor(), equalTo(typeDescriptor));
+  }
+}

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

@@ -17,11 +17,17 @@
  */
 package org.apache.beam.sdk.io.gcp.pubsub;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import org.apache.beam.sdk.io.GenerateSequence;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
 import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
@@ -46,11 +52,12 @@ public class PubsubWriteIT {
   private PubsubClient pubsubClient;
 
   private TopicPath testTopic;
+  private String project;
 
   @Before
   public void setup() throws IOException {
     PubsubOptions options = TestPipeline.testingPipelineOptions().as(PubsubOptions.class);
-    String project = options.getProject();
+    project = options.getProject();
     pubsubClient = PubsubGrpcClient.FACTORY.newClient(null, null, options);
     testTopic =
         PubsubClient.topicPathFromName(project, "pubsub-write-" + Instant.now().getMillis());
@@ -102,4 +109,55 @@ public void testBoundedWriteMessageWithAttributes() {
         .apply(PubsubIO.writeMessages().to(testTopic.getPath()));
     pipeline.run();
   }
+
+  @Test
+  public void testBoundedWriteMessageWithAttributesAndMessageIdAndOrderingKey() throws IOException {
+    TopicPath testTopicPath =
+        PubsubClient.topicPathFromName(
+            project, "pubsub-write-ordering-key-" + Instant.now().getMillis());
+    pubsubClient.createTopic(testTopicPath);
+    SubscriptionPath testSubscriptionPath =
+        pubsubClient.createRandomSubscription(
+            PubsubClient.projectPathFromId(project), testTopicPath, 10);
+
+    byte[] payload = RandomStringUtils.randomAscii(1_000_000).getBytes(StandardCharsets.UTF_8);
+    Map<String, String> attributes =
+        ImmutableMap.<String, String>builder()
+            .put("id", "1")
+            .put("description", RandomStringUtils.randomAscii(100))
+            .build();
+
+    PubsubMessage outgoingMessage =
+        new PubsubMessage(payload, attributes, "test_message", "111222");
+
+    pipeline
+        .apply(Create.of(outgoingMessage).withCoder(PubsubMessageSchemaCoder.getSchemaCoder()))
+        .apply(PubsubIO.writeMessagesWithOrderingKey().to(testTopicPath.getPath()));
+    pipeline.run().waitUntilFinish();
+
+    List<PubsubClient.IncomingMessage> incomingMessages =
+        pubsubClient.pull(Instant.now().getMillis(), testSubscriptionPath, 1, true);
+
+    // sometimes the first pull comes up short. try 4 pulls to avoid flaky false-negatives
+    int numPulls = 1;
+    while (incomingMessages.isEmpty()) {
+      if (numPulls >= 4) {
+        throw new RuntimeException(
+            String.format("Pulled %s times from PubSub but retrieved no elements.", numPulls));
+      }
+      incomingMessages =
+          pubsubClient.pull(Instant.now().getMillis(), testSubscriptionPath, 1, true);
+      numPulls++;
+    }
+    assertEquals(1, incomingMessages.size());
+
+    com.google.pubsub.v1.PubsubMessage incomingMessage = incomingMessages.get(0).message();
+    assertTrue(
+        Arrays.equals(outgoingMessage.getPayload(), incomingMessage.getData().toByteArray()));
+    assertEquals(outgoingMessage.getAttributeMap(), incomingMessage.getAttributesMap());
+    assertEquals(outgoingMessage.getOrderingKey(), incomingMessage.getOrderingKey());
+
+    pubsubClient.deleteSubscription(testSubscriptionPath);
+    pubsubClient.deleteTopic(testTopicPath);
+  }
 }