feat: Add topic-partition-key record grouper

build.gradle

@@ -181,6 +181,7 @@ dependencies {
     testImplementation "io.confluent:kafka-connect-avro-converter:$confluentPlatformVersion"
 
     testRuntimeOnly "org.junit.jupiter:junit-jupiter-engine:$junitVersion"
+    testRuntimeOnly "ch.qos.logback:logback-classic:1.4.11"
 }
 
 test {

src/main/java/io/aiven/kafka/connect/common/grouper/KeyAndTopicPartitionRecordGrouper.java

@@ -75,7 +75,7 @@ private String generateRecordKey(final SinkRecord record) {
         final Supplier<String> setKey = () -> {
             if (record.key() == null) {
                 return "null";
-            } else if (record.keySchema().type() == Schema.Type.STRING) {
+            } else if (record.keySchema() != null && record.keySchema().type() == Schema.Type.STRING) {
                 return (String) record.key();
             } else {
                 return record.key().toString();

src/main/java/io/aiven/kafka/connect/common/grouper/KeyRecordGrouper.java

@@ -72,7 +72,7 @@ private String generateRecordKey(final SinkRecord record) {
         final Supplier<String> setKey = () -> {
             if (record.key() == null) {
                 return "null";
-            } else if (record.keySchema().type() == Schema.Type.STRING) {
+            } else if (record.keySchema() != null && record.keySchema().type() == Schema.Type.STRING) {
                 return (String) record.key();
             } else {
                 return record.key().toString();

src/main/java/io/aiven/kafka/connect/common/grouper/RecordGrouperFactory.java

@@ -17,10 +17,12 @@
 package io.aiven.kafka.connect.common.grouper;
 
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import io.aiven.kafka.connect.common.config.AivenCommonConfig;
@@ -36,23 +38,31 @@ public final class RecordGrouperFactory {
     public static final String KEY_RECORD = KeyRecordGrouper.class.getName();
 
     public static final String TOPIC_PARTITION_RECORD = TopicPartitionRecordGrouper.class.getName();
+    public static final String TOPIC_PARTITION_KEY_RECORD = TopicPartitionKeyRecordGrouper.class.getName();
 
     public static final String KEY_TOPIC_PARTITION_RECORD = KeyAndTopicPartitionRecordGrouper.class.getName();
 
     public static final Map<String, List<Pair<String, Boolean>>> SUPPORTED_VARIABLES = new LinkedHashMap<>() {{
-            put(TOPIC_PARTITION_RECORD, List.of(
-                Pair.of(FilenameTemplateVariable.TOPIC.name, true),
-                Pair.of(FilenameTemplateVariable.PARTITION.name, true),
-                Pair.of(FilenameTemplateVariable.START_OFFSET.name, true),
-                Pair.of(FilenameTemplateVariable.TIMESTAMP.name, false)
-            ));
-            put(KEY_RECORD, List.of(Pair.of(FilenameTemplateVariable.KEY.name, true)));
-            put(KEY_TOPIC_PARTITION_RECORD, List.of(
-                Pair.of(FilenameTemplateVariable.KEY.name, true),    
-                Pair.of(FilenameTemplateVariable.TOPIC.name, false),
-                Pair.of(FilenameTemplateVariable.PARTITION.name, false)    
-            ));
-        }};
+                put(TOPIC_PARTITION_RECORD, List.of(
+                    Pair.of(FilenameTemplateVariable.TOPIC.name, true),
+                    Pair.of(FilenameTemplateVariable.PARTITION.name, true),
+                    Pair.of(FilenameTemplateVariable.START_OFFSET.name, true),
+                    Pair.of(FilenameTemplateVariable.TIMESTAMP.name, false)
+                ));
+                put(TOPIC_PARTITION_KEY_RECORD, List.of(
+                    Pair.of(FilenameTemplateVariable.TOPIC.name, true),
+                    Pair.of(FilenameTemplateVariable.PARTITION.name, true),
+                    Pair.of(FilenameTemplateVariable.KEY.name, true),
+                    Pair.of(FilenameTemplateVariable.START_OFFSET.name, true),
+                    Pair.of(FilenameTemplateVariable.TIMESTAMP.name, false)
+                ));
+                put(KEY_RECORD, List.of(Pair.of(FilenameTemplateVariable.KEY.name, true)));
+                put(KEY_TOPIC_PARTITION_RECORD, List.of(
+                    Pair.of(FilenameTemplateVariable.KEY.name, true),
+                    Pair.of(FilenameTemplateVariable.TOPIC.name, false),
+                    Pair.of(FilenameTemplateVariable.PARTITION.name, false)
+                ));
+            }};
 
     public static final List<String> ALL_SUPPORTED_VARIABLES =
         SUPPORTED_VARIABLES.values()
@@ -73,6 +83,13 @@ public final class RecordGrouperFactory {
             .map(Pair::getLeft)
             .collect(Collectors.toSet());
 
+
+    private static final Set<String> TOPIC_PARTITION_KEY_RECORD_REQUIRED_VARS =
+        SUPPORTED_VARIABLES.get(TOPIC_PARTITION_KEY_RECORD).stream()
+            .filter(Pair::getRight)
+            .map(Pair::getLeft)
+            .collect(Collectors.toSet());
+
     private static final Set<String> KEY_TOPIC_PARTITION_RECORD_REQUIRED_VARS =
         SUPPORTED_VARIABLES.get(KEY_TOPIC_PARTITION_RECORD).stream()
             .filter(Pair::getRight)
@@ -85,6 +102,12 @@ public final class RecordGrouperFactory {
             .map(Pair::getLeft)
             .collect(Collectors.toSet());
 
+    private static final Set<String> TOPIC_PARTITION_KEY_RECORD_OPT_VARS =
+        SUPPORTED_VARIABLES.get(TOPIC_PARTITION_KEY_RECORD).stream()
+            .filter(p -> !p.getRight())
+            .map(Pair::getLeft)
+            .collect(Collectors.toSet());
+
     private static final Set<String> KEY_TOPIC_PARTITION_RECORD_OPT_VARS =
         SUPPORTED_VARIABLES.get(KEY_TOPIC_PARTITION_RECORD).stream()
             .filter(p -> !p.getRight())
@@ -102,13 +125,15 @@ private RecordGrouperFactory() {
     }
 
     public static String resolveRecordGrouperType(final Template template) {
-        final Set<String> variables = template.variablesSet();
-        if (isByKeyRecord(variables)) {
+        final Supplier<Set<String>> variables = () -> new HashSet<>(template.variablesSet());
+        if (isByTopicPartitionKeyRecord(variables.get())) {
+            return TOPIC_PARTITION_KEY_RECORD;
+        } else if (isByTopicPartitionRecord(variables.get())) {
+            return TOPIC_PARTITION_RECORD;
+        } else if (isByKeyRecord(variables.get())) {
             return KEY_RECORD;
-        } else if (isByTopicPartitionKeyRecord(variables)) {
+        } else if (isByKeyTopicPartitionRecord(variables.get())) {
             return KEY_TOPIC_PARTITION_RECORD;
-        } else if (isByTopicPartitionRecord(variables)) {
-            return TOPIC_PARTITION_RECORD;
         } else {
             throw new IllegalArgumentException(
                 String.format(
@@ -131,11 +156,19 @@ public static RecordGrouper newRecordGrouper(final AivenCommonConfig config) {
                 config.getMaxRecordsPerFile() != 0
                     ? config.getMaxRecordsPerFile()
                     : null;
-            return config.getFormatType() == FormatType.PARQUET || config.getFormatType() == FormatType.AVRO
+            if (TOPIC_PARTITION_KEY_RECORD.equals(grType)) {
+                return config.getFormatType() == FormatType.PARQUET || config.getFormatType() == FormatType.AVRO
+                    ? new SchemaBasedTopicPartitionKeyRecordGrouper(
+                    fileNameTemplate, maxRecordsPerFile, config.getFilenameTimestampSource())
+                    : new TopicPartitionKeyRecordGrouper(
+                    fileNameTemplate, maxRecordsPerFile, config.getFilenameTimestampSource());
+            } else {
+                return config.getFormatType() == FormatType.PARQUET || config.getFormatType() == FormatType.AVRO
                     ? new SchemaBasedTopicPartitionRecordGrouper(
-                            fileNameTemplate, maxRecordsPerFile, config.getFilenameTimestampSource())
+                    fileNameTemplate, maxRecordsPerFile, config.getFilenameTimestampSource())
                     : new TopicPartitionRecordGrouper(
-                            fileNameTemplate, maxRecordsPerFile, config.getFilenameTimestampSource());
+                    fileNameTemplate, maxRecordsPerFile, config.getFilenameTimestampSource());
+            }
         }
     }
 
@@ -155,6 +188,17 @@ private static boolean isByTopicPartitionRecord(final Set<String> vars) {
     }
 
     private static boolean isByTopicPartitionKeyRecord(final Set<String> vars) {
+        final Set<String> requiredVars =
+            Sets.intersection(TOPIC_PARTITION_KEY_RECORD_REQUIRED_VARS, vars)
+                .immutableCopy();
+        vars.removeAll(requiredVars);
+        final boolean containsRequiredVars = TOPIC_PARTITION_KEY_RECORD_REQUIRED_VARS.equals(requiredVars);
+        final boolean containsOptionalVars =
+            vars.isEmpty() || !Collections.disjoint(TOPIC_PARTITION_KEY_RECORD_OPT_VARS, vars);
+        return containsRequiredVars && containsOptionalVars;
+    }
+
+    private static boolean isByKeyTopicPartitionRecord(final Set<String> vars) {
         final Set<String> requiredVars =
             Sets.intersection(KEY_TOPIC_PARTITION_RECORD_REQUIRED_VARS, vars)
                 .immutableCopy();

src/main/java/io/aiven/kafka/connect/common/grouper/SchemaBasedTopicPartitionKeyRecordGrouper.java

@@ -0,0 +1,130 @@
+/*
+ * Copyright 2023 Aiven Oy
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.aiven.kafka.connect.common.grouper;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.errors.SchemaProjectorException;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+import io.aiven.kafka.connect.common.config.TimestampSource;
+import io.aiven.kafka.connect.common.templating.Template;
+
+final class SchemaBasedTopicPartitionKeyRecordGrouper extends TopicPartitionKeyRecordGrouper {
+
+    private final SchemaBasedRotator schemaBasedRotator = new SchemaBasedRotator();
+
+    SchemaBasedTopicPartitionKeyRecordGrouper(final Template filenameTemplate,
+                                              final Integer maxRecordsPerFile,
+                                              final TimestampSource tsSource) {
+        super(filenameTemplate, maxRecordsPerFile, tsSource);
+    }
+
+    @Override
+    protected String resolveRecordKeyFor(final SinkRecord record) {
+        if (schemaBasedRotator.rotate(record)) {
+            return generateNewRecordKey(record);
+        } else {
+            return super.resolveRecordKeyFor(record);
+        }
+    }
+
+    @Override
+    public void clear() {
+        schemaBasedRotator.clear();
+        super.clear();
+    }
+
+    private static final class SchemaBasedRotator implements Rotator<SinkRecord> {
+
+        private final Map<TopicPartitionKey, KeyValueSchema> keyValueSchemas = new HashMap<>();
+
+        @Override
+        public boolean rotate(final SinkRecord record) {
+            if (Objects.isNull(record.valueSchema()) || Objects.isNull(record.keySchema())) {
+                throw new SchemaProjectorException("Record must have schemas for key and value");
+            }
+            final var key = recordKey(record);
+            final var tpk = new TopicPartitionKey(new TopicPartition(record.topic(), record.kafkaPartition()), key);
+            final var keyValueVersion =
+                keyValueSchemas.computeIfAbsent(tpk, ignored -> new KeyValueSchema(
+                    record.keySchema(),
+                    record.valueSchema()));
+            final var schemaChanged =
+                !keyValueVersion.keySchema.equals(record.keySchema())
+                    || !keyValueVersion.valueSchema.equals(record.valueSchema());
+            if (schemaChanged) {
+                keyValueSchemas.put(tpk,
+                    new KeyValueSchema(record.keySchema(), record.valueSchema())
+                );
+            }
+            return schemaChanged;
+        }
+
+        private String recordKey(final SinkRecord record) {
+            final String key;
+            if (record.key() == null) {
+                key = "null";
+            } else if (record.keySchema().type() == Schema.Type.STRING) {
+                key = (String) record.key();
+            } else {
+                key = record.key().toString();
+            }
+            return key;
+        }
+
+
+        private static class KeyValueSchema {
+
+            final Schema keySchema;
+
+            final Schema valueSchema;
+
+            KeyValueSchema(final Schema keySchema, final Schema valueSchema) {
+                this.keySchema = keySchema;
+                this.valueSchema = valueSchema;
+            }
+
+            @Override
+            public boolean equals(final Object o) {
+                if (this == o) {
+                    return true;
+                }
+                if (o == null || getClass() != o.getClass()) {
+                    return false;
+                }
+                final var that = (KeyValueSchema) o;
+                return keySchema.equals(that.keySchema) && valueSchema.equals(that.valueSchema);
+            }
+
+            @Override
+            public int hashCode() {
+                return Objects.hash(keySchema, valueSchema);
+            }
+        }
+
+        void clear() {
+            keyValueSchemas.clear();
+        }
+
+    }
+
+}