KAFKA-13431: Expose the original pre-transform topic partition and offset in sink records

checkstyle/import-control.xml

@@ -448,6 +448,7 @@
     <subpackage name="sink">
       <allow pkg="org.apache.kafka.clients.consumer" />
       <allow pkg="org.apache.kafka.connect.connector" />
+      <allow pkg="org.apache.kafka.connect.transforms" />
       <allow pkg="org.apache.kafka.connect.storage" />
     </subpackage>
 

connect/api/src/main/java/org/apache/kafka/connect/sink/SinkRecord.java

@@ -20,18 +20,26 @@
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.transforms.Transformation;
+
+import java.util.Map;
+import java.util.Objects;
 
 /**
- * SinkRecord is a {@link ConnectRecord} that has been read from Kafka and includes the kafkaOffset of
- * the record in the Kafka topic-partition in addition to the standard fields. This information
- * should be used by the {@link SinkTask} to coordinate kafkaOffset commits.
+ * SinkRecord is a {@link ConnectRecord} that has been read from Kafka and includes the original Kafka record's
+ * topic, partition and offset (before any {@link Transformation transformations} have been applied)
+ * in addition to the standard fields. This information should be used by the {@link SinkTask} to coordinate
+ * offset commits.
  * <p>
  * It also includes the {@link TimestampType}, which may be {@link TimestampType#NO_TIMESTAMP_TYPE}, and the relevant
  * timestamp, which may be {@code null}.
  */
 public class SinkRecord extends ConnectRecord<SinkRecord> {
     private final long kafkaOffset;
     private final TimestampType timestampType;
+    private final String originalTopic;
+    private final Integer originalKafkaPartition;
+    private final long originalKafkaOffset;
 
     public SinkRecord(String topic, int partition, Schema keySchema, Object key, Schema valueSchema, Object value, long kafkaOffset) {
         this(topic, partition, keySchema, key, valueSchema, value, kafkaOffset, null, TimestampType.NO_TIMESTAMP_TYPE);
@@ -44,9 +52,22 @@ public SinkRecord(String topic, int partition, Schema keySchema, Object key, Sch
 
     public SinkRecord(String topic, int partition, Schema keySchema, Object key, Schema valueSchema, Object value, long kafkaOffset,
                       Long timestamp, TimestampType timestampType, Iterable<Header> headers) {
+        this(topic, partition, keySchema, key, valueSchema, value, kafkaOffset, timestamp, timestampType, headers, topic, partition, kafkaOffset);
+    }
+
+    /**
+     * This constructor is intended for use by the Connect runtime only and plugins (sink connectors or transformations)
+     * should not use this directly outside testing code.
+     */
+    public SinkRecord(String topic, int partition, Schema keySchema, Object key, Schema valueSchema, Object value, long kafkaOffset,
+                      Long timestamp, TimestampType timestampType, Iterable<Header> headers, String originalTopic,
+                      Integer originalKafkaPartition, long originalKafkaOffset) {
         super(topic, partition, keySchema, key, valueSchema, value, timestamp, headers);
         this.kafkaOffset = kafkaOffset;
         this.timestampType = timestampType;
+        this.originalTopic = originalTopic;
+        this.originalKafkaPartition = originalKafkaPartition;
+        this.originalKafkaOffset = originalKafkaOffset;
     }
 
     public long kafkaOffset() {
@@ -57,6 +78,105 @@ public TimestampType timestampType() {
         return timestampType;
     }
 
+    /**
+     * Get the original topic for this sink record, before any {@link Transformation transformations} were applied.
+     * In order to be compatible with transformations that mutate topic names, this method should be used
+     * by sink tasks instead of {@link #topic()} for any internal offset tracking purposes (for instance, reporting
+     * offsets to the Connect runtime via {@link SinkTask#preCommit(Map)}).
+     * <p>
+     * This method was added in Apache Kafka 3.6. Sink connectors that use this method but want to maintain backward
+     * compatibility in order to be able to be deployed on older Connect runtimes should guard the call to this method
+     * with a try-catch block, since calling this method will result in a {@link NoSuchMethodError} when the sink
+     * connector is deployed to Connect runtimes older than Kafka 3.6.
+     * For example:
+     * <pre>{@code
+     * String originalTopic;
+     * try {
+     *     originalTopic = record.originalTopic();
+     * } catch (NoSuchMethodError e) {
+     *     log.warn("This connector is not compatible with SMTs that mutate topic names, topic partitions or offset values on this version of Kafka Connect");
+     *     originalTopic = record.topic();
+     * }
+     * }
+     * </pre>
+     * <p>
+     * Note that sink connectors that do their own offset tracking will be incompatible with SMTs that mutate topic
+     * names when deployed to older Connect runtimes that do not support this method.
+     *
+     * @return the topic for this record before any transformations were applied
+     *
+     * @since 3.6
+     */
+    public String originalTopic() {
+        return originalTopic;
+    }
+
+    /**
+     * Get the original topic partition for this sink record, before any {@link Transformation transformations} were applied.
+     * In order to be compatible with transformations that mutate topic partitions, this method should be used
+     * by sink tasks instead of {@link #kafkaPartition()} for any internal offset tracking purposes (for instance, reporting
+     * offsets to the Connect runtime via {@link SinkTask#preCommit(Map)}).
+     * <p>
+     * This method was added in Apache Kafka 3.6. Sink connectors that use this method but want to maintain backward
+     * compatibility in order to be able to be deployed on older Connect runtimes should guard the call to this method
+     * with a try-catch block, since calling this method will result in a {@link NoSuchMethodError} when the sink
+     * connector is deployed to Connect runtimes older than Kafka 3.6.
+     * For example:
+     * <pre>{@code
+     * String originalKafkaPartition;
+     * try {
+     *     originalKafkaPartition = record.originalKafkaPartition();
+     * } catch (NoSuchMethodError e) {
+     *     log.warn("This connector is not compatible with SMTs that mutate topic names, topic partitions or offset values on this version of Kafka Connect");
+     *     originalKafkaPartition = record.kafkaPartition();
+     * }
+     * }
+     * </pre>
+     * <p>
+     * Note that sink connectors that do their own offset tracking will be incompatible with SMTs that mutate topic
+     * partitions when deployed to older Connect runtimes that do not support this method.
+     *
+     * @return the topic partition for this record before any transformations were applied
+     *
+     * @since 3.6
+     */
+    public Integer originalKafkaPartition() {
+        return originalKafkaPartition;
+    }
+
+    /**
+     * Get the original offset for this sink record, before any {@link Transformation transformations} were applied.
+     * In order to be compatible with transformations that mutate offset values, this method should be used
+     * by sink tasks instead of {@link #kafkaOffset()} for any internal offset tracking purposes (for instance, reporting
+     * offsets to the Connect runtime via {@link SinkTask#preCommit(Map)}).
+     * <p>
+     * This method was added in Apache Kafka 3.6. Sink connectors that use this method but want to maintain backward
+     * compatibility in order to be able to be deployed on older Connect runtimes should guard the call to this method
+     * with a try-catch block, since calling this method will result in a {@link NoSuchMethodError} when the sink
+     * connector is deployed to Connect runtimes older than Kafka 3.6.
+     * For example:
+     * <pre>{@code
+     * String originalKafkaOffset;
+     * try {
+     *     originalKafkaOffset = record.originalKafkaOffset();
+     * } catch (NoSuchMethodError e) {
+     *     log.warn("This connector is not compatible with SMTs that mutate topic names, topic partitions or offset values on this version of Kafka Connect");
+     *     originalKafkaOffset = record.kafkaOffset();
+     * }
+     * }
+     * </pre>
+     * <p>
+     * Note that sink connectors that do their own offset tracking will be incompatible with SMTs that mutate offset
+     * values when deployed to older Connect runtimes that do not support this method.
+     *
+     * @return the offset for this record before any transformations were applied
+     *
+     * @since 3.6
+     */
+    public long originalKafkaOffset() {
+        return originalKafkaOffset;
+    }
+
     @Override
     public SinkRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value, Long timestamp) {
         return newRecord(topic, kafkaPartition, keySchema, key, valueSchema, value, timestamp, headers().duplicate());
@@ -65,7 +185,8 @@ public SinkRecord newRecord(String topic, Integer kafkaPartition, Schema keySche
     @Override
     public SinkRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value,
                                 Long timestamp, Iterable<Header> headers) {
-        return new SinkRecord(topic, kafkaPartition, keySchema, key, valueSchema, value, kafkaOffset, timestamp, timestampType, headers);
+        return new SinkRecord(topic, kafkaPartition, keySchema, key, valueSchema, value, kafkaOffset, timestamp, timestampType, headers,
+                originalTopic, originalKafkaPartition, originalKafkaOffset);
     }
 
     @Override
@@ -79,17 +200,21 @@ public boolean equals(Object o) {
 
         SinkRecord that = (SinkRecord) o;
 
-        if (kafkaOffset != that.kafkaOffset)
-            return false;
-
-        return timestampType == that.timestampType;
+        return kafkaOffset == that.kafkaOffset &&
+                timestampType == that.timestampType &&
+                Objects.equals(originalTopic, that.originalTopic) &&
+                Objects.equals(originalKafkaPartition, that.originalKafkaPartition) &&
+                originalKafkaOffset == that.originalKafkaOffset;
     }
 
     @Override
     public int hashCode() {
         int result = super.hashCode();
         result = 31 * result + Long.hashCode(kafkaOffset);
         result = 31 * result + timestampType.hashCode();
+        result = 31 * result + originalTopic.hashCode();
+        result = 31 * result + originalKafkaPartition.hashCode();
+        result = 31 * result + Long.hashCode(originalKafkaOffset);
         return result;
     }
 
@@ -98,6 +223,9 @@ public String toString() {
         return "SinkRecord{" +
                 "kafkaOffset=" + kafkaOffset +
                 ", timestampType=" + timestampType +
+                ", originalTopic=" + originalTopic +
+                ", originalKafkaPartition=" + originalKafkaPartition +
+                ", originalKafkaOffset=" + originalKafkaOffset +
                 "} " + super.toString();
     }
 }

connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTask.java

@@ -19,6 +19,7 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.transforms.Transformation;
 
 import java.util.Collection;
 import java.util.Map;
@@ -105,9 +106,13 @@ public void initialize(SinkTaskContext context) {
     /**
      * Flush all records that have been {@link #put(Collection)} for the specified topic-partitions.
      *
-     * @param currentOffsets the current offset state as of the last call to {@link #put(Collection)}},
-     *                       provided for convenience but could also be determined by tracking all offsets included in the {@link SinkRecord}s
-     *                       passed to {@link #put}.
+     * @param currentOffsets the current offset state as of the last call to {@link #put(Collection)}, provided for
+     *                       convenience but could also be determined by tracking all offsets included in the
+     *                       {@link SinkRecord}s passed to {@link #put}. Note that the topic, partition and offset
+     *                       here correspond to the original Kafka topic partition and offset, before any
+     *                       {@link Transformation transformations} have been applied. These can be tracked by the task
+     *                       through the {@link SinkRecord#originalTopic()}, {@link SinkRecord#originalKafkaPartition()}
+     *                       and {@link SinkRecord#originalKafkaOffset()} methods.
      */
     public void flush(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
     }
@@ -118,11 +123,17 @@ public void flush(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
      * The default implementation simply invokes {@link #flush(Map)} and is thus able to assume all {@code currentOffsets}
      * are safe to commit.
      *
-     * @param currentOffsets the current offset state as of the last call to {@link #put(Collection)}},
-     *                       provided for convenience but could also be determined by tracking all offsets included in the {@link SinkRecord}s
-     *                       passed to {@link #put}.
+     * @param currentOffsets the current offset state as of the last call to {@link #put(Collection)}, provided for
+     *                       convenience but could also be determined by tracking all offsets included in the
+     *                       {@link SinkRecord}s passed to {@link #put}. Note that the topic, partition and offset
+     *                       here correspond to the original Kafka topic partition and offset, before any
+     *                       {@link Transformation transformations} have been applied. These can be tracked by the task
+     *                       through the {@link SinkRecord#originalTopic()}, {@link SinkRecord#originalKafkaPartition()}
+     *                       and {@link SinkRecord#originalKafkaOffset()} methods.
      *
-     * @return an empty map if Connect-managed offset commit is not desired, otherwise a map of offsets by topic-partition that are safe to commit.
+     * @return an empty map if Connect-managed offset commit is not desired, otherwise a map of offsets by topic-partition that are
+     *         safe to commit. Note that the returned topic-partition to offsets map should use the original Kafka
+     *         topic partitions and offsets instead of the transformed values.
      */
     public Map<TopicPartition, OffsetAndMetadata> preCommit(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
         flush(currentOffsets);
@@ -132,7 +143,11 @@ public Map<TopicPartition, OffsetAndMetadata> preCommit(Map<TopicPartition, Offs
     /**
      * The SinkTask uses this method to create writers for newly assigned partitions in case of partition
      * rebalance. This method will be called after partition re-assignment completes and before the SinkTask starts
-     * fetching data. Note that any errors raised from this method will cause the task to stop.
+     * fetching data. Any errors raised from this method will cause the task to stop.
+     * <p>
+     * Note that the topic partitions here correspond to the original Kafka topic partitions, before any
+     * {@link Transformation transformations} have been applied.
+     *
      * @param partitions The list of partitions that are now assigned to the task (may include
      *                   partitions previously assigned to the task)
      */
@@ -151,8 +166,12 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
      * The SinkTask uses this method to close writers for partitions that are no
      * longer assigned to the SinkTask. This method will be called before a rebalance operation starts
      * and after the SinkTask stops fetching data. After being closed, Connect will not write
-     * any records to the task until a new set of partitions has been opened. Note that any errors raised
+     * any records to the task until a new set of partitions has been opened. Any errors raised
      * from this method will cause the task to stop.
+     * <p>
+     * Note that the topic partitions here correspond to the original Kafka topic partitions, before any
+     * {@link Transformation transformations} have been applied.
+     *
      * @param partitions The list of partitions that should be closed
      */
     public void close(Collection<TopicPartition> partitions) {

connect/api/src/test/java/org/apache/kafka/connect/sink/SinkRecordTest.java

@@ -46,7 +46,7 @@ public class SinkRecordTest {
     @BeforeEach
     public void beforeEach() {
         record = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET,
-                                KAFKA_TIMESTAMP, TS_TYPE, null);
+                                KAFKA_TIMESTAMP, TS_TYPE, null, TOPIC_NAME, PARTITION_NUMBER, KAFKA_OFFSET);
     }
 
     @Test
@@ -125,4 +125,36 @@ public void shouldModifyRecordHeader() {
         Header header = record.headers().lastWithName("intHeader");
         assertEquals(100, (int) Values.convertToInteger(header.schema(), header.value()));
     }
-}
+
+    @Test
+    public void shouldRetainOriginalTopicPartition() {
+        SinkRecord transformed = record.newRecord("transformed-topic", PARTITION_NUMBER + 1, Schema.STRING_SCHEMA, "key",
+                Schema.BOOLEAN_SCHEMA, false, KAFKA_TIMESTAMP);
+
+        assertEquals(TOPIC_NAME, transformed.originalTopic());
+        assertEquals(PARTITION_NUMBER, transformed.originalKafkaPartition());
+
+        SinkRecord transformed2 = transformed.newRecord("transformed-topic-2", PARTITION_NUMBER + 2, Schema.STRING_SCHEMA, "key",
+                Schema.BOOLEAN_SCHEMA, false, KAFKA_TIMESTAMP);
+
+        assertEquals(TOPIC_NAME, transformed2.originalTopic());
+        assertEquals(PARTITION_NUMBER, transformed2.originalKafkaPartition());
+    }
+
+    @Test
+    public void shouldRetainOriginalTopicPartitionWithOlderConstructor() {
+        SinkRecord record = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA,
+                false, KAFKA_OFFSET, KAFKA_TIMESTAMP, TS_TYPE, null);
+        SinkRecord transformed = record.newRecord("transformed-topic", PARTITION_NUMBER + 1, Schema.STRING_SCHEMA, "key",
+                Schema.BOOLEAN_SCHEMA, false, KAFKA_TIMESTAMP);
+
+        assertEquals(TOPIC_NAME, transformed.originalTopic());
+        assertEquals(PARTITION_NUMBER, transformed.originalKafkaPartition());
+
+        SinkRecord transformed2 = transformed.newRecord("transformed-topic-2", PARTITION_NUMBER + 2, Schema.STRING_SCHEMA, "key",
+                Schema.BOOLEAN_SCHEMA, false, KAFKA_TIMESTAMP);
+
+        assertEquals(TOPIC_NAME, transformed2.originalTopic());
+        assertEquals(PARTITION_NUMBER, transformed2.originalKafkaPartition());
+    }
+}