Merge pull request #25 from aiven/timestamp_extractor_key

ivanyu · web-flow · commit dfd70445add5 · 2020-11-14T08:41:24.000+02:00
Add support for extracting timestamp from message key
diff --git a/README.md b/README.md
@@ -17,8 +17,9 @@ The transformation:
 - expects it to have a specified field;
 - expects the value of the field to be either `INT64` or `org.apache.kafka.connect.data.Timestamp` and not be `null`.
 
-Exists in one variant:
-- `io.aiven.kafka.connect.transforms.ExtractTimestamp$Value` - works on values.
+Exists in two variants:
+ - `io.aiven.kafka.connect.transforms.ExtractTimestamp$Key` - works on keys;
+ - `io.aiven.kafka.connect.transforms.ExtractTimestamp$Value` - works on values.
 
 The transformation defines the following configurations:
 
@@ -28,6 +29,7 @@ Here's an example of this transformation configuration:
 
 ```properties
 transforms=ExtractTimestampFromValueField
+transforms.ExtractTimestampFromValueField.type=io.aiven.kafka.connect.transforms.ExtractTimestamp$Value
 transforms.ExtractTimestampFromValueField.field.name=inner_field_name
 ```
 
diff --git a/src/main/java/io/aiven/kafka/connect/transforms/ExtractTimestamp.java b/src/main/java/io/aiven/kafka/connect/transforms/ExtractTimestamp.java
@@ -21,15 +21,12 @@
 
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.data.SchemaAndValue;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.transforms.Transformation;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 public abstract class ExtractTimestamp<R extends ConnectRecord<R>> implements Transformation<R> {
-    private static final Logger log = LoggerFactory.getLogger(ExtractTimestamp.class);
 
     private ExtractTimestampConfig config;
 
@@ -45,23 +42,25 @@ public void configure(final Map<String, ?> configs) {
 
     @Override
     public R apply(final R record) {
-        if (record.value() == null) {
-            throw new DataException("Value can't be null: " + record);
+        final SchemaAndValue schemaAndValue = getSchemaAndValue(record);
+        
+        if (schemaAndValue.value() == null) {
+            throw new DataException(keyOrValue() + " can't be null: " + record);
         }
 
         final Object fieldValue;
-        if (record.value() instanceof Struct) {
-            final Struct struct = (Struct) record.value();
+        if (schemaAndValue.value() instanceof Struct) {
+            final Struct struct = (Struct) schemaAndValue.value();
             if (struct.schema().field(config.fieldName()) == null) {
                 throw new DataException(config.fieldName() + " field must be present and its value can't be null: "
                     + record);
             }
             fieldValue = struct.get(config.fieldName());
-        } else if (record.value() instanceof Map) {
-            final Map map = (Map) record.value();
+        } else if (schemaAndValue.value() instanceof Map) {
+            final Map<?, ?> map = (Map<?, ?>) schemaAndValue.value();
             fieldValue = map.get(config.fieldName());
         } else {
-            throw new DataException("Value type must be STRUCT or MAP: " + record);
+            throw new DataException(keyOrValue() + " type must be STRUCT or MAP: " + record);
         }
 
         if (fieldValue == null) {
@@ -91,13 +90,36 @@ public R apply(final R record) {
         );
     }
 
-    public static final class Value<R extends ConnectRecord<R>> extends ExtractTimestamp<R> {
-        // There's an implementation only for value, not for key.
-        // We provide $Value class anyway for the consistency sake
-        // and in case we need a $Key version in the future as well.
-    }
-
     @Override
     public void close() {
     }
+
+    protected abstract String keyOrValue();
+
+    protected abstract SchemaAndValue getSchemaAndValue(final R record);
+
+    public static final class Key<R extends ConnectRecord<R>> extends ExtractTimestamp<R> {
+        @Override
+        protected SchemaAndValue getSchemaAndValue(final R record) {
+            return new SchemaAndValue(record.keySchema(), record.key());
+        }
+
+        @Override
+        protected String keyOrValue() {
+            return "key";
+        }
+    }
+
+    public static final class Value<R extends ConnectRecord<R>> extends ExtractTimestamp<R> {
+        @Override
+        protected SchemaAndValue getSchemaAndValue(final R record) {
+            return new SchemaAndValue(record.valueSchema(), record.value());
+        }
+
+        @Override
+        protected String keyOrValue() {
+            return "value";
+        }
+    }
+    
 }
diff --git a/src/main/java/io/aiven/kafka/connect/transforms/ExtractTimestampConfig.java b/src/main/java/io/aiven/kafka/connect/transforms/ExtractTimestampConfig.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Aiven Oy
+ * Copyright 2020 Aiven Oy
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/src/test/java/io/aiven/kafka/connect/transforms/ExtractTimestampKeyTest.java b/src/test/java/io/aiven/kafka/connect/transforms/ExtractTimestampKeyTest.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright 2020 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.transforms;
+
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+public class ExtractTimestampKeyTest extends ExtractTimestampTest {
+
+    @Override
+    protected String keyOrValue() {
+        return "key";
+    }
+
+    @Override
+    protected ExtractTimestamp<SinkRecord> createTransformationObject() {
+        return new ExtractTimestamp.Key<>();
+    }
+
+    @Override
+    protected SinkRecord record(final Schema schema, final Object data) {
+        return record(schema, data, null, null);
+    }
+}
diff --git a/src/test/java/io/aiven/kafka/connect/transforms/ExtractTimestampTest.java b/src/test/java/io/aiven/kafka/connect/transforms/ExtractTimestampTest.java
@@ -0,0 +1,224 @@
+/*
+ * Copyright 2020 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.transforms;
+
+import java.util.Date;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaBuilder;
+import org.apache.kafka.connect.data.Struct;
+import org.apache.kafka.connect.data.Timestamp;
+import org.apache.kafka.connect.errors.DataException;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+abstract class ExtractTimestampTest {
+    private static final String FIELD = "test_field";
+
+    @Test
+    void recordNotStructOrMap() {
+        final SinkRecord originalRecord = record(SchemaBuilder.INT8_SCHEMA, (byte) 123);
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(keyOrValue() + " type must be STRUCT or MAP: " + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void recordStructNull() {
+        final Schema schema = SchemaBuilder.struct().schema();
+        final SinkRecord originalRecord = record(schema, null);
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(keyOrValue() + " can't be null: " + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void recordMapNull() {
+        final SinkRecord originalRecord = record(null, null);
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(keyOrValue() + " can't be null: " + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void structWithMissingField() {
+        final Schema schema = SchemaBuilder.struct()
+            .field(FIELD, Schema.INT64_SCHEMA)
+            .build();
+        final SinkRecord originalRecord = record(null, new Struct(schema));
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(FIELD + " field must be present and its value can't be null: " + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void mapWithMissingField() {
+        final SinkRecord originalRecord = record(null, new HashMap<>());
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(FIELD + " field must be present and its value can't be null: " + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void structWithNullField() {
+        final Schema schema = SchemaBuilder.struct()
+            .field(FIELD, Schema.OPTIONAL_INT64_SCHEMA)
+            .build();
+        final SinkRecord originalRecord = record(null, new Struct(schema).put(FIELD, null));
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(FIELD + " field must be present and its value can't be null: " + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void mapWithNullField() {
+        final HashMap<Object, Object> valueMap = new HashMap<>();
+        valueMap.put(FIELD, null);
+        final SinkRecord originalRecord = record(null, valueMap);
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(FIELD + " field must be present and its value can't be null: " + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void structWithFieldOfIncorrectType() {
+        final Schema schema = SchemaBuilder.struct()
+            .field(FIELD, Schema.STRING_SCHEMA)
+            .build();
+        final SinkRecord originalRecord = record(null, new Struct(schema).put(FIELD, "aaa"));
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(FIELD + " field must be INT64 or org.apache.kafka.connect.data.Timestamp: "
+                + originalRecord,
+            e.getMessage());
+    }
+
+    @Test
+    void mapWithFieldOfIncorrectType() {
+        final HashMap<Object, Object> valueMap = new HashMap<>();
+        valueMap.put(FIELD, "aaa");
+        final SinkRecord originalRecord = record(null, valueMap);
+        final Throwable e = assertThrows(DataException.class,
+            () -> transformation().apply(originalRecord));
+        assertEquals(FIELD + " field must be INT64 or org.apache.kafka.connect.data.Timestamp: "
+                + originalRecord,
+            e.getMessage());
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = { true, false })
+    void structWithIntField(final boolean optional) {
+        final Schema schema;
+        if (optional) {
+            schema = SchemaBuilder.struct()
+                .field(FIELD, Schema.OPTIONAL_INT64_SCHEMA)
+                .build();
+        } else {
+            schema = SchemaBuilder.struct()
+                .field(FIELD, Schema.INT64_SCHEMA)
+                .build();
+        }
+        final long timestamp = 11363151277L;
+        final SinkRecord originalRecord = record(null, new Struct(schema).put(FIELD, timestamp));
+        final SinkRecord transformedRecord = transformation().apply(originalRecord);
+        assertEquals(setNewTimestamp(originalRecord, timestamp), transformedRecord);
+    }
+
+    @Test
+    void mapWithIntField() {
+        final long timestamp = 11363151277L;
+        final HashMap<Object, Object> valueMap = new HashMap<>();
+        valueMap.put(FIELD, timestamp);
+        final SinkRecord originalRecord = record(null, valueMap);
+        final SinkRecord transformedRecord = transformation().apply(originalRecord);
+        assertEquals(setNewTimestamp(originalRecord, timestamp), transformedRecord);
+    }
+
+    @Test
+    void structWithTimestampField() {
+        final Schema schema = SchemaBuilder.struct()
+            .field(FIELD, Timestamp.SCHEMA)
+            .build();
+        final long timestamp = 11363151277L;
+        final SinkRecord originalRecord = record(null, new Struct(schema).put(FIELD, new Date(timestamp)));
+        final SinkRecord transformedRecord = transformation().apply(originalRecord);
+        assertEquals(setNewTimestamp(originalRecord, timestamp), transformedRecord);
+    }
+
+    @Test
+    void mapWithTimestampField() {
+        final long timestamp = 11363151277L;
+        final HashMap<Object, Object> valueMap = new HashMap<>();
+        valueMap.put(FIELD, new Date(timestamp));
+        final SinkRecord originalRecord = record(null, valueMap);
+        final SinkRecord transformedRecord = transformation().apply(originalRecord);
+        assertEquals(setNewTimestamp(originalRecord, timestamp), transformedRecord);
+    }
+
+    private ExtractTimestamp<SinkRecord> transformation() {
+        final Map<String, String> props = new HashMap<>();
+        props.put("field.name", FIELD);
+        final ExtractTimestamp<SinkRecord> transform = createTransformationObject();
+        transform.configure(props);
+        return transform;
+    }
+
+    protected abstract SinkRecord record(final Schema schema, final Object data);
+
+    protected SinkRecord record(final Schema keySchema,
+                                final Object key,
+                                final Schema valueSchema,
+                                final Object value) {
+        return new SinkRecord("original_topic", 0,
+            keySchema, key,
+            valueSchema, value,
+            123L,
+            456L, TimestampType.CREATE_TIME);
+    }
+
+    private SinkRecord setNewTimestamp(final SinkRecord record, final long newTimestamp) {
+        return record.newRecord(record.topic(),
+            record.kafkaPartition(),
+            record.keySchema(),
+            record.key(),
+            record.valueSchema(),
+            record.value(),
+            newTimestamp,
+            record.headers()
+        );
+    }
+
+    protected abstract String keyOrValue();
+
+    protected abstract ExtractTimestamp<SinkRecord> createTransformationObject();
+}
diff --git a/src/test/java/io/aiven/kafka/connect/transforms/ExtractTimestampValueTest.java b/src/test/java/io/aiven/kafka/connect/transforms/ExtractTimestampValueTest.java

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`/*`
`2`		`- * Copyright 2019 Aiven Oy`
	`2`	`+ * Copyright 2020 Aiven Oy`
`3`	`3`	`*`
`4`	`4`	`* Licensed under the Apache License, Version 2.0 (the "License");`
`5`	`5`	`* you may not use this file except in compliance with the License.`