[DBZ-PGYB][yugabyte/yugabyte-db#21591] Support to process update records with CHANGE

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/PostgresChangeRecordEmitter.java

@@ -17,6 +17,8 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import io.debezium.connector.postgresql.connection.ReplicaIdentityInfo;
+import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.postgresql.core.BaseConnection;
@@ -57,10 +59,11 @@ public class PostgresChangeRecordEmitter extends RelationalChangeRecordEmitter<P
     private final PostgresConnection connection;
     private final TableId tableId;
     private final Map<String, Object> cachedOldToastedValues = new HashMap<>();
+    private final ReplicaIdentityInfo.ReplicaIdentity replicaIdentity;
 
     public PostgresChangeRecordEmitter(PostgresPartition partition, OffsetContext offset, Clock clock, PostgresConnectorConfig connectorConfig, PostgresSchema schema,
                                        PostgresConnection connection, TableId tableId,
-                                       ReplicationMessage message) {
+                                       ReplicationMessage message, ReplicaIdentityInfo.ReplicaIdentity replicaIdentity) {
         super(partition, offset, clock, connectorConfig);
 
         this.schema = schema;
@@ -70,6 +73,14 @@ public PostgresChangeRecordEmitter(PostgresPartition partition, OffsetContext of
 
         this.tableId = tableId;
         Objects.requireNonNull(this.tableId);
+
+        this.replicaIdentity = replicaIdentity;
+    }
+
+    public PostgresChangeRecordEmitter(PostgresPartition partition, OffsetContext offset, Clock clock, PostgresConnectorConfig connectorConfig, PostgresSchema schema,
+                                       PostgresConnection connection, TableId tableId,
+                                       ReplicationMessage message) {
+        this(partition, offset, clock, connectorConfig, schema, connection, tableId, message, ReplicaIdentityInfo.ReplicaIdentity.CHANGE);
     }
 
     @Override
@@ -107,6 +118,11 @@ protected Object[] getOldColumnValues() {
                 case CREATE:
                     return null;
                 case UPDATE:
+                    // YB Note: For replica identity CHANGE, there is no old column value available.
+                    if (replicaIdentity == ReplicaIdentityInfo.ReplicaIdentity.CHANGE) {
+                        return null;
+                    }
+
                     return columnValues(message.getOldTupleList(), tableId, true, true, true);
                 default:
                     return columnValues(message.getOldTupleList(), tableId, true, false, true);
@@ -151,7 +167,7 @@ private DataCollectionSchema synchronizeTableSchema(DataCollectionSchema tableSc
         return schema.schemaFor(tableId);
     }
 
-    private Object[] columnValues(List<ReplicationMessage.Column> columns, TableId tableId, boolean refreshSchemaIfChanged,
+    protected Object[] columnValues(List<ReplicationMessage.Column> columns, TableId tableId, boolean refreshSchemaIfChanged,
                                   boolean sourceOfToasted, boolean oldValues)
             throws SQLException {
         if (columns == null || columns.isEmpty()) {
@@ -188,12 +204,62 @@ private Object[] columnValues(List<ReplicationMessage.Column> columns, TableId t
                         }
                     }
                 }
-                values[position] = value;
+
+                if (replicaIdentity == ReplicaIdentityInfo.ReplicaIdentity.CHANGE) {
+                    if (!UnchangedToastedReplicationMessageColumn.isUnchangedToastedValue(value)) {
+                        values[position] = new Object[]{value, Boolean.TRUE};
+                    }
+                } else {
+                    LOGGER.info("Replica identity is NOT CHANGE");
+                    values[position] = value;
+                }
             }
         }
         return values;
     }
 
+    @Override
+    protected void emitUpdateRecord(Receiver<PostgresPartition> receiver, TableSchema tableSchema) throws InterruptedException {
+        Object[] oldColumnValues = getOldColumnValues();
+        Object[] newColumnValues = getNewColumnValues();
+
+        Struct oldKey = tableSchema.keyFromColumnData(oldColumnValues);
+        Struct newKey = tableSchema.keyFromColumnData(newColumnValues);
+
+        Struct newValue = tableSchema.valueFromColumnData(newColumnValues);
+        Struct oldValue = tableSchema.valueFromColumnData(oldColumnValues);
+
+        if (skipEmptyMessages() && (newColumnValues == null || newColumnValues.length == 0)) {
+            LOGGER.debug("no new values found for table '{}' from update message at '{}'; skipping record", tableSchema, getOffset().getSourceInfo());
+            return;
+        }
+
+        /*
+         * If skip.messages.without.change is configured true,
+         * Skip Publishing the message in case there is no change in monitored columns
+         * (Postgres) Only works if REPLICA IDENTITY is set to FULL - as oldValues won't be available
+         */
+        if (skipMessagesWithoutChange() && Objects.nonNull(newValue) && newValue.equals(oldValue)) {
+            LOGGER.debug("No new values found for table '{}' in included columns from update message at '{}'; skipping record", tableSchema,
+              getOffset().getSourceInfo());
+            return;
+        }
+        // some configurations does not provide old values in case of updates
+        // in this case we handle all updates as regular ones
+
+        // YB Note: If replica identity is change, one hack is that we always know there will be no
+        // oldKey present so we should simply go ahead with this block. Also, oldKey would be null
+        // at this stage if replica identity is CHANGE
+        if (oldKey == null || replicaIdentity == ReplicaIdentityInfo.ReplicaIdentity.CHANGE || Objects.equals(oldKey, newKey)) {
+            Struct envelope = tableSchema.getEnvelopeSchema().update(oldValue, newValue, getOffset().getSourceInfo(), getClock().currentTimeAsInstant());
+            receiver.changeRecord(getPartition(), tableSchema, Operation.UPDATE, newKey, envelope, getOffset(), null);
+        }
+        // PK update -> emit as delete and re-insert with new key
+        else {
+            emitUpdateAsPrimaryKeyChangeRecord(receiver, tableSchema, oldKey, newKey, oldValue, newValue);
+        }
+    }
+
     private int getPosition(String columnName, Table table, Object[] values) {
         final Column tableColumn = table.columnWithName(columnName);
 

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/PostgresConnectorConfig.java

@@ -644,6 +644,11 @@ public static AutoCreateMode parse(String value, String defaultValue) {
                             "FULL - Records the old values of all columns in the row." +
                             "NOTHING - Records no information about the old row. This is the default for system tables.");
 
+    public static final Field REPLICA_IDENTITY = Field.create("replica.identity")
+           .withDisplayName("Internal: Replica Identity")
+           .withType(Type.STRING)
+           .withImportance(Importance.LOW);
+
     public static final Field STREAM_PARAMS = Field.create("slot.stream.params")
             .withDisplayName("Optional parameters to pass to the logical decoder when the stream is started.")
             .withType(Type.STRING)

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/PostgresSchema.java

@@ -66,9 +66,7 @@ protected PostgresSchema(PostgresConnectorConfig config, PostgresDefaultValueCon
 
     private static TableSchemaBuilder getTableSchemaBuilder(PostgresConnectorConfig config, PostgresValueConverter valueConverter,
                                                             PostgresDefaultValueConverter defaultValueConverter) {
-        return new TableSchemaBuilder(valueConverter, defaultValueConverter, config.schemaNameAdjuster(),
-                config.customConverterRegistry(), config.getSourceInfoStructMaker().schema(),
-                config.getFieldNamer(), false);
+        return new PGTableSchemaBuilder(valueConverter, defaultValueConverter, config, false /* multiPartitionMode */);
     }
 
     /**

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/PostgresSnapshotChangeEventSource.java

@@ -7,12 +7,18 @@
 
 import java.sql.SQLException;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
+import io.debezium.connector.postgresql.connection.ReplicaIdentityInfo;
+import io.debezium.connector.postgresql.connection.YBReplicaIdentity;
+import io.debezium.connector.postgresql.connection.pgoutput.PgOutputReplicationMessage;
+import io.debezium.pipeline.spi.ChangeRecordEmitter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -46,6 +52,7 @@ public class PostgresSnapshotChangeEventSource extends RelationalSnapshotChangeE
     private final Snapshotter blockingSnapshotter;
     private final SlotCreationResult slotCreatedInfo;
     private final SlotState startingSlotInfo;
+    private final Map<TableId, YBReplicaIdentity> replicaIdentityMap;
 
     public PostgresSnapshotChangeEventSource(PostgresConnectorConfig connectorConfig, Snapshotter snapshotter,
                                              MainConnectionProvidingConnectionFactory<PostgresConnection> connectionFactory, PostgresSchema schema,
@@ -60,6 +67,7 @@ public PostgresSnapshotChangeEventSource(PostgresConnectorConfig connectorConfig
         this.slotCreatedInfo = slotCreatedInfo;
         this.startingSlotInfo = startingSlotInfo;
         this.blockingSnapshotter = new AlwaysSnapshotter();
+        this.replicaIdentityMap = new ConcurrentHashMap<>();
     }
 
     @Override
@@ -88,6 +96,22 @@ protected SnapshotContext<PostgresPartition, PostgresOffsetContext> prepare(Post
         return new PostgresSnapshotContext(partition, connectorConfig.databaseName(), onDemand);
     }
 
+    @Override
+    protected ChangeRecordEmitter<PostgresPartition> getChangeRecordEmitter(
+      PostgresPartition partition, PostgresOffsetContext offset, TableId tableId, Object[] row,
+      Instant timestamp) {
+        offset.event(tableId, timestamp);
+
+        YBReplicaIdentity ybReplicaIdentity = replicaIdentityMap.get(tableId);
+        if (ybReplicaIdentity == null) {
+            ybReplicaIdentity = new YBReplicaIdentity(connectorConfig, tableId);
+            replicaIdentityMap.put(tableId, ybReplicaIdentity);
+        }
+
+        return new YBSnapshotChangeRecordEmitter<>(partition, offset, row, getClock(),
+                                                   connectorConfig, ybReplicaIdentity.getReplicaIdentity());
+    }
+
     @Override
     protected void connectionCreated(RelationalSnapshotContext<PostgresPartition, PostgresOffsetContext> snapshotContext)
             throws Exception {

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/PostgresStreamingChangeEventSource.java

@@ -9,22 +9,17 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.OptionalLong;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 
+import io.debezium.connector.postgresql.connection.*;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.postgresql.core.BaseConnection;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import io.debezium.DebeziumException;
-import io.debezium.connector.postgresql.connection.LogicalDecodingMessage;
-import io.debezium.connector.postgresql.connection.Lsn;
-import io.debezium.connector.postgresql.connection.PostgresConnection;
-import io.debezium.connector.postgresql.connection.ReplicationConnection;
-import io.debezium.connector.postgresql.connection.ReplicationMessage;
 import io.debezium.connector.postgresql.connection.ReplicationMessage.Operation;
-import io.debezium.connector.postgresql.connection.ReplicationStream;
-import io.debezium.connector.postgresql.connection.WalPositionLocator;
 import io.debezium.connector.postgresql.spi.Snapshotter;
 import io.debezium.heartbeat.Heartbeat;
 import io.debezium.pipeline.ErrorHandler;
@@ -81,6 +76,7 @@ public class PostgresStreamingChangeEventSource implements StreamingChangeEventS
     private long numberOfEventsSinceLastEventSentOrWalGrowingWarning = 0;
     private Lsn lastCompletelyProcessedLsn;
     private PostgresOffsetContext effectiveOffset;
+    private final Map<TableId, YBReplicaIdentity> replicaIdentityMap;
 
     public PostgresStreamingChangeEventSource(PostgresConnectorConfig connectorConfig, Snapshotter snapshotter,
                                               PostgresConnection connection, PostgresEventDispatcher<TableId> dispatcher, ErrorHandler errorHandler, Clock clock,
@@ -96,7 +92,7 @@ public PostgresStreamingChangeEventSource(PostgresConnectorConfig connectorConfi
         this.snapshotter = snapshotter;
         this.replicationConnection = replicationConnection;
         this.connectionProbeTimer = ElapsedTimeStrategy.constant(Clock.system(), connectorConfig.statusUpdateInterval());
-
+        this.replicaIdentityMap = new ConcurrentHashMap<>();
     }
 
     @Override
@@ -316,6 +312,15 @@ else if (message.getOperation() == Operation.MESSAGE) {
                     tableId,
                     message.getOperation());
 
+            LOGGER.info("Received record of type {}", message.getOperation());
+
+            // YB Note: Get the cached replica identity.
+            YBReplicaIdentity ybReplicaIdentity = replicaIdentityMap.get(tableId);
+            if (ybReplicaIdentity == null) {
+                ybReplicaIdentity = new YBReplicaIdentity(connectorConfig, tableId);
+                replicaIdentityMap.put(tableId, ybReplicaIdentity);
+            }
+
             boolean dispatched = message.getOperation() != Operation.NOOP && dispatcher.dispatchDataChangeEvent(
                     partition,
                     tableId,
@@ -327,7 +332,8 @@ else if (message.getOperation() == Operation.MESSAGE) {
                             schema,
                             connection,
                             tableId,
-                            message));
+                            message,
+                            ybReplicaIdentity.getReplicaIdentity()));
 
             maybeWarnAboutGrowingWalBacklog(dispatched);
         }

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/YBSnapshotChangeRecordEmitter.java

@@ -0,0 +1,52 @@
+package io.debezium.connector.postgresql;
+
+import io.debezium.connector.postgresql.connection.ReplicaIdentityInfo;
+import io.debezium.data.Envelope;
+import io.debezium.pipeline.spi.OffsetContext;
+import io.debezium.relational.RelationalChangeRecordEmitter;
+import io.debezium.relational.RelationalDatabaseConnectorConfig;
+import io.debezium.util.Clock;
+
+/**
+ * Custom snapshot change record emitter for YugabyteDB which forms column values object based on
+ * the replica identity type
+ * @param <P> instance of {@link io.debezium.pipeline.spi.Partition}
+ * @author Vaibhav Kushwaha (vkushwaha@yugabyte.com)
+ */
+public class YBSnapshotChangeRecordEmitter<P extends PostgresPartition> extends RelationalChangeRecordEmitter<P> {
+  private final Object[] row;
+  private final ReplicaIdentityInfo.ReplicaIdentity replicaIdentity;
+
+  public YBSnapshotChangeRecordEmitter(P partition, OffsetContext offset, Object[] row, Clock clock, RelationalDatabaseConnectorConfig connectorConfig,
+                                       ReplicaIdentityInfo.ReplicaIdentity replicaIdentity) {
+    super(partition, offset, clock, connectorConfig);
+
+    this.row = row;
+    this.replicaIdentity = replicaIdentity;
+  }
+
+  @Override
+  public Envelope.Operation getOperation() {
+    return Envelope.Operation.READ;
+  }
+
+  @Override
+  protected Object[] getOldColumnValues() {
+    throw new UnsupportedOperationException("Can't get old row values for READ record");
+  }
+
+  @Override
+  protected Object[] getNewColumnValues() {
+    Object[] values = new Object[row.length];
+
+    for (int position = 0; position < values.length; ++position) {
+      if (replicaIdentity == ReplicaIdentityInfo.ReplicaIdentity.CHANGE) {
+        values[position] = new Object[]{row[position], Boolean.TRUE};
+      } else {
+        values[position] = row[position];
+      }
+    }
+
+    return values;
+  }
+}

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/connection/PostgresConnection.java

@@ -65,6 +65,7 @@ public class PostgresConnection extends JdbcConnection {
     public static final String CONNECTION_VALIDATE_CONNECTION = "Debezium Validate Connection";
     public static final String CONNECTION_HEARTBEAT = "Debezium Heartbeat";
     public static final String CONNECTION_GENERAL = "Debezium General";
+    public static final String CONNECTION_FETCH_REPLICA_IDENTITY = "Debezium YB Fetch Replica Identity";
 
     private static final Pattern FUNCTION_DEFAULT_PATTERN = Pattern.compile("^[(]?[A-Za-z0-9_.]+\\((?:.+(?:, ?.+)*)?\\)");
     private static final Pattern EXPRESSION_DEFAULT_PATTERN = Pattern.compile("\\(+(?:.+(?:[+ - * / < > = ~ ! @ # % ^ & | ` ?] ?.+)+)+\\)");

debezium-connector-postgres/src/main/java/io/debezium/connector/postgresql/connection/ReplicaIdentityInfo.java

@@ -50,7 +50,8 @@ public enum ReplicaIdentity {
         FULL("UPDATE AND DELETE events will contain the previous values of all the columns"),
         DEFAULT("UPDATE and DELETE events will contain previous values only for PK columns"),
         INDEX("UPDATE and DELETE events will contain previous values only for columns present in the REPLICA IDENTITY index"),
-        UNKNOWN("Unknown REPLICA IDENTITY");
+        UNKNOWN("Unknown REPLICA IDENTITY"),
+        CHANGE("UPDATE events will contain values only for changed columns");
 
         private final String description;
 
@@ -77,6 +78,8 @@ static ReplicaIdentityInfo.ReplicaIdentity parseFromDB(String s) {
                     return INDEX;
                 case "f":
                     return FULL;
+                case "c":
+                    return CHANGE;
                 default:
                     return UNKNOWN;
             }