feat: add processing time on gcs partition

Author: lavkesh

src/main/java/com/gotocompany/firehose/config/BlobSinkConfig.java

@@ -49,6 +49,10 @@ public interface BlobSinkConfig extends AppConfig {
     @Key("SINK_BLOB_FILE_PARTITION_PROTO_TIMESTAMP_FIELD_NAME")
     String getFilePartitionProtoTimestampFieldName();
 
+    @Key("SINK_BLOB_FILE_PARTITION_PROCESSING_TIME_ENABLED")
+    @DefaultValue("false")
+    boolean getFilePartitionProcessingTimeEnabled();
+
     @Key("SINK_BLOB_FILE_PARTITION_TIME_GRANULARITY_TYPE")
     @DefaultValue("day")
     @ConverterClass(BlobSinkFilePartitionTypeConverter.class)

src/main/java/com/gotocompany/firehose/sink/blob/message/Record.java

@@ -2,11 +2,14 @@
 
 import com.google.protobuf.Descriptors;
 import com.google.protobuf.DynamicMessage;
+import com.gotocompany.firehose.config.BlobSinkConfig;
 import com.gotocompany.firehose.sink.blob.proto.KafkaMetadataProtoMessage;
 import lombok.AllArgsConstructor;
 import lombok.Data;
 
 import java.time.Instant;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
 
 @AllArgsConstructor
 @Data
@@ -34,4 +37,14 @@ public Instant getTimestamp(String fieldName) {
         int nanos = (int) timestamp.getField(timestamp.getDescriptorForType().findFieldByName("nanos"));
         return Instant.ofEpochSecond(seconds, nanos);
     }
+
+    public LocalDateTime getLocalDateTime(BlobSinkConfig config) {
+        if (config.getFilePartitionProcessingTimeEnabled()) {
+            return LocalDateTime.now();
+        } else {
+            return LocalDateTime.ofInstant(
+                    getTimestamp(config.getFilePartitionProtoTimestampFieldName()),
+                    ZoneId.of(config.getFilePartitionProtoTimestampTimezone()));
+        }
+    }
 }

src/main/java/com/gotocompany/firehose/sink/blob/writer/local/path/TimePartitionedPathUtils.java

@@ -7,11 +7,7 @@
 
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.time.Instant;
-import java.time.LocalDate;
 import java.time.LocalDateTime;
-import java.time.LocalTime;
-import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
 
 /**
@@ -25,22 +21,21 @@ public class TimePartitionedPathUtils {
 
     public static Path getTimePartitionedPath(Record record, BlobSinkConfig sinkConfig) {
         String topic = record.getTopic(sinkConfig.getOutputKafkaMetadataColumnName());
-        Instant timestamp = record.getTimestamp(sinkConfig.getFilePartitionProtoTimestampFieldName());
+        Path path = Paths.get(topic);
         if (sinkConfig.getFilePartitionTimeGranularityType() == Constants.FilePartitionType.NONE) {
-            return Paths.get(topic);
+            return path;
         }
-        LocalDate localDate = LocalDateTime.ofInstant(timestamp, ZoneId.of(sinkConfig.getFilePartitionProtoTimestampTimezone())).toLocalDate();
-        String datePart = DATE_FORMATTER.format(localDate);
-        LocalTime localTime = LocalDateTime.ofInstant(timestamp, ZoneId.of(sinkConfig.getFilePartitionProtoTimestampTimezone())).toLocalTime();
-        String hourPart = HOUR_FORMATTER.format(localTime);
+        LocalDateTime dateTime = record.getLocalDateTime(sinkConfig);
+        String datePart = DATE_FORMATTER.format(dateTime.toLocalDate());
+        String hourPart = HOUR_FORMATTER.format(dateTime.toLocalTime());
 
         String dateSegment = String.format("%s%s", sinkConfig.getFilePartitionTimeDatePrefix(), datePart);
         String hourSegment = String.format("%s%s", sinkConfig.getFilePartitionTimeHourPrefix(), hourPart);
 
         String dateTimePartition;
         switch (sinkConfig.getFilePartitionTimeGranularityType()) {
             case NONE:
-                return Paths.get(topic);
+                return path;
             case DAY:
                 dateTimePartition = String.format("%s", dateSegment);
                 break;

src/test/java/com/gotocompany/firehose/sink/blob/message/RecordTest.java

@@ -1,11 +1,15 @@
 package com.gotocompany.firehose.sink.blob.message;
 
 import com.google.protobuf.DynamicMessage;
+import com.gotocompany.firehose.config.BlobSinkConfig;
 import com.gotocompany.firehose.sink.blob.TestUtils;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.time.Instant;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
 
 public class RecordTest {
 
@@ -39,4 +43,33 @@ public void shouldGetTimeStampFromMessage() {
         Record record = new Record(message, metadata);
         Assert.assertEquals(defaultTimestamp, record.getTimestamp("created_time"));
     }
+
+    @Test
+    public void shouldGetDateTimeLocally() throws InterruptedException {
+        BlobSinkConfig config = Mockito.mock(BlobSinkConfig.class);
+        Mockito.when(config.getFilePartitionProcessingTimeEnabled()).thenReturn(true);
+        DynamicMessage message = TestUtils.createMessage(defaultTimestamp, defaultOrderNumber);
+        DynamicMessage metadata = TestUtils.createMetadata("nested_field", defaultTimestamp, defaultOffset, defaultPartition, defaultTopic);
+        Record record = new Record(message, metadata);
+        LocalDateTime before = LocalDateTime.now();
+        Thread.sleep(1000);
+        LocalDateTime localDateTime = record.getLocalDateTime(config);
+        Thread.sleep(1000);
+        LocalDateTime after = LocalDateTime.now();
+        Assert.assertTrue(localDateTime.isAfter(before));
+        Assert.assertTrue(localDateTime.isBefore(after));
+    }
+
+    @Test
+    public void shouldGetDateTimeFromMessage() throws InterruptedException {
+        BlobSinkConfig config = Mockito.mock(BlobSinkConfig.class);
+        Mockito.when(config.getFilePartitionProcessingTimeEnabled()).thenReturn(false);
+        Mockito.when(config.getFilePartitionProtoTimestampFieldName()).thenReturn("created_time");
+        Mockito.when(config.getFilePartitionProtoTimestampTimezone()).thenReturn("UTC");
+        DynamicMessage message = TestUtils.createMessage(defaultTimestamp, defaultOrderNumber);
+        DynamicMessage metadata = TestUtils.createMetadata("nested_field", defaultTimestamp, defaultOffset, defaultPartition, defaultTopic);
+        Record record = new Record(message, metadata);
+        LocalDateTime localDateTime = record.getLocalDateTime(config);
+        Assert.assertEquals(LocalDateTime.ofInstant(defaultTimestamp, ZoneId.of("UTC")), localDateTime);
+    }
 }