db-sources: disable counts for state messages for FULL_REFRESH streams (#38208)

Co-authored-by: Xiaohan Song <[email protected]>

Author: subodh1810

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

@@ -1 +1 @@
-version=0.35.2
+version=0.35.4

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

@@ -7,6 +7,7 @@ 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 java.time.Duration
 import java.time.Instant
 import java.time.OffsetDateTime
@@ -41,9 +42,11 @@ open class SourceStateIterator<T>(
             ) {
                 val stateMessage =
                     sourceStateMessageProducer.generateStateMessageAtCheckpoint(stream)
-                stateMessage!!.withSourceStats(
-                    AirbyteStateStats().withRecordCount(recordCount.toDouble())
-                )
+                if (shouldAttachCountWithState()) {
+                    stateMessage!!.withSourceStats(
+                        AirbyteStateStats().withRecordCount(recordCount.toDouble())
+                    )
+                }
 
                 recordCount = 0L
                 lastCheckpoint = Instant.now()
@@ -64,9 +67,11 @@ open class SourceStateIterator<T>(
             hasEmittedFinalState = true
             val finalStateMessageForStream =
                 sourceStateMessageProducer.createFinalStateMessage(stream)
-            finalStateMessageForStream!!.withSourceStats(
-                AirbyteStateStats().withRecordCount(recordCount.toDouble())
-            )
+            if (shouldAttachCountWithState()) {
+                finalStateMessageForStream!!.withSourceStats(
+                    AirbyteStateStats().withRecordCount(recordCount.toDouble())
+                )
+            }
             recordCount = 0L
             return AirbyteMessage()
                 .withType(AirbyteMessage.Type.STATE)
@@ -76,6 +81,14 @@ 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/test/kotlin/io/airbyte/cdk/integrations/source/relationaldb/state/CursorStateMessageProducerTest.kt

@@ -446,6 +446,7 @@ internal class CursorStateMessageProducerTest {
                     NAMESPACE,
                     Field.of(UUID_FIELD_NAME, JsonSchemaType.STRING)
                 )
+                .withSyncMode(SyncMode.INCREMENTAL)
                 .withCursorField(List.of(UUID_FIELD_NAME))
 
         private val EMPTY_STATE_MESSAGE = createEmptyStateMessage(0.0)

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

@@ -665,10 +665,6 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
                 modelsSchema(),
             )
         } else {
-            assertExpectedStateMessageCountMatches(
-                stateMessages1,
-                MODEL_RECORDS.size.toLong() + MODEL_RECORDS_2.size.toLong(),
-            )
             assertExpectedRecords(
                 Streams.concat(MODEL_RECORDS_2.stream(), MODEL_RECORDS.stream())
                     .collect(Collectors.toSet()),
@@ -690,7 +686,6 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
             val recordMessages2 = extractRecordMessages(actualRecords2)
             val stateMessages2 = extractStateMessages(actualRecords2)
 
-            assertExpectedStateMessageCountMatches(stateMessages2, 7)
             assertExpectedRecords(
                 Streams.concat(MODEL_RECORDS_2.stream(), Stream.of(puntoRecord))
                     .collect(Collectors.toSet()),
@@ -1134,7 +1129,6 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
         val recordsFromFirstBatch = extractRecordMessages(dataFromFirstBatch)
         val stateAfterFirstBatch = extractStateMessages(dataFromFirstBatch)
         assertExpectedStateMessagesForFullRefresh(stateAfterFirstBatch)
-        assertExpectedStateMessageCountMatches(stateAfterFirstBatch, MODEL_RECORDS.size.toLong())
 
         val stateMessageEmittedAfterFirstSyncCompletion =
             stateAfterFirstBatch[stateAfterFirstBatch.size - 1]
@@ -1244,10 +1238,6 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
 
         Assertions.assertEquals(12, recordsFromFirstBatch.size)
 
-        assertExpectedStateMessageCountMatches(
-            stateAfterFirstBatch,
-            MODEL_RECORDS.size.toLong() + MODEL_RECORDS_2.size.toLong(),
-        )
         stateAfterFirstBatch.map { state -> assertStateDoNotHaveDuplicateStreams(state) }
     }
 

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

@@ -3,7 +3,7 @@ plugins {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.33.2'
+    cdkVersionRequired = '0.35.4'
     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.22
+  dockerImageTag: 4.0.23
   dockerRepository: airbyte/source-mssql
   documentationUrl: https://docs.airbyte.com/integrations/sources/mssql
   githubIssueLabel: source-mssql

airbyte-integrations/connectors/source-mssql/src/test/java/io/airbyte/integrations/source/mssql/CdcMssqlSourceTest.java

@@ -58,6 +58,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import javax.sql.DataSource;
 import org.junit.jupiter.api.*;
@@ -110,6 +111,13 @@ protected JsonNode config() {
         .build();
   }
 
+  @Override
+  protected void assertExpectedStateMessageCountMatches(final List<? extends AirbyteStateMessage> stateMessages, long totalCount) {
+    AtomicLong count = new AtomicLong(0L);
+    stateMessages.stream().forEach(stateMessage -> count.addAndGet(stateMessage.getSourceStats().getRecordCount().longValue()));
+    assertEquals(totalCount, count.get());
+  }
+
   @Override
   @BeforeEach
   protected void setup() {

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

@@ -6,7 +6,7 @@ plugins {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.35.2'
+    cdkVersionRequired = '0.35.4'
     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.3
+  dockerImageTag: 3.4.4
   dockerRepository: airbyte/source-mysql
   documentationUrl: https://docs.airbyte.com/integrations/sources/mysql
   githubIssueLabel: source-mysql

airbyte-integrations/connectors/source-mysql/src/test/java/io/airbyte/integrations/source/mysql/CdcMysqlSourceTest.java

@@ -74,6 +74,7 @@
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -92,6 +93,13 @@ public class CdcMysqlSourceTest extends CdcSourceTest<MySqlSource, MySQLTestData
   private static final List<JsonNode> DATE_TIME_RECORDS = ImmutableList.of(
       Jsons.jsonNode(ImmutableMap.of(COL_ID, 120, COL_DATE_TIME, "'2023-00-00 20:37:47'")));
 
+  @Override
+  protected void assertExpectedStateMessageCountMatches(final List<? extends AirbyteStateMessage> stateMessages, long totalCount) {
+    AtomicLong count = new AtomicLong(0L);
+    stateMessages.stream().forEach(stateMessage -> count.addAndGet(stateMessage.getSourceStats().getRecordCount().longValue()));
+    assertEquals(totalCount, count.get());
+  }
+
   @Override
   protected MySQLTestDatabase createTestDatabase() {
     return MySQLTestDatabase.in(BaseImage.MYSQL_8, ContainerModifier.INVALID_TIMEZONE_CEST).withCdcPermissions();

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

@@ -12,7 +12,7 @@ java {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.35.2'
+    cdkVersionRequired = '0.35.4'
     features = ['db-sources', 'datastore-postgres']
     useLocalCdk = false
 }

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

@@ -9,7 +9,7 @@ data:
   connectorSubtype: database
   connectorType: source
   definitionId: decd338e-5647-4c0b-adf4-da0e75f5a750
-  dockerImageTag: 3.4.3
+  dockerImageTag: 3.4.4
   dockerRepository: airbyte/source-postgres
   documentationUrl: https://docs.airbyte.com/integrations/sources/postgres
   githubIssueLabel: source-postgres