[source-postgres] State counter on postgres

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

@@ -45,6 +45,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.AfterEach;
@@ -141,6 +142,15 @@ protected ConfiguredAirbyteCatalog getConfiguredCatalog() {
 
   protected abstract void assertExpectedStateMessages(final List<AirbyteStateMessage> stateMessages);
 
+  // TODO: this assertion should be added into test cases in this class, we will need to implement
+  // corresponding iterator for other connectors before
+  // doing so.
+  protected void assertExpectedStateMessageCountMatches(final List<AirbyteStateMessage> stateMessages, long totalCount) {
+    AtomicLong count = new AtomicLong(0L);
+    stateMessages.stream().forEach(stateMessage -> count.addAndGet(stateMessage.getSourceStats().getRecordCount().longValue()));
+    assertEquals(totalCount, count.get());
+  }
+
   @BeforeEach
   protected void setup() {
     testdb = createTestDatabase();

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

@@ -12,9 +12,9 @@ java {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.20.6'
+    cdkVersionRequired = '0.20.8'
     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.3.13
+  dockerImageTag: 3.3.14
   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/ctid/CtidStateManager.java

@@ -5,18 +5,32 @@
 package io.airbyte.integrations.source.postgres.ctid;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import io.airbyte.cdk.integrations.source.relationaldb.state.SourceStateMessageProducer;
 import io.airbyte.integrations.source.postgres.internal.models.CtidStatus;
+import io.airbyte.integrations.source.postgres.internal.models.InternalModels.StateType;
 import io.airbyte.protocol.models.AirbyteStreamNameNamespacePair;
+import io.airbyte.protocol.models.v0.AirbyteMessage;
 import io.airbyte.protocol.models.v0.AirbyteStateMessage;
+import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-public abstract class CtidStateManager {
+public abstract class CtidStateManager implements SourceStateMessageProducer<AirbyteMessageWithCtid> {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(CtidStateManager.class);
 
   public static final long CTID_STATUS_VERSION = 2;
   public static final String STATE_TYPE_KEY = "state_type";
 
   protected final Map<AirbyteStreamNameNamespacePair, CtidStatus> pairToCtidStatus;
+  private Function<AirbyteStreamNameNamespacePair, JsonNode> streamStateForIncrementalRunSupplier;
+
+  private String lastCtid;
+  private FileNodeHandler fileNodeHandler;
 
   protected CtidStateManager(final Map<AirbyteStreamNameNamespacePair, CtidStatus> pairToCtidStatus) {
     this.pairToCtidStatus = pairToCtidStatus;
@@ -41,4 +55,65 @@ public static boolean validateRelationFileNode(final CtidStatus ctidstatus,
 
   public abstract AirbyteStateMessage createFinalStateMessage(final AirbyteStreamNameNamespacePair pair, final JsonNode streamStateForIncrementalRun);
 
+  public void setStreamStateIteratorFields(Function<AirbyteStreamNameNamespacePair, JsonNode> streamStateForIncrementalRunSupplier,
+                                           FileNodeHandler fileNodeHandler) {
+    this.streamStateForIncrementalRunSupplier = streamStateForIncrementalRunSupplier;
+    this.fileNodeHandler = fileNodeHandler;
+  }
+
+  @Override
+  public AirbyteStateMessage generateStateMessageAtCheckpoint(final ConfiguredAirbyteStream stream) {
+    final AirbyteStreamNameNamespacePair pair = new AirbyteStreamNameNamespacePair(stream.getStream().getName(),
+        stream.getStream().getNamespace());
+    final Long fileNode = fileNodeHandler.getFileNode(pair);
+    assert fileNode != null;
+    final CtidStatus ctidStatus = new CtidStatus()
+        .withVersion(CTID_STATUS_VERSION)
+        .withStateType(StateType.CTID)
+        .withCtid(lastCtid)
+        .withIncrementalState(getStreamState(pair))
+        .withRelationFilenode(fileNode);
+    LOGGER.info("Emitting ctid state for stream {}, state is {}", pair, ctidStatus);
+    return createCtidStateMessage(pair, ctidStatus);
+  }
+
+  /**
+   * Stores the latest CTID.
+   */
+  @Override
+  public AirbyteMessage processRecordMessage(final ConfiguredAirbyteStream stream, AirbyteMessageWithCtid message) {
+    if (Objects.nonNull(message.ctid())) {
+      this.lastCtid = message.ctid();
+    }
+    return message.recordMessage();
+  }
+
+  /**
+   * Creates a final state message for the stream.
+   */
+  @Override
+  public AirbyteStateMessage createFinalStateMessage(final ConfiguredAirbyteStream stream) {
+    final AirbyteStreamNameNamespacePair pair = new AirbyteStreamNameNamespacePair(stream.getStream().getName(),
+        stream.getStream().getNamespace());
+
+    final AirbyteStateMessage finalStateMessage = createFinalStateMessage(pair, getStreamState(pair));
+    LOGGER.info("Finished initial sync of stream {}, Emitting final state, state is {}", pair, finalStateMessage);
+    return finalStateMessage;
+  }
+
+  /**
+   * Extra criteria(besides checking frequency) to check if we should emit state message.
+   */
+  @Override
+  public boolean shouldEmitStateMessage(final ConfiguredAirbyteStream stream) {
+    return Objects.nonNull(lastCtid)
+        && StringUtils.isNotBlank(lastCtid);
+  }
+
+  private JsonNode getStreamState(final AirbyteStreamNameNamespacePair pair) {
+    final CtidStatus currentCtidStatus = getCtidStatus(pair);
+    return (currentCtidStatus == null || currentCtidStatus.getIncrementalState() == null) ? streamStateForIncrementalRunSupplier.apply(pair)
+        : currentCtidStatus.getIncrementalState();
+  }
+
 }

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

@@ -10,15 +10,17 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import io.airbyte.cdk.db.jdbc.JdbcDatabase;
+import io.airbyte.cdk.integrations.debezium.DebeziumIteratorConstants;
 import io.airbyte.cdk.integrations.source.relationaldb.DbSourceDiscoverUtil;
 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.commons.stream.AirbyteStreamUtils;
 import io.airbyte.commons.util.AutoCloseableIterator;
 import io.airbyte.commons.util.AutoCloseableIterators;
 import io.airbyte.integrations.source.postgres.PostgresQueryUtils.TableBlockSize;
 import io.airbyte.integrations.source.postgres.PostgresType;
 import io.airbyte.integrations.source.postgres.ctid.CtidPostgresSourceOperations.RowDataWithCtid;
-import io.airbyte.integrations.source.postgres.internal.models.CtidStatus;
 import io.airbyte.protocol.models.AirbyteStreamNameNamespacePair;
 import io.airbyte.protocol.models.CommonField;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
@@ -109,7 +111,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> getInitialSyncCtidIterator(
             tablesMaxTuple.orElseGet(() -> Map.of(pair, -1)).get(pair));
         final AutoCloseableIterator<AirbyteMessageWithCtid> recordIterator =
             getRecordIterator(queryStream, streamName, namespace, emmitedAt.toEpochMilli());
-        final AutoCloseableIterator<AirbyteMessage> recordAndMessageIterator = augmentWithState(recordIterator, pair);
+        final AutoCloseableIterator<AirbyteMessage> recordAndMessageIterator = augmentWithState(recordIterator, airbyteStream);
         final AutoCloseableIterator<AirbyteMessage> logAugmented = augmentWithLogs(recordAndMessageIterator, pair, streamName);
         iteratorList.add(logAugmented);
 
@@ -165,21 +167,20 @@ private AutoCloseableIterator<AirbyteMessage> augmentWithLogs(final AutoCloseabl
   }
 
   private AutoCloseableIterator<AirbyteMessage> augmentWithState(final AutoCloseableIterator<AirbyteMessageWithCtid> recordIterator,
-                                                                 final AirbyteStreamNameNamespacePair pair) {
+                                                                 final ConfiguredAirbyteStream airbyteStream) {
 
-    final CtidStatus currentCtidStatus = ctidStateManager.getCtidStatus(pair);
-    final JsonNode incrementalState =
-        (currentCtidStatus == null || currentCtidStatus.getIncrementalState() == null) ? streamStateForIncrementalRunSupplier.apply(pair)
-            : currentCtidStatus.getIncrementalState();
     final Duration syncCheckpointDuration =
         config.get(SYNC_CHECKPOINT_DURATION_PROPERTY) != null ? Duration.ofSeconds(config.get(SYNC_CHECKPOINT_DURATION_PROPERTY).asLong())
-            : CtidStateIterator.SYNC_CHECKPOINT_DURATION;
+            : DebeziumIteratorConstants.SYNC_CHECKPOINT_DURATION;
     final Long syncCheckpointRecords = config.get(SYNC_CHECKPOINT_RECORDS_PROPERTY) != null ? config.get(SYNC_CHECKPOINT_RECORDS_PROPERTY).asLong()
-        : CtidStateIterator.SYNC_CHECKPOINT_RECORDS;
+        : DebeziumIteratorConstants.SYNC_CHECKPOINT_RECORDS;
 
+    ctidStateManager.setStreamStateIteratorFields(streamStateForIncrementalRunSupplier, fileNodeHandler);
+
+    final AirbyteStreamNameNamespacePair pair =
+        new AirbyteStreamNameNamespacePair(airbyteStream.getStream().getName(), airbyteStream.getStream().getNamespace());
     return AutoCloseableIterators.transformIterator(
-        r -> new CtidStateIterator(r, pair, fileNodeHandler, ctidStateManager, incrementalState,
-            syncCheckpointDuration, syncCheckpointRecords),
+        r -> new SourceStateIterator(r, airbyteStream, ctidStateManager, new StateEmitFrequency(syncCheckpointRecords, syncCheckpointDuration)),
         recordIterator, pair);
   }
 

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

@@ -99,7 +99,7 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(
         final AutoCloseableIterator<JsonNode> queryStream = queryTableXmin(selectedDatabaseFields, table.getNameSpace(), table.getName());
         final AutoCloseableIterator<AirbyteMessage> recordIterator =
             getRecordIterator(queryStream, streamName, namespace, emittedAt.toEpochMilli());
-        final AutoCloseableIterator<AirbyteMessage> recordAndMessageIterator = augmentWithState(recordIterator, pair);
+        final AutoCloseableIterator<AirbyteMessage> recordAndMessageIterator = augmentWithState(recordIterator, airbyteStream, pair);
 
         iteratorList.add(augmentWithLogs(recordAndMessageIterator, pair, streamName));
       }
@@ -233,12 +233,14 @@ private AutoCloseableIterator<AirbyteMessage> augmentWithLogs(final AutoCloseabl
   }
 
   private AutoCloseableIterator<AirbyteMessage> augmentWithState(final AutoCloseableIterator<AirbyteMessage> recordIterator,
+                                                                 final ConfiguredAirbyteStream airbyteStream,
                                                                  final AirbyteStreamNameNamespacePair pair) {
+    xminStateManager.setCurrentXminStatus(currentXminStatus);
     return AutoCloseableIterators.transform(
         autoCloseableIterator -> new XminStateIterator(
             autoCloseableIterator,
-            pair,
-            currentXminStatus),
+            airbyteStream,
+            xminStateManager),
         recordIterator,
         AirbyteStreamUtils.convertFromNameAndNamespace(pair.getName(), pair.getNamespace()));
   }