add logic to swtich between xmin and ctid sync

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

@@ -27,10 +27,12 @@
 import static java.util.stream.Collectors.toSet;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
 import datadog.trace.api.Trace;
 import io.airbyte.commons.exceptions.ConfigErrorException;
 import io.airbyte.commons.features.EnvVariableFeatureFlags;
@@ -70,11 +72,13 @@
 import io.airbyte.protocol.models.v0.AirbyteConnectionStatus.Status;
 import io.airbyte.protocol.models.v0.AirbyteEstimateTraceMessage.Type;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
+import io.airbyte.protocol.models.v0.AirbyteStateMessage;
 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.ConfiguredAirbyteCatalog;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream;
+import io.airbyte.protocol.models.v0.StreamDescriptor;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.Path;
@@ -85,7 +89,9 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -94,6 +100,7 @@
 import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -455,17 +462,78 @@ public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(final
               AutoCloseableIterators.lazyIterator(incrementalIteratorSupplier, null)));
 
     } else if (PostgresUtils.isXmin(sourceConfig) && isIncrementalSyncMode(catalog)) {
-      final XminStateManager xminStateManager = new XminStateManager(stateManager.getRawStateMessages());
-      final PostgresXminHandler handler = new PostgresXminHandler(database, sourceOperations, getQuoteString(), xminStatus, xminStateManager);
-      return handler.getIncrementalIterators(catalog, tableNameToTable, emittedAt);
+      final List<AirbyteStateMessage> rawStateMessages = stateManager.getRawStateMessages();
+
+      final List<AirbyteStateMessage> statesFromCtidSync = new ArrayList<>();
+      final List<AirbyteStateMessage> statesFromXminSync = new ArrayList<>();
+
+      final Set<AirbyteStreamNameNamespacePair> alreadySeenStreams = new HashSet<>();
+      final Set<AirbyteStreamNameNamespacePair> streamsStillInCtidSync = new HashSet<>();
+      rawStateMessages.forEach(s -> {
+        final JsonNode streamState = s.getStream().getStreamState();
+        final StreamDescriptor streamDescriptor = s.getStream().getStreamDescriptor();
+        final AirbyteStateMessage clonedState = Jsons.clone(s);
+        if (streamState.has("type") && streamState.get("type").asText().equalsIgnoreCase("ctid")) {
+          statesFromCtidSync.add(clonedState);
+          streamsStillInCtidSync.add(new AirbyteStreamNameNamespacePair(streamDescriptor.getName(), streamDescriptor.getNamespace()));
+        } else {
+          statesFromXminSync.add(clonedState);
+        }
+        alreadySeenStreams.add(new AirbyteStreamNameNamespacePair(streamDescriptor.getName(), streamDescriptor.getNamespace()));
+      });
+
+      final List<ConfiguredAirbyteStream> newlyAddedStreams = identifyNewlyAddedStreams(catalog, alreadySeenStreams);
+
+      final List<ConfiguredAirbyteStream> configuredStreamsStillInCtidSync = catalog.getStreams().stream()
+          .filter(stream -> streamsStillInCtidSync.contains(AirbyteStreamNameNamespacePair.fromAirbyteStream(stream.getStream())))
+          .map(Jsons::clone)
+          .toList();
+
+      configuredStreamsStillInCtidSync.addAll(newlyAddedStreams);
+
+      final List<ConfiguredAirbyteStream> streamsInXminSync = catalog.getStreams().stream()
+          .filter(stream -> !configuredStreamsStillInCtidSync.contains(stream))
+          .map(Jsons::clone)
+          .toList();
+
+      final XminStateManager xminStateManager = new XminStateManager(statesFromXminSync);
+      final PostgresXminHandler xminHandler = new PostgresXminHandler(database, sourceOperations, getQuoteString(), xminStatus, xminStateManager);
+
+      final List<AutoCloseableIterator<AirbyteMessage>> xminIterator = xminHandler.getIncrementalIterators(
+          new ConfiguredAirbyteCatalog().withStreams(streamsInXminSync), tableNameToTable, emittedAt);
+
+      if (configuredStreamsStillInCtidSync.isEmpty()) {
+        return xminIterator;
+      }
+
+      final CtidStateManager ctidStateManager = new CtidStateManager(statesFromCtidSync);
+      final PostgresCtidHandler ctidHandler = new PostgresCtidHandler(database, sourceOperations, getQuoteString(), ctidStateManager,
+          x -> new ObjectMapper().valueToTree(xminStatus),
+          (pair, jsonState) -> XminStateManager.getAirbyteStateMessage(pair, Jsons.object(jsonState, XminStatus.class)));
+      final List<AutoCloseableIterator<AirbyteMessage>> ctidIterator = ctidHandler.getIncrementalIterators(
+          new ConfiguredAirbyteCatalog().withStreams(configuredStreamsStillInCtidSync), tableNameToTable, emittedAt);
+      return Stream
+          .of(ctidIterator, xminIterator)
+          .flatMap(Collection::stream)
+          .collect(Collectors.toList());
     } else {
-      final CtidStateManager ctidStateManager = new CtidStateManager(stateManager.getRawStateMessages());
-      final PostgresCtidHandler handler = new PostgresCtidHandler(database, sourceOperations, getQuoteString(), null, ctidStateManager);
-      return handler.getIncrementalIterators(catalog, tableNameToTable, emittedAt);
-//      return super.getIncrementalIterators(database, catalog, tableNameToTable, stateManager, emittedAt);
+      return super.getIncrementalIterators(database, catalog, tableNameToTable, stateManager, emittedAt);
     }
   }
 
+  protected List<ConfiguredAirbyteStream> identifyNewlyAddedStreams(final ConfiguredAirbyteCatalog catalog, final Set<AirbyteStreamNameNamespacePair> alreadySeenStreams) {
+    final Set<AirbyteStreamNameNamespacePair> allStreams = AirbyteStreamNameNamespacePair.fromConfiguredCatalog(catalog);
+
+    final Set<AirbyteStreamNameNamespacePair> newlyAddedStreams = new HashSet<>(Sets.difference(allStreams, alreadySeenStreams));
+
+    return catalog.getStreams().stream()
+        .filter(stream -> newlyAddedStreams.contains(AirbyteStreamNameNamespacePair.fromAirbyteStream(stream.getStream())))
+        .map(Jsons::clone)
+        .collect(Collectors.toList());
+  }
+
+
+
   @Override
   public Set<JdbcPrivilegeDto> getPrivilegesTableForCurrentUser(final JdbcDatabase database,
                                                                 final String schema)

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

@@ -3,12 +3,18 @@
 import static io.airbyte.integrations.source.postgres.ctid.CtidStateManager.CTID_STATUS_TYPE;
 import static io.airbyte.integrations.source.postgres.ctid.CtidStateManager.CTID_STATUS_VERSION;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.collect.AbstractIterator;
 import io.airbyte.integrations.source.postgres.internal.models.CtidStatus;
 import io.airbyte.protocol.models.AirbyteStreamNameNamespacePair;
 import io.airbyte.protocol.models.v0.AirbyteMessage;
+import io.airbyte.protocol.models.v0.AirbyteMessage.Type;
+import io.airbyte.protocol.models.v0.AirbyteStateMessage;
 import java.util.Iterator;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.annotation.CheckForNull;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
@@ -22,11 +28,18 @@ public class CtidStateIterator extends AbstractIterator<AirbyteMessage> implemen
   private boolean hasEmittedFinalState;
   private boolean hasCaughtException = false;
   private String lastCtid;
+  private final JsonNode streamStateForIncrementalRun;
+  final BiFunction<AirbyteStreamNameNamespacePair, JsonNode, AirbyteStateMessage> finalStateMessageSupplier;
   final AtomicLong recordCount = new AtomicLong();
+
   public CtidStateIterator(final Iterator<AirbyteMessage> messageIterator,
-      final AirbyteStreamNameNamespacePair pair) {
+      final AirbyteStreamNameNamespacePair pair,
+      final JsonNode streamStateForIncrementalRun,
+      final BiFunction<AirbyteStreamNameNamespacePair, JsonNode, AirbyteStateMessage> finalStateMessageSupplier) {
     this.messageIterator = messageIterator;
     this.pair = pair;
+    this.streamStateForIncrementalRun = streamStateForIncrementalRun;
+    this.finalStateMessageSupplier = finalStateMessageSupplier;
   }
 
   @CheckForNull
@@ -49,7 +62,8 @@ protected AirbyteMessage computeNext() {
             new CtidStatus()
                 .withVer(CTID_STATUS_VERSION)
                 .withType(CTID_STATUS_TYPE)
-                .withCtid(lastCtid));
+                .withCtid(lastCtid)
+                .withIncrementalState(streamStateForIncrementalRun));
       }
       // Use try-catch to catch Exception that could occur when connection to the database fails
       try {
@@ -68,11 +82,9 @@ protected AirbyteMessage computeNext() {
       }
     } else if (!hasEmittedFinalState) {
       hasEmittedFinalState = true;
-      return CtidStateManager.createStateMessage(pair,
-          new CtidStatus()
-              .withVer(CTID_STATUS_VERSION)
-              .withType(CTID_STATUS_TYPE)
-              .withCtid(lastCtid));
+      return new AirbyteMessage()
+          .withType(Type.STATE)
+          .withState(finalStateMessageSupplier.apply(pair, streamStateForIncrementalRun));
     } else {
       return endOfData();
     }

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

@@ -19,6 +19,7 @@
 import io.airbyte.protocol.models.v0.AirbyteMessage;
 import io.airbyte.protocol.models.v0.AirbyteMessage.Type;
 import io.airbyte.protocol.models.v0.AirbyteRecordMessage;
+import io.airbyte.protocol.models.v0.AirbyteStateMessage;
 import io.airbyte.protocol.models.v0.AirbyteStream;
 import io.airbyte.protocol.models.v0.CatalogHelpers;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
@@ -32,6 +33,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -42,18 +46,22 @@ public class PostgresCtidHandler {
   private final JdbcDatabase database;
   private final JdbcCompatibleSourceOperations<?> sourceOperations;
   private final String quoteString;
-  private final CtidStatus ctidStatus;
   private final CtidStateManager ctidStateManager;
+  private final Function<AirbyteStreamNameNamespacePair, JsonNode> streamStateForIncrementalRunSupplier;
+  private final BiFunction<AirbyteStreamNameNamespacePair, JsonNode, AirbyteStateMessage> finalStateMessageSupplier;
+
   public PostgresCtidHandler(final JdbcDatabase database,
       final JdbcCompatibleSourceOperations<?> sourceOperations,
       final String quoteString,
-      final CtidStatus ctidStatus,
-      final CtidStateManager ctidStateManager) {
+      final CtidStateManager ctidStateManager,
+      final Function<AirbyteStreamNameNamespacePair, JsonNode> streamStateForIncrementalRunSupplier,
+      final BiFunction<AirbyteStreamNameNamespacePair, JsonNode, AirbyteStateMessage> finalStateMessageSupplier) {
     this.database = database;
     this.sourceOperations = sourceOperations;
     this.quoteString = quoteString;
-    this.ctidStatus = ctidStatus;
     this.ctidStateManager = ctidStateManager;
+    this.streamStateForIncrementalRunSupplier = streamStateForIncrementalRunSupplier;
+    this.finalStateMessageSupplier = finalStateMessageSupplier;
   }
   public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(
       final ConfiguredAirbyteCatalog catalog,
@@ -174,8 +182,13 @@ private AutoCloseableIterator<AirbyteMessage> augmentWithLogs(final AutoCloseabl
 
   private AutoCloseableIterator<AirbyteMessage> augmentWithState(final AutoCloseableIterator<AirbyteMessage> recordIterator,
       final AirbyteStreamNameNamespacePair pair) {
+
+    final CtidStatus currentCtidStatus = ctidStateManager.getCtidStatus(pair);
+    final JsonNode incrementalState = currentCtidStatus.getIncrementalState() != null ? currentCtidStatus.getIncrementalState() : streamStateForIncrementalRunSupplier.apply(pair);
+
+
     return AutoCloseableIterators.transform(
-        autoClosableIterator -> new CtidStateIterator(recordIterator, pair), recordIterator, pair);
+        autoClosableIterator -> new CtidStateIterator(recordIterator, pair, incrementalState, finalStateMessageSupplier ), recordIterator, pair);
   }
 
   private AutoCloseableIterator<AirbyteMessage> swallowCtid(final AutoCloseableIterator<AirbyteMessage> iterator,

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

@@ -73,6 +73,14 @@ public XminStatus getXminStatus(final AirbyteStreamNameNamespacePair pair) {
    * @return AirbyteMessage which includes information on state of records read so far
    */
   public static AirbyteMessage createStateMessage(final AirbyteStreamNameNamespacePair pair, final XminStatus xminStatus) {
+    final AirbyteStateMessage stateMessage = getAirbyteStateMessage(pair, xminStatus);
+
+    return new AirbyteMessage()
+        .withType(Type.STATE)
+        .withState(stateMessage);
+  }
+
+  public static AirbyteStateMessage getAirbyteStateMessage(final AirbyteStreamNameNamespacePair pair, final XminStatus xminStatus) {
     final AirbyteStreamState airbyteStreamState =
         new AirbyteStreamState()
             .withStreamDescriptor(
@@ -86,10 +94,7 @@ public static AirbyteMessage createStateMessage(final AirbyteStreamNameNamespace
         new AirbyteStateMessage()
             .withType(AirbyteStateType.STREAM)
             .withStream(airbyteStreamState);
-
-    return new AirbyteMessage()
-        .withType(Type.STATE)
-        .withState(stateMessage);
+    return stateMessage;
   }
 
 }

airbyte-integrations/connectors/source-postgres/src/main/resources/internal_models/ctid_internal_models.yaml

@@ -19,3 +19,7 @@ definitions:
       ctid:
         description: ctid bookmark
         type: string
+      incremental_state:
+        description: "State to switch to after completion of ctid snapsho"
+        type: object
+        existingJavaType: com.fasterxml.jackson.databind.JsonNode