[BEAM-4610] Add SSL support for RedisIO

apache · Mar 28, 2019 · a36af54 · a36af54
1 parent cea9c32
commit a36af54
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 1 deletion.
diff --git a/...ava/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisConnectionConfiguration.java b/...ava/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisConnectionConfiguration.java
@@ -42,6 +42,8 @@ public abstract class RedisConnectionConfiguration implements Serializable {
 
   abstract int timeout();
 
+  abstract boolean ssl();
+
   abstract Builder builder();
 
   @AutoValue.Builder
@@ -54,6 +56,8 @@ abstract static class Builder {
 
     abstract Builder setTimeout(int timeout);
 
+    abstract Builder setSsl(boolean ssl);
+
     abstract RedisConnectionConfiguration build();
   }
 
@@ -62,6 +66,7 @@ public static RedisConnectionConfiguration create() {
         .setHost(Protocol.DEFAULT_HOST)
         .setPort(Protocol.DEFAULT_PORT)
         .setTimeout(Protocol.DEFAULT_TIMEOUT)
+        .setSsl(false)
         .build();
   }
 
@@ -70,6 +75,7 @@ public static RedisConnectionConfiguration create(String host, int port) {
         .setHost(host)
         .setPort(port)
         .setTimeout(Protocol.DEFAULT_TIMEOUT)
+        .setSsl(false)
         .build();
   }
 
@@ -99,9 +105,18 @@ public RedisConnectionConfiguration withTimeout(int timeout) {
     return builder().setTimeout(timeout).build();
   }
 
+  /**
+   * Enable SSL connection to Redis server.
+   *
+   * @return a {@link RedisConnectionConfiguration} object
+   */
+  public RedisConnectionConfiguration enableSSL() {
+    return builder().setSsl(true).build();
+  }
+
   /** Connect to the Redis instance. */
   public Jedis connect() {
-    Jedis jedis = new Jedis(host(), port(), timeout());
+    Jedis jedis = new Jedis(host(), port(), timeout(), ssl());
     if (auth() != null) {
       jedis.auth(auth());
     }
@@ -113,5 +128,6 @@ public void populateDisplayData(DisplayData.Builder builder) {
     builder.add(DisplayData.item("host", host()));
     builder.add(DisplayData.item("port", port()));
     builder.addIfNotNull(DisplayData.item("timeout", timeout()));
+    builder.add(DisplayData.item("ssl", ssl()));
   }
 }
diff --git a/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java b/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
@@ -79,6 +79,16 @@
  *
  * }</pre>
  *
+ * <p>It's also possible to enable Redis SSL connection with the corresponding methods:
+ *
+ * <pre>{@code
+ * pipeline.apply(RedisIO.read()
+ *   .withEndpoint("::1", 6379)
+ *   .withKeyPattern("foo*"))
+ *   .enableSSL()
+ *
+ * }</pre>
+ *
  * <p>{@link #readAll()} can be used to request Redis server using input PCollection elements as key
  * pattern (as String).
  *
@@ -199,6 +209,10 @@ public Read withBatchSize(int batchSize) {
       return builder().setBatchSize(batchSize).build();
     }
 
+    public Read enableSSL() {
+      return builder().setConnectionConfiguration(connectionConfiguration().enableSSL()).build();
+    }
+
     @Override
     public void populateDisplayData(DisplayData.Builder builder) {
       connectionConfiguration().populateDisplayData(builder);

diff --git a/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java b/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
@@ -249,6 +249,18 @@ public void testReadBuildsCorrectly() {
     Assert.assertEquals(111, read.connectionConfiguration().port());
     Assert.assertEquals("pass", read.connectionConfiguration().auth());
     Assert.assertEquals(5, read.connectionConfiguration().timeout());
+    Assert.assertEquals(false, read.connectionConfiguration().ssl());
+  }
+
+  @Test
+  public void testReadBuildsWithSslCorrectly() {
+    RedisIO.Read read =
+        RedisIO.read().withEndpoint("test", 111).withAuth("pass").withTimeout(5).enableSSL();
+    Assert.assertEquals("test", read.connectionConfiguration().host());
+    Assert.assertEquals(111, read.connectionConfiguration().port());
+    Assert.assertEquals("pass", read.connectionConfiguration().auth());
+    Assert.assertEquals(5, read.connectionConfiguration().timeout());
+    Assert.assertEquals(true, read.connectionConfiguration().ssl());
   }
 
   @Test
@@ -258,6 +270,7 @@ public void testWriteBuildsCorrectly() {
     Assert.assertEquals(111, write.connectionConfiguration().port());
     Assert.assertEquals("pass", write.connectionConfiguration().auth());
     Assert.assertEquals(5, write.connectionConfiguration().timeout());
+    Assert.assertEquals(false, write.connectionConfiguration().ssl());
     Assert.assertEquals(Method.APPEND, write.method());
   }