Stream status on postgres

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

@@ -111,7 +111,11 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
         expectedStreamStatus: AirbyteStreamStatusTraceMessage
     ) {
         var actualMessage = allMessages[idx]
-        Assertions.assertEquals(AirbyteMessage.Type.TRACE, actualMessage.type)
+        Assertions.assertEquals(
+            AirbyteMessage.Type.TRACE,
+            actualMessage.type,
+            "[Debug] all Message: $allMessages"
+        )
         var traceMessage = actualMessage.trace
         Assertions.assertNotNull(traceMessage.streamStatus)
         Assertions.assertEquals(expectedStreamStatus, traceMessage.streamStatus)
@@ -576,6 +580,25 @@ abstract class CdcSourceTest<S : Source, T : TestDatabase<*, T, *>> {
         val secondBatchIterator = source().read(config()!!, configuredCatalog, state)
         val dataFromSecondBatch = AutoCloseableIterators.toListAndClose(secondBatchIterator)
 
+        assertStreamStatusTraceMessageIndex(
+            0,
+            dataFromSecondBatch,
+            createAirbteStreanStatusTraceMessage(
+                modelsSchema(),
+                MODELS_STREAM_NAME,
+                AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.STARTED
+            )
+        )
+        assertStreamStatusTraceMessageIndex(
+            dataFromSecondBatch.size - 1,
+            dataFromSecondBatch,
+            createAirbteStreanStatusTraceMessage(
+                modelsSchema(),
+                MODELS_STREAM_NAME,
+                AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.COMPLETE
+            )
+        )
+
         val stateAfterSecondBatch = extractStateMessages(dataFromSecondBatch)
         assertExpectedStateMessagesFromIncrementalSync(stateAfterSecondBatch)
 

airbyte-cdk/java/airbyte-cdk/db-sources/src/testFixtures/kotlin/io/airbyte/cdk/integrations/source/jdbc/test/JdbcSourceAcceptanceTest.kt

@@ -424,7 +424,7 @@ abstract class JdbcSourceAcceptanceTest<S : Source, T : TestDatabase<*, T, *>> {
                 defaultNamespace,
                 Field.of(COL_ID, JsonSchemaType.NUMBER),
             )
-        val actualMessages = MoreIterators.toList(source().read(config(), catalog, null))
+        var actualMessages = MoreIterators.toList(source().read(config(), catalog, null))
 
         assertStreamStatusTraceMessageIndex(
             0,
@@ -469,11 +469,19 @@ abstract class JdbcSourceAcceptanceTest<S : Source, T : TestDatabase<*, T, *>> {
         setTraceEmittedAtToNull(actualMessages)
         setTraceEmittedAtToNull(expectedMessages)
 
+        actualMessages = removeStateMessage(actualMessages)
+
         Assertions.assertEquals(expectedMessages.size, actualMessages.size)
         Assertions.assertTrue(expectedMessages.containsAll(actualMessages))
         Assertions.assertTrue(actualMessages.containsAll(expectedMessages))
     }
 
+    private fun removeStateMessage(airbyteMessages: List<AirbyteMessage>): List<AirbyteMessage> {
+        var mutableListMessages = airbyteMessages.toMutableList()
+        mutableListMessages.removeIf { message -> message.type == AirbyteMessage.Type.STATE }
+        return mutableListMessages
+    }
+
     protected open val airbyteMessagesReadOneColumn: MutableList<AirbyteMessage>
         get() {
             val expectedMessages =

airbyte-integrations/connectors/source-postgres/acceptance-test-config.yml

@@ -1,7 +1,6 @@
 # See [Connector Acceptance Tests](https://docs.airbyte.com/connector-development/testing-connectors/connector-acceptance-tests-reference)
 # for more information about how to configure these tests
 connector_image: airbyte/source-postgres:dev
-test_strictness_level: high
 custom_environment_variables:
   USE_STREAM_CAPABLE_STATE: true
 acceptance_tests:
@@ -32,16 +31,6 @@ acceptance_tests:
       - config_path: "secrets/config_cdc.json"
         backward_compatibility_tests_config:
           disable_for_version: "2.1.1"
-  basic_read:
-    tests:
-      - config_path: "secrets/config.json"
-        expect_records:
-          path: "integration_tests/expected_records.txt"
-        validate_state_messages: false
-      - config_path: "secrets/config_cdc.json"
-        expect_records:
-          path: "integration_tests/expected_records.txt"
-        validate_state_messages: false
   full_refresh:
     tests:
       - config_path: "secrets/config.json"

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

@@ -12,9 +12,9 @@ java {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.35.4'
+    cdkVersionRequired = '0.35.12'
     features = ['db-sources', 'datastore-postgres']
-    useLocalCdk = false
+    useLocalCdk = true
 }
 
 application {

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.7
+  dockerImageTag: 3.4.8
   dockerRepository: airbyte/source-postgres
   documentationUrl: https://docs.airbyte.com/integrations/sources/postgres
   githubIssueLabel: source-postgres

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/PostgresSource.java

@@ -66,12 +66,15 @@
 import io.airbyte.cdk.integrations.source.relationaldb.InitialLoadHandler;
 import io.airbyte.cdk.integrations.source.relationaldb.TableInfo;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateManager;
+import io.airbyte.cdk.integrations.source.relationaldb.streamstatus.StreamStatusTraceEmitterIterator;
 import io.airbyte.commons.exceptions.ConfigErrorException;
 import io.airbyte.commons.functional.CheckedConsumer;
 import io.airbyte.commons.functional.CheckedFunction;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.map.MoreMaps;
+import io.airbyte.commons.stream.AirbyteStreamStatusHolder;
 import io.airbyte.commons.util.AutoCloseableIterator;
+import io.airbyte.commons.util.AutoCloseableIterators;
 import io.airbyte.integrations.source.postgres.PostgresQueryUtils.ResultWithFailed;
 import io.airbyte.integrations.source.postgres.cdc.PostgresReplicationConnection;
 import io.airbyte.integrations.source.postgres.ctid.CtidGlobalStateManager;
@@ -96,6 +99,7 @@
 import io.airbyte.protocol.models.v0.AirbyteStateMessage.AirbyteStateType;
 import io.airbyte.protocol.models.v0.AirbyteStream;
 import io.airbyte.protocol.models.v0.AirbyteStreamNameNamespacePair;
+import io.airbyte.protocol.models.v0.AirbyteStreamStatusTraceMessage;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
 import io.airbyte.protocol.models.v0.ConnectorSpecification;
@@ -524,7 +528,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(final
       final PostgresXminHandler xminHandler = new PostgresXminHandler(database, sourceOperations, getQuoteString(), xminStatus, xminStateManager);
 
       final List<AutoCloseableIterator<AirbyteMessage>> initialSyncCtidIterators = new ArrayList<>(ctidHandler.getInitialSyncCtidIterator(
-          new ConfiguredAirbyteCatalog().withStreams(finalListOfStreamsToBeSyncedViaCtid), tableNameToTable, emittedAt));
+          new ConfiguredAirbyteCatalog().withStreams(finalListOfStreamsToBeSyncedViaCtid), tableNameToTable, emittedAt, true, true));
       final List<AutoCloseableIterator<AirbyteMessage>> xminIterators = new ArrayList<>(xminHandler.getIncrementalIterators(
           new ConfiguredAirbyteCatalog().withStreams(xminStreams.streamsForXminSync()), tableNameToTable, emittedAt));
 
@@ -568,7 +572,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(final
       final List<AutoCloseableIterator<AirbyteMessage>> initialSyncCtidIterators = new ArrayList<>(
           cursorBasedCtidHandler.getInitialSyncCtidIterator(new ConfiguredAirbyteCatalog().withStreams(finalListOfStreamsToBeSyncedViaCtid),
               tableNameToTable,
-              emittedAt));
+              emittedAt, true, true));
       final List<AutoCloseableIterator<AirbyteMessage>> cursorBasedIterators = new ArrayList<>(super.getIncrementalIterators(database,
           new ConfiguredAirbyteCatalog().withStreams(
               cursorBasedStreamsCategorised.remainingStreams()
@@ -937,6 +941,18 @@ protected void estimateFullRefreshSyncSize(final JdbcDatabase database,
     }
   }
 
+  @Override
+  public AutoCloseableIterator<AirbyteMessage> augmentWithStreamStatus(final ConfiguredAirbyteStream airbyteStream,
+                                                                       final AutoCloseableIterator<AirbyteMessage> streamIterator) {
+    final var pair =
+        new io.airbyte.protocol.models.AirbyteStreamNameNamespacePair(airbyteStream.getStream().getName(), airbyteStream.getStream().getNamespace());
+    final var starterStatus =
+        new StreamStatusTraceEmitterIterator(new AirbyteStreamStatusHolder(pair, AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.STARTED));
+    final var completeStatus =
+        new StreamStatusTraceEmitterIterator(new AirbyteStreamStatusHolder(pair, AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.COMPLETE));
+    return AutoCloseableIterators.concatWithEagerClose(starterStatus, streamIterator, completeStatus);
+  }
+
   private List<JsonNode> getFullTableEstimate(final JdbcDatabase database,
                                               final String fullTableName,
                                               final String schemaName,

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/cdc/PostgresCdcCtidInitializer.java

@@ -22,8 +22,10 @@
 import io.airbyte.cdk.integrations.source.relationaldb.TableInfo;
 import io.airbyte.cdk.integrations.source.relationaldb.models.CdcState;
 import io.airbyte.cdk.integrations.source.relationaldb.state.StateManager;
+import io.airbyte.cdk.integrations.source.relationaldb.streamstatus.StreamStatusTraceEmitterIterator;
 import io.airbyte.commons.exceptions.ConfigErrorException;
 import io.airbyte.commons.json.Jsons;
+import io.airbyte.commons.stream.AirbyteStreamStatusHolder;
 import io.airbyte.commons.util.AutoCloseableIterator;
 import io.airbyte.commons.util.AutoCloseableIterators;
 import io.airbyte.integrations.source.postgres.PostgresQueryUtils;
@@ -36,8 +38,10 @@
 import io.airbyte.protocol.models.CommonField;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
 import io.airbyte.protocol.models.v0.AirbyteStreamNameNamespacePair;
+import io.airbyte.protocol.models.v0.AirbyteStreamStatusTraceMessage;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
+import io.airbyte.protocol.models.v0.SyncMode;
 import java.sql.SQLException;
 import java.time.Duration;
 import java.time.Instant;
@@ -186,11 +190,12 @@ public static List<AutoCloseableIterator<AirbyteMessage>> cdcCtidIteratorsCombin
         savedOffsetAfterReplicationSlotLSN);
     final List<AutoCloseableIterator<AirbyteMessage>> initialSyncCtidIterators = new ArrayList<>();
     final List<AirbyteStreamNameNamespacePair> streamsUnderVacuum = new ArrayList<>();
+    List<ConfiguredAirbyteStream> finalListOfStreamsToBeSyncedViaCtid = new ArrayList<>();
     if (!ctidStreams.streamsForCtidSync().isEmpty()) {
       streamsUnderVacuum.addAll(streamsUnderVacuum(database,
           ctidStreams.streamsForCtidSync(), quoteString).result());
 
-      List<ConfiguredAirbyteStream> finalListOfStreamsToBeSyncedViaCtid =
+      finalListOfStreamsToBeSyncedViaCtid =
           streamsUnderVacuum.isEmpty() ? ctidStreams.streamsForCtidSync()
               : ctidStreams.streamsForCtidSync().stream()
                   .filter(c -> !streamsUnderVacuum.contains(AirbyteStreamNameNamespacePair.fromConfiguredAirbyteSteam(c)))
@@ -217,7 +222,7 @@ public static List<AutoCloseableIterator<AirbyteMessage>> cdcCtidIteratorsCombin
       }
 
       initialSyncCtidIterators.addAll(ctidHandler.getInitialSyncCtidIterator(
-          new ConfiguredAirbyteCatalog().withStreams(finalListOfStreamsToBeSyncedViaCtid), tableNameToTable, emittedAt));
+          new ConfiguredAirbyteCatalog().withStreams(finalListOfStreamsToBeSyncedViaCtid), tableNameToTable, emittedAt, true, false));
     } else {
       LOGGER.info("No streams will be synced via ctid");
     }
@@ -238,8 +243,27 @@ public static List<AutoCloseableIterator<AirbyteMessage>> cdcCtidIteratorsCombin
     final Supplier<AutoCloseableIterator<AirbyteMessage>> incrementalIteratorSupplier = () -> handler.getIncrementalIterators(
         propertiesManager, eventConverter, new PostgresCdcSavedInfoFetcher(stateToBeUsed), postgresCdcStateHandler);
 
+    final List<ConfiguredAirbyteStream> finalListOfStreamsToBeSyncedViaCtidInLambda = finalListOfStreamsToBeSyncedViaCtid;
+    final List<AutoCloseableIterator<AirbyteMessage>> cdcStreamsStartStatusEmitters = catalog.getStreams().stream()
+        .filter(stream -> !finalListOfStreamsToBeSyncedViaCtidInLambda.contains(stream))
+        .map(stream -> (AutoCloseableIterator<AirbyteMessage>) new StreamStatusTraceEmitterIterator(
+            new AirbyteStreamStatusHolder(
+                new io.airbyte.protocol.models.AirbyteStreamNameNamespacePair(stream.getStream().getName(), stream.getStream().getNamespace()),
+                AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.STARTED)))
+        .toList();
+
+    final List<AutoCloseableIterator<AirbyteMessage>> allStreamsCompleteStatusEmitters = catalog.getStreams().stream()
+        .filter(stream -> stream.getSyncMode() == SyncMode.INCREMENTAL)
+        .map(stream -> (AutoCloseableIterator<AirbyteMessage>) new StreamStatusTraceEmitterIterator(
+            new AirbyteStreamStatusHolder(
+                new io.airbyte.protocol.models.AirbyteStreamNameNamespacePair(stream.getStream().getName(), stream.getStream().getNamespace()),
+                AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.COMPLETE)))
+        .toList();
+
     if (initialSyncCtidIterators.isEmpty()) {
-      return Collections.singletonList(incrementalIteratorSupplier.get());
+      return Stream.of(cdcStreamsStartStatusEmitters, Collections.singletonList(incrementalIteratorSupplier.get()), allStreamsCompleteStatusEmitters)
+          .flatMap(Collection::stream)
+          .collect(Collectors.toList());
     }
 
     if (streamsUnderVacuum.isEmpty()) {
@@ -248,12 +272,15 @@ public static List<AutoCloseableIterator<AirbyteMessage>> cdcCtidIteratorsCombin
       // We finish the current CDC once the initial snapshot is complete and the next sync starts
       // processing the WAL
       return Stream
-          .of(initialSyncCtidIterators, Collections.singletonList(AutoCloseableIterators.lazyIterator(incrementalIteratorSupplier, null)))
+          .of(initialSyncCtidIterators, cdcStreamsStartStatusEmitters,
+              Collections.singletonList(AutoCloseableIterators.lazyIterator(incrementalIteratorSupplier, null)),
+              allStreamsCompleteStatusEmitters)
           .flatMap(Collection::stream)
           .collect(Collectors.toList());
     } else {
       LOGGER.warn("Streams are under vacuuming, not going to process WAL");
-      return initialSyncCtidIterators;
+      return Stream.of(initialSyncCtidIterators, cdcStreamsStartStatusEmitters, allStreamsCompleteStatusEmitters).flatMap(Collection::stream)
+          .collect(Collectors.toList());
     }
   }
 

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/ctid/PostgresCtidHandler.java

@@ -16,6 +16,8 @@
 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.StateEmitFrequency;
+import io.airbyte.cdk.integrations.source.relationaldb.streamstatus.StreamStatusTraceEmitterIterator;
+import io.airbyte.commons.stream.AirbyteStreamStatusHolder;
 import io.airbyte.commons.stream.AirbyteStreamUtils;
 import io.airbyte.commons.util.AutoCloseableIterator;
 import io.airbyte.commons.util.AutoCloseableIterators;
@@ -29,6 +31,7 @@
 import io.airbyte.protocol.models.v0.AirbyteRecordMessage;
 import io.airbyte.protocol.models.v0.AirbyteRecordMessageMeta;
 import io.airbyte.protocol.models.v0.AirbyteStream;
+import io.airbyte.protocol.models.v0.AirbyteStreamStatusTraceMessage;
 import io.airbyte.protocol.models.v0.CatalogHelpers;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
@@ -107,7 +110,9 @@ public AutoCloseableIterator<AirbyteMessage> getIteratorForStream(@NotNull Confi
   public List<AutoCloseableIterator<AirbyteMessage>> getInitialSyncCtidIterator(
                                                                                 final ConfiguredAirbyteCatalog catalog,
                                                                                 final Map<String, TableInfo<CommonField<PostgresType>>> tableNameToTable,
-                                                                                final Instant emmitedAt) {
+                                                                                final Instant emmitedAt,
+                                                                                final boolean decorateWithStartedStatus,
+                                                                                final boolean decorateWithCompletedStatus) {
     final List<AutoCloseableIterator<AirbyteMessage>> iteratorList = new ArrayList<>();
     for (final ConfiguredAirbyteStream airbyteStream : catalog.getStreams()) {
       final AirbyteStream stream = airbyteStream.getStream();
@@ -119,11 +124,23 @@ public List<AutoCloseableIterator<AirbyteMessage>> getInitialSyncCtidIterator(
         continue;
       }
       if (airbyteStream.getSyncMode().equals(SyncMode.INCREMENTAL)) {
+
+        final AirbyteStreamNameNamespacePair pair = new AirbyteStreamNameNamespacePair(streamName, namespace);
+        if (decorateWithStartedStatus) {
+          iteratorList.add(
+              new StreamStatusTraceEmitterIterator(new AirbyteStreamStatusHolder(pair, AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.STARTED)));
+        }
+
         // Grab the selected fields to sync
         final TableInfo<CommonField<PostgresType>> table = tableNameToTable
             .get(fullyQualifiedTableName);
         final var iterator = getIteratorForStream(airbyteStream, table, emmitedAt);
         iteratorList.add(iterator);
+
+        if (decorateWithCompletedStatus) {
+          iteratorList.add(new StreamStatusTraceEmitterIterator(
+              new AirbyteStreamStatusHolder(pair, AirbyteStreamStatusTraceMessage.AirbyteStreamStatus.COMPLETE)));
+        }
       }
     }
     return iteratorList;

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/xmin/PostgresXminHandler.java

@@ -16,6 +16,8 @@
 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.StateEmitFrequency;
+import io.airbyte.cdk.integrations.source.relationaldb.streamstatus.StreamStatusTraceEmitterIterator;
+import io.airbyte.commons.stream.AirbyteStreamStatusHolder;
 import io.airbyte.commons.stream.AirbyteStreamUtils;
 import io.airbyte.commons.util.AutoCloseableIterator;
 import io.airbyte.commons.util.AutoCloseableIterators;
@@ -28,6 +30,7 @@
 import io.airbyte.protocol.models.v0.AirbyteRecordMessage;
 import io.airbyte.protocol.models.v0.AirbyteRecordMessageMeta;
 import io.airbyte.protocol.models.v0.AirbyteStream;
+import io.airbyte.protocol.models.v0.AirbyteStreamStatusTraceMessage.AirbyteStreamStatus;
 import io.airbyte.protocol.models.v0.CatalogHelpers;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
@@ -92,6 +95,9 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(
       }
 
       if (airbyteStream.getSyncMode().equals(SyncMode.INCREMENTAL)) {
+        iteratorList.add(
+            new StreamStatusTraceEmitterIterator(new AirbyteStreamStatusHolder(pair, AirbyteStreamStatus.STARTED)));
+
         // Grab the selected fields to sync
         final TableInfo<CommonField<PostgresType>> table = tableNameToTable
             .get(fullyQualifiedTableName);
@@ -107,6 +113,8 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(
         final AutoCloseableIterator<AirbyteMessage> recordAndMessageIterator = augmentWithState(recordIterator, airbyteStream, pair);
 
         iteratorList.add(augmentWithLogs(recordAndMessageIterator, pair, streamName));
+        iteratorList.add(
+            new StreamStatusTraceEmitterIterator(new AirbyteStreamStatusHolder(pair, AirbyteStreamStatus.COMPLETE)));
       }
     }