/*
* 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.common.UncaughtExceptionLogger;
import com.google.cloud.teleport.v2.templates.io.RedisHashIO;
import com.google.cloud.teleport.v2.templates.transforms.MessageTransformation;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
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.PipelineOptions;
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.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.TypeDescriptors;
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.",
"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 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 {@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 PipelineOptions {
@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);
}
/** 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);
PCollection<PubsubMessage> input;
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) Extract PubSubMessage message to PCollection<String>.
* 3) Transform PCollection<String> to PCollection<KV<String, String>> so it can be consumed by RedisIO
* 4) Write to Redis using SET
*
*/
LOG.info(
"Starting PubSub-To-Redis Pipeline. Reading from subscription: {}",
options.getInputSubscription());
input =
pipeline.apply(
"Read PubSub Events",
MessageTransformation.readFromPubSub(options.getInputSubscription()));
if (options.getRedisSinkType().equals(STRING_SINK)) {
PCollection<String> pCollectionString =
input.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 =
input.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 =
input.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 =
input.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();
}
}
