Pre-Work for adding Checkpointing (#3163)

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/base/FailureTrackingAirbyteMessageConsumer.java

@@ -25,7 +25,6 @@
 package io.airbyte.integrations.base;
 
 import io.airbyte.protocol.models.AirbyteMessage;
-import io.airbyte.protocol.models.AirbyteRecordMessage;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -59,15 +58,12 @@ public void start() throws Exception {
     }
   }
 
-  protected abstract void acceptTracked(AirbyteRecordMessage msg) throws Exception;
+  protected abstract void acceptTracked(AirbyteMessage msg) throws Exception;
 
   @Override
   public void accept(AirbyteMessage msg) throws Exception {
     try {
-      // ignore all other message types
-      if (msg.getType() == AirbyteMessage.Type.RECORD) {
-        acceptTracked(msg.getRecord());
-      }
+      acceptTracked(msg);
     } catch (Exception e) {
       hasFailed = true;
       throw e;

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/destination/buffered_stream_consumer/BufferedStreamConsumer.java

@@ -28,14 +28,14 @@
 import com.google.common.base.Preconditions;
 import io.airbyte.commons.concurrency.GracefulShutdownHandler;
 import io.airbyte.commons.concurrency.VoidCallable;
-import io.airbyte.commons.functional.CheckedBiConsumer;
 import io.airbyte.commons.functional.CheckedConsumer;
 import io.airbyte.commons.functional.CheckedFunction;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.lang.CloseableQueue;
 import io.airbyte.commons.lang.Queues;
 import io.airbyte.integrations.base.AirbyteStreamNameNamespacePair;
 import io.airbyte.integrations.base.FailureTrackingAirbyteMessageConsumer;
+import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
 import io.airbyte.queue.OnDiskQueue;
@@ -50,7 +50,6 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -146,11 +145,15 @@ protected void startTracked() throws Exception {
   }
 
   @Override
-  protected void acceptTracked(AirbyteRecordMessage message) throws Exception {
+  protected void acceptTracked(AirbyteMessage message) throws Exception {
     Preconditions.checkState(hasStarted, "Cannot accept records until consumer has started");
 
+    if (message.getType() != AirbyteMessage.Type.RECORD) {
+      return;
+    }
+
     // ignore other message types.
-    final AirbyteStreamNameNamespacePair pair = AirbyteStreamNameNamespacePair.fromRecordMessage(message);
+    final AirbyteStreamNameNamespacePair pair = AirbyteStreamNameNamespacePair.fromRecordMessage(message.getRecord());
     if (!pairs.contains(pair)) {
       throw new IllegalArgumentException(
           String.format("Message contained record from a stream that was not in the catalog. \ncatalog: %s , \nmessage: %s",
@@ -209,28 +212,12 @@ private static void writeStreamsWithNRecords(int minRecords,
 
           LOGGER.info("Writing stream {}. Max batch size: {}, Actual batch size: {}, Remaining buffered records: {}",
               pair, BufferedStreamConsumer.BATCH_SIZE, records.size(), writeBuffer.size());
-          recordWriter.accept(pair, records.stream());
+          recordWriter.accept(pair, records);
         } catch (Exception e) {
           throw new RuntimeException(e);
         }
       }
     }
   }
 
-  public interface OnStartFunction extends VoidCallable {}
-
-  public interface RecordWriter extends CheckedBiConsumer<AirbyteStreamNameNamespacePair, Stream<AirbyteRecordMessage>, Exception> {
-
-    @Override
-    void accept(AirbyteStreamNameNamespacePair pair, Stream<AirbyteRecordMessage> recordStream) throws Exception;
-
-  }
-
-  public interface OnCloseFunction extends CheckedConsumer<Boolean, Exception> {
-
-    @Override
-    void accept(Boolean hasFailed) throws Exception;
-
-  }
-
 }

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/destination/buffered_stream_consumer/OnCloseFunction.java

@@ -0,0 +1,34 @@
+/*
+ * MIT License
+ *
+ * Copyright (c) 2020 Airbyte
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package io.airbyte.integrations.destination.buffered_stream_consumer;
+
+import io.airbyte.commons.functional.CheckedConsumer;
+
+public interface OnCloseFunction extends CheckedConsumer<Boolean, Exception> {
+
+  @Override
+  void accept(Boolean hasFailed) throws Exception;
+
+}

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/destination/buffered_stream_consumer/OnStartFunction.java

@@ -0,0 +1,31 @@
+/*
+ * MIT License
+ *
+ * Copyright (c) 2020 Airbyte
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package io.airbyte.integrations.destination.buffered_stream_consumer;
+
+import io.airbyte.commons.concurrency.VoidCallable;
+
+public interface OnStartFunction extends VoidCallable {
+
+}

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/destination/buffered_stream_consumer/RecordWriter.java

@@ -0,0 +1,37 @@
+/*
+ * MIT License
+ *
+ * Copyright (c) 2020 Airbyte
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package io.airbyte.integrations.destination.buffered_stream_consumer;
+
+import io.airbyte.commons.functional.CheckedBiConsumer;
+import io.airbyte.integrations.base.AirbyteStreamNameNamespacePair;
+import io.airbyte.protocol.models.AirbyteRecordMessage;
+import java.util.List;
+
+public interface RecordWriter extends CheckedBiConsumer<AirbyteStreamNameNamespacePair, List<AirbyteRecordMessage>, Exception> {
+
+  @Override
+  void accept(AirbyteStreamNameNamespacePair pair, List<AirbyteRecordMessage> recordStream) throws Exception;
+
+}

airbyte-integrations/bases/base-java/src/test/java/io/airbyte/integrations/base/FailureTrackingAirbyteMessageConsumerTest.java

@@ -34,7 +34,6 @@
 
 import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.AirbyteMessage.Type;
-import io.airbyte.protocol.models.AirbyteRecordMessage;
 import org.junit.jupiter.api.Test;
 
 class FailureTrackingAirbyteMessageConsumerTest {
@@ -93,7 +92,7 @@ protected void startTracked() {
     }
 
     @Override
-    protected void acceptTracked(AirbyteRecordMessage s) {
+    protected void acceptTracked(AirbyteMessage s) {
 
     }
 

airbyte-integrations/connectors/destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/DefaultSqlOperations.java

@@ -29,8 +29,6 @@
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import java.sql.SQLException;
 import java.util.List;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -64,9 +62,7 @@ public String createTableQuery(String schemaName, String tableName) {
   }
 
   @Override
-  public void insertRecords(JdbcDatabase database, Stream<AirbyteRecordMessage> recordsStream, String schemaName, String tmpTableName)
-      throws SQLException {
-    final List<AirbyteRecordMessage> records = recordsStream.collect(Collectors.toList());
+  public void insertRecords(JdbcDatabase database, List<AirbyteRecordMessage> records, String schemaName, String tmpTableName) throws SQLException {
 
     // todo (cgardens) - move this into a postgres version of this. this syntax is VERY postgres
     // specific.

airbyte-integrations/connectors/destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/JdbcBufferedConsumerFactory.java

@@ -35,7 +35,7 @@
 import io.airbyte.integrations.destination.buffered_stream_consumer.BufferedStreamConsumer;
 import io.airbyte.integrations.destination.buffered_stream_consumer.BufferedStreamConsumer.OnCloseFunction;
 import io.airbyte.integrations.destination.buffered_stream_consumer.BufferedStreamConsumer.OnStartFunction;
-import io.airbyte.integrations.destination.buffered_stream_consumer.BufferedStreamConsumer.RecordWriter;
+import io.airbyte.integrations.destination.buffered_stream_consumer.RecordWriter;
 import io.airbyte.protocol.models.AirbyteStream;
 import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.ConfiguredAirbyteStream;
@@ -140,14 +140,14 @@ private static RecordWriter recordWriterFunction(JdbcDatabase database,
     final Map<AirbyteStreamNameNamespacePair, WriteConfig> pairToWriteConfig = writeConfigs.stream()
         .collect(Collectors.toUnmodifiableMap(JdbcBufferedConsumerFactory::toNameNamespacePair, Function.identity()));
 
-    return (pair, recordStream) -> {
+    return (pair, records) -> {
       if (!pairToWriteConfig.containsKey(pair)) {
         throw new IllegalArgumentException(
             String.format("Message contained record from a stream that was not in the catalog. \ncatalog: %s", Jsons.serialize(catalog)));
       }
 
       final WriteConfig writeConfig = pairToWriteConfig.get(pair);
-      sqlOperations.insertRecords(database, recordStream, writeConfig.getOutputSchemaName(), writeConfig.getTmpTableName());
+      sqlOperations.insertRecords(database, records, writeConfig.getOutputSchemaName(), writeConfig.getTmpTableName());
     };
   }
 

airbyte-integrations/connectors/destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/SqlOperations.java

@@ -26,7 +26,7 @@
 
 import io.airbyte.db.jdbc.JdbcDatabase;
 import io.airbyte.protocol.models.AirbyteRecordMessage;
-import java.util.stream.Stream;
+import java.util.List;
 
 // todo (cgardens) - is it necessary to expose so much configurability in this interface. review if
 // we can narrow the surface area.
@@ -84,7 +84,7 @@ public interface SqlOperations {
    * @param tableName name of table
    * @throws Exception exception
    */
-  void insertRecords(JdbcDatabase database, Stream<AirbyteRecordMessage> records, String schemaName, String tableName) throws Exception;
+  void insertRecords(JdbcDatabase database, List<AirbyteRecordMessage> records, String schemaName, String tableName) throws Exception;
 
   /**
    * Query to copy all records from source table to destination table. Both tables must be in the

airbyte-integrations/connectors/destination-meilisearch/src/main/java/io/airbyte/integrations/destination/meilisearch/MeiliSearchDestination.java

@@ -37,7 +37,7 @@
 import io.airbyte.integrations.base.Destination;
 import io.airbyte.integrations.base.IntegrationRunner;
 import io.airbyte.integrations.destination.buffered_stream_consumer.BufferedStreamConsumer;
-import io.airbyte.integrations.destination.buffered_stream_consumer.BufferedStreamConsumer.RecordWriter;
+import io.airbyte.integrations.destination.buffered_stream_consumer.RecordWriter;
 import io.airbyte.protocol.models.AirbyteConnectionStatus;
 import io.airbyte.protocol.models.AirbyteConnectionStatus.Status;
 import io.airbyte.protocol.models.AirbyteRecordMessage;
@@ -139,7 +139,7 @@ private static boolean indexExists(Client client, String indexName) throws Excep
   }
 
   private static RecordWriter recordWriterFunction(final Map<String, Index> indexNameToWriteConfig) {
-    return (namePair, recordStream) -> {
+    return (namePair, records) -> {
       final String resolvedIndexName = getIndexName(namePair.getName());
       if (!indexNameToWriteConfig.containsKey(resolvedIndexName)) {
         throw new IllegalArgumentException(
@@ -153,7 +153,8 @@ private static RecordWriter recordWriterFunction(final Map<String, Index> indexN
       // destinations work. There is not really a viable way to "transform" data after it is MeiliSearch.
       // Tools like DBT do not apply. Therefore, we need to try to write data in the most usable format
       // possible that does not require alteration.
-      final String json = Jsons.serialize(recordStream
+      final String json = Jsons.serialize(records
+          .stream()
           .map(AirbyteRecordMessage::getData)
           .peek(o -> ((ObjectNode) o).put(AB_PK_COLUMN, Names.toAlphanumericAndUnderscore(UUID.randomUUID().toString())))
           .collect(Collectors.toList()));

airbyte-integrations/connectors/destination-redshift/src/main/java/io/airbyte/integrations/destination/redshift/RedshiftSqlOperations.java

@@ -32,8 +32,6 @@
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import java.sql.SQLException;
 import java.util.List;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -54,9 +52,7 @@ public String createTableQuery(String schemaName, String tableName) {
   }
 
   @Override
-  public void insertRecords(JdbcDatabase database, Stream<AirbyteRecordMessage> recordsStream, String schemaName, String tmpTableName)
-      throws SQLException {
-    final List<AirbyteRecordMessage> records = recordsStream.collect(Collectors.toList());
+  public void insertRecords(JdbcDatabase database, List<AirbyteRecordMessage> records, String schemaName, String tmpTableName) throws SQLException {
     LOGGER.info("actual size of batch: {}", records.size());
 
     // query syntax:

airbyte-integrations/connectors/destination-snowflake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeSqlOperations.java

@@ -32,8 +32,6 @@
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import java.sql.SQLException;
 import java.util.List;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -54,9 +52,7 @@ public void createTableIfNotExists(JdbcDatabase database, String schemaName, Str
   }
 
   @Override
-  public void insertRecords(JdbcDatabase database, Stream<AirbyteRecordMessage> recordsStream, String schemaName, String tableName)
-      throws SQLException {
-    final List<AirbyteRecordMessage> records = recordsStream.collect(Collectors.toList());
+  public void insertRecords(JdbcDatabase database, List<AirbyteRecordMessage> records, String schemaName, String tableName) throws SQLException {
     LOGGER.info("actual size of batch: {}", records.size());
 
     // snowflake query syntax: