Skip to content

Fix nullness for PubsubIO #39444

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kennknowles wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix nullness for PubsubIO #39444

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions ...cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/AddTimestampAttribute.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down
15 changes: 8 additions & 7 deletions .../google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/ExternalWrite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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() {}

Expand All @@ -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;
Expand Down Expand Up @@ -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) {
Expand All @@ -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"));
}
}
}
8 changes: 5 additions & 3 deletions ...le-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/NestedRowToMessage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand Down
46 changes: 24 additions & 22 deletions ...ud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubSubPayloadTranslation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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.
Expand All @@ -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)) {
Expand Down Expand Up @@ -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());
Expand All @@ -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(
Expand All @@ -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()))
Expand All @@ -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()))
Expand Down
59 changes: 32 additions & 27 deletions ...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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>>
Expand Down Expand Up @@ -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.");
}
Expand Down Expand Up @@ -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. */
Expand Down Expand Up @@ -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(
Expand All @@ -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);
}
Expand All @@ -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(
Expand Down Expand Up @@ -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;
Expand All @@ -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}. */
Expand All @@ -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());
}

Expand Down
Loading
Loading