[source-postgres/mssql/mysql] Send state and count for full refresh (#39349)

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: btkcodedev <[email protected]>
Co-authored-by: Anjay Goel <[email protected]>
Co-authored-by: marcosmarxm <[email protected]>
Co-authored-by: Baz <[email protected]>
Co-authored-by: Natalie Kwong <[email protected]>
Co-authored-by: Audrey Maldonado <[email protected]>
Co-authored-by: gosusnp <[email protected]>
Co-authored-by: Augustin <[email protected]>
Co-authored-by: Abdul Rahman Zantout <[email protected]>
Co-authored-by: Anatolii Yatsuk <[email protected]>
Co-authored-by: Dhroov Makwana <[email protected]>
Co-authored-by: Alexandre Girard <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Serhii Lazebnyi <[email protected]>
Co-authored-by: Marcos Marx <[email protected]>
Co-authored-by: Ben Church <[email protected]>
Co-authored-by: Justin Flannery <[email protected]>
Co-authored-by: Mal Hancock <[email protected]>
Co-authored-by: Enrique Alcázar Garzás <[email protected]>
Co-authored-by: Daniela García Nistor <[email protected]>
Co-authored-by: Carlo Nuccio <[email protected]>
Co-authored-by: Natik Gadzhi <[email protected]>
Co-authored-by: Charlie Duong <[email protected]>
Co-authored-by: Yue Li <[email protected]>
Co-authored-by: Cristina Mariscal <[email protected]>
Co-authored-by: cristina.mariscal <[email protected]>
Co-authored-by: Gonzalo Villafañe Tapia <[email protected]>
Co-authored-by: Jérémy Denquin <[email protected]>
Co-authored-by: Maxime Carbonneau-Leclerc <[email protected]>
Co-authored-by: williammcguinness <[email protected]>
Co-authored-by: Marius Posta <[email protected]>
Co-authored-by: Stephane Geneix <[email protected]>
Co-authored-by: Gireesh Sreepathi <[email protected]>
Co-authored-by: Akash Kulkarni <[email protected]>
Co-authored-by: Danylo Jablonski <[email protected]>
Co-authored-by: Serhii Lazebnyi <[email protected]>
Co-authored-by: alafanechere <[email protected]>

Author: 35 people

airbyte-cdk/java/airbyte-cdk/README.md

@@ -173,7 +173,8 @@ corresponds to that version.
 ### Java CDK
 
 | Version | Date       | Pull Request                                               | Subject                                                                                                                                                        |
-| :------ | :--------- | :--------------------------------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+|:--------| :--------- | :--------------------------------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 0.40.1  | 2024-06-14 | [\#39349](https://github.com/airbytehq/airbyte/pull/39349) | Source stats for full refresh streams                                                                                                                          |
 | 0.40.0  | 2024-06-17 | [\#38622](https://github.com/airbytehq/airbyte/pull/38622) | Destinations: Implement refreshes logic in AbstractStreamOperation                                                                                             |
 | 0.39.0  | 2024-06-17 | [\#38067](https://github.com/airbytehq/airbyte/pull/38067) | Destinations: Breaking changes for refreshes (fail on INCOMPLETE stream status; ignore OVERWRITE sync mode)                                                    |
 | 0.38.2  | 2024-06-14 | [\#39460](https://github.com/airbytehq/airbyte/pull/39460) | Bump postgres JDBC driver version                                                                                                                              |

airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties

@@ -1 +1 @@
-version=0.40.0
+version=0.40.1

airbyte-cdk/java/airbyte-cdk/db-sources/src/main/kotlin/io/airbyte/cdk/integrations/source/jdbc/AbstractJdbcSource.kt

@@ -46,6 +46,9 @@ import io.airbyte.cdk.integrations.source.relationaldb.InitialLoadHandler
 import io.airbyte.cdk.integrations.source.relationaldb.RelationalDbQueryUtils
 import io.airbyte.cdk.integrations.source.relationaldb.RelationalDbQueryUtils.enquoteIdentifier
 import io.airbyte.cdk.integrations.source.relationaldb.TableInfo
+import io.airbyte.cdk.integrations.source.relationaldb.state.SourceStateIterator
+import io.airbyte.cdk.integrations.source.relationaldb.state.SourceStateMessageProducer
+import io.airbyte.cdk.integrations.source.relationaldb.state.StateEmitFrequency
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateManager
 import io.airbyte.commons.exceptions.ConfigErrorException
 import io.airbyte.commons.functional.CheckedConsumer
@@ -70,6 +73,7 @@ import java.sql.Connection
 import java.sql.PreparedStatement
 import java.sql.ResultSet
 import java.sql.SQLException
+import java.time.Duration
 import java.time.Instant
 import java.util.*
 import java.util.function.Consumer
@@ -173,12 +177,40 @@ abstract class AbstractJdbcSource<Datatype>(
                 cursorField,
             )
 
+        if (airbyteStream.syncMode == FULL_REFRESH) {
+            var defaultProducer = getSourceStateProducerForNonResumableFullRefreshStream(database)
+            if (defaultProducer != null) {
+                iterator =
+                    AutoCloseableIterators.transform(
+                        { autoCloseableIterator: AutoCloseableIterator<AirbyteMessage> ->
+                            SourceStateIterator(
+                                autoCloseableIterator,
+                                airbyteStream,
+                                defaultProducer,
+                                StateEmitFrequency(stateEmissionFrequency.toLong(), Duration.ZERO)
+                            )
+                        },
+                        iterator,
+                        AirbyteStreamUtils.convertFromNameAndNamespace(
+                            airbyteStream.stream.name,
+                            airbyteStream.stream.namespace
+                        )
+                    )
+            }
+        }
+
         return when (airbyteStream.syncMode) {
             FULL_REFRESH -> augmentWithStreamStatus(airbyteStream, iterator)
             else -> iterator
         }
     }
 
+    protected open fun getSourceStateProducerForNonResumableFullRefreshStream(
+        database: JdbcDatabase
+    ): SourceStateMessageProducer<AirbyteMessage>? {
+        return null
+    }
+
     open fun augmentWithStreamStatus(
         airbyteStream: ConfiguredAirbyteStream,
         streamItrator: AutoCloseableIterator<AirbyteMessage>

airbyte-cdk/java/airbyte-cdk/db-sources/src/main/kotlin/io/airbyte/cdk/integrations/source/relationaldb/state/NonResumableStateMessageProducer.kt

@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.integrations.source.relationaldb.state
+
+import io.airbyte.protocol.models.v0.AirbyteMessage
+import io.airbyte.protocol.models.v0.AirbyteStateMessage
+import io.airbyte.protocol.models.v0.AirbyteStreamState
+import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream
+import io.airbyte.protocol.models.v0.StreamDescriptor
+
+class NonResumableStateMessageProducer<T>(
+    private val isCdc: Boolean,
+    private val sourceStateMessageProducer: SourceStateMessageProducer<T>
+) : SourceStateMessageProducer<AirbyteMessage> {
+    override fun generateStateMessageAtCheckpoint(
+        stream: ConfiguredAirbyteStream?
+    ): AirbyteStateMessage? {
+        return null
+    }
+
+    override fun processRecordMessage(
+        stream: ConfiguredAirbyteStream?,
+        message: AirbyteMessage
+    ): AirbyteMessage {
+        return message
+    }
+
+    override fun createFinalStateMessage(stream: ConfiguredAirbyteStream?): AirbyteStateMessage? {
+        if (isCdc) {
+            return sourceStateMessageProducer.createFinalStateMessage(stream)
+        } else {
+            val airbyteStreamState =
+                AirbyteStreamState()
+                    .withStreamDescriptor(
+                        StreamDescriptor()
+                            .withName(stream!!.stream.name)
+                            .withNamespace(stream.stream.namespace),
+                    )
+
+            return AirbyteStateMessage()
+                .withType(AirbyteStateMessage.AirbyteStateType.STREAM)
+                .withStream(airbyteStreamState)
+        }
+    }
+
+    // no intermediate state message.
+    override fun shouldEmitStateMessage(stream: ConfiguredAirbyteStream?): Boolean {
+        return false
+    }
+}

airbyte-cdk/java/airbyte-cdk/db-sources/src/main/kotlin/io/airbyte/cdk/integrations/source/relationaldb/state/SourceStateIterator.kt

@@ -7,11 +7,13 @@ import com.google.common.collect.AbstractIterator
 import io.airbyte.protocol.models.v0.AirbyteMessage
 import io.airbyte.protocol.models.v0.AirbyteStateStats
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream
-import io.airbyte.protocol.models.v0.SyncMode
+import io.github.oshai.kotlinlogging.KotlinLogging
 import java.time.Duration
 import java.time.Instant
 import java.time.OffsetDateTime
 
+private val LOGGER = KotlinLogging.logger {}
+
 open class SourceStateIterator<T>(
     private val messageIterator: Iterator<T>,
     private val stream: ConfiguredAirbyteStream?,
@@ -40,11 +42,9 @@ open class SourceStateIterator<T>(
             ) {
                 val stateMessage =
                     sourceStateMessageProducer.generateStateMessageAtCheckpoint(stream)
-                if (shouldAttachCountWithState()) {
-                    stateMessage!!.withSourceStats(
-                        AirbyteStateStats().withRecordCount(recordCount.toDouble())
-                    )
-                }
+                stateMessage!!.withSourceStats(
+                    AirbyteStateStats().withRecordCount(recordCount.toDouble())
+                )
 
                 recordCount = 0L
                 lastCheckpoint = Instant.now()
@@ -65,11 +65,10 @@ open class SourceStateIterator<T>(
             hasEmittedFinalState = true
             val finalStateMessageForStream =
                 sourceStateMessageProducer.createFinalStateMessage(stream)
-            if (shouldAttachCountWithState()) {
-                finalStateMessageForStream!!.withSourceStats(
-                    AirbyteStateStats().withRecordCount(recordCount.toDouble())
-                )
-            }
+            finalStateMessageForStream!!.withSourceStats(
+                AirbyteStateStats().withRecordCount(recordCount.toDouble())
+            )
+
             recordCount = 0L
             return AirbyteMessage()
                 .withType(AirbyteMessage.Type.STATE)
@@ -79,14 +78,6 @@ open class SourceStateIterator<T>(
         }
     }
 
-    /**
-     * We are disabling counts for FULL_REFRESH streams cause there is are issues with it. We should
-     * re-enable it once we do the work for project Counts: Emit Counts in Full Refresh
-     */
-    private fun shouldAttachCountWithState(): Boolean {
-        return stream?.syncMode != SyncMode.FULL_REFRESH
-    }
-
     // This method is used to check if we should emit a state message. If the record count is set to
     // 0,
     // we should not emit a state message.

airbyte-cdk/java/airbyte-cdk/db-sources/src/testFixtures/kotlin/io/airbyte/cdk/integrations/debezium/CdcSourceTest.kt

@@ -781,7 +781,11 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
                 modelsSchema(),
             )
         } else {
-            assertExpectedStateMessageCountMatches(stateMessages1, MODEL_RECORDS.size.toLong())
+            // We are expecting count match for all streams, including non RFR streams.
+            assertExpectedStateMessageCountMatches(
+                stateMessages1,
+                MODEL_RECORDS.size.toLong() + MODEL_RECORDS_2.size.toLong()
+            )
 
             // Expect state and record message from MODEL_RECORDS_2.
             assertStreamStatusTraceMessageIndex(
@@ -923,8 +927,11 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
         waitForCdcRecords(modelsSchema(), MODELS_STREAM_NAME, 1)
 
         // assertExpectedStateMessages(stateMessages1)
-        // Non resumeable full refresh does not get any state messages.
-        assertExpectedStateMessageCountMatches(stateMessages1, MODEL_RECORDS.size.toLong())
+        // Non resumeable full refresh will also get state messages with count.
+        assertExpectedStateMessageCountMatches(
+            stateMessages1,
+            MODEL_RECORDS.size.toLong() + MODEL_RECORDS_2.size.toLong()
+        )
         assertExpectedRecords(
             (MODEL_RECORDS_2 + MODEL_RECORDS).toSet(),
             recordMessages1,
@@ -933,15 +940,18 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
             modelsSchema(),
         )
 
+        // Platform will remove non RFR streams before each new sync.
         val state = Jsons.jsonNode(listOf(stateMessages1[stateMessages1.size - 1]))
+        val streamStates = state.get(0).get("global").get("stream_states") as ArrayNode
+        removeStreamState(MODELS_STREAM_NAME_2, streamStates)
+
         val read2 = source().read(config()!!, configuredCatalog, state)
         val actualRecords2 = AutoCloseableIterators.toListAndClose(read2)
 
         val recordMessages2 = extractRecordMessages(actualRecords2)
         val stateMessages2 = extractStateMessages(actualRecords2)
 
-        assertExpectedStateMessagesFromIncrementalSync(stateMessages2)
-        assertExpectedStateMessageCountMatches(stateMessages2, 1)
+        assertExpectedStateMessageCountMatches(stateMessages2, 1 + MODEL_RECORDS_2.size.toLong())
         assertExpectedRecords(
             (MODEL_RECORDS_2 + listOf(puntoRecord)).toSet(),
             recordMessages2,

airbyte-integrations/connectors/source-mssql/build.gradle

@@ -3,7 +3,7 @@ plugins {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.38.1'
+    cdkVersionRequired = '0.40.1'
     features = ['db-sources']
     useLocalCdk = false
 }

airbyte-integrations/connectors/source-mssql/metadata.yaml

@@ -9,7 +9,7 @@ data:
   connectorSubtype: database
   connectorType: source
   definitionId: b5ea17b1-f170-46dc-bc31-cc744ca984c1
-  dockerImageTag: 4.0.29
+  dockerImageTag: 4.0.30
   dockerRepository: airbyte/source-mssql
   documentationUrl: https://docs.airbyte.com/integrations/sources/mssql
   githubIssueLabel: source-mssql

airbyte-integrations/connectors/source-mssql/src/main/java/io/airbyte/integrations/source/mssql/MssqlSource.java

@@ -9,6 +9,7 @@
 import static io.airbyte.cdk.integrations.debezium.internals.DebeziumEventConverter.CDC_UPDATED_AT;
 import static io.airbyte.cdk.integrations.source.relationaldb.RelationalDbQueryUtils.*;
 import static io.airbyte.cdk.integrations.source.relationaldb.RelationalDbReadUtil.identifyStreamsForCursorBased;
+import static io.airbyte.integrations.source.mssql.MssqlCdcHelper.isCdc;
 import static io.airbyte.integrations.source.mssql.MssqlQueryUtils.getCursorBasedSyncStatusForStreams;
 import static io.airbyte.integrations.source.mssql.MssqlQueryUtils.getTableSizeInfoForStreams;
 import static io.airbyte.integrations.source.mssql.initialsync.MssqlInitialReadUtil.*;
@@ -37,6 +38,8 @@
 import io.airbyte.cdk.integrations.source.relationaldb.InitialLoadHandler;
 import io.airbyte.cdk.integrations.source.relationaldb.TableInfo;
 import io.airbyte.cdk.integrations.source.relationaldb.models.CursorBasedStatus;
+import io.airbyte.cdk.integrations.source.relationaldb.state.NonResumableStateMessageProducer;
+import io.airbyte.cdk.integrations.source.relationaldb.state.SourceStateMessageProducer;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateGeneratorUtils;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateManager;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateManagerFactory;
@@ -639,7 +642,7 @@ protected void initializeForStateManager(final JdbcDatabase database,
       return;
     }
     var sourceConfig = database.getSourceConfig();
-    if (MssqlCdcHelper.isCdc(sourceConfig)) {
+    if (isCdc(sourceConfig)) {
       initialLoadStateManager = getMssqlInitialLoadGlobalStateManager(database, catalog, stateManager, tableNameToTable, getQuoteString());
     } else {
       final MssqlCursorBasedStateManager cursorBasedStateManager = new MssqlCursorBasedStateManager(stateManager.getRawStateMessages(), catalog);
@@ -656,7 +659,7 @@ public InitialLoadHandler<JDBCType> getInitialLoadHandler(final JdbcDatabase dat
                                                             final ConfiguredAirbyteCatalog catalog,
                                                             final StateManager stateManager) {
     var sourceConfig = database.getSourceConfig();
-    if (MssqlCdcHelper.isCdc(sourceConfig)) {
+    if (isCdc(sourceConfig)) {
       return getMssqlFullRefreshInitialLoadHandler(database, catalog, initialLoadStateManager, stateManager, airbyteStream, Instant.now(),
           getQuoteString())
               .get();
@@ -677,6 +680,11 @@ public boolean supportResumableFullRefresh(final JdbcDatabase database, final Co
     return false;
   }
 
+  @Override
+  protected SourceStateMessageProducer<AirbyteMessage> getSourceStateProducerForNonResumableFullRefreshStream(final JdbcDatabase database) {
+    return new NonResumableStateMessageProducer<>(isCdc(database.getSourceConfig()), initialLoadStateManager);
+  }
+
   @NotNull
   @Override
   public AutoCloseableIterator<AirbyteMessage> augmentWithStreamStatus(@NotNull final ConfiguredAirbyteStream airbyteStream,

airbyte-integrations/connectors/source-mssql/src/main/java/io/airbyte/integrations/source/mssql/initialsync/MssqlInitialLoadGlobalStateManager.java

@@ -30,6 +30,7 @@ public class MssqlInitialLoadGlobalStateManager extends MssqlInitialLoadStateMan
 
   // No special handling for resumable full refresh streams. We will report the cursor as it is.
   private Set<AirbyteStreamNameNamespacePair> resumableFullRefreshStreams;
+  private Set<AirbyteStreamNameNamespacePair> nonResumableFullRefreshStreams;
 
   public MssqlInitialLoadGlobalStateManager(final InitialLoadStreams initialLoadStreams,
                                             final Map<AirbyteStreamNameNamespacePair, OrderedColumnInfo> pairToOrderedColInfo,
@@ -60,16 +61,21 @@ private void initStreams(final InitialLoadStreams initialLoadStreams,
                            final ConfiguredAirbyteCatalog catalog) {
     this.streamsThatHaveCompletedSnapshot = new HashSet<>();
     this.resumableFullRefreshStreams = new HashSet<>();
+    this.nonResumableFullRefreshStreams = new HashSet<>();
+
     catalog.getStreams().forEach(configuredAirbyteStream -> {
+      var pairInStream =
+          new AirbyteStreamNameNamespacePair(configuredAirbyteStream.getStream().getName(), configuredAirbyteStream.getStream().getNamespace());
       if (!initialLoadStreams.streamsForInitialLoad().contains(configuredAirbyteStream)
           && configuredAirbyteStream.getSyncMode() == SyncMode.INCREMENTAL) {
-        this.streamsThatHaveCompletedSnapshot.add(
-            new AirbyteStreamNameNamespacePair(configuredAirbyteStream.getStream().getName(), configuredAirbyteStream.getStream().getNamespace()));
+        this.streamsThatHaveCompletedSnapshot.add(pairInStream);
       }
-      if (initialLoadStreams.streamsForInitialLoad().contains(configuredAirbyteStream)
-          && configuredAirbyteStream.getSyncMode() == SyncMode.FULL_REFRESH) {
-        this.resumableFullRefreshStreams.add(
-            new AirbyteStreamNameNamespacePair(configuredAirbyteStream.getStream().getName(), configuredAirbyteStream.getStream().getNamespace()));
+      if (configuredAirbyteStream.getSyncMode() == SyncMode.FULL_REFRESH) {
+        if (initialLoadStreams.streamsForInitialLoad().contains(configuredAirbyteStream)) {
+          this.resumableFullRefreshStreams.add(pairInStream);
+        } else {
+          this.nonResumableFullRefreshStreams.add(pairInStream);
+        }
       }
     });
   }
@@ -128,6 +134,12 @@ public AirbyteStateMessage createFinalStateMessage(final ConfiguredAirbyteStream
       streamStates.add(getAirbyteStreamState(stream, Jsons.jsonNode(ocStatus)));
     });
 
+    nonResumableFullRefreshStreams.forEach(stream -> {
+      streamStates.add(new AirbyteStreamState()
+          .withStreamDescriptor(
+              new StreamDescriptor().withName(stream.getName()).withNamespace(stream.getNamespace())));
+    });
+
     return new AirbyteStateMessage()
         .withType(AirbyteStateType.GLOBAL)
         .withGlobal(generateGlobalState(streamStates));

airbyte-integrations/connectors/source-mysql/build.gradle

@@ -6,7 +6,7 @@ plugins {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.38.0'
+    cdkVersionRequired = '0.40.1'
     features = ['db-sources']
     useLocalCdk = false
 }

airbyte-integrations/connectors/source-mysql/metadata.yaml

@@ -9,7 +9,7 @@ data:
   connectorSubtype: database
   connectorType: source
   definitionId: 435bb9a5-7887-4809-aa58-28c27df0d7ad
-  dockerImageTag: 3.4.9
+  dockerImageTag: 3.4.10
   dockerRepository: airbyte/source-mysql
   documentationUrl: https://docs.airbyte.com/integrations/sources/mysql
   githubIssueLabel: source-mysql

airbyte-integrations/connectors/source-mysql/src/main/java/io/airbyte/integrations/source/mysql/MySqlSource.java

@@ -46,6 +46,8 @@
 import io.airbyte.cdk.integrations.source.relationaldb.DbSourceDiscoverUtil;
 import io.airbyte.cdk.integrations.source.relationaldb.InitialLoadHandler;
 import io.airbyte.cdk.integrations.source.relationaldb.TableInfo;
+import io.airbyte.cdk.integrations.source.relationaldb.state.NonResumableStateMessageProducer;
+import io.airbyte.cdk.integrations.source.relationaldb.state.SourceStateMessageProducer;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateGeneratorUtils;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateManager;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateManagerFactory;
@@ -228,6 +230,11 @@ public InitialLoadHandler<MysqlType> getInitialLoadHandler(final JdbcDatabase da
     }
   }
 
+  @Override
+  protected SourceStateMessageProducer<AirbyteMessage> getSourceStateProducerForNonResumableFullRefreshStream(final JdbcDatabase database) {
+    return new NonResumableStateMessageProducer<>(isCdc(database.getSourceConfig()), initialLoadStateManager);
+  }
+
   private static AirbyteStream overrideSyncModes(final AirbyteStream stream) {
     return stream.withSupportedSyncModes(Lists.newArrayList(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL));
   }