Per-stream state support for Postgres source

airbyte-integrations/bases/debezium-v1-4-2/src/testFixtures/java/io/airbyte/integrations/debezium/CdcSourceTest.java

@@ -316,7 +316,7 @@ void testDelete() throws Exception {
         .format("DELETE FROM %s.%s WHERE %s = %s", MODELS_SCHEMA, MODELS_STREAM_NAME, COL_ID,
             11));
 
-    final JsonNode state = stateMessages1.get(0).getData();
+    final JsonNode state = Jsons.jsonNode(stateMessages1);
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);
     final List<AirbyteMessage> actualRecords2 = AutoCloseableIterators.toListAndClose(read2);
@@ -347,7 +347,7 @@ void testUpdate() throws Exception {
         .format("UPDATE %s.%s SET %s = '%s' WHERE %s = %s", MODELS_SCHEMA, MODELS_STREAM_NAME,
             COL_MODEL, updatedModel, COL_ID, 11));
 
-    final JsonNode state = stateMessages1.get(0).getData();
+    final JsonNode state = Jsons.jsonNode(stateMessages1);
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);
     final List<AirbyteMessage> actualRecords2 = AutoCloseableIterators.toListAndClose(read2);
@@ -403,7 +403,7 @@ void testRecordsProducedDuringAndAfterSync() throws Exception {
       recordsCreated[0]++;
     }
 
-    final JsonNode state = stateAfterFirstBatch.get(0).getData();
+    final JsonNode state = Jsons.jsonNode(stateAfterFirstBatch);
     final AutoCloseableIterator<AirbyteMessage> secondBatchIterator = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);
     final List<AirbyteMessage> dataFromSecondBatch = AutoCloseableIterators
@@ -492,7 +492,7 @@ void testCdcAndFullRefreshInSameSync() throws Exception {
         .jsonNode(ImmutableMap.of(COL_ID, 100, COL_MAKE_ID, 3, COL_MODEL, "Punto"));
     writeModelRecord(puntoRecord);
 
-    final JsonNode state = extractStateMessages(actualRecords1).get(0).getData();
+    final JsonNode state = Jsons.jsonNode(extractStateMessages(actualRecords1));
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), configuredCatalog, state);
     final List<AirbyteMessage> actualRecords2 = AutoCloseableIterators.toListAndClose(read2);
@@ -535,7 +535,7 @@ void testNoDataOnSecondSync() throws Exception {
     final AutoCloseableIterator<AirbyteMessage> read1 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, null);
     final List<AirbyteMessage> actualRecords1 = AutoCloseableIterators.toListAndClose(read1);
-    final JsonNode state = extractStateMessages(actualRecords1).get(0).getData();
+    final JsonNode state = Jsons.jsonNode(extractStateMessages(actualRecords1));
 
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);

airbyte-integrations/bases/debezium-v1-9-2/src/testFixtures/java/io/airbyte/integrations/debezium/CdcSourceTest.java

@@ -316,7 +316,7 @@ void testDelete() throws Exception {
         .format("DELETE FROM %s.%s WHERE %s = %s", MODELS_SCHEMA, MODELS_STREAM_NAME, COL_ID,
             11));
 
-    final JsonNode state = stateMessages1.get(0).getData();
+    final JsonNode state = Jsons.jsonNode(stateMessages1);
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);
     final List<AirbyteMessage> actualRecords2 = AutoCloseableIterators.toListAndClose(read2);
@@ -347,7 +347,7 @@ void testUpdate() throws Exception {
         .format("UPDATE %s.%s SET %s = '%s' WHERE %s = %s", MODELS_SCHEMA, MODELS_STREAM_NAME,
             COL_MODEL, updatedModel, COL_ID, 11));
 
-    final JsonNode state = stateMessages1.get(0).getData();
+    final JsonNode state = Jsons.jsonNode(stateMessages1);
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);
     final List<AirbyteMessage> actualRecords2 = AutoCloseableIterators.toListAndClose(read2);
@@ -399,7 +399,7 @@ protected void testRecordsProducedDuringAndAfterSync() throws Exception {
       writeModelRecord(record);
     }
 
-    final JsonNode state = stateAfterFirstBatch.get(0).getData();
+    final JsonNode state = Jsons.jsonNode(stateAfterFirstBatch);
     final AutoCloseableIterator<AirbyteMessage> secondBatchIterator = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);
     final List<AirbyteMessage> dataFromSecondBatch = AutoCloseableIterators
@@ -488,7 +488,7 @@ void testCdcAndFullRefreshInSameSync() throws Exception {
         .jsonNode(ImmutableMap.of(COL_ID, 100, COL_MAKE_ID, 3, COL_MODEL, "Punto"));
     writeModelRecord(puntoRecord);
 
-    final JsonNode state = extractStateMessages(actualRecords1).get(0).getData();
+    final JsonNode state = Jsons.jsonNode(extractStateMessages(actualRecords1));
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), configuredCatalog, state);
     final List<AirbyteMessage> actualRecords2 = AutoCloseableIterators.toListAndClose(read2);
@@ -531,7 +531,7 @@ void testNoDataOnSecondSync() throws Exception {
     final AutoCloseableIterator<AirbyteMessage> read1 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, null);
     final List<AirbyteMessage> actualRecords1 = AutoCloseableIterators.toListAndClose(read1);
-    final JsonNode state = extractStateMessages(actualRecords1).get(0).getData();
+    final JsonNode state = Jsons.jsonNode(extractStateMessages(actualRecords1));
 
     final AutoCloseableIterator<AirbyteMessage> read2 = getSource()
         .read(getConfig(), CONFIGURED_CATALOG, state);

airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/SourceAcceptanceTest.java

@@ -13,6 +13,7 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.config.StandardCheckConnectionOutput.Status;
@@ -106,6 +107,18 @@ public abstract class SourceAcceptanceTest extends AbstractSourceConnectorTest {
    */
   protected abstract JsonNode getState() throws Exception;
 
+  /**
+   * Tests whether the connector under test supports the per-stream state format or should use the
+   * legacy format for data generated by this test.
+   *
+   * @return {@code true} if the connector supports the per-stream state format or {@code false} if it
+   *         does not support the per-stream state format (e.g. legacy format supported). Default
+   *         value is {@code false}.
+   */
+  protected boolean supportsPerStream() {
+    return false;
+  }
+
   /**
    * Verify that a spec operation issued to the connector returns a valid spec.
    */
@@ -236,7 +249,7 @@ public void testIncrementalSyncWithState() throws Exception {
 
     // when we run incremental sync again there should be no new records. Run a sync with the latest
     // state message and assert no records were emitted.
-    final JsonNode latestState = stateMessages.get(stateMessages.size() - 1).getData();
+    final JsonNode latestState = Jsons.jsonNode(supportsPerStream() ? stateMessages : List.of(Iterables.getLast(stateMessages)));
     final List<AirbyteRecordMessage> secondSyncRecords = filterRecords(runRead(configuredCatalog, latestState));
     assertTrue(
         secondSyncRecords.isEmpty(),

airbyte-integrations/connectors/source-jdbc/src/test/java/io/airbyte/integrations/source/jdbc/AbstractJdbcSourceAcceptanceTest.java

@@ -15,8 +15,14 @@
 import io.airbyte.integrations.base.IntegrationRunner;
 import io.airbyte.integrations.base.Source;
 import io.airbyte.integrations.source.jdbc.test.JdbcSourceAcceptanceTest;
+import io.airbyte.integrations.source.relationaldb.models.CdcState;
+import io.airbyte.protocol.models.AirbyteGlobalState;
+import io.airbyte.protocol.models.AirbyteStateMessage;
+import io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType;
+import io.airbyte.protocol.models.AirbyteStreamState;
 import io.airbyte.test.utils.PostgreSQLContainerHelper;
 import java.sql.JDBCType;
+import java.util.List;
 import java.util.Set;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
@@ -82,6 +88,11 @@ public String getDriverClass() {
     return PostgresTestSource.DRIVER_CLASS;
   }
 
+  @Override
+  protected boolean supportsPerStream() {
+    return true;
+  }
+
   @AfterAll
   static void cleanUp() {
     PSQL_DB.close();
@@ -118,6 +129,27 @@ public Set<String> getExcludedInternalNameSpaces() {
       return Set.of("information_schema", "pg_catalog", "pg_internal", "catalog_history");
     }
 
+    // TODO This is a temporary override so that the Postgres source can take advantage of per-stream
+    // state
+    @Override
+    protected List<AirbyteStateMessage> generateEmptyInitialState(final JsonNode config) {
+      if (getSupportedStateType(config) == AirbyteStateType.GLOBAL) {
+        final AirbyteGlobalState globalState = new AirbyteGlobalState()
+            .withSharedState(Jsons.jsonNode(new CdcState()))
+            .withStreamStates(List.of());
+        return List.of(new AirbyteStateMessage().withStateType(AirbyteStateType.GLOBAL).withGlobal(globalState));
+      } else {
+        return List.of(new AirbyteStateMessage()
+            .withStateType(AirbyteStateType.STREAM)
+            .withStream(new AirbyteStreamState()));
+      }
+    }
+
+    @Override
+    protected AirbyteStateType getSupportedStateType(final JsonNode config) {
+      return AirbyteStateType.STREAM;
+    }
+
     public static void main(final String[] args) throws Exception {
       final Source source = new PostgresTestSource();
       LOGGER.info("starting source: {}", PostgresTestSource.class);