Merge pull request #39444: Fix nullness for PubsubIO

Author: kennknowles

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

@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.io.gcp.pubsub;
 
 import static org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageToRow.TIMESTAMP_FIELD;
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
 
 import org.apache.beam.sdk.schemas.transforms.DropFields;
 import org.apache.beam.sdk.transforms.PTransform;
@@ -31,9 +32,6 @@
  * Adds a timestamp attribute if desired and filters it out of the underlying row if no timestamp
  * attribute exists.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 class AddTimestampAttribute extends PTransform<PCollection<Row>, PCollection<Row>> {
   private static final Logger LOG = LoggerFactory.getLogger(AddTimestampAttribute.class);
   private final boolean useTimestampAttribute;
@@ -48,7 +46,14 @@ public PCollection<Row> expand(PCollection<Row> input) {
     // element's event time. PubSubIO will populate the attribute from there.
     PCollection<Row> withTimestamp =
         useTimestampAttribute
-            ? input.apply(WithTimestamps.of((row) -> row.getDateTime(TIMESTAMP_FIELD).toInstant()))
+            ? input.apply(
+                WithTimestamps.of(
+                    (row) ->
+                        checkArgumentNotNull(
+                                row.getDateTime(TIMESTAMP_FIELD),
+                                "Field '%s' must be present and non-null in the input row when writing to PubSub with a timestamp attribute.",
+                                TIMESTAMP_FIELD)
+                            .toInstant()))
             : input;
 
     PCollection<Row> rows;

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

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.io.gcp.pubsub;
 
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+
 import com.google.auto.service.AutoService;
 import java.util.Map;
 import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
@@ -35,9 +37,6 @@
 
 /** Exposes {@link PubsubIO.Write} as an external transform for cross-language usage. */
 @AutoService(ExternalTransformRegistrar.class)
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public final class ExternalWrite implements ExternalTransformRegistrar {
   public ExternalWrite() {}
 
@@ -50,7 +49,7 @@ public ExternalWrite() {}
 
   /** Parameters class to expose the transform to an external SDK. */
   public static class Configuration {
-    private String topic;
+    private @Nullable String topic;
     private @Nullable String idAttribute;
     private @Nullable String timestampAttribute;
     private boolean publishWithOrderingKey = false;
@@ -80,8 +79,9 @@ public WriteBuilder() {}
     public PTransform<PCollection<byte[]>, PDone> buildExternal(Configuration config) {
       PubsubIO.Write.Builder<byte[]> writeBuilder =
           PubsubIO.Write.newBuilder(new ParsePubsubMessageProtoAsPayloadFromWindowedValue());
-      if (config.topic != null) {
-        StaticValueProvider<String> topic = StaticValueProvider.of(config.topic);
+      String topicStr = config.topic;
+      if (topicStr != null) {
+        StaticValueProvider<String> topic = StaticValueProvider.of(topicStr);
         writeBuilder.setTopicProvider(NestedValueProvider.of(topic, PubsubTopic::fromPath));
       }
       if (config.idAttribute != null) {
@@ -102,7 +102,8 @@ public static class ParsePubsubMessageProtoAsPayloadFromWindowedValue
 
     @Override
     public PubsubMessage apply(ValueInSingleWindow<byte[]> input) {
-      return INNER.apply(input.getValue());
+      return INNER.apply(
+          checkArgumentNotNull(input.getValue(), "Windowed value input cannot be null"));
     }
   }
 }

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

@@ -22,6 +22,7 @@
 import static org.apache.beam.sdk.io.gcp.pubsub.PubsubSchemaIOProvider.ATTRIBUTE_ARRAY_FIELD_TYPE;
 import static org.apache.beam.sdk.io.gcp.pubsub.PubsubSchemaIOProvider.ATTRIBUTE_MAP_FIELD_TYPE;
 import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 
 import java.util.Collection;
@@ -34,16 +35,17 @@
 import org.apache.beam.sdk.transforms.SimpleFunction;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 class NestedRowToMessage extends SimpleFunction<Row, PubsubMessage> {
   private static final long serialVersionUID = 65176815766314684L;
 
-  private final PayloadSerializer serializer;
+  private final @Nullable PayloadSerializer serializer;
   private final SerializableFunction<Row, Map<String, String>> attributesExtractor;
   private final SerializableFunction<Row, byte[]> payloadExtractor;
 
   @SuppressWarnings("methodref.receiver.bound")
-  NestedRowToMessage(PayloadSerializer serializer, Schema schema) {
+  NestedRowToMessage(@Nullable PayloadSerializer serializer, Schema schema) {
     this.serializer = serializer;
     if (schema.getField(ATTRIBUTES_FIELD).getType().equals(ATTRIBUTE_MAP_FIELD_TYPE)) {
       attributesExtractor = NestedRowToMessage::getAttributesFromMap;
@@ -81,7 +83,7 @@ private static byte[] getPayloadFromBytes(Row row) {
   }
 
   private byte[] getPayloadFromNested(Row row) {
-    return serializer.serialize(checkArgumentNotNull(row.getRow(PAYLOAD_FIELD)));
+    return checkStateNotNull(serializer).serialize(checkArgumentNotNull(row.getRow(PAYLOAD_FIELD)));
   }
 
   @Override

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

@@ -19,6 +19,7 @@
 
 import static org.apache.beam.sdk.io.gcp.pubsub.PubsubIO.ENABLE_CUSTOM_PUBSUB_SINK;
 import static org.apache.beam.sdk.io.gcp.pubsub.PubsubIO.ENABLE_CUSTOM_PUBSUB_SOURCE;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 
 import com.google.auto.service.AutoService;
 import java.util.Collections;
@@ -37,16 +38,13 @@
 import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
 import org.apache.beam.sdk.runners.AppliedPTransform;
 import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.util.Preconditions;
 import org.apache.beam.sdk.util.construction.PTransformTranslation;
 import org.apache.beam.sdk.util.construction.PTransformTranslation.TransformPayloadTranslator;
 import org.apache.beam.sdk.util.construction.SdkComponents;
 import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 /**
  * Utility methods for translating a {@link Unbounded} which reads from {@link
  * PubsubUnboundedSource} to {@link RunnerApi} representations.
@@ -61,7 +59,7 @@ public String getUrn() {
     }
 
     @Override
-    public RunnerApi.FunctionSpec translate(
+    public RunnerApi.@Nullable FunctionSpec translate(
         AppliedPTransform<?, ?, Unbounded<?>> transform, SdkComponents components) {
       if (ExperimentalOptions.hasExperiment(
           transform.getPipeline().getOptions(), ENABLE_CUSTOM_PUBSUB_SOURCE)) {
@@ -93,11 +91,13 @@ public RunnerApi.FunctionSpec translate(
         }
       }
 
-      if (pubsubUnboundedSource.getTimestampAttribute() != null) {
-        payloadBuilder.setTimestampAttribute(pubsubUnboundedSource.getTimestampAttribute());
+      String timestampAttribute = pubsubUnboundedSource.getTimestampAttribute();
+      if (timestampAttribute != null) {
+        payloadBuilder.setTimestampAttribute(timestampAttribute);
       }
-      if (pubsubUnboundedSource.getIdAttribute() != null) {
-        payloadBuilder.setIdAttribute(pubsubUnboundedSource.getIdAttribute());
+      String idAttribute = pubsubUnboundedSource.getIdAttribute();
+      if (idAttribute != null) {
+        payloadBuilder.setIdAttribute(idAttribute);
       }
       payloadBuilder.setWithAttributes(
           pubsubUnboundedSource.getNeedsAttributes() || pubsubUnboundedSource.getNeedsMessageId());
@@ -116,7 +116,7 @@ public String getUrn() {
     }
 
     @Override
-    public RunnerApi.FunctionSpec translate(
+    public RunnerApi.@Nullable FunctionSpec translate(
         AppliedPTransform<?, ?, PubsubUnboundedSink.PubsubSink> transform,
         SdkComponents components) {
       if (ExperimentalOptions.hasExperiment(
@@ -125,19 +125,20 @@ public RunnerApi.FunctionSpec translate(
       }
       PubSubWritePayload.Builder payloadBuilder = PubSubWritePayload.newBuilder();
       ValueProvider<TopicPath> topicProvider =
-          Preconditions.checkStateNotNull(transform.getTransform().outer.getTopicProvider());
+          checkStateNotNull(transform.getTransform().outer.getTopicProvider());
       if (topicProvider.isAccessible()) {
         payloadBuilder.setTopic(topicProvider.get().getFullPath());
       } else {
         payloadBuilder.setTopicRuntimeOverridden(
             ((NestedValueProvider) topicProvider).propertyName());
       }
-      if (transform.getTransform().outer.getTimestampAttribute() != null) {
-        payloadBuilder.setTimestampAttribute(
-            transform.getTransform().outer.getTimestampAttribute());
+      String timestampAttribute = transform.getTransform().outer.getTimestampAttribute();
+      if (timestampAttribute != null) {
+        payloadBuilder.setTimestampAttribute(timestampAttribute);
       }
-      if (transform.getTransform().outer.getIdAttribute() != null) {
-        payloadBuilder.setIdAttribute(transform.getTransform().outer.getIdAttribute());
+      String idAttribute = transform.getTransform().outer.getIdAttribute();
+      if (idAttribute != null) {
+        payloadBuilder.setIdAttribute(idAttribute);
       }
       return FunctionSpec.newBuilder()
           .setUrn(getUrn(transform.getTransform()))
@@ -154,20 +155,21 @@ public String getUrn() {
     }
 
     @Override
-    public RunnerApi.FunctionSpec translate(
+    public RunnerApi.@Nullable FunctionSpec translate(
         AppliedPTransform<?, ?, PubsubUnboundedSink.PubsubDynamicSink> transform,
         SdkComponents components) {
       if (ExperimentalOptions.hasExperiment(
           transform.getPipeline().getOptions(), ENABLE_CUSTOM_PUBSUB_SINK)) {
         return null;
       }
       PubSubWritePayload.Builder payloadBuilder = PubSubWritePayload.newBuilder();
-      if (transform.getTransform().outer.getTimestampAttribute() != null) {
-        payloadBuilder.setTimestampAttribute(
-            transform.getTransform().outer.getTimestampAttribute());
+      String timestampAttribute = transform.getTransform().outer.getTimestampAttribute();
+      if (timestampAttribute != null) {
+        payloadBuilder.setTimestampAttribute(timestampAttribute);
       }
-      if (transform.getTransform().outer.getIdAttribute() != null) {
-        payloadBuilder.setIdAttribute(transform.getTransform().outer.getIdAttribute());
+      String idAttribute = transform.getTransform().outer.getIdAttribute();
+      if (idAttribute != null) {
+        payloadBuilder.setIdAttribute(idAttribute);
       }
       return FunctionSpec.newBuilder()
           .setUrn(getUrn(transform.getTransform()))

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

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.io.gcp.pubsub;
 
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
@@ -40,13 +42,11 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.checkerframework.dataflow.qual.Pure;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /** An (abstract) helper class for talking to Pubsub via an underlying transport. */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public abstract class PubsubClient implements Closeable {
   private static final Logger LOG = LoggerFactory.getLogger(PubsubClient.class);
   private static final Map<String, SerializableFunction<String, Schema>>
@@ -82,7 +82,7 @@ PubsubClient newClient(
    * Return timestamp as ms-since-unix-epoch corresponding to {@code timestamp}. Throw {@link
    * IllegalArgumentException} if timestamp cannot be recognized.
    */
-  protected static Long parseTimestampAsMsSinceEpoch(String timestamp) {
+  protected static long parseTimestampAsMsSinceEpoch(String timestamp) {
     if (timestamp.isEmpty()) {
       throw new IllegalArgumentException("Empty timestamp.");
     }
@@ -113,17 +113,12 @@ protected static long extractTimestampAttribute(
       String timestampAttribute, @Nullable Map<String, String> attributes) {
     Preconditions.checkState(!timestampAttribute.isEmpty());
     String value = attributes == null ? null : attributes.get(timestampAttribute);
-    checkArgument(
-        value != null,
-        "PubSub message is missing a value for timestamp attribute %s",
-        timestampAttribute);
-    Long timestampMsSinceEpoch = parseTimestampAsMsSinceEpoch(value);
-    checkArgument(
-        timestampMsSinceEpoch != null,
-        "Cannot interpret value of attribute %s as timestamp: %s",
-        timestampAttribute,
-        value);
-    return timestampMsSinceEpoch;
+    String nonNullValue =
+        checkArgumentNotNull(
+            value,
+            "PubSub message is missing a value for timestamp attribute %s",
+            timestampAttribute);
+    return parseTimestampAsMsSinceEpoch(nonNullValue);
   }
 
   /** Path representing a cloud project id. */
@@ -386,17 +381,21 @@ public static TopicPath topicPathFromName(String projectId, String topicName) {
   public abstract static class OutgoingMessage implements Serializable {
 
     /** Underlying Message. May not have publish timestamp set. */
+    @Pure
     public abstract PubsubMessage getMessage();
 
     /** Timestamp for element (ms since epoch). */
+    @Pure
     public abstract long getTimestampMsSinceEpoch();
 
     /**
      * If using an id attribute, the record id to associate with this record's metadata so the
      * receiver can reject duplicates. Otherwise {@literal null}.
      */
+    @Pure
     public abstract @Nullable String recordId();
 
+    @Pure
     public abstract @Nullable String topic();
 
     public static OutgoingMessage of(
@@ -415,11 +414,13 @@ public static OutgoingMessage of(
         @Nullable String topic) {
       PubsubMessage.Builder builder =
           PubsubMessage.newBuilder().setData(ByteString.copyFrom(message.getPayload()));
-      if (message.getAttributeMap() != null) {
-        builder.putAllAttributes(message.getAttributeMap());
+      Map<String, String> attributeMap = message.getAttributeMap();
+      if (attributeMap != null) {
+        builder.putAllAttributes(attributeMap);
       }
-      if (message.getOrderingKey() != null) {
-        builder.setOrderingKey(message.getOrderingKey());
+      String orderingKey = message.getOrderingKey();
+      if (orderingKey != null) {
+        builder.setOrderingKey(orderingKey);
       }
       return of(builder.build(), timestampMsSinceEpoch, recordId, topic);
     }
@@ -435,21 +436,23 @@ public static OutgoingMessage of(
   public abstract static class IncomingMessage implements Serializable {
 
     /** Underlying Message. */
+    @Pure
     public abstract PubsubMessage message();
 
-    /**
-     * Timestamp for element (ms since epoch). Either Pubsub's processing time, or the custom
-     * timestamp associated with the message.
-     */
+    /** Either Pubsub's processing time, or the custom timestamp associated with the message. */
+    @Pure
     public abstract long timestampMsSinceEpoch();
 
     /** Timestamp (in system time) at which we requested the message (ms since epoch). */
+    @Pure
     public abstract long requestTimeMsSinceEpoch();
 
     /** Id to pass back to Pubsub to acknowledge receipt of this message. */
+    @Pure
     public abstract String ackId();
 
     /** Id to pass to the runner to distinguish this message from all others. */
+    @Pure
     public abstract String recordId();
 
     public static IncomingMessage of(
@@ -554,7 +557,7 @@ public abstract void createSchema(
   public abstract void deleteSchema(SchemaPath schemaPath) throws IOException;
 
   /** Return {@link SchemaPath} from {@link TopicPath} if exists. */
-  public abstract SchemaPath getSchemaPath(TopicPath topicPath) throws IOException;
+  public abstract @Nullable SchemaPath getSchemaPath(TopicPath topicPath) throws IOException;
 
   /** Return a Beam {@link Schema} from the Pub/Sub schema resource, if exists. */
   public abstract Schema getSchema(SchemaPath schemaPath) throws IOException;
@@ -567,8 +570,10 @@ static Schema fromPubsubSchema(com.google.api.services.pubsub.model.Schema pubsu
               "Pub/Sub schema type %s is not supported at this time", pubsubSchema.getType()));
     }
     SerializableFunction<String, Schema> definitionToSchemaFn =
-        schemaTypeToConversionFnMap.get(pubsubSchema.getType());
-    return definitionToSchemaFn.apply(pubsubSchema.getDefinition());
+        checkStateNotNull(schemaTypeToConversionFnMap.get(pubsubSchema.getType()));
+    String definition =
+        checkNotNull(pubsubSchema.getDefinition(), "Pub/Sub schema definition is null");
+    return definitionToSchemaFn.apply(definition);
   }
 
   /** Convert a {@link com.google.pubsub.v1.Schema} to a Beam {@link Schema}. */
@@ -579,7 +584,7 @@ static Schema fromPubsubSchema(com.google.pubsub.v1.Schema pubsubSchema) {
           String.format("Pub/Sub schema type %s is not supported at this time", typeName));
     }
     SerializableFunction<String, Schema> definitionToSchemaFn =
-        schemaTypeToConversionFnMap.get(typeName);
+        checkStateNotNull(schemaTypeToConversionFnMap.get(typeName));
     return definitionToSchemaFn.apply(pubsubSchema.getDefinition());
   }