[FLINK-15571] Add stream command

Author: eskabetxe

src/main/java/org/apache/flink/connector/redis/shared/JedisConnector.java

@@ -17,7 +17,6 @@
 
 package org.apache.flink.connector.redis.shared;
 
-import org.apache.flink.connector.redis.sink2.RedisAction;
 import redis.clients.jedis.JedisCluster;
 import redis.clients.jedis.JedisPool;
 import redis.clients.jedis.JedisSentinelPool;
@@ -50,16 +49,6 @@ public JedisCommands getJedisCommands() {
 
     }
 
-    public void execute(RedisAction action) {
-        switch (action.command) {
-            case SET:
-                getJedisCommands().set(action.key, action.value);
-                break;
-            default:
-                throw new IllegalArgumentException("Cannot process such data type: " + action.command);
-        }
-    }
-
     @Override
     public void close() {
         if (jedisCluster != null) jedisCluster.close();

src/main/java/org/apache/flink/connector/redis/shared/command/RedisCommand.java

@@ -0,0 +1,10 @@
+package org.apache.flink.connector.redis.shared.command;
+
+import org.apache.flink.connector.redis.shared.JedisConnector;
+
+import java.io.Serializable;
+
+public interface RedisCommand extends Serializable {
+    void send(JedisConnector connector);
+    long getMessageSize();
+}

src/main/java/org/apache/flink/connector/redis/shared/command/StreamRedisCommand.java

@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.flink.connector.redis.shared.command;
+
+import org.apache.flink.connector.redis.shared.JedisConnector;
+import redis.clients.jedis.StreamEntryID;
+
+import java.util.Map;
+
+public class StreamRedisCommand implements RedisCommand {
+    public final String key;
+    public final Map<String, String> value;
+
+    private StreamRedisCommand(String key, Map<String, String> value) {
+        this.key = key;
+        this.value = value;
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    @Override
+    public void send(JedisConnector connector) {
+        connector.getJedisCommands().xadd(key, StreamEntryID.NEW_ENTRY, value);
+    }
+
+    @Override
+    public long getMessageSize() {
+        return value.entrySet().stream()
+                .map(k -> k.getKey().length()+k.getValue().length())
+                .reduce(Integer::sum)
+                .orElse(0);
+    }
+
+    public static class Builder {
+        private String key;
+        private Map<String, String> value;
+
+
+        public Builder withKey(String key) {
+            this.key = key;
+            return this;
+        }
+
+        public Builder withValue(Map<String, String> value) {
+            this.value = value;
+            return this;
+        }
+
+        public StreamRedisCommand build() {
+            return new StreamRedisCommand(key, value);
+        }
+    }
+}

src/main/java/org/apache/flink/connector/redis/sink2/RedisAction.java renamed to src/main/java/org/apache/flink/connector/redis/shared/command/StringRedisCommand.java

@@ -14,17 +14,15 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.flink.connector.redis.sink2;
+package org.apache.flink.connector.redis.shared.command;
 
-import java.io.Serializable;
+import org.apache.flink.connector.redis.shared.JedisConnector;
 
-public class RedisAction implements Serializable {
-    public final RedisCommand command;
+public class StringRedisCommand implements RedisCommand {
     public final String key;
     public final String value;
 
-    private RedisAction(RedisCommand command, String key, String value) {
-        this.command = command;
+    private StringRedisCommand(String key, String value) {
         this.key = key;
         this.value = value;
     }
@@ -33,15 +31,20 @@ public static Builder builder() {
         return new Builder();
     }
 
+    @Override
+    public void send(JedisConnector connector) {
+        connector.getJedisCommands().set(key, value);
+    }
+
+    @Override
+    public long getMessageSize() {
+        return value.length();
+    }
+
     public static class Builder {
-        private RedisCommand command;
         private String key;
         private String value;
 
-        public Builder withCommand(RedisCommand command) {
-            this.command = command;
-            return this;
-        }
 
         public Builder withKey(String key) {
             this.key = key;
@@ -53,10 +56,8 @@ public Builder withValue(String value) {
             return this;
         }
 
-        public RedisAction build() {
-            return new RedisAction(command, key, value);
+        public StringRedisCommand build() {
+            return new StringRedisCommand(key, value);
         }
     }
-
-
 }

src/main/java/org/apache/flink/connector/redis/shared/config/JedisSentinelConfig.java

@@ -17,8 +17,6 @@
 package org.apache.flink.connector.redis.shared.config;
 
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import redis.clients.jedis.Protocol;
 
 import java.util.HashSet;

src/main/java/org/apache/flink/connector/redis/sink2/RedisCommand.java

@@ -17,9 +17,9 @@
 package org.apache.flink.connector.redis.sink2;
 
 import org.apache.flink.api.common.functions.Function;
+import org.apache.flink.connector.redis.shared.command.RedisCommand;
 
 import java.io.Serializable;
-import java.util.Optional;
 
 /**
  * Function that creates the description how the input data should be mapped to redis type.
@@ -28,48 +28,5 @@
  */
 public interface RedisSerializer<T> extends Function, Serializable {
 
-    /**
-     * Returns a redis command.
-     *
-     * @return RedisCommand
-     */
-    RedisCommand getCommand(T data);
-
-    /**
-     * Extracts key from data.
-     *
-     * @param data source data
-     * @return key
-     */
-    String getKeyFromData(T data);
-
-    /**
-     * Extracts value from data.
-     *
-     * @param data source data
-     * @return value
-     */
-    String getValueFromData(T data);
-
-    /**
-     * Extracts the additional key from data as an {@link Optional<String>}.
-     * The default implementation returns an empty Optional.
-     *
-     * @param data
-     * @return Optional
-     */
-    default Optional<String> getAdditionalKey(T data) {
-        return Optional.empty();
-    }
-
-    /**
-     * Extracts the additional time to live (TTL) for data as an {@link Optional<Integer>}.
-     * The default implementation returns an empty Optional.
-     *
-     * @param data
-     * @return Optional
-     */
-    default Optional<Integer> getAdditionalTTL(T data) {
-        return Optional.empty();
-    }
+    RedisCommand getMessage(T input);
 }

src/main/java/org/apache/flink/connector/redis/sink2/RedisSerializer.java

@@ -18,25 +18,16 @@
 package org.apache.flink.connector.redis.sink2.async;
 
 import org.apache.flink.api.connector.sink2.Sink;
-import org.apache.flink.api.connector.sink2.SinkWriter;
 import org.apache.flink.connector.base.sink.writer.AsyncSinkWriter;
-import org.apache.flink.connector.base.sink.writer.BufferedRequestState;
-import org.apache.flink.connector.base.sink.writer.ElementConverter;
 import org.apache.flink.connector.redis.shared.JedisConnector;
-import org.apache.flink.connector.redis.sink2.RedisAction;
+import org.apache.flink.connector.redis.shared.command.RedisCommand;
 import org.apache.flink.connector.redis.sink2.RedisSinkConfig;
-import redis.clients.jedis.Protocol.Command;
-import redis.clients.jedis.commands.JedisCommands;
 
-import java.io.IOException;
-import java.util.*;
-import java.util.concurrent.CompletableFuture;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.Consumer;
-import java.util.function.Supplier;
 
-import static redis.clients.jedis.Protocol.Command.SET;
-
-public class RedisAsyncWriter<T> extends AsyncSinkWriter<T, RedisAction> {
+public class RedisAsyncWriter<T> extends AsyncSinkWriter<T, RedisCommand> {
 
     private final JedisConnector jedisConnector;
 
@@ -51,11 +42,11 @@ public RedisAsyncWriter(JedisConnector jedisConnector,
     }
 
     @Override
-    protected void submitRequestEntries(List<RedisAction> actions, Consumer<List<RedisAction>> toRetry) {
-        List<RedisAction> errors = new ArrayList<>();
-        for(RedisAction action : actions) {
+    protected void submitRequestEntries(List<RedisCommand> actions, Consumer<List<RedisCommand>> toRetry) {
+        List<RedisCommand> errors = new ArrayList<>();
+        for(RedisCommand action : actions) {
             try {
-                this.jedisConnector.execute(action);
+                action.send(this.jedisConnector);
             } catch (Exception e) {
                 errors.add(action);
             }
@@ -64,8 +55,8 @@ protected void submitRequestEntries(List<RedisAction> actions, Consumer<List<Red
     }
 
     @Override
-    protected long getSizeInBytes(RedisAction redisAction) {
-        return redisAction.value.length();
+    protected long getSizeInBytes(RedisCommand message) {
+        return message.getMessageSize();
     }
 
     public void close() {

src/main/java/org/apache/flink/connector/redis/sink2/async/RedisAsyncWriter.java

@@ -18,10 +18,10 @@
 
 import org.apache.flink.api.connector.sink2.SinkWriter;
 import org.apache.flink.connector.base.sink.writer.ElementConverter;
-import org.apache.flink.connector.redis.sink2.RedisAction;
+import org.apache.flink.connector.redis.shared.command.RedisCommand;
 import org.apache.flink.connector.redis.sink2.RedisSerializer;
 
-public class RedisConverter<T> implements ElementConverter<T, RedisAction> {
+public class RedisConverter<T> implements ElementConverter<T, RedisCommand> {
 
     private final RedisSerializer<T> serializer;
 
@@ -30,11 +30,7 @@ public RedisConverter(RedisSerializer<T> serializer) {
     }
 
     @Override
-    public RedisAction apply(T input, SinkWriter.Context context) {
-        return RedisAction.builder()
-                .withCommand(serializer.getCommand(input))
-                .withKey(serializer.getKeyFromData(input))
-                .withValue(serializer.getValueFromData(input))
-                .build();
+    public RedisCommand apply(T input, SinkWriter.Context context) {
+        return serializer.getMessage(input);
     }
 }

src/main/java/org/apache/flink/connector/redis/sink2/async/RedisConverter.java

@@ -19,7 +19,7 @@
 
 import org.apache.flink.api.connector.sink2.SinkWriter;
 import org.apache.flink.connector.redis.shared.JedisConnector;
-import org.apache.flink.connector.redis.sink2.RedisAction;
+import org.apache.flink.connector.redis.shared.command.RedisCommand;
 import org.apache.flink.connector.redis.sink2.RedisSerializer;
 
 import java.io.IOException;
@@ -30,21 +30,17 @@ public class RedisWriter<T> implements SinkWriter<T> {
     private final JedisConnector jedisConnector;
     private final RedisSerializer<T> serializer;
 
-    private final Queue<RedisAction> queue = new ArrayDeque<>();
+    private final Queue<RedisCommand> queue = new ArrayDeque<>();
     public RedisWriter(JedisConnector jedisConnector, RedisSerializer<T> serializer) {
         this.jedisConnector = jedisConnector;
         this.serializer = serializer;
     }
 
     @Override
     public void write(T input, Context context) throws IOException, InterruptedException {
-        RedisAction action = RedisAction.builder()
-                .withCommand(serializer.getCommand(input))
-                .withKey(serializer.getKeyFromData(input))
-                .withValue(serializer.getValueFromData(input))
-                .build();
+        RedisCommand message = serializer.getMessage(input);
 
-        queue.add(action);
+        queue.add(message);
     }
 
     @Override
@@ -54,8 +50,8 @@ public void flush(boolean endOfInput) throws IOException, InterruptedException {
 
     private void flush() {
         while(!this.queue.isEmpty()) {
-            RedisAction element = this.queue.remove();
-            this.jedisConnector.execute(element);
+            RedisCommand element = this.queue.remove();
+            element.send(this.jedisConnector);
         }
     }