changes

Author: eskabetxe

src/main/java/org/apache/flink/connector/redis/source/RedisRecordEmitter.java

@@ -0,0 +1,21 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.api.connector.source.SourceOutput;
+import org.apache.flink.connector.base.source.reader.RecordEmitter;
+
+import java.util.function.Function;
+
+public class RedisRecordEmitter<T> implements RecordEmitter<String, T, RedisSourceSplit> {
+
+    final Function<String, T> deserializer;
+
+    public RedisRecordEmitter(Function<String, T> deserializer) {
+        this.deserializer = deserializer;
+    }
+
+    @Override
+    public void emitRecord(String element, SourceOutput<T> output, RedisSourceSplit testSourceSplit) throws Exception {
+        output.collect(this.deserializer.apply(element));
+//        output.collect(this.deserializer.apply(element), element.getTimestamp());
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisSource.java

@@ -0,0 +1,53 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.api.connector.source.*;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.io.SimpleVersionedSerializer;
+
+import java.util.List;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+public class RedisSource<T> implements Source<T, RedisSourceSplit, RedisStateEnumerator> {
+
+    private final Configuration configuration;
+    private final Function<String, T> deserializer;
+
+    public RedisSource(Function<String, T> deserializer) {
+        this.deserializer = deserializer;
+        this.configuration = new Configuration();
+    }
+
+    @Override
+    public Boundedness getBoundedness() {
+        return Boundedness.CONTINUOUS_UNBOUNDED;
+    }
+
+    @Override
+    public SourceReader<T, RedisSourceSplit> createReader(SourceReaderContext sourceReaderContext) throws Exception {
+        final Supplier<RedisSplitReader> splitReaderSupplier = () -> new RedisSplitReader(events);
+        RedisRecordEmitter<T> recordEmitter = new RedisRecordEmitter<>(this.deserializer);
+
+        return new RedisSourceReader<>(splitReaderSupplier, recordEmitter, this.configuration, sourceReaderContext);
+    }
+
+    @Override
+    public SplitEnumerator<RedisSourceSplit, RedisStateEnumerator> createEnumerator(SplitEnumeratorContext<RedisSourceSplit> splitEnumeratorContext) throws Exception {
+        return new RedisSplitEnumerator(splitEnumeratorContext);
+    }
+
+    @Override
+    public SplitEnumerator<RedisSourceSplit, RedisStateEnumerator> restoreEnumerator(SplitEnumeratorContext<RedisSourceSplit> splitEnumeratorContext, RedisStateEnumerator c) throws Exception {
+        return new RedisSplitEnumerator(splitEnumeratorContext);
+    }
+
+    @Override
+    public SimpleVersionedSerializer<RedisSourceSplit> getSplitSerializer() {
+        return new RedisSplitSerializer();
+    }
+
+    @Override
+    public SimpleVersionedSerializer<RedisStateEnumerator> getEnumeratorCheckpointSerializer() {
+        return new RedisStateSerializer();
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisSourceReader.java

@@ -0,0 +1,48 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.api.connector.source.SourceReaderContext;
+import org.apache.flink.configuration.ConfigOptions;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.connector.base.source.reader.RecordEmitter;
+import org.apache.flink.connector.base.source.reader.SingleThreadMultiplexSourceReaderBase;
+
+import java.util.Map;
+import java.util.function.Supplier;
+
+public class RedisSourceReader<T> extends SingleThreadMultiplexSourceReaderBase<String, T, RedisSourceSplit, RedisSourceSplit> {
+
+    public RedisSourceReader(Supplier<RedisSplitReader> splitReaderSupplier,
+                             RecordEmitter<String, T, RedisSourceSplit> recordEmitter,
+                             Configuration config,
+                             SourceReaderContext context) {
+        super(splitReaderSupplier::get, recordEmitter, config, context);
+    }
+
+    @Override
+    public void start() {
+        // we request a split only if we did not get splits during the checkpoint restore
+        if (this.getNumberOfCurrentlyAssignedSplits() == 0) {
+            this.context.sendSplitRequest();
+        }
+    }
+
+    @Override
+    protected void onSplitFinished(Map<String, RedisSourceSplit> map) {
+        boolean finishSplits = map.isEmpty();
+        if (finishSplits) {
+            this.notifyNoMoreSplits();
+        } else {
+            this.context.sendSplitRequest();
+        }
+    }
+
+    @Override
+    protected RedisSourceSplit initializedState(RedisSourceSplit testSourceSplit) {
+        return testSourceSplit;
+    }
+
+    @Override
+    protected RedisSourceSplit toSplitType(String s, RedisSourceSplit testSourceSplit) {
+        return testSourceSplit;
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisSourceSplit.java

@@ -0,0 +1,17 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.api.connector.source.SourceSplit;
+
+public class RedisSourceSplit implements SourceSplit {
+
+    final long id;
+
+    public RedisSourceSplit(long id) {
+        this.id = id;
+    }
+
+    @Override
+    public String splitId() {
+        return String.valueOf(this.id);
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisSplitEnumerator.java

@@ -0,0 +1,49 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.api.connector.source.SplitEnumerator;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.util.List;
+
+public class RedisSplitEnumerator implements SplitEnumerator<RedisSourceSplit, RedisStateEnumerator> {
+
+    private final SplitEnumeratorContext<RedisSourceSplit> context;
+    private boolean firstEmitted = false;
+
+    public RedisSplitEnumerator(SplitEnumeratorContext<RedisSourceSplit> context) {
+        this.context = context;
+    }
+
+    @Override
+    public void start() {
+    }
+
+    @Override
+    public void handleSplitRequest(int subtaskId, @Nullable String s) {
+        if (!this.firstEmitted) {
+            this.context.assignSplit(new RedisSourceSplit(subtaskId), subtaskId);
+            this.firstEmitted = true;
+        } else {
+            this.context.assignSplit(new RedisSourceSplit(System.currentTimeMillis()), subtaskId);
+        }
+    }
+
+    @Override
+    public void addSplitsBack(List<RedisSourceSplit> list, int i) {
+    }
+
+    @Override
+    public void addReader(int i) {
+    }
+
+    @Override
+    public RedisStateEnumerator snapshotState(long l) throws Exception {
+        return new RedisStateEnumerator();
+    }
+
+    @Override
+    public void close() throws IOException {
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisSplitReader.java

@@ -0,0 +1,45 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
+import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
+import org.apache.flink.connector.base.source.reader.splitreader.SplitsChange;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+public class RedisSplitReader implements SplitReader<String, RedisSourceSplit> {
+
+    private final List<String> events;
+    private RedisSourceSplit split;
+
+    public RedisSplitReader(List<String> events) {
+        this.events = events;
+    }
+
+    @Override
+    public RecordsWithSplitIds<String> fetch() throws IOException {
+        if (this.split == null) {
+            return null;
+        }
+        if (split.splitId().equals("0")) {
+            return new RedisSplitRecords(this.split.splitId(), events);
+        } else {
+            return new RedisSplitRecords(this.split.splitId(), new ArrayList<>());
+        }
+    }
+
+    @Override
+    public void handleSplitsChanges(SplitsChange<RedisSourceSplit> splitsChange) {
+        if (splitsChange.splits().isEmpty()) return;
+        this.split = splitsChange.splits().get(0);
+    }
+
+    @Override
+    public void wakeUp() {
+    }
+
+    @Override
+    public void close() throws Exception {
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisSplitRecords.java

@@ -0,0 +1,43 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
+
+import javax.annotation.Nullable;
+import java.util.*;
+
+public class RedisSplitRecords implements RecordsWithSplitIds<String> {
+
+    private final String splitId;
+    private final Iterator<String> recordIterator;
+
+    public RedisSplitRecords(String splitId, List<String> records) {
+        this.splitId = splitId;
+        this.recordIterator = records.iterator();
+    }
+
+    @Nullable
+    @Override
+    public String nextSplit() {
+        if (this.recordIterator.hasNext()) {
+            return this.splitId;
+        }
+        return null;
+    }
+
+    @Nullable
+    @Override
+    public String nextRecordFromSplit() {
+        if (this.recordIterator.hasNext()) {
+            return this.recordIterator.next();
+        } else {
+            return null;
+        }
+    }
+
+    @Override
+    public Set<String> finishedSplits() {
+        HashSet<String> ids = new HashSet<>();
+        ids.add(splitId);
+        return ids;
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisSplitSerializer.java

@@ -0,0 +1,29 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.core.io.SimpleVersionedSerializer;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public class RedisSplitSerializer implements SimpleVersionedSerializer<RedisSourceSplit> {
+
+    private static final int CURRENT_VERSION = 0;
+
+    @Override
+    public int getVersion() {
+        return CURRENT_VERSION;
+    }
+
+    @Override
+    public byte[] serialize(RedisSourceSplit testSplitReader) throws IOException {
+        final ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
+        buffer.putLong(0, testSplitReader.id);
+        return buffer.array();
+    }
+
+    @Override
+    public RedisSourceSplit deserialize(int version, byte[] serialized) throws IOException {
+        final ByteBuffer buffer = ByteBuffer.wrap(serialized);
+        return new RedisSourceSplit(buffer.getLong());
+    }
+}

src/main/java/org/apache/flink/connector/redis/source/RedisStateEnumerator.java

@@ -0,0 +1,4 @@
+package org.apache.flink.connector.redis.source;
+
+public class RedisStateEnumerator {
+}

src/main/java/org/apache/flink/connector/redis/source/RedisStateSerializer.java

@@ -0,0 +1,22 @@
+package org.apache.flink.connector.redis.source;
+
+import org.apache.flink.core.io.SimpleVersionedSerializer;
+
+import java.io.IOException;
+
+public class RedisStateSerializer implements SimpleVersionedSerializer<RedisStateEnumerator> {
+    @Override
+    public int getVersion() {
+        return 0;
+    }
+
+    @Override
+    public byte[] serialize(RedisStateEnumerator testStateEnumerator) throws IOException {
+        return new byte[0];
+    }
+
+    @Override
+    public RedisStateEnumerator deserialize(int i, byte[] bytes) throws IOException {
+        return new RedisStateEnumerator();
+    }
+}