/** * Returns A {@link PTransform} that continuously reads from a Google Cloud Pub/Sub stream. The * messages will only contain a {@link PubsubMessage#getPayload() payload}, but no {@link * PubsubMessage#getAttributeMap() attributes}. */ public static Read<PubsubMessage> readMessages() { return new AutoValue_PubsubIO_Read.Builder<PubsubMessage>() .setCoder(PubsubMessagePayloadOnlyCoder.of()) .setParseFn(new IdentityMessageFn()) .setNeedsAttributes(false) .build(); }
private void initializePubsub(Description description) throws IOException { pubsub = PubsubGrpcClient.FACTORY.newClient( NO_TIMESTAMP_ATTRIBUTE, NO_ID_ATTRIBUTE, pipelineOptions); TopicPath eventsTopicPathTmp = PubsubClient.topicPathFromName( pipelineOptions.getProject(), createTopicName(description, EVENTS_TOPIC_NAME)); pubsub.createTopic(eventsTopicPathTmp); eventsTopicPath = eventsTopicPathTmp; }
/** * Writes to Pub/Sub are limited by 10mb in general. This attribute controls the maximum allowed * bytes to be sent to Pub/Sub in a single batched message. */ public Write<T> withMaxBatchBytesSize(int maxBatchBytesSize) { return toBuilder().setMaxBatchBytesSize(maxBatchBytesSize).build(); }
private static Subscription buildSubscription(int i) { Subscription subscription = new Subscription(); subscription.setName( PubsubClient.subscriptionPathFromName(PROJECT.getId(), "Subscription" + i).getPath()); subscription.setTopic(PubsubClient.topicPathFromName(PROJECT.getId(), "Topic" + i).getPath()); return subscription; } }
/** * Returns A {@link PTransform} that continuously reads from a Google Cloud Pub/Sub stream. The * messages will contain both a {@link PubsubMessage#getPayload() payload} and {@link * PubsubMessage#getAttributeMap() attributes}. */ public static Read<PubsubMessage> readMessagesWithAttributes() { return new AutoValue_PubsubIO_Read.Builder<PubsubMessage>() .setCoder(PubsubMessageWithAttributesCoder.of()) .setParseFn(new IdentityMessageFn()) .setNeedsAttributes(true) .build(); }
@Override public Coder<PubsubMessage> getOutputCoder() { return outer.getNeedsAttributes() ? PubsubMessageWithAttributesCoder.of() : PubsubMessagePayloadOnlyCoder.of(); }
/** * Create a random subscription for {@code topic}. Return the {@link SubscriptionPath}. It is the * responsibility of the caller to later delete the subscription. */ public SubscriptionPath createRandomSubscription( ProjectPath project, TopicPath topic, int ackDeadlineSeconds) throws IOException { // Create a randomized subscription derived from the topic name. String subscriptionName = topic.getName() + "_beam_" + ThreadLocalRandom.current().nextLong(); SubscriptionPath subscription = PubsubClient.subscriptionPathFromName(project.getId(), subscriptionName); createSubscription(topic, subscription, ackDeadlineSeconds); return subscription; }
private static Topic buildTopic(int i) { Topic topic = new Topic(); topic.setName(PubsubClient.topicPathFromName(PROJECT.getId(), "Topic" + i).getPath()); return topic; }
/** Tests that checkpoints finalized after the reader is closed succeed. */ @Test public void closeWithActiveCheckpoints() throws Exception { setupOneMessage(); PubsubReader reader = primSource.createReader(p.getOptions(), null); reader.start(); PubsubCheckpoint checkpoint = reader.getCheckpointMark(); reader.close(); checkpoint.finalizeCheckpoint(); } }
/** Like {@code topic()} but with a {@link ValueProvider}. */ public Read<T> fromTopic(ValueProvider<String> topic) { if (topic.isAccessible()) { // Validate. PubsubTopic.fromPath(topic.get()); } return toBuilder() .setTopicProvider(NestedValueProvider.of(topic, new TopicTranslator())) .build(); }
/** * Causes the source to return a PubsubMessage that includes Pubsub attributes, and uses the * given parsing function to transform the PubsubMessage into an output type. A Coder for the * output type T must be registered or set on the output via {@link * PCollection#setCoder(Coder)}. */ private Read<T> withCoderAndParseFn(Coder<T> coder, SimpleFunction<PubsubMessage, T> parseFn) { return toBuilder().setCoder(coder).setParseFn(parseFn).build(); }
/** Like {@code topic()} but with a {@link ValueProvider}. */ public Write<T> to(ValueProvider<String> topic) { return toBuilder() .setTopicProvider(NestedValueProvider.of(topic, new TopicTranslator())) .build(); }
/** * When reading from Cloud Pub/Sub where unique record identifiers are provided as Pub/Sub * message attributes, specifies the name of the attribute containing the unique identifier. The * value of the attribute can be any string that uniquely identifies this record. * * <p>Pub/Sub cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream. * If {@code idAttribute} is not provided, Beam cannot guarantee that no duplicate data will be * delivered, and deduplication of the stream will be strictly best effort. */ public Read<T> withIdAttribute(String idAttribute) { return toBuilder().setIdAttribute(idAttribute).build(); }
/** * Used to write a PubSub message together with PubSub attributes. The user-supplied format * function translates the input type T to a PubsubMessage object, which is used by the sink to * separately set the PubSub message's payload and attributes. */ private Write<T> withFormatFn(SimpleFunction<T, PubsubMessage> formatFn) { return toBuilder().setFormatFn(formatFn).build(); }
/** * Writes to Pub/Sub, adding each record's unique identifier to the published messages in an * attribute with the specified name. The value of the attribute is an opaque string. * * <p>If the the output from this sink is being read by another Beam pipeline, then {@link * PubsubIO.Read#withIdAttribute(String)} can be used to ensure that* the other source reads * these unique identifiers from the appropriate attribute. */ public Write<T> withIdAttribute(String idAttribute) { return toBuilder().setIdAttribute(idAttribute).build(); }
/** * Returns A {@link PTransform} that writes UTF-8 encoded strings to a Google Cloud Pub/Sub * stream. */ public static Write<String> writeStrings() { return PubsubIO.<String>write().withFormatFn(new FormatPayloadAsUtf8()); }
void publish(String pubSubTopicUrl, FeatureRow row) throws IOException { TopicPath topicPath = PubsubClient.topicPathFromPath(pubSubTopicUrl); pubsubHelper.publish(topicPath, row); }
@Override public TopicPath apply(PubsubTopic from) { return PubsubClient.topicPathFromName(from.project, from.topic); } }
@FinishBundle public void finishBundle() throws Exception { pubsubClient.close(); pubsubClient = null; }
/** * Returns A {@link PTransform} that continuously reads binary encoded protobuf messages of the * given type from a Google Cloud Pub/Sub stream. */ public static <T extends Message> Read<T> readProtos(Class<T> messageClass) { // TODO: Stop using ProtoCoder and instead parse the payload directly. // We should not be relying on the fact that ProtoCoder's wire format is identical to // the protobuf wire format, as the wire format is not part of a coder's API. ProtoCoder<T> coder = ProtoCoder.of(messageClass); return PubsubIO.<T>read().withCoderAndParseFn(coder, new ParsePayloadUsingCoder<>(coder)); }