Merge Redis cluster connector with Redis connector

docs/coverage/java/pom.xml

@@ -59,12 +59,6 @@
       <version>${project.version}</version>
     </dependency>
 
-    <dependency>
-      <groupId>dev.feast</groupId>
-      <artifactId>feast-storage-connector-redis-cluster</artifactId>
-      <version>${project.version}</version>
-    </dependency>
-
     <dependency>
       <groupId>dev.feast</groupId>
       <artifactId>feast-ingestion</artifactId>

ingestion/pom.xml

@@ -126,12 +126,6 @@
       <version>${project.version}</version>
     </dependency>
 
-    <dependency>
-      <groupId>dev.feast</groupId>
-      <artifactId>feast-storage-connector-redis-cluster</artifactId>
-      <version>${project.version}</version>
-    </dependency>
-
     <dependency>
       <groupId>dev.feast</groupId>
       <artifactId>feast-storage-connector-bigquery</artifactId>

ingestion/src/main/java/feast/ingestion/utils/StoreUtil.java

@@ -26,7 +26,6 @@
 import feast.storage.api.writer.FeatureSink;
 import feast.storage.connectors.bigquery.writer.BigQueryFeatureSink;
 import feast.storage.connectors.redis.writer.RedisFeatureSink;
-import feast.storage.connectors.rediscluster.writer.RedisClusterFeatureSink;
 import java.util.HashMap;
 import java.util.Map;
 import org.slf4j.Logger;
@@ -84,7 +83,7 @@ public static FeatureSink getFeatureSink(
     StoreType storeType = store.getType();
     switch (storeType) {
       case REDIS_CLUSTER:
-        return RedisClusterFeatureSink.fromConfig(store.getRedisClusterConfig(), featureSetSpecs);
+        return RedisFeatureSink.fromConfig(store.getRedisClusterConfig(), featureSetSpecs);
       case REDIS:
         return RedisFeatureSink.fromConfig(store.getRedisConfig(), featureSetSpecs);
       case BIGQUERY:

serving/pom.xml

@@ -110,12 +110,6 @@
       <version>${project.version}</version>
     </dependency>
 
-    <dependency>
-      <groupId>dev.feast</groupId>
-      <artifactId>feast-storage-connector-redis-cluster</artifactId>
-      <version>${project.version}</version>
-    </dependency>
-
     <dependency>
       <groupId>dev.feast</groupId>
       <artifactId>feast-storage-connector-bigquery</artifactId>

serving/src/main/java/feast/serving/config/ServingServiceConfig.java

@@ -28,8 +28,8 @@
 import feast.storage.api.retriever.HistoricalRetriever;
 import feast.storage.api.retriever.OnlineRetriever;
 import feast.storage.connectors.bigquery.retriever.BigQueryHistoricalRetriever;
+import feast.storage.connectors.redis.retriever.RedisClusterOnlineRetriever;
 import feast.storage.connectors.redis.retriever.RedisOnlineRetriever;
-import feast.storage.connectors.rediscluster.retriever.RedisClusterOnlineRetriever;
 import io.opentracing.Tracer;
 import java.util.Map;
 import org.slf4j.Logger;

storage/connectors/pom.xml

@@ -16,7 +16,6 @@
 
     <modules>
         <module>redis</module>
-        <module>rediscluster</module>
         <module>bigquery</module>
     </modules>
 

storage/connectors/redis/pom.xml

@@ -70,6 +70,13 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>net.ishiis.redis</groupId>
+            <artifactId>redis-unit</artifactId>
+            <version>1.0.3</version>
+            <scope>test</scope>
+        </dependency>
+
 
         <dependency>
             <groupId>junit</groupId>

storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetriever.java → storage/connectors/redis/src/main/java/feast/storage/connectors/redis/retriever/RedisClusterOnlineRetriever.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package feast.storage.connectors.rediscluster.retriever;
+package feast.storage.connectors.redis.retriever;
 
 import com.google.protobuf.AbstractMessageLite;
 import com.google.protobuf.InvalidProtocolBufferException;

storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterIngestionClient.java → storage/connectors/redis/src/main/java/feast/storage/connectors/redis/writer/RedisClusterIngestionClient.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package feast.storage.connectors.rediscluster.writer;
+package feast.storage.connectors.redis.writer;
 
 import com.google.common.collect.Lists;
 import feast.proto.core.StoreProto;

storage/connectors/redis/src/main/java/feast/storage/connectors/redis/writer/RedisCustomIO.java

@@ -19,7 +19,7 @@
 import feast.proto.core.FeatureSetProto.EntitySpec;
 import feast.proto.core.FeatureSetProto.FeatureSetSpec;
 import feast.proto.core.FeatureSetProto.FeatureSpec;
-import feast.proto.core.StoreProto.Store.RedisConfig;
+import feast.proto.core.StoreProto.Store.*;
 import feast.proto.storage.RedisProto.RedisKey;
 import feast.proto.storage.RedisProto.RedisKey.Builder;
 import feast.proto.types.FeatureRowProto.FeatureRow;
@@ -28,7 +28,7 @@
 import feast.storage.api.writer.FailedElement;
 import feast.storage.api.writer.WriteResult;
 import feast.storage.common.retry.Retriable;
-import io.lettuce.core.RedisConnectionException;
+import io.lettuce.core.RedisException;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -63,21 +63,22 @@ public class RedisCustomIO {
 
   private RedisCustomIO() {}
 
-  public static Write write(RedisConfig redisConfig, Map<String, FeatureSetSpec> featureSetSpecs) {
-    return new Write(redisConfig, featureSetSpecs);
+  public static Write write(
+      RedisIngestionClient redisIngestionClient, Map<String, FeatureSetSpec> featureSetSpecs) {
+    return new Write(redisIngestionClient, featureSetSpecs);
   }
 
   /** ServingStoreWrite data to a Redis server. */
   public static class Write extends PTransform<PCollection<FeatureRow>, WriteResult> {
 
     private Map<String, FeatureSetSpec> featureSetSpecs;
-    private RedisConfig redisConfig;
+    private RedisIngestionClient redisIngestionClient;
     private int batchSize;
     private int timeout;
 
-    public Write(RedisConfig redisConfig, Map<String, FeatureSetSpec> featureSetSpecs) {
-
-      this.redisConfig = redisConfig;
+    public Write(
+        RedisIngestionClient redisIngestionClient, Map<String, FeatureSetSpec> featureSetSpecs) {
+      this.redisIngestionClient = redisIngestionClient;
       this.featureSetSpecs = featureSetSpecs;
     }
 
@@ -95,7 +96,7 @@ public Write withTimeout(int timeout) {
     public WriteResult expand(PCollection<FeatureRow> input) {
       PCollectionTuple redisWrite =
           input.apply(
-              ParDo.of(new WriteDoFn(redisConfig, featureSetSpecs))
+              ParDo.of(new WriteDoFn(redisIngestionClient, featureSetSpecs))
                   .withOutputTags(successfulInsertsTag, TupleTagList.of(failedInsertsTupleTag)));
       return WriteResult.in(
           input.getPipeline(),
@@ -111,9 +112,10 @@ public static class WriteDoFn extends DoFn<FeatureRow, FeatureRow> {
       private int timeout = DEFAULT_TIMEOUT;
       private RedisIngestionClient redisIngestionClient;
 
-      WriteDoFn(RedisConfig config, Map<String, FeatureSetSpec> featureSetSpecs) {
+      WriteDoFn(
+          RedisIngestionClient redisIngestionClient, Map<String, FeatureSetSpec> featureSetSpecs) {
 
-        this.redisIngestionClient = new RedisStandaloneIngestionClient(config);
+        this.redisIngestionClient = redisIngestionClient;
         this.featureSetSpecs = featureSetSpecs;
       }
 
@@ -140,7 +142,7 @@ public void setup() {
       public void startBundle() {
         try {
           redisIngestionClient.connect();
-        } catch (RedisConnectionException e) {
+        } catch (RedisException e) {
           log.error("Connection to redis cannot be established ", e);
         }
         featureRows.clear();
@@ -165,7 +167,7 @@ public void execute() throws ExecutionException, InterruptedException {
 
                   @Override
                   public Boolean isExceptionRetriable(Exception e) {
-                    return e instanceof RedisConnectionException;
+                    return e instanceof RedisException;
                   }
 
                   @Override

storage/connectors/redis/src/main/java/feast/storage/connectors/redis/writer/RedisFeatureSink.java

@@ -20,6 +20,7 @@
 import feast.proto.core.FeatureSetProto.FeatureSet;
 import feast.proto.core.FeatureSetProto.FeatureSetSpec;
 import feast.proto.core.StoreProto;
+import feast.proto.core.StoreProto.Store.RedisClusterConfig;
 import feast.proto.core.StoreProto.Store.RedisConfig;
 import feast.proto.types.FeatureRowProto.FeatureRow;
 import feast.storage.api.writer.FeatureSink;
@@ -28,6 +29,7 @@
 import io.lettuce.core.RedisConnectionException;
 import io.lettuce.core.RedisURI;
 import java.util.Map;
+import javax.annotation.Nullable;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.values.PCollection;
 
@@ -46,8 +48,17 @@ public static FeatureSink fromConfig(
     return builder().setFeatureSetSpecs(featureSetSpecs).setRedisConfig(redisConfig).build();
   }
 
+  public static FeatureSink fromConfig(
+      RedisClusterConfig redisConfig, Map<String, FeatureSetSpec> featureSetSpecs) {
+    return builder().setFeatureSetSpecs(featureSetSpecs).setRedisClusterConfig(redisConfig).build();
+  }
+
+  @Nullable
   public abstract RedisConfig getRedisConfig();
 
+  @Nullable
+  public abstract RedisClusterConfig getRedisClusterConfig();
+
   public abstract Map<String, FeatureSetSpec> getFeatureSetSpecs();
 
   public abstract Builder toBuilder();
@@ -60,29 +71,45 @@ public static Builder builder() {
   public abstract static class Builder {
     public abstract Builder setRedisConfig(RedisConfig redisConfig);
 
+    public abstract Builder setRedisClusterConfig(RedisClusterConfig redisConfig);
+
     public abstract Builder setFeatureSetSpecs(Map<String, FeatureSetSpec> featureSetSpecs);
 
     public abstract RedisFeatureSink build();
   }
 
   @Override
   public void prepareWrite(FeatureSet featureSet) {
-
-    RedisClient redisClient =
-        RedisClient.create(RedisURI.create(getRedisConfig().getHost(), getRedisConfig().getPort()));
-    try {
-      redisClient.connect();
-    } catch (RedisConnectionException e) {
+    if (getRedisConfig() != null) {
+      RedisClient redisClient =
+          RedisClient.create(
+              RedisURI.create(getRedisConfig().getHost(), getRedisConfig().getPort()));
+      try {
+        redisClient.connect();
+      } catch (RedisConnectionException e) {
+        throw new RuntimeException(
+            String.format(
+                "Failed to connect to Redis at host: '%s' port: '%d'. Please check that your Redis is running and accessible from Feast.",
+                getRedisConfig().getHost(), getRedisConfig().getPort()));
+      }
+      redisClient.shutdown();
+    } else if (getRedisClusterConfig() == null) {
       throw new RuntimeException(
-          String.format(
-              "Failed to connect to Redis at host: '%s' port: '%d'. Please check that your Redis is running and accessible from Feast.",
-              getRedisConfig().getHost(), getRedisConfig().getPort()));
+          "At least one RedisConfig or RedisClusterConfig must be provided to Redis Sink");
     }
-    redisClient.shutdown();
   }
 
   @Override
   public PTransform<PCollection<FeatureRow>, WriteResult> writer() {
-    return new RedisCustomIO.Write(getRedisConfig(), getFeatureSetSpecs());
+    if (getRedisClusterConfig() != null) {
+      return new RedisCustomIO.Write(
+          new RedisClusterIngestionClient(getRedisClusterConfig()), getFeatureSetSpecs());
+    } else if (getRedisConfig() != null) {
+      return new RedisCustomIO.Write(
+          new RedisStandaloneIngestionClient(getRedisConfig()), getFeatureSetSpecs());
+    } else {
+      throw new RuntimeException(
+          "At least one RedisConfig or RedisClusterConfig must be provided to Redis Sink");
+    }
   }
 }

storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetrieverTest.java → storage/connectors/redis/src/test/java/feast/storage/connectors/redis/retriever/RedisClusterOnlineRetrieverTest.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package feast.storage.connectors.rediscluster.retriever;
+package feast.storage.connectors.redis.retriever;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;

storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/writer/RedisClusterFeatureSinkTest.java → storage/connectors/redis/src/test/java/feast/storage/connectors/redis/writer/RedisClusterFeatureSinkTest.java

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package feast.storage.connectors.rediscluster.writer;
+package feast.storage.connectors.redis.writer;
 
 import static feast.storage.common.testing.TestUtil.field;
 import static org.hamcrest.CoreMatchers.equalTo;
@@ -67,7 +67,7 @@ public class RedisClusterFeatureSinkTest {
   private RedisClusterClient redisClusterClient;
   private RedisClusterCommands<byte[], byte[]> redisClusterCommands;
 
-  private RedisClusterFeatureSink redisClusterFeatureSink;
+  private RedisFeatureSink redisClusterFeatureSink;
 
   @Before
   public void setUp() throws IOException {
@@ -123,7 +123,7 @@ public void setUp() throws IOException {
             .build();
 
     redisClusterFeatureSink =
-        RedisClusterFeatureSink.builder()
+        RedisFeatureSink.builder()
             .setFeatureSetSpecs(specMap)
             .setRedisClusterConfig(redisClusterConfig)
             .build();
@@ -141,8 +141,8 @@ static boolean deleteDirectory(File directoryToBeDeleted) {
 
   @After
   public void teardown() {
-    redisClusterClient.shutdown();
     redisCluster.stop();
+    redisClusterClient.shutdown();
     deleteDirectory(new File(String.valueOf(Paths.get(System.getProperty("user.dir"), ".redis"))));
   }
 
@@ -259,7 +259,7 @@ public void shouldProduceFailedElementIfRetryExceeded() {
             .build();
     Map<String, FeatureSetSpec> specMap = ImmutableMap.of("myproject/fs", spec1);
     redisClusterFeatureSink =
-        RedisClusterFeatureSink.builder()
+        RedisFeatureSink.builder()
             .setFeatureSetSpecs(specMap)
             .setRedisClusterConfig(redisClusterConfig)
             .build();