[source-postgres] State counter on postgres (#34724)

Co-authored-by: Aaron ("AJ") Steers <[email protected]> Co-authored-by: Joe Bell <[email protected]> Co-authored-by: Marcos Marx <[email protected]> Co-authored-by: SatishChGit <[email protected]> Co-authored-by: evantahler <[email protected]> Co-authored-by: Rodi Reich Zilberman <[email protected]> Co-authored-by: Joe Reuter <[email protected]> Co-authored-by: Catherine Noll <[email protected]> Co-authored-by: Anton Karpets <[email protected]> Co-authored-by: Christo Grabowski <[email protected]> Co-authored-by: Akash Kulkarni <[email protected]> Co-authored-by: Akash Kulkarni <[email protected]> Co-authored-by: Gireesh Sreepathi <[email protected]> Co-authored-by: Artem Inzhyyants <[email protected]> Co-authored-by: Ella Rohm-Ensing <[email protected]> Co-authored-by: Daryna Ishchenko <[email protected]> Co-authored-by: Baz <[email protected]> Co-authored-by: Patrick Nilan <[email protected]> Co-authored-by: Serhii Lazebnyi <[email protected]> Co-authored-by: Marius Posta <[email protected]> Co-authored-by: Chandler Prall <[email protected]> Co-authored-by: pmossman <[email protected]> Co-authored-by: Anatolii Yatsuk <[email protected]> Co-authored-by: Cole Snodgrass <[email protected]> Co-authored-by: Anatolii Yatsuk <[email protected]> Co-authored-by: bgroff <[email protected]> Co-authored-by: Maxime Carbonneau-Leclerc <[email protected]> Co-authored-by: Ryan Waskewich <[email protected]> Co-authored-by: maxi297 <[email protected]> Co-authored-by: Pedro S. Lopez <[email protected]> Co-authored-by: Håkon Åmdal <[email protected]> Co-authored-by: Roman Yermilov [GL] <[email protected]> Co-authored-by: Augustin <[email protected]> Co-authored-by: Lake Mossman <[email protected]> Co-authored-by: lmossman <[email protected]> Co-authored-by: Brian Lai <[email protected]> Co-authored-by: Subodh Kant Chaturvedi <[email protected]> Co-authored-by: Stephane Geneix <[email protected]> Co-authored-by: Sajarin <[email protected]> Co-authored-by: Edward Gao <[email protected]> Co-authored-by: nguyenaiden <[email protected]> Co-authored-by: Natalie Kwong <[email protected]> Co-authored-by: terencecho <[email protected]> Co-authored-by: Alex Birdsall <[email protected]>
airbytehq · Mar 7, 2024 · 1bb7a1c · 1bb7a1c
1 parent be6849f
commit 1bb7a1c
Show file tree

Hide file tree

Showing 16 changed files with 206 additions and 225 deletions.
diff --git a/airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties b/airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties
@@ -1 +1 @@
-version=0.23.16
+version=0.23.17
diff --git a/...sources/src/main/java/io/airbyte/cdk/integrations/debezium/DebeziumIteratorConstants.java b/...sources/src/main/java/io/airbyte/cdk/integrations/debezium/DebeziumIteratorConstants.java
@@ -11,6 +11,8 @@ public class DebeziumIteratorConstants {
   public static final String SYNC_CHECKPOINT_DURATION_PROPERTY = "sync_checkpoint_seconds";
   public static final String SYNC_CHECKPOINT_RECORDS_PROPERTY = "sync_checkpoint_records";
 
+  // TODO: Move these variables to a separate class IteratorConstants, as they will be used in state
+  // iterators for non debezium cases too.
   public static final Duration SYNC_CHECKPOINT_DURATION = Duration.ofMinutes(15);
   public static final Integer SYNC_CHECKPOINT_RECORDS = 10_000;
 

diff --git a/.../db-sources/src/testFixtures/java/io/airbyte/cdk/integrations/debezium/CdcSourceTest.java b/.../db-sources/src/testFixtures/java/io/airbyte/cdk/integrations/debezium/CdcSourceTest.java
@@ -141,6 +141,13 @@ 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) {
+    // Do nothing.
+  }
+
   @BeforeEach
   protected void setup() {
     testdb = createTestDatabase();
@@ -350,6 +357,7 @@ void testExistingData() throws Exception {
 
     assertExpectedRecords(new HashSet<>(MODEL_RECORDS), recordMessages);
     assertExpectedStateMessages(stateMessages);
+    assertExpectedStateMessageCountMatches(stateMessages, MODEL_RECORDS.size());
   }
 
   protected void compareTargetPositionFromTheRecordsWithTargetPostionGeneratedBeforeSync(final CdcTargetPosition targetPosition,
@@ -377,6 +385,7 @@ public void testDelete() throws Exception {
         extractRecordMessages(actualRecords2));
     final List<AirbyteStateMessage> stateMessages2 = extractStateMessages(actualRecords2);
     assertExpectedStateMessagesFromIncrementalSync(stateMessages2);
+    assertExpectedStateMessageCountMatches(stateMessages2, 1);
     assertEquals(1, recordMessages2.size());
     assertEquals(11, recordMessages2.get(0).getData().get(COL_ID).asInt());
     assertCdcMetaData(recordMessages2.get(0).getData(), false);
@@ -411,6 +420,7 @@ public void testUpdate() throws Exception {
     assertEquals(11, recordMessages2.get(0).getData().get(COL_ID).asInt());
     assertEquals(updatedModel, recordMessages2.get(0).getData().get(COL_MODEL).asText());
     assertCdcMetaData(recordMessages2.get(0).getData(), true);
+    assertExpectedStateMessageCountMatches(stateMessages2, 1);
   }
 
   @SuppressWarnings({"BusyWait", "CodeBlock2Expr"})
@@ -534,6 +544,8 @@ public void testCdcAndFullRefreshInSameSync() throws Exception {
     final HashSet<String> names = new HashSet<>(STREAM_NAMES);
     names.add(MODELS_STREAM_NAME + "_2");
     assertExpectedStateMessages(stateMessages1);
+    // Full refresh does not get any state messages.
+    assertExpectedStateMessageCountMatches(stateMessages1, MODEL_RECORDS_2.size());
     assertExpectedRecords(Streams.concat(MODEL_RECORDS_2.stream(), MODEL_RECORDS.stream())
         .collect(Collectors.toSet()),
         recordMessages1,
@@ -554,6 +566,7 @@ public void testCdcAndFullRefreshInSameSync() throws Exception {
     final Set<AirbyteRecordMessage> recordMessages2 = extractRecordMessages(actualRecords2);
     final List<AirbyteStateMessage> stateMessages2 = extractStateMessages(actualRecords2);
     assertExpectedStateMessagesFromIncrementalSync(stateMessages2);
+    assertExpectedStateMessageCountMatches(stateMessages2, 1);
     assertExpectedRecords(
         Streams.concat(MODEL_RECORDS_2.stream(), Stream.of(puntoRecord))
             .collect(Collectors.toSet()),
@@ -576,6 +589,7 @@ public void testNoData() throws Exception {
     final List<AirbyteStateMessage> stateMessages = extractStateMessages(actualRecords);
     assertExpectedRecords(Collections.emptySet(), recordMessages);
     assertExpectedStateMessagesForNoData(stateMessages);
+    assertExpectedStateMessageCountMatches(stateMessages, 0);
   }
 
   protected void assertExpectedStateMessagesForNoData(final List<AirbyteStateMessage> stateMessages) {
@@ -600,6 +614,7 @@ public void testNoDataOnSecondSync() throws Exception {
 
     assertExpectedRecords(Collections.emptySet(), recordMessages2);
     assertExpectedStateMessagesFromIncrementalSync(stateMessages2);
+    assertExpectedStateMessageCountMatches(stateMessages2, 0);
   }
 
   @Test
@@ -630,6 +645,7 @@ public void newTableSnapshotTest() throws Exception {
         dataFromFirstBatch);
     final List<AirbyteStateMessage> stateAfterFirstBatch = extractStateMessages(dataFromFirstBatch);
     assertExpectedStateMessages(stateAfterFirstBatch);
+    assertExpectedStateMessageCountMatches(stateAfterFirstBatch, MODEL_RECORDS.size());
 
     final AirbyteStateMessage stateMessageEmittedAfterFirstSyncCompletion = stateAfterFirstBatch.get(stateAfterFirstBatch.size() - 1);
     assertEquals(AirbyteStateMessage.AirbyteStateType.GLOBAL, stateMessageEmittedAfterFirstSyncCompletion.getType());

diff --git a/airbyte-integrations/connectors/source-postgres/build.gradle b/airbyte-integrations/connectors/source-postgres/build.gradle
@@ -12,7 +12,7 @@ java {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.23.16'
+    cdkVersionRequired = '0.23.17'
     features = ['db-sources', 'datastore-postgres']
     useLocalCdk = false
 }

diff --git a/airbyte-integrations/connectors/source-postgres/metadata.yaml b/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.14
+  dockerImageTag: 3.3.15
   dockerRepository: airbyte/source-postgres
   documentationUrl: https://docs.airbyte.com/integrations/sources/postgres
   githubIssueLabel: source-postgres

diff --git a/...ostgres/src/main/java/io/airbyte/integrations/source/postgres/ctid/CtidStateIterator.java b/...ostgres/src/main/java/io/airbyte/integrations/source/postgres/ctid/CtidStateIterator.java
diff --git a/...postgres/src/main/java/io/airbyte/integrations/source/postgres/ctid/CtidStateManager.java b/...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();
+  }
+
 }
diff --git a/...tgres/src/main/java/io/airbyte/integrations/source/postgres/ctid/PostgresCtidHandler.java b/...tgres/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);
   }