source-postgres-cdc: allow LSN to be acknowledged after sync

airbyte-integrations/bases/debezium-v1-9-6/build.gradle

@@ -13,12 +13,16 @@ dependencies {
     implementation 'io.debezium:debezium-embedded:1.9.6.Final'
     implementation 'io.debezium:debezium-connector-sqlserver:1.9.6.Final'
     implementation 'io.debezium:debezium-connector-mysql:1.9.6.Final'
-    implementation 'io.debezium:debezium-connector-postgres:1.9.6.Final'
+    implementation files('debezium-connector-postgres-1.9.6.Final.jar')
     implementation 'org.codehaus.plexus:plexus-utils:3.4.2'
 
     testFixturesImplementation project(':airbyte-db:db-lib')
     testFixturesImplementation project(':airbyte-integrations:bases:base-java')
 
+    testImplementation project(':airbyte-test-utils')
+    testImplementation libs.connectors.testcontainers.jdbc
+    testImplementation libs.connectors.testcontainers.postgresql
+
     testFixturesImplementation 'org.junit.jupiter:junit-jupiter-engine:5.4.2'
     testFixturesImplementation 'org.junit.jupiter:junit-jupiter-api:5.4.2'
     testFixturesImplementation 'org.junit.jupiter:junit-jupiter-params:5.4.2'

airbyte-integrations/bases/debezium-v1-9-6/src/main/java/io/airbyte/integrations/debezium/internals/PostgresDebeziumStateUtil.java

@@ -8,6 +8,7 @@
 import static io.debezium.connector.postgresql.SourceInfo.LSN_KEY;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import io.airbyte.db.jdbc.JdbcUtils;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
 import io.debezium.config.Configuration;
 import io.debezium.connector.common.OffsetReader;
@@ -16,10 +17,18 @@
 import io.debezium.connector.postgresql.PostgresOffsetContext.Loader;
 import io.debezium.connector.postgresql.PostgresPartition;
 import io.debezium.connector.postgresql.connection.Lsn;
+import io.debezium.jdbc.JdbcConnection.ResultSetMapper;
+import io.debezium.jdbc.JdbcConnection.StatementFactory;
 import io.debezium.pipeline.spi.Offsets;
 import io.debezium.pipeline.spi.Partition;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.Properties;
@@ -29,6 +38,10 @@
 import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
 import org.apache.kafka.connect.storage.FileOffsetBackingStore;
 import org.apache.kafka.connect.storage.OffsetStorageReaderImpl;
+import org.postgresql.core.BaseConnection;
+import org.postgresql.replication.LogSequenceNumber;
+import org.postgresql.replication.PGReplicationStream;
+import org.postgresql.replication.fluent.logical.ChainedLogicalStreamBuilder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -40,33 +53,131 @@ public class PostgresDebeziumStateUtil {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(PostgresDebeziumStateUtil.class);
 
-  public boolean isSavedOffsetAfterReplicationSlotLSN(final Properties baseProperties,
-                                                      final ConfiguredAirbyteCatalog catalog,
-                                                      final JsonNode cdcState,
-                                                      final JsonNode replicationSlot,
-                                                      final JsonNode config) {
+  public boolean isSavedOffsetAfterReplicationSlotLSN(final JsonNode replicationSlot,
+                                                      final OptionalLong savedOffset) {
 
+    if (Objects.isNull(savedOffset) || savedOffset.isEmpty()) {
+      return true;
+    }
+
+    if (replicationSlot.has("confirmed_flush_lsn")) {
+      final long confirmedFlushLsnOnServerSide = Lsn.valueOf(replicationSlot.get("confirmed_flush_lsn").asText()).asLong();
+      LOGGER.info("Replication slot confirmed_flush_lsn : " + confirmedFlushLsnOnServerSide + " Saved offset LSN : " + savedOffset.getAsLong());
+      return savedOffset.getAsLong() >= confirmedFlushLsnOnServerSide;
+    } else if (replicationSlot.has("restart_lsn")) {
+      final long restartLsn = Lsn.valueOf(replicationSlot.get("restart_lsn").asText()).asLong();
+      LOGGER.info("Replication slot restart_lsn : " + restartLsn + " Saved offset LSN : " + savedOffset.getAsLong());
+      return savedOffset.getAsLong() >= restartLsn;
+    }
+
+    // We return true when saved offset is not present cause using an empty offset would result in sync
+    // from scratch anyway
+    return true;
+  }
+
+  public OptionalLong savedOffset(final Properties baseProperties,
+                                  final ConfiguredAirbyteCatalog catalog,
+                                  final JsonNode cdcState,
+                                  final JsonNode config) {
     final DebeziumPropertiesManager debeziumPropertiesManager = new DebeziumPropertiesManager(baseProperties, config, catalog,
         AirbyteFileOffsetBackingStore.initializeState(cdcState),
         Optional.empty());
     final Properties debeziumProperties = debeziumPropertiesManager.getDebeziumProperties();
-    final OptionalLong savedOffset = parseSavedOffset(debeziumProperties);
-
-    if (savedOffset.isPresent()) {
-      if (replicationSlot.has("confirmed_flush_lsn")) {
-        final long confirmedFlushLsnOnServerSide = Lsn.valueOf(replicationSlot.get("confirmed_flush_lsn").asText()).asLong();
-        LOGGER.info("Replication slot confirmed_flush_lsn : " + confirmedFlushLsnOnServerSide + " Saved offset LSN : " + savedOffset.getAsLong());
-        return savedOffset.getAsLong() >= confirmedFlushLsnOnServerSide;
-      } else if (replicationSlot.has("restart_lsn")) {
-        final long restartLsn = Lsn.valueOf(replicationSlot.get("restart_lsn").asText()).asLong();
-        LOGGER.info("Replication slot restart_lsn : " + restartLsn + " Saved offset LSN : " + savedOffset.getAsLong());
-        return savedOffset.getAsLong() >= restartLsn;
+    return parseSavedOffset(debeziumProperties);
+  }
+
+  private Connection connection(final JsonNode jdbcConfig) throws SQLException {
+    Properties properties = new Properties();
+    properties.setProperty("user", jdbcConfig.has(JdbcUtils.USERNAME_KEY) ? jdbcConfig.get(JdbcUtils.USERNAME_KEY).asText()
+        : null);
+    properties.setProperty("password", jdbcConfig.has(JdbcUtils.PASSWORD_KEY) ? jdbcConfig.get(JdbcUtils.PASSWORD_KEY).asText()
+        : null);
+    properties.setProperty("assumeMinServerVersion", "9.4");
+    properties.setProperty("ApplicationName", "Airbyte Debezium Streaming");
+    properties.setProperty("replication", "database");
+    properties.setProperty("preferQueryMode", "simple"); // replication protocol only supports simple query mode
+
+    return DriverManager.getConnection(jdbcConfig.get(JdbcUtils.JDBC_URL_KEY).asText(), properties);
+  }
+
+  public void commitLSNToPostgresDatabase(final JsonNode jdbcConfig,
+                                          final OptionalLong savedOffset,
+                                          final String slotName,
+                                          final String publicationName,
+                                          final String plugin) {
+    if (Objects.isNull(savedOffset) || savedOffset.isEmpty()) {
+      return;
+    }
+
+    final LogSequenceNumber logSequenceNumber = LogSequenceNumber.valueOf(savedOffset.getAsLong());
+
+    try (final BaseConnection pgConnection = ((BaseConnection) connection(jdbcConfig))) {
+      validateReplicationConnection(pgConnection);
+
+      ChainedLogicalStreamBuilder streamBuilder = pgConnection
+          .getReplicationAPI()
+          .replicationStream()
+          .logical()
+          .withSlotName("\"" + slotName + "\"")
+          .withStartPosition(logSequenceNumber);
+
+      streamBuilder = addSlotOption(publicationName, plugin, pgConnection, streamBuilder);
+
+      try (final PGReplicationStream stream = streamBuilder.start()) {
+        stream.forceUpdateStatus();
+
+        stream.setFlushedLSN(logSequenceNumber);
+        stream.setAppliedLSN(logSequenceNumber);
+
+        stream.forceUpdateStatus();
       }
+    } catch (SQLException e) {
+      throw new RuntimeException(e);
     }
+  }
 
-    // We return true when saved offset is not present cause using an empty offset would result in sync
-    // from scratch anyway
-    return true;
+  private ChainedLogicalStreamBuilder addSlotOption(final String publicationName,
+                                                    final String plugin,
+                                                    final BaseConnection pgConnection,
+                                                    ChainedLogicalStreamBuilder streamBuilder) {
+    if (plugin.equalsIgnoreCase("pgoutput")) {
+      streamBuilder = streamBuilder.withSlotOption("proto_version", 1)
+          .withSlotOption("publication_names", publicationName);
+
+      if (pgConnection.haveMinimumServerVersion(140000)) {
+        streamBuilder = streamBuilder.withSlotOption("messages", true);
+      }
+    } else if (plugin.equalsIgnoreCase("wal2json")) {
+      streamBuilder = streamBuilder
+          .withSlotOption("pretty-print", 1)
+          .withSlotOption("write-in-chunks", 1)
+          .withSlotOption("include-xids", 1)
+          .withSlotOption("include-timestamp", 1)
+          .withSlotOption("include-not-null", "true");
+    } else {
+      throw new RuntimeException("Unknown plugin value : " + plugin);
+    }
+    return streamBuilder;
+  }
+
+  private void validateReplicationConnection(final BaseConnection pgConnection) throws SQLException {
+    final Lsn xlogStart = queryAndMap(pgConnection, "IDENTIFY_SYSTEM", Connection::createStatement, rs -> {
+      if (!rs.next()) {
+        throw new IllegalStateException("The DB connection is not a valid replication connection");
+      }
+      String xlogpos = rs.getString("xlogpos");
+      return Lsn.valueOf(xlogpos);
+    });
+  }
+
+  private <T> T queryAndMap(final Connection conn, final String query, final StatementFactory statementFactory, final ResultSetMapper<T> mapper)
+      throws SQLException {
+    Objects.requireNonNull(mapper, "Mapper must be provided");
+    try (Statement statement = statementFactory.createStatement(conn)) {
+      try (ResultSet resultSet = statement.executeQuery(query);) {
+        return mapper.apply(resultSet);
+      }
+    }
   }
 
   /**

airbyte-integrations/bases/debezium-v1-9-6/src/test/java/io/airbyte/integrations/debezium/internals/PostgresDebeziumStateUtilTest.java

@@ -6,13 +6,30 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.collect.ImmutableMap;
+import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.json.Jsons;
+import io.airbyte.commons.string.Strings;
+import io.airbyte.db.PostgresUtils;
+import io.airbyte.db.factory.DataSourceFactory;
+import io.airbyte.db.factory.DatabaseDriver;
+import io.airbyte.db.jdbc.DefaultJdbcDatabase;
+import io.airbyte.db.jdbc.JdbcDatabase;
 import io.airbyte.db.jdbc.JdbcUtils;
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog;
+import io.airbyte.test.utils.PostgreSQLContainerHelper;
+import io.debezium.connector.postgresql.connection.Lsn;
+import java.sql.SQLException;
 import java.util.List;
+import java.util.Map;
+import java.util.OptionalLong;
 import java.util.Properties;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.testcontainers.containers.PostgreSQLContainer;
+import org.testcontainers.utility.DockerImageName;
+import org.testcontainers.utility.MountableFile;
 
 public class PostgresDebeziumStateUtilTest {
 
@@ -47,9 +64,12 @@ public void stateGeneratedAfterSnapshotCompletionAfterReplicationSlot() {
     final JsonNode cdcState = Jsons.deserialize(
         "{\"{\\\"schema\\\":null,\\\"payload\\\":[\\\"db_jagkjrgxhw\\\",{\\\"server\\\":\\\"db_jagkjrgxhw\\\"}]}\":\"{\\\"last_snapshot_record\\\":true,\\\"lsn\\\":23897640,\\\"txId\\\":505,\\\"ts_usec\\\":1659422332985000,\\\"snapshot\\\":true}\"}");
 
-    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(new Properties(),
-        new ConfiguredAirbyteCatalog(), cdcState, REPLICATION_SLOT, CONFIG);
+    final OptionalLong savedOffset = postgresDebeziumStateUtil.savedOffset(new Properties(),
+        new ConfiguredAirbyteCatalog(), cdcState, CONFIG);
+    Assertions.assertTrue(savedOffset.isPresent());
+    Assertions.assertEquals(savedOffset.getAsLong(), 23897640L);
 
+    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(REPLICATION_SLOT, savedOffset);
     Assertions.assertTrue(savedOffsetAfterReplicationSlotLSN);
   }
 
@@ -58,9 +78,12 @@ public void stateGeneratedAfterSnapshotCompletionBeforeReplicationSlot() {
     final JsonNode cdcState = Jsons.deserialize(
         "{\"{\\\"schema\\\":null,\\\"payload\\\":[\\\"db_jagkjrgxhw\\\",{\\\"server\\\":\\\"db_jagkjrgxhw\\\"}]}\":\"{\\\"last_snapshot_record\\\":true,\\\"lsn\\\":23896935,\\\"txId\\\":505,\\\"ts_usec\\\":1659422332985000,\\\"snapshot\\\":true}\"}");
 
-    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(new Properties(),
-        new ConfiguredAirbyteCatalog(), cdcState, REPLICATION_SLOT, CONFIG);
+    final OptionalLong savedOffset = postgresDebeziumStateUtil.savedOffset(new Properties(),
+        new ConfiguredAirbyteCatalog(), cdcState, CONFIG);
+    Assertions.assertTrue(savedOffset.isPresent());
+    Assertions.assertEquals(savedOffset.getAsLong(), 23896935L);
 
+    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(REPLICATION_SLOT, savedOffset);
     Assertions.assertFalse(savedOffsetAfterReplicationSlotLSN);
   }
 
@@ -69,9 +92,12 @@ public void stateGeneratedFromWalStreamingAfterReplicationSlot() {
     final JsonNode cdcState = Jsons.deserialize(
         "{\"{\\\"schema\\\":null,\\\"payload\\\":[\\\"db_jagkjrgxhw\\\",{\\\"server\\\":\\\"db_jagkjrgxhw\\\"}]}\":\"{\\\"transaction_id\\\":null,\\\"lsn_proc\\\":23901120,\\\"lsn_commit\\\":23901120,\\\"lsn\\\":23901120,\\\"txId\\\":525,\\\"ts_usec\\\":1659422649959099}\"}");
 
-    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(new Properties(),
-        new ConfiguredAirbyteCatalog(), cdcState, REPLICATION_SLOT, CONFIG);
+    final OptionalLong savedOffset = postgresDebeziumStateUtil.savedOffset(new Properties(),
+        new ConfiguredAirbyteCatalog(), cdcState, CONFIG);
+    Assertions.assertTrue(savedOffset.isPresent());
+    Assertions.assertEquals(savedOffset.getAsLong(), 23901120L);
 
+    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(REPLICATION_SLOT, savedOffset);
     Assertions.assertTrue(savedOffsetAfterReplicationSlotLSN);
   }
 
@@ -80,18 +106,92 @@ public void stateGeneratedFromWalStreamingBeforeReplicationSlot() {
     final JsonNode cdcState = Jsons.deserialize(
         "{\"{\\\"schema\\\":null,\\\"payload\\\":[\\\"db_jagkjrgxhw\\\",{\\\"server\\\":\\\"db_jagkjrgxhw\\\"}]}\":\"{\\\"transaction_id\\\":null,\\\"lsn_proc\\\":23896935,\\\"lsn_commit\\\":23896935,\\\"lsn\\\":23896935,\\\"txId\\\":525,\\\"ts_usec\\\":1659422649959099}\"}");
 
-    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(new Properties(),
-        new ConfiguredAirbyteCatalog(), cdcState, REPLICATION_SLOT, CONFIG);
+    final OptionalLong savedOffset = postgresDebeziumStateUtil.savedOffset(new Properties(),
+        new ConfiguredAirbyteCatalog(), cdcState, CONFIG);
+    Assertions.assertTrue(savedOffset.isPresent());
+    Assertions.assertEquals(savedOffset.getAsLong(), 23896935L);
 
+    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(REPLICATION_SLOT, savedOffset);
     Assertions.assertFalse(savedOffsetAfterReplicationSlotLSN);
   }
 
   @Test
-  public void nullState() {
-    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(new Properties(),
-        new ConfiguredAirbyteCatalog(), null, REPLICATION_SLOT, CONFIG);
+  public void nullOffset() {
+    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(REPLICATION_SLOT, null);
+    Assertions.assertTrue(savedOffsetAfterReplicationSlotLSN);
+  }
+
+  @Test
+  public void emptyState() {
+    final OptionalLong savedOffset = postgresDebeziumStateUtil.savedOffset(new Properties(),
+        new ConfiguredAirbyteCatalog(), null, CONFIG);
+    Assertions.assertTrue(savedOffset.isEmpty());
 
+    final boolean savedOffsetAfterReplicationSlotLSN = postgresDebeziumStateUtil.isSavedOffsetAfterReplicationSlotLSN(REPLICATION_SLOT, savedOffset);
     Assertions.assertTrue(savedOffsetAfterReplicationSlotLSN);
   }
 
+  @ParameterizedTest
+  @ValueSource(strings = {"pgoutput", "wal2json"})
+  public void LsnCommitTest(final String plugin) throws SQLException {
+    final DockerImageName myImage = DockerImageName.parse("debezium/postgres:13-alpine").asCompatibleSubstituteFor("postgres");
+    final String dbName = Strings.addRandomSuffix("db", "_", 10).toLowerCase();
+    final String fullReplicationSlot = "debezium_slot" + "_" + dbName;
+    final String publication = "publication";
+    try (final PostgreSQLContainer<?> container = new PostgreSQLContainer<>(myImage)) {
+      container.start();
+
+      final String initScriptName = "init_" + dbName.concat(".sql");
+      final String tmpFilePath = IOs.writeFileToRandomTmpDir(initScriptName, "CREATE DATABASE " + dbName + ";");
+      PostgreSQLContainerHelper.runSqlScript(MountableFile.forHostPath(tmpFilePath), container);
+
+      final Map<String, String> databaseConfig = Map.of(JdbcUtils.USERNAME_KEY, container.getUsername(),
+          JdbcUtils.PASSWORD_KEY, container.getPassword(),
+          JdbcUtils.JDBC_URL_KEY, String.format(DatabaseDriver.POSTGRESQL.getUrlFormatString(),
+              container.getHost(),
+              container.getFirstMappedPort(),
+              dbName));
+
+      final JdbcDatabase database = new DefaultJdbcDatabase(
+          DataSourceFactory.create(
+              databaseConfig.get(JdbcUtils.USERNAME_KEY),
+              databaseConfig.get(JdbcUtils.PASSWORD_KEY),
+              DatabaseDriver.POSTGRESQL.getDriverClassName(),
+              databaseConfig.get(JdbcUtils.JDBC_URL_KEY)));
+
+      database.execute("SELECT pg_create_logical_replication_slot('" + fullReplicationSlot + "', '" + plugin + "');");
+      database.execute("CREATE PUBLICATION " + publication + " FOR ALL TABLES;");
+
+      database.execute("CREATE TABLE public.test_table (id int primary key, name varchar(256));");
+      database.execute("insert into public.test_table values (1, 'foo');");
+      database.execute("insert into public.test_table values (2, 'bar');");
+
+      final Lsn lsnAtTheBeginning = Lsn.valueOf(
+          getReplicationSlot(database, fullReplicationSlot, plugin, dbName).get("confirmed_flush_lsn").asText());
+
+      final long targetLsn = PostgresUtils.getLsn(database).asLong();
+      postgresDebeziumStateUtil.commitLSNToPostgresDatabase(Jsons.jsonNode(databaseConfig),
+          OptionalLong.of(targetLsn),
+          fullReplicationSlot,
+          publication,
+          plugin);
+
+      final Lsn lsnAfterCommit = Lsn.valueOf(
+          getReplicationSlot(database, fullReplicationSlot, plugin, dbName).get("confirmed_flush_lsn").asText());
+
+      Assertions.assertEquals(1, lsnAfterCommit.compareTo(lsnAtTheBeginning));
+      Assertions.assertEquals(targetLsn, lsnAfterCommit.asLong());
+      container.stop();
+    }
+  }
+
+  private JsonNode getReplicationSlot(final JdbcDatabase database, String slotName, String plugin, String dbName) {
+    try {
+      return database.queryJsons("SELECT * FROM pg_replication_slots WHERE slot_name = ? AND plugin = ? AND database = ?", slotName, plugin, dbName)
+          .get(0);
+    } catch (final SQLException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
 }