/*
* Copyright (C) 2023 Google LLC
*
* Licensed 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 com.google.cloud.teleport.v2.templates;
import static com.google.cloud.teleport.v2.templates.PubSubToRedis.RedisSinkType.HASH_SINK;
import static com.google.cloud.teleport.v2.templates.PubSubToRedis.RedisSinkType.LOGGING_SINK;
import static com.google.cloud.teleport.v2.templates.PubSubToRedis.RedisSinkType.STREAMS_SINK;
import static com.google.cloud.teleport.v2.templates.PubSubToRedis.RedisSinkType.STRING_SINK;
import com.google.cloud.teleport.metadata.Template;
import com.google.cloud.teleport.metadata.TemplateCategory;
import com.google.cloud.teleport.metadata.TemplateParameter;
import com.google.cloud.teleport.metadata.TemplateParameter.TemplateEnumOption;
import com.google.cloud.teleport.v2.coders.FailsafeElementCoder;
import com.google.cloud.teleport.v2.common.UncaughtExceptionLogger;
import com.google.cloud.teleport.v2.templates.io.RedisHashIO;
import com.google.cloud.teleport.v2.templates.transforms.MessageTransformation;
import com.google.cloud.teleport.v2.transforms.FailsafeElementTransforms.ConvertFailsafeElementToPubsubMessage;
import com.google.cloud.teleport.v2.transforms.JavascriptTextTransformer;
import com.google.cloud.teleport.v2.values.FailsafeElement;
import com.google.common.base.Strings;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageWithAttributesCoder;
import org.apache.beam.sdk.io.redis.RedisConnectionConfiguration;
import org.apache.beam.sdk.io.redis.RedisIO;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.Validation;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.commons.lang3.ArrayUtils;
import org.checkerframework.checker.initialization.qual.Initialized;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.UnknownKeyFor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* The {@link PubSubToRedis} pipeline is a streaming pipeline which ingests data in Bytes from
* PubSub, and inserts resulting records as KV in Redis.
*
* <p><b>Pipeline Requirements</b>
*
* <ul>
* <li>The PubSub topic and subscriptions exist
* <li>The Redis is up and running
* </ul>
*
* <p><b>Example Usage</b>
*
* <pre>
* # Set the pipeline vars
* PROJECT_NAME=my-project
* BUCKET_NAME=my-bucket
* INPUT_SUBSCRIPTION=my-subscription
* REDIS_HOST=my-host
* REDIS_PORT=my-port
* REDIS_PASSWORD=my-pwd
*
* mvn compile exec:java \
* -Dexec.mainClass=com.google.cloud.teleport.v2.templates.PubSubToRedis \
* -Dexec.cleanupDaemonThreads=false \
* -Dexec.args=" \
* --project=${PROJECT_NAME} \
* --stagingLocation=gs://${BUCKET_NAME}/staging \
* --tempLocation=gs://${BUCKET_NAME}/temp \
* --runner=DataflowRunner \
* --inputSubscription=${INPUT_SUBSCRIPTION} \
* --redisHost=${REDIS_HOST}
* --redisPort=${REDIS_PORT}
* --redisPassword=${REDIS_PASSWORD}"
* </pre>
*/
@Template(
name = "Cloud_PubSub_to_Redis",
category = TemplateCategory.STREAMING,
displayName = "Pub/Sub to Redis",
description = {
"The Pub/Sub to Redis template is a streaming pipeline that reads messages from a Pub/Sub subscription and "
+ "writes the message payload to Redis. The most common use case of this template is to export logs to Redis "
+ "Enterprise for advanced search-based log analysis in real time.",
"Before writing to Redis, you can apply a JavaScript user-defined function to the message payload. Any "
+ "messages that experience processing failures are forwarded to a Pub/Sub unprocessed topic for further "
+ "troubleshooting and reprocessing.",
"For added security, enable an SSL connection when setting up your database endpoint connection."
},
optionsClass = PubSubToRedis.PubSubToRedisOptions.class,
flexContainerName = "pubsub-to-redis",
contactInformation = "https://github.com/GoogleCloudPlatform/DataflowTemplates/issues",
documentation =
"https://cloud.google.com/dataflow/docs/guides/templates/provided/pubsub-to-redis",
requirements = {
"The source Pub/Sub subscription must exist prior to running the pipeline.",
"The Pub/Sub unprocessed (dead-letter) topic must exist prior to running the pipeline if using a JavaScript UDF.",
"The Redis database endpoint must be accessible from the Dataflow workers' subnetwork.",
},
preview = true,
streaming = true,
supportsAtLeastOnce = true)
public class PubSubToRedis {
/*
* Options supported by {@link PubSubToRedis}
*
* <p>Inherits standard configuration options.
*/
/** The log to output status messages to. */
private static final Logger LOG = LoggerFactory.getLogger(PubSubToRedis.class);
/** The tag for the main output for the UDF. */
public static final TupleTag<FailsafeElement<PubsubMessage, String>> UDF_OUT =
new TupleTag<FailsafeElement<PubsubMessage, String>>() {};
/** The tag for the dead-letter output of the udf. */
public static final TupleTag<FailsafeElement<PubsubMessage, String>> UDF_DEADLETTER_OUT =
new TupleTag<FailsafeElement<PubsubMessage, String>>() {};
/** Pubsub message/string coder for FailsafeElement. */
public static final FailsafeElementCoder<PubsubMessage, String> FAILSAFE_ELEMENT_CODER =
FailsafeElementCoder.of(PubsubMessageWithAttributesCoder.of(), StringUtf8Coder.of());
/**
* The {@link PubSubToRedisOptions} class provides the custom execution options passed by the
* executor at the command-line.
*
* <p>Inherits standard configuration options.
*/
public interface PubSubToRedisOptions
extends JavascriptTextTransformer.JavascriptTextTransformerOptions {
@TemplateParameter.PubsubSubscription(
order = 1,
groupName = "Source",
description = "Pub/Sub input subscription",
helpText = "The Pub/Sub subscription to read the input from.",
example = "projects/<PROJECT_ID>/subscriptions/<SUBSCRIPTION_ID>")
String getInputSubscription();
void setInputSubscription(String value);
@TemplateParameter.Text(
order = 2,
groupName = "Target",
description = "Redis DB Host",
helpText = "The Redis database host.",
example = "your.cloud.db.redislabs.com")
@Default.String("127.0.0.1")
@Validation.Required
String getRedisHost();
void setRedisHost(String redisHost);
@TemplateParameter.Integer(
order = 3,
groupName = "Target",
description = "Redis DB Port",
helpText = "The Redis database port.",
example = "12345")
@Default.Integer(6379)
@Validation.Required
int getRedisPort();
void setRedisPort(int redisPort);
@TemplateParameter.Password(
order = 4,
groupName = "Target",
description = "Redis DB Password",
helpText = "The Redis database password. Defaults to `empty`.")
@Default.String("")
@Validation.Required
String getRedisPassword();
void setRedisPassword(String redisPassword);
@TemplateParameter.Boolean(
order = 5,
optional = true,
description = "Redis ssl enabled",
helpText = "The Redis database SSL parameter.")
@Default.Boolean(false)
@UnknownKeyFor
@NonNull
@Initialized
ValueProvider<@UnknownKeyFor @NonNull @Initialized Boolean> getSslEnabled();
void setSslEnabled(ValueProvider<Boolean> sslEnabled);
@TemplateParameter.Enum(
order = 6,
optional = true,
enumOptions = {
@TemplateEnumOption("STRING_SINK"),
@TemplateEnumOption("HASH_SINK"),
@TemplateEnumOption("STREAMS_SINK"),
@TemplateEnumOption("LOGGING_SINK")
},
description = "Redis sink to write",
helpText =
"The Redis sink. Supported values are `STRING_SINK, HASH_SINK, STREAMS_SINK, and LOGGING_SINK`.",
example = "STRING_SINK")
@Default.Enum("STRING_SINK")
RedisSinkType getRedisSinkType();
void setRedisSinkType(RedisSinkType redisSinkType);
@TemplateParameter.Integer(
order = 7,
optional = true,
description = "Redis connection timeout in milliseconds",
helpText = "The Redis connection timeout in milliseconds. ",
example = "2000")
@Default.Integer(2000)
int getConnectionTimeout();
void setConnectionTimeout(int timeout);
@TemplateParameter.Long(
order = 8,
optional = true,
parentName = "redisSinkType",
parentTriggerValues = {"HASH_SINK", "LOGGING_SINK"},
description =
"Hash key expiration time in sec (ttl), supported only for HASH_SINK and LOGGING_SINK",
helpText =
"The key expiration time in seconds. The `ttl` default for `HASH_SINK` is -1, which means it never expires.")
@Default.Long(-1L)
Long getTtl();
void setTtl(Long ttl);
@TemplateParameter.PubsubTopic(
order = 9,
optional = true,
description = "Output deadletter Pub/Sub topic",
helpText =
"The Pub/Sub topic to forward unprocessable messages to. Messages that fail UDF transformation are forwarded here, Required if using a JavaScript UDF.",
example = "projects/<PROJECT_ID>/topics/<TOPIC_NAME>")
String getOutputDeadletterTopic();
void setOutputDeadletterTopic(String outputDeadletterTopic);
}
/** Allowed list of sink types. */
public enum RedisSinkType {
HASH_SINK,
LOGGING_SINK,
STREAMS_SINK,
STRING_SINK
}
/**
* Main entry point for executing the pipeline.
*
* @param args The command-line arguments to the pipeline.
*/
public static void main(String[] args) {
UncaughtExceptionLogger.register();
// Parse the user options passed from the command-line.
PubSubToRedisOptions options =
PipelineOptionsFactory.fromArgs(args).withValidation().as(PubSubToRedisOptions.class);
run(options);
}
/**
* Runs the pipeline with the supplied options.
*
* @param options The execution parameters to the pipeline.
* @return The result of the pipeline execution.
*/
public static PipelineResult run(PubSubToRedisOptions options) {
// Create the pipeline
Pipeline pipeline = Pipeline.create(options);
// Register the coders for pipeline
CoderRegistry coderRegistry = pipeline.getCoderRegistry();
coderRegistry.registerCoderForType(
FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor(), FAILSAFE_ELEMENT_CODER);
RedisConnectionConfiguration redisConnectionConfiguration =
RedisConnectionConfiguration.create()
.withHost(options.getRedisHost())
.withPort(options.getRedisPort())
.withAuth(options.getRedisPassword())
.withTimeout(options.getConnectionTimeout())
.withSSL(options.getSslEnabled());
/*
* Steps: 1) Read PubSubMessage with attributes and messageId from input PubSub subscription.
* 2) Apply JavaScript UDF transformation to message payload (if configured).
* 3) Extract PubSubMessage message to appropriate Redis format.
* 4) Write to Redis using the configured sink type.
*
*/
LOG.info(
"Starting PubSub-To-Redis Pipeline. Reading from subscription: {}",
options.getInputSubscription());
PCollection<PubsubMessage> messages =
pipeline.apply(
"Read PubSub Events",
MessageTransformation.readFromPubSub(options.getInputSubscription()));
PCollection<PubsubMessage> maybeTransformed = applyUdf(messages, options);
if (options.getRedisSinkType().equals(STRING_SINK)) {
PCollection<String> pCollectionString =
maybeTransformed.apply(
"Map to Redis String", ParDo.of(new MessageTransformation.MessageToRedisString()));
PCollection<KV<String, String>> kvStringCollection =
pCollectionString.apply(
"Transform to String KV",
MapElements.into(
TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.strings()))
.via(record -> KV.of(MessageTransformation.key, record)));
kvStringCollection.apply(
"Write to " + STRING_SINK.name(),
RedisIO.write()
.withMethod(RedisIO.Write.Method.SET)
.withConnectionConfiguration(redisConnectionConfiguration));
}
if (options.getRedisSinkType().equals(HASH_SINK)) {
PCollection<KV<String, KV<String, String>>> pCollectionHash =
maybeTransformed.apply(
"Map to Redis Hash", ParDo.of(new MessageTransformation.MessageToRedisHash()));
pCollectionHash.apply(
"Write to " + HASH_SINK.name(),
RedisHashIO.write()
.withConnectionConfiguration(redisConnectionConfiguration)
.withTtl(options.getTtl()));
}
if (options.getRedisSinkType().equals(LOGGING_SINK)) {
PCollection<KV<String, KV<String, String>>> pCollectionHash =
maybeTransformed.apply(
"Map to Redis Logs", ParDo.of(new MessageTransformation.MessageToRedisLogs()));
pCollectionHash.apply(
"Write to " + LOGGING_SINK.name(),
RedisHashIO.write()
.withConnectionConfiguration(redisConnectionConfiguration)
.withTtl(options.getTtl()));
}
if (options.getRedisSinkType().equals(STREAMS_SINK)) {
PCollection<KV<String, Map<String, String>>> pCollectionStreams =
maybeTransformed.apply(
"Map to Redis Streams", ParDo.of(new MessageTransformation.MessageToRedisStreams()));
pCollectionStreams.apply(
"Write to " + STREAMS_SINK.name(),
RedisIO.writeStreams().withConnectionConfiguration(redisConnectionConfiguration));
}
// Execute the pipeline and return the result.
return pipeline.run();
}
/**
* Applies the JavaScript UDF to messages if configured, and writes UDF failures to the
* dead-letter topic.
*
* <p>If no UDF is configured, returns the input messages unchanged.
*
* @param messages the input PubSub messages
* @param options the pipeline options
* @return the (possibly transformed) PubSub messages
*/
// This follows the same pattern as PubsubProtoToBigQuery.runUdf
static PCollection<PubsubMessage> applyUdf(
PCollection<PubsubMessage> messages, PubSubToRedisOptions options) {
boolean useJavascriptUdf = !Strings.isNullOrEmpty(options.getJavascriptTextTransformGcsPath());
if (!useJavascriptUdf) {
return messages;
}
if (Strings.isNullOrEmpty(options.getJavascriptTextTransformFunctionName())) {
throw new IllegalArgumentException(
"JavaScript function name cannot be null or empty if file is set");
}
if (Strings.isNullOrEmpty(options.getOutputDeadletterTopic())) {
throw new IllegalArgumentException(
"A dead-letter Pub/Sub topic (--outputDeadletterTopic) must be provided when using a JavaScript UDF.");
}
// Map incoming messages to FailsafeElement so we can recover from failures
// across multiple transforms.
PCollection<FailsafeElement<PubsubMessage, String>> failsafeElements =
messages.apply("MapToRecord", ParDo.of(new PubsubMessageToFailsafeElementFn()));
PCollectionTuple udfResult =
failsafeElements.apply(
"InvokeUDF",
JavascriptTextTransformer.FailsafeJavascriptUdf.<PubsubMessage>newBuilder()
.setFileSystemPath(options.getJavascriptTextTransformGcsPath())
.setFunctionName(options.getJavascriptTextTransformFunctionName())
.setReloadIntervalMinutes(options.getJavascriptTextTransformReloadIntervalMinutes())
.setSuccessTag(UDF_OUT)
.setFailureTag(UDF_DEADLETTER_OUT)
.build());
// Write UDF failures to the dead-letter topic using the shared
// ConvertFailsafeElementToPubsubMessage transform, following the pattern
// from PubsubProtoToBigQuery.
udfResult
.get(UDF_DEADLETTER_OUT)
.setCoder(FAILSAFE_ELEMENT_CODER)
.apply(
"Get UDF Failures",
ConvertFailsafeElementToPubsubMessage.<PubsubMessage, String>builder()
.setOriginalPayloadSerializeFn(msg -> ArrayUtils.toObject(msg.getPayload()))
.setErrorMessageAttributeKey("udfErrorMessage")
.build())
.apply("Write Failed UDF", writeUdfFailures(options));
// Extract the successfully transformed messages
return udfResult
.get(UDF_OUT)
.apply(
"Extract Transformed Payload",
ParDo.of(
new DoFn<FailsafeElement<PubsubMessage, String>, PubsubMessage>() {
@ProcessElement
public void processElement(ProcessContext c) {
FailsafeElement<PubsubMessage, String> element = c.element();
c.output(
new PubsubMessage(
element.getPayload().getBytes(StandardCharsets.UTF_8),
element.getOriginalPayload().getAttributeMap(),
element.getOriginalPayload().getMessageId()));
}
}));
}
/**
* Returns a {@link PubsubIO.Write} configured to write UDF failures to the dead-letter topic.
*
* <p>Follows the same pattern as {@code PubsubProtoToBigQuery.writeUdfFailures}.
*/
private static PubsubIO.Write<PubsubMessage> writeUdfFailures(PubSubToRedisOptions options) {
return PubsubIO.writeMessages().to(options.getOutputDeadletterTopic());
}
/**
* The {@link PubsubMessageToFailsafeElementFn} wraps an incoming {@link PubsubMessage} with the
* {@link FailsafeElement} class so errors can be recovered from and the original message can be
* output to a dead-letter output.
*/
static class PubsubMessageToFailsafeElementFn
extends DoFn<PubsubMessage, FailsafeElement<PubsubMessage, String>> {
@ProcessElement
public void processElement(ProcessContext context) {
PubsubMessage message = context.element();
context.output(
FailsafeElement.of(message, new String(message.getPayload(), StandardCharsets.UTF_8)));
}
}
}
