KAFKA-13431: Sink Connectors: Add support for topic-mutating SMTs to async users

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

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.sink;
 
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.data.Schema;
@@ -32,6 +33,7 @@
 public class SinkRecord extends ConnectRecord<SinkRecord> {
     private final long kafkaOffset;
     private final TimestampType timestampType;
+    private final TopicPartition originalTopicPartition;
 
     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 +46,15 @@ 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, null);
+    }
+
+    public SinkRecord(String topic, int partition, Schema keySchema, Object key, Schema valueSchema, Object value, long kafkaOffset,
+                      Long timestamp, TimestampType timestampType, Iterable<Header> headers, TopicPartition originalTopicPartition) {
         super(topic, partition, keySchema, key, valueSchema, value, timestamp, headers);
         this.kafkaOffset = kafkaOffset;
         this.timestampType = timestampType;
+        this.originalTopicPartition = originalTopicPartition;
     }
 
     public long kafkaOffset() {
@@ -57,6 +65,14 @@ public TimestampType timestampType() {
         return timestampType;
     }
 
+    /**
+     * @return topic and partition corresponding to the kafka record before transformations were applied. This is
+     * necessary for internal offset tracking, to be compatible with SMTs that mutate the topic name.
+     */
+    public TopicPartition originalTopicPartition() {
+        return originalTopicPartition;
+    }
+
     @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 +81,17 @@ 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,
+                originalTopicPartition);
     }
 
     @Override
@@ -98,6 +124,8 @@ public String toString() {
         return "SinkRecord{" +
                 "kafkaOffset=" + kafkaOffset +
                 ", timestampType=" + timestampType +
+                ", originalTopicPartition=" + originalTopicPartition +
                 "} " + super.toString();
     }
+
 }

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

@@ -29,7 +29,7 @@
  * from those partitions. As records are fetched from Kafka, they will be passed to the sink task using the
  * {@link #put(Collection)} API, which should either write them to the downstream system or batch them for
  * later writing. Periodically, Connect will call {@link #flush(Map)} to ensure that batched records are
- * actually pushed to the downstream system..
+ * actually pushed to the downstream system.
  *
  * Below we describe the lifecycle of a SinkTask.
  *
@@ -96,6 +96,9 @@ public void initialize(SinkTaskContext context) {
      * be stopped immediately. {@link SinkTaskContext#timeout(long)} can be used to set the maximum time before the
      * batch will be retried.
      *
+     * In case the connector does his own offset tracking (e.g. used by preCommit), the {@link SinkRecord#originalTopicPartition()}
+     * will have to be used, otherwise, the behavior will not be compatible with SMTs that mutate the topic name.
+     *
      * @param records the set of records to send
      */
     public abstract void put(Collection<SinkRecord> records);
@@ -104,8 +107,8 @@ 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}.
+     *                       provided for convenience but could also be determined by tracking all offsets included in the
+     *                       {@link SinkRecord#originalTopicPartition()}s passed to {@link #put}.
      */
     public void flush(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
     }
@@ -116,8 +119,8 @@ 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}.
+     *                       provided for convenience but could also be determined by tracking all offsets included in the
+     *                       {@link SinkRecord#originalTopicPartition()}s passed to {@link #put}.
      *
      * @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.
      */
@@ -171,4 +174,5 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
      */
     @Override
     public abstract void stop();
+
 }

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

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.sink;
 
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.Values;
@@ -46,7 +47,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, new TopicPartition(TOPIC_NAME, PARTITION_NUMBER));
     }
 
     @Test
@@ -125,4 +126,37 @@ public void shouldModifyRecordHeader() {
         Header header = record.headers().lastWithName("intHeader");
         assertEquals(100, (int) Values.convertToInteger(header.schema(), header.value()));
     }
+
+    @Test
+    public void shouldCreateSinkRecordWithOriginalTopicPartition() {
+        TopicPartition tp = new TopicPartition("originalTopic", 100);
+        record = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET,
+                KAFKA_TIMESTAMP, TS_TYPE, null, tp);
+        assertSame(tp, record.originalTopicPartition());
+    }
+
+    @Test
+    public void shouldKeepOriginalTopicPartition() {
+        TopicPartition tp = new TopicPartition("originalTopic", 100);
+        record = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET,
+                KAFKA_TIMESTAMP, TS_TYPE, null, tp);
+        record = record.newRecord("otherTopic", 200, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET);
+
+        assertSame(tp, record.originalTopicPartition());
+    }
+
+    @Test
+    public void shouldNotConsiderOriginalTopicPartitionForEquality() {
+        // For backwards compatibility
+        TopicPartition tp1 = new TopicPartition("originalTopic1", 1);
+        TopicPartition tp2 = new TopicPartition("originalTopic2", 1);
+        SinkRecord record1 = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET,
+                KAFKA_TIMESTAMP, TS_TYPE, null, tp1);
+
+        SinkRecord record2 = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET,
+                KAFKA_TIMESTAMP, TS_TYPE, null, tp2);
+
+        assertTrue(record1.equals(record2));
+    }
+
 }

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/InternalSinkRecord.java

@@ -18,6 +18,7 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.header.Header;
@@ -32,26 +33,26 @@ public class InternalSinkRecord extends SinkRecord {
 
     private final ConsumerRecord<byte[], byte[]> originalRecord;
 
-    public InternalSinkRecord(ConsumerRecord<byte[], byte[]> originalRecord, SinkRecord record) {
+    public InternalSinkRecord(ConsumerRecord<byte[], byte[]> originalRecord, TopicPartition originalTopicPartition, SinkRecord record) {
         super(record.topic(), record.kafkaPartition(), record.keySchema(), record.key(),
             record.valueSchema(), record.value(), record.kafkaOffset(), record.timestamp(),
-            record.timestampType(), record.headers());
+            record.timestampType(), record.headers(), originalTopicPartition);
         this.originalRecord = originalRecord;
     }
 
-    protected InternalSinkRecord(ConsumerRecord<byte[], byte[]> originalRecord, String topic,
+    protected InternalSinkRecord(ConsumerRecord<byte[], byte[]> originalRecord, TopicPartition originalTopicPartition, String topic,
                                  int partition, Schema keySchema, Object key, Schema valueSchema,
                                  Object value, long kafkaOffset, Long timestamp,
                                  TimestampType timestampType, Iterable<Header> headers) {
-        super(topic, partition, keySchema, key, valueSchema, value, kafkaOffset, timestamp, timestampType, headers);
+        super(topic, partition, keySchema, key, valueSchema, value, kafkaOffset, timestamp, timestampType, headers, originalTopicPartition);
         this.originalRecord = originalRecord;
     }
 
     @Override
     public SinkRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key,
                                 Schema valueSchema, Object value, Long timestamp,
                                 Iterable<Header> headers) {
-        return new InternalSinkRecord(originalRecord, topic, kafkaPartition, keySchema, key,
+        return new InternalSinkRecord(originalRecord, originalTopicPartition(), topic, kafkaPartition, keySchema, key,
             valueSchema, value, kafkaOffset(), timestamp, timestampType(), headers());
     }
 

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java

@@ -503,14 +503,17 @@ private SinkRecord convertAndTransformRecord(final ConsumerRecord<byte[], byte[]
             return null;
         }
 
+        TopicPartition originalTopicPartition = new TopicPartition(msg.topic(), msg.partition());
         Long timestamp = ConnectUtils.checkAndConvertTimestamp(msg.timestamp());
         SinkRecord origRecord = new SinkRecord(msg.topic(), msg.partition(),
                 keyAndSchema.schema(), keyAndSchema.value(),
                 valueAndSchema.schema(), valueAndSchema.value(),
                 msg.offset(),
                 timestamp,
                 msg.timestampType(),
-                headers);
+                headers,
+                originalTopicPartition);
+
         log.trace("{} Applying transformations to record in topic '{}' partition {} at offset {} and timestamp {} with key {} and value {}",
                 this, msg.topic(), msg.partition(), msg.offset(), timestamp, keyAndSchema.value(), valueAndSchema.value());
         if (isTopicTrackingEnabled) {
@@ -523,7 +526,7 @@ private SinkRecord convertAndTransformRecord(final ConsumerRecord<byte[], byte[]
             return null;
         }
         // Error reporting will need to correlate each sink record with the original consumer record
-        return new InternalSinkRecord(msg, transformedRecord);
+        return new InternalSinkRecord(msg, originalTopicPartition, transformedRecord);
     }
 
     private Headers convertHeadersFor(ConsumerRecord<byte[], byte[]> record) {

connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSinkTask.java

@@ -72,7 +72,10 @@ public void put(Collection<SinkRecord> records) {
             data.put("topic", record.topic());
             data.put("time_ms", nowMs);
             data.put("seqno", record.value());
+            data.put("partition", record.kafkaPartition());
             data.put("offset", record.kafkaOffset());
+            data.put("originalTopic", record.originalTopicPartition().topic());
+            data.put("originalPartition", record.originalTopicPartition().partition());
             String dataJson;
             try {
                 dataJson = JSON_SERDE.writeValueAsString(data);

connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.java

@@ -17,13 +17,10 @@
 
 package org.apache.kafka.connect.integration;
 
-import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.connect.connector.Task;
 import org.apache.kafka.connect.sink.ErrantRecordReporter;
 import org.apache.kafka.connect.sink.SinkRecord;
 
-import java.util.Collection;
-import java.util.HashMap;
 import java.util.Map;
 
 public class ErrantRecordSinkConnector extends MonitorableSinkConnector {
@@ -47,15 +44,10 @@ public void start(Map<String, String> props) {
         }
 
         @Override
-        public void put(Collection<SinkRecord> records) {
-            for (SinkRecord rec : records) {
-                taskHandle.record();
-                TopicPartition tp = cachedTopicPartitions
-                    .computeIfAbsent(rec.topic(), v -> new HashMap<>())
-                    .computeIfAbsent(rec.kafkaPartition(), v -> new TopicPartition(rec.topic(), rec.kafkaPartition()));
-                committedOffsets.put(tp, committedOffsets.getOrDefault(tp, 0L) + 1);
-                reporter.report(rec, new Throwable());
-            }
+        protected void processRecord(SinkRecord rec) {
+            super.processRecord(rec);
+            reporter.report(rec, new Throwable());
         }
+
     }
 }

connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java

@@ -88,17 +88,14 @@ public ConfigDef config() {
 
     public static class MonitorableSinkTask extends SinkTask {
 
-        private String connectorName;
+        private final Set<TopicPartition> assignments;
+        private final Map<TopicPartition, Long> committedOffsets;
         private String taskId;
-        TaskHandle taskHandle;
-        Set<TopicPartition> assignments;
-        Map<TopicPartition, Long> committedOffsets;
-        Map<String, Map<Integer, TopicPartition>> cachedTopicPartitions;
+        private TaskHandle taskHandle;
 
         public MonitorableSinkTask() {
             this.assignments = new HashSet<>();
             this.committedOffsets = new HashMap<>();
-            this.cachedTopicPartitions = new HashMap<>();
         }
 
         @Override
@@ -109,7 +106,7 @@ public String version() {
         @Override
         public void start(Map<String, String> props) {
             taskId = props.get("task.id");
-            connectorName = props.get("connector.name");
+            String connectorName = props.get("connector.name");
             taskHandle = RuntimeHandles.get().connectorHandle(connectorName).taskHandle(taskId);
             log.debug("Starting task {}", taskId);
             taskHandle.recordTaskStart();
@@ -125,28 +122,29 @@ public void open(Collection<TopicPartition> partitions) {
         @Override
         public void put(Collection<SinkRecord> records) {
             for (SinkRecord rec : records) {
-                taskHandle.record(rec);
-                TopicPartition tp = cachedTopicPartitions
-                        .computeIfAbsent(rec.topic(), v -> new HashMap<>())
-                        .computeIfAbsent(rec.kafkaPartition(), v -> new TopicPartition(rec.topic(), rec.kafkaPartition()));
-                committedOffsets.put(tp, committedOffsets.getOrDefault(tp, 0L) + 1);
-                log.trace("Task {} obtained record (key='{}' value='{}')", taskId, rec.key(), rec.value());
+                processRecord(rec);
             }
         }
 
+        protected void processRecord(SinkRecord rec) {
+            taskHandle.record(rec);
+            committedOffsets.put(rec.originalTopicPartition(), rec.kafkaOffset() + 1);
+            log.trace("Task {} obtained record (key='{}' value='{}')", taskId, rec.key(), rec.value());
+        }
+
         @Override
-        public Map<TopicPartition, OffsetAndMetadata> preCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {
+        public Map<TopicPartition, OffsetAndMetadata> preCommit(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
+            Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
             for (TopicPartition tp : assignments) {
                 Long recordsSinceLastCommit = committedOffsets.get(tp);
                 if (recordsSinceLastCommit == null) {
                     log.warn("preCommit was called with topic-partition {} that is not included "
                             + "in the assignments of this task {}", tp, assignments);
                 } else {
-                    taskHandle.commit(recordsSinceLastCommit.intValue());
                     log.error("Forwarding to framework request to commit additional {} for {}",
                             recordsSinceLastCommit, tp);
-                    taskHandle.commit((int) (long) recordsSinceLastCommit);
-                    committedOffsets.put(tp, 0L);
+                    taskHandle.commit(recordsSinceLastCommit.intValue());
+                    offsets.put(tp, new OffsetAndMetadata(committedOffsets.remove(tp)));
                 }
             }
             return offsets;