added a timestamp unit to partitioner config

confluentinc · Feb 3, 2020 · 21b0fb4 · 21b0fb4
1 parent 8080189
commit 21b0fb4
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 17 deletions.
diff --git a/partitioner/src/main/java/io/confluent/connect/storage/partitioner/PartitionerConfig.java b/partitioner/src/main/java/io/confluent/connect/storage/partitioner/PartitionerConfig.java
@@ -101,6 +101,13 @@ public class PartitionerConfig extends AbstractConfig implements ComposableConfi
   public static final String TIMESTAMP_FIELD_NAME_DEFAULT = "timestamp";
   public static final String TIMESTAMP_FIELD_NAME_DISPLAY = "Record Field for Timestamp Extractor";
 
+  public static final String TIMESTAMP_UNIT_NAME_CONFIG = "timestamp.unit";
+  public static final String TIMESTAMP_UNIT_NAME_DOC =
+          "The time unit to be used for timestamp by the timestamp extractor.";
+  public static final String TIMESTAMP_UNIT_NAME_DEFAULT = "ms";
+  public static final String TIMESTAMP_UNIT_NAME_DISPLAY =
+          "Time Unit of specified Record Field for Timestamp Extractor";
+
   /**
    * Create a new configuration definition.
    *
@@ -208,6 +215,16 @@ public static ConfigDef newConfigDef(ConfigDef.Recommender partitionerClassRecom
           ++orderInGroup,
           Width.LONG,
           TIMESTAMP_FIELD_NAME_DISPLAY);
+
+      configDef.define(TIMESTAMP_UNIT_NAME_CONFIG,
+          Type.STRING,
+          TIMESTAMP_UNIT_NAME_DEFAULT,
+          Importance.MEDIUM,
+          TIMESTAMP_UNIT_NAME_DOC,
+          group,
+          ++orderInGroup,
+          Width.LONG,
+          TIMESTAMP_UNIT_NAME_DISPLAY);
     }
 
     return configDef;

diff --git a/partitioner/src/main/java/io/confluent/connect/storage/partitioner/TimeBasedPartitioner.java b/partitioner/src/main/java/io/confluent/connect/storage/partitioner/TimeBasedPartitioner.java
@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import io.confluent.connect.storage.common.SchemaGenerator;
 import io.confluent.connect.storage.common.StorageCommonConfig;
@@ -266,15 +267,32 @@ public Long extract(ConnectRecord<?> record) {
   public static class RecordFieldTimestampExtractor implements TimestampExtractor {
     private String fieldName;
     private DateTimeFormatter dateTime;
+    private TimeUnit timeUnit;
 
     @Override
     public void configure(Map<String, Object> config) {
       fieldName = (String) config.get(PartitionerConfig.TIMESTAMP_FIELD_NAME_CONFIG);
       dateTime = ISODateTimeFormat.dateTimeParser();
+      timeUnit = getTimeUnit(config);
+    }
+
+    private TimeUnit getTimeUnit(Map<String, Object> config) {
+      String timestampTimeUnit = (String) config.get(PartitionerConfig.TIMESTAMP_UNIT_NAME_CONFIG);
+
+      if ("s".equalsIgnoreCase(timestampTimeUnit)) {
+        return TimeUnit.SECONDS;
+      }
+
+      return TimeUnit.MILLISECONDS;
     }
 
     @Override
     public Long extract(ConnectRecord<?> record) {
+      Long timestamp = extractTimestamp(record);
+      return timeUnit == TimeUnit.SECONDS ? timestamp * 1000 : timestamp;
+    }
+
+    private Long extractTimestamp(ConnectRecord<?> record) {
       Object value = record.value();
       if (value instanceof Struct) {
         Struct struct = (Struct) value;

diff --git a/...oner/src/test/java/io/confluent/connect/storage/partitioner/TimeBasedPartitionerTest.java b/...oner/src/test/java/io/confluent/connect/storage/partitioner/TimeBasedPartitionerTest.java
@@ -15,6 +15,10 @@
 
 package io.confluent.connect.storage.partitioner;
 
+import io.confluent.connect.storage.StorageSinkTestBase;
+import io.confluent.connect.storage.common.StorageCommonConfig;
+import io.confluent.connect.storage.errors.PartitionException;
+import io.confluent.connect.storage.util.DateTimeUtils;
 import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.utils.Time;
@@ -35,26 +39,12 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
-import java.util.Collections;
-import java.util.Date;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Locale;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.TimeUnit;
 
-import io.confluent.connect.storage.StorageSinkTestBase;
-import io.confluent.connect.storage.common.StorageCommonConfig;
-import io.confluent.connect.storage.errors.PartitionException;
-import io.confluent.connect.storage.util.DateTimeUtils;
-
-import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.CoreMatchers.instanceOf;
-import static org.hamcrest.CoreMatchers.notNullValue;
-import static org.hamcrest.CoreMatchers.startsWith;
+import static org.hamcrest.CoreMatchers.*;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.fail;
 
 public class TimeBasedPartitionerTest extends StorageSinkTestBase {
 
@@ -473,6 +463,17 @@ public void testNumericRecordFieldTimeMap() {
     ts = DATE_TIME.getMillis();
     testMapNumericTimestampPartitionEncoding(
           partitioner, timeField, ts, Schema.INT64_SCHEMA, DATE_TIME);
+
+    Map<String, Object> configOverride = new HashMap<>();
+    configOverride.put(PartitionerConfig.TIMESTAMP_UNIT_NAME_CONFIG, "s");
+    TimeBasedPartitioner<String> secondsPartitioner = configurePartitioner(
+            new TimeBasedPartitioner<>(), timeField, configOverride);
+
+    assertThat(secondsPartitioner.getTimestampExtractor(),
+            instanceOf(TimeBasedPartitioner.RecordFieldTimestampExtractor.class));
+
+    testMapNumericTimestampPartitionEncoding(
+            secondsPartitioner, timeField, ((long) ts)/1000, Schema.INT64_SCHEMA, DATE_TIME);
   }
 
   @Test
@@ -484,6 +485,14 @@ public void testRecordFieldTimeDateExtractor() {
     assertThat(partitioner.getTimestampExtractor(),
           instanceOf(TimeBasedPartitioner.RecordFieldTimestampExtractor.class));
 
+    Map<String, Object> configOverride = new HashMap<>();
+    configOverride.put(PartitionerConfig.TIMESTAMP_UNIT_NAME_CONFIG, "s");
+    TimeBasedPartitioner<String> secondsPartitioner = configurePartitioner(
+            new TimeBasedPartitioner<>(), timeField, configOverride);
+
+    assertThat(secondsPartitioner.getTimestampExtractor(),
+            instanceOf(TimeBasedPartitioner.RecordFieldTimestampExtractor.class));
+
     DateTime moment = new DateTime(2015, 4, 2, 1, 0, 0, 0, DateTimeZone.forID(TIME_ZONE));
     String expectedPartition = "year=2015/month=4/day=2/hour=1";
 
@@ -492,6 +501,11 @@ public void testRecordFieldTimeDateExtractor() {
     String encodedPartition = partitioner.encodePartition(sinkRecord);
     assertEquals(expectedPartition, encodedPartition);
 
+    long rawTimestampSeconds = moment.getMillis()/1000;
+    sinkRecord = createSinkRecord(rawTimestampSeconds);
+    encodedPartition = secondsPartitioner.encodePartition(sinkRecord);
+    assertEquals(expectedPartition, encodedPartition);
+
     String timestamp = ISODateTimeFormat.dateTimeNoMillis().print(moment);
     sinkRecord = createSinkRecord(Schema.STRING_SCHEMA, timestamp);
     encodedPartition = partitioner.encodePartition(sinkRecord);
@@ -511,6 +525,11 @@ public void testRecordFieldTimeDateExtractor() {
     encodedPartition = partitioner.encodePartition(sinkRecord);
     assertEquals("year=1970/month=1/day=1/hour=0", encodedPartition);
 
+    int shortTimestampSeconds = shortTimestamp/1000;
+    sinkRecord = createSinkRecord(Schema.INT32_SCHEMA, shortTimestampSeconds);
+    encodedPartition = secondsPartitioner.encodePartition(sinkRecord);
+    assertEquals("year=1970/month=1/day=1/hour=0", encodedPartition);
+
     // Struct - Date extraction
     sinkRecord = createSinkRecord(rawTimestamp);
     String structEncodedPartition = partitioner.encodePartition(sinkRecord);
@@ -525,8 +544,9 @@ public void testRecordFieldTimeDateExtractor() {
 
   @Test
   public void testNestedRecordFieldTimeExtractor() throws Exception {
+    String timeField = "nested.timestamp";
     TimeBasedPartitioner<String> partitioner = configurePartitioner(
-          new TimeBasedPartitioner<>(), "nested.timestamp", null);
+          new TimeBasedPartitioner<>(), timeField, null);
 
     assertThat(partitioner.getTimestampExtractor(),
           instanceOf(TimeBasedPartitioner.RecordFieldTimestampExtractor.class));
@@ -537,6 +557,20 @@ public void testNestedRecordFieldTimeExtractor() throws Exception {
     String encodedPartition = partitioner.encodePartition(sinkRecord);
 
     validateEncodedPartition(encodedPartition);
+
+    Map<String, Object> configOverride = new HashMap<>();
+    configOverride.put(PartitionerConfig.TIMESTAMP_UNIT_NAME_CONFIG, "s");
+    TimeBasedPartitioner<String> secondsPartitioner = configurePartitioner(
+            new TimeBasedPartitioner<>(), timeField, configOverride);
+
+    assertThat(secondsPartitioner.getTimestampExtractor(),
+            instanceOf(TimeBasedPartitioner.RecordFieldTimestampExtractor.class));
+
+    sinkRecord = createSinkRecordWithNestedTimestampField(timestamp/1000);
+
+    encodedPartition = secondsPartitioner.encodePartition(sinkRecord);
+
+    validateEncodedPartition(encodedPartition);
   }
 
   @Test