feat: add the DeserializingSplitReader

pom.xml

@@ -47,6 +47,10 @@
             <groupId>com.google.protobuf</groupId>
             <artifactId>protobuf-java</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-core</artifactId>
+        </dependency>
         <dependency>
             <groupId>com.google.cloud</groupId>
             <artifactId>google-cloud-pubsublite</artifactId>
@@ -65,6 +69,10 @@
             <groupId>com.google.errorprone</groupId>
             <artifactId>error_prone_annotations</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.google.code.findbugs</groupId>
+            <artifactId>jsr305</artifactId>
+        </dependency>
         <dependency>
             <groupId>com.google.guava</groupId>
             <artifactId>guava</artifactId>

src/main/java/com/google/cloud/pubsublite/flink/MessageTimestampExtractor.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2021 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.pubsublite.flink;
+
+import com.google.cloud.Timestamp;
+import com.google.cloud.pubsublite.SequencedMessage;
+import java.io.Serializable;
+import java.time.Instant;
+
+public interface MessageTimestampExtractor extends Serializable {
+  static MessageTimestampExtractor publishTimeExtractor() {
+    return (MessageTimestampExtractor)
+        m -> Timestamp.fromProto(m.publishTime()).toDate().toInstant();
+  }
+
+  static MessageTimestampExtractor eventTimeExtractor() {
+    return (MessageTimestampExtractor)
+        m -> {
+          if (m.message().eventTime().isPresent()) {
+            return Timestamp.fromProto(m.message().eventTime().get()).toDate().toInstant();
+          }
+          return Timestamp.fromProto(m.publishTime()).toDate().toInstant();
+        };
+  }
+
+  Instant timestamp(SequencedMessage m) throws Exception;
+}

src/main/java/com/google/cloud/pubsublite/flink/PubsubLiteDeserializationSchema.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2021 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.pubsublite.flink;
+
+import com.google.cloud.pubsublite.SequencedMessage;
+import java.io.Serializable;
+import javax.annotation.Nullable;
+import org.apache.flink.api.common.serialization.DeserializationSchema;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+
+public interface PubsubLiteDeserializationSchema<T> extends Serializable {
+
+  static <T> PubsubLiteDeserializationSchema<T> dataOnly(DeserializationSchema<T> schema) {
+    return new PubsubLiteDeserializationSchema<T>() {
+      @Override
+      public void open(DeserializationSchema.InitializationContext context) throws Exception {
+        schema.open(context);
+      }
+
+      @Override
+      public T deserialize(SequencedMessage message) throws Exception {
+        return schema.deserialize(message.message().data().toByteArray());
+      }
+
+      @Override
+      public TypeInformation<T> getProducedType() {
+        return schema.getProducedType();
+      }
+    };
+  }
+
+  void open(DeserializationSchema.InitializationContext context) throws Exception;
+
+  /**
+   * Deserialize a Pub/Sub Lite message
+   *
+   * <p>If a message cannot be deserialized, the schema can either thrown a exception which will
+   * fail the source, or it can return null in which case the source will skip the message and
+   * proceed.
+   *
+   * @param message The pub/sub lite message
+   * @return The deserialized message as an object (null if the message cannot be deserialized).
+   */
+  @Nullable
+  T deserialize(SequencedMessage message) throws Exception;
+
+  TypeInformation<T> getProducedType();
+}

src/main/java/com/google/cloud/pubsublite/flink/reader/DeserializingSplitReader.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2021 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.pubsublite.flink.reader;
+
+import com.google.cloud.pubsublite.SequencedMessage;
+import com.google.cloud.pubsublite.flink.MessageTimestampExtractor;
+import com.google.cloud.pubsublite.flink.PubsubLiteDeserializationSchema;
+import com.google.cloud.pubsublite.flink.split.SubscriptionPartitionSplit;
+import com.google.common.collect.Multimap;
+import java.io.IOException;
+import java.time.Instant;
+import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import org.apache.flink.connector.base.source.reader.RecordsBySplits;
+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;
+
+public class DeserializingSplitReader<T>
+    implements SplitReader<Record<T>, SubscriptionPartitionSplit> {
+  final SplitReader<SequencedMessage, SubscriptionPartitionSplit> underlying;
+  final PubsubLiteDeserializationSchema<T> schema;
+  final MessageTimestampExtractor extractor;
+
+  public DeserializingSplitReader(
+      SplitReader<SequencedMessage, SubscriptionPartitionSplit> underlying,
+      PubsubLiteDeserializationSchema<T> schema,
+      MessageTimestampExtractor extractor) {
+    this.underlying = underlying;
+    this.schema = schema;
+    this.extractor = extractor;
+  }
+
+  @Override
+  public RecordsBySplits<Record<T>> fetch() throws IOException {
+    RecordsBySplits.Builder<Record<T>> builder = new RecordsBySplits.Builder<>();
+    RecordsWithSplitIds<SequencedMessage> fetch = underlying.fetch();
+    Multimap<String, SequencedMessage> messageMap = ReaderUtils.recordWithSplitsToMap(fetch);
+    for (Map.Entry<String, SequencedMessage> entry : messageMap.entries()) {
+      String split = entry.getKey();
+      SequencedMessage message = entry.getValue();
+      try {
+        @Nullable T value = schema.deserialize(message);
+        Instant timestamp = extractor.timestamp(message);
+        builder.add(split, Record.create(Optional.ofNullable(value), message.offset(), timestamp));
+      } catch (Exception e) {
+        throw new IOException(e);
+      }
+    }
+    builder.addFinishedSplits(fetch.finishedSplits());
+    return builder.build();
+  }
+
+  @Override
+  public void handleSplitsChanges(SplitsChange<SubscriptionPartitionSplit> splitsChange) {
+    underlying.handleSplitsChanges(splitsChange);
+  }
+
+  @Override
+  public void wakeUp() {
+    underlying.wakeUp();
+  }
+
+  @Override
+  public void close() throws Exception {
+    underlying.close();
+  }
+}

src/main/java/com/google/cloud/pubsublite/flink/reader/ReaderUtils.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2021 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.pubsublite.flink.reader;
+
+import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.Multimap;
+import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
+
+public class ReaderUtils {
+  public static <T> Multimap<String, T> recordWithSplitsToMap(RecordsWithSplitIds<T> records) {
+    ImmutableListMultimap.Builder<String, T> builder = ImmutableListMultimap.builder();
+    for (String split = records.nextSplit(); split != null; split = records.nextSplit()) {
+      for (T m = records.nextRecordFromSplit(); m != null; m = records.nextRecordFromSplit()) {
+        builder.put(split, m);
+      }
+    }
+    return builder.build();
+  }
+}

src/main/java/com/google/cloud/pubsublite/flink/reader/Record.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2021 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.pubsublite.flink.reader;
+
+import com.google.auto.value.AutoValue;
+import com.google.cloud.pubsublite.Offset;
+import java.time.Instant;
+import java.util.Optional;
+
+@AutoValue
+public abstract class Record<T> {
+  public abstract Optional<T> value();
+
+  public abstract Offset offset();
+
+  public abstract Instant timestamp();
+
+  public static <T> Record<T> create(Optional<T> value, Offset offset, Instant timestamp) {
+    return new AutoValue_Record<>(value, offset, timestamp);
+  }
+}

src/test/java/com/google/cloud/pubsublite/flink/TestUtilities.java

@@ -18,9 +18,6 @@
 import com.google.cloud.pubsublite.Offset;
 import com.google.cloud.pubsublite.SequencedMessage;
 import com.google.cloud.pubsublite.proto.Cursor;
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.Multimap;
-import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
 
 public class TestUtilities {
   public static SequencedMessage messageFromOffset(Offset offset) {
@@ -29,14 +26,4 @@ public static SequencedMessage messageFromOffset(Offset offset) {
             .setCursor(Cursor.newBuilder().setOffset(offset.value()))
             .build());
   }
-
-  public static <T> Multimap<String, T> recordWithSplitsToMap(RecordsWithSplitIds<T> records) {
-    Multimap<String, T> map = HashMultimap.create();
-    for (String split = records.nextSplit(); split != null; split = records.nextSplit()) {
-      for (T m = records.nextRecordFromSplit(); m != null; m = records.nextRecordFromSplit()) {
-        map.put(split, m);
-      }
-    }
-    return map;
-  }
 }