CDAP-19939 Fix handling of high precision datetime (#190)

* Fix handling of high precision datetime

* Review fixes for handling of high precision datetime

* CDAP-19939 - Minor refactoring

* Review fixes for handling of high precision datetime

Author: sumitjnn

src/main/java/io/cdap/delta/bigquery/StructuredRecordToJson.java

@@ -20,6 +20,8 @@
 import io.cdap.cdap.api.common.Bytes;
 import io.cdap.cdap.api.data.format.StructuredRecord;
 import io.cdap.cdap.api.data.schema.Schema;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.math.BigDecimal;
@@ -39,22 +41,35 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
 /**
  * Util class to convert structured record into json.
  */
 public final class StructuredRecordToJson {
+  private static final Logger LOG = LoggerFactory.getLogger(StructuredRecordToJson.class);
+
   private static final DateTimeFormatter DATETIME_FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");
   private static final DateTimeFormatter TIME_FORMATTER = DateTimeFormatter.ofPattern("HH:mm:ss.SSSSSS");
   // array of arrays and map of arrays are not supported by big query
   private static final Set<Schema.Type> UNSUPPORTED_ARRAY_TYPES = ImmutableSet.of(Schema.Type.ARRAY, Schema.Type.MAP);
 
+  private static final int MAX_LOGICAL_DATE_TIME_FRACTION_PRECISION = 6;
+  /* BigQuery format for DateTime: YYYY-[M]M-[D]D[( |T)[H]H:[M]M:[S]S[.F]]
+   * [.F]: Up to six fractional digits (microsecond precision)
+   */
+  private static final Pattern LOGICAL_DATE_PATTERN =
+    Pattern.compile("\\d{4}-\\d{1,2}-\\d{1,2}([ T]\\d{1,2}:\\d{1,2}:\\d{1,2}(\\.(\\d+))?)?");
+  private static final int TIME_FRACTION_GROUP = 3;
+
   /**
    * Writes object and writes to json writer.
-   * @param writer json writer to write the object to
-   * @param name name of the field to be written
-   * @param object object to be written
+   *
+   * @param writer      json writer to write the object to
+   * @param name        name of the field to be written
+   * @param object      object to be written
    * @param fieldSchema field schema to be written
    */
   public static void write(JsonWriter writer, String name, Object object, Schema fieldSchema) throws IOException {
@@ -143,8 +158,14 @@ private static void writeSimpleTypes(JsonWriter writer, String name, boolean isA
           writer.value(Objects.requireNonNull(getDecimal((byte[]) object, schema)).toPlainString());
           break;
         case DATETIME:
-          //datetime should be already an ISO-8601 string
-          writer.value(Objects.requireNonNull(object.toString()));
+          String strValue = object.toString();
+          // Datetime should be already an ISO-8601 string
+          // But BigQuery format is stricter than ISO-8601 and does not support Zone and Offset
+          // Hence it is more closer to DateTimeFormatter.ISO_LOCAL_DATE_TIME but with microsecond precision
+          // Check if the value matches expected format for DateTime and trim time fraction to
+          // MAX_TIME_FRACTION_PRECISION if it exceeds it
+          strValue = checkAndTrimToMaxSupportedPrecision(strValue);
+          writer.value(Objects.requireNonNull(strValue));
           break;
         default:
           throw new IllegalStateException(
@@ -185,6 +206,27 @@ private static void writeSimpleTypes(JsonWriter writer, String name, boolean isA
     }
   }
 
+  private static String checkAndTrimToMaxSupportedPrecision(String strValue) {
+    Matcher matcher = LOGICAL_DATE_PATTERN.matcher(strValue);
+    if (matcher.matches()) {
+      String timeFraction = matcher.group(TIME_FRACTION_GROUP);
+      //matcher.group returns null for a group if an optional group did not exist in the string
+      if (timeFraction != null && timeFraction.length() > MAX_LOGICAL_DATE_TIME_FRACTION_PRECISION) {
+        //Trim the time fraction to max supported precision
+        String trimmedTimeFraction = timeFraction.substring(0, MAX_LOGICAL_DATE_TIME_FRACTION_PRECISION);
+        strValue = new StringBuilder(strValue)
+          .replace(matcher.start(TIME_FRACTION_GROUP), matcher.end(TIME_FRACTION_GROUP), trimmedTimeFraction)
+          .toString();
+      }
+    } else {
+      //Don't throw exception for now as we might be missing some scenario in the format
+      //Let it fail during BigQuery insert in case of wrong format
+      LOG.warn("Invalid value {} for DATETIME type, it should match the " +
+                 "format YYYY-[M]M-[D]D[( |T)[H]H:[M]M:[S]S[.F]]", strValue);
+    }
+    return strValue;
+  }
+
   private static void writeArray(JsonWriter writer,
                                  String name,
                                  @Nullable Object value,

src/test/java/io/cdap/delta/bigquery/BigQueryEventConsumerTest.java

@@ -63,8 +63,10 @@
 import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.time.LocalDate;
+import java.time.LocalDateTime;
 import java.time.ZoneId;
 import java.time.ZonedDateTime;
+import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -84,7 +86,7 @@
  * Tests for BigQueryEventConsumer. In order to run these tests, service account credentials must be set in the system
  * properties. The service account must have permission to create and write to BigQuery datasets and tables,
  * as well as permission to write to GCS.
- *
+ * <p>
  * The tests create real resources in GCP and will cost some small amount of money for each run.
  */
 public class BigQueryEventConsumerTest {
@@ -93,6 +95,7 @@ public class BigQueryEventConsumerTest {
     Schema.Field.of("id", Schema.of(Schema.Type.INT)),
     Schema.Field.of("name", Schema.of(Schema.Type.STRING)),
     Schema.Field.of("created", Schema.of(Schema.LogicalType.TIMESTAMP_MILLIS)),
+    Schema.Field.of("updated", Schema.nullableOf(Schema.of(Schema.LogicalType.DATETIME))),
     Schema.Field.of("bday", Schema.of(Schema.LogicalType.DATE)),
     Schema.Field.of("score", Schema.nullableOf(Schema.of(Schema.Type.DOUBLE))),
     Schema.Field.of("partition", Schema.nullableOf(Schema.of(Schema.Type.INT))));
@@ -794,6 +797,8 @@ private void insertTruncate(BigQueryEventConsumer eventConsumer, String dataset)
       .set("id", 0)
       .set("name", "alice")
       .setTimestamp("created", ZonedDateTime.ofInstant(Instant.EPOCH, ZoneId.of("UTC")))
+      //check nano precision value is correctly handled, minusNanos is required otherwise zeroes in nanos are truncated
+      .setDateTime("updated", LocalDateTime.ofInstant(Instant.now(), ZoneId.of("UTC")).minusNanos(1))
       .setDate("bday", LocalDate.ofEpochDay(0))
       .set("score", 0.0d)
       .build();
@@ -814,6 +819,8 @@ private void insertTruncate(BigQueryEventConsumer eventConsumer, String dataset)
       .set("id", 1)
       .set("name", "bob")
       .setTimestamp("created", ZonedDateTime.ofInstant(Instant.ofEpochSecond(86400), ZoneId.of("UTC")))
+      .setDateTime("updated", LocalDateTime.ofInstant(Instant.now().truncatedTo(ChronoUnit.SECONDS),
+                                                      ZoneId.of("UTC")))
       .setDate("bday", LocalDate.ofEpochDay(1))
       .set("score", 1.0d)
       .build();

src/test/java/io/cdap/delta/bigquery/StructuredRecordToJsonTest.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright © 2022 Cask Data, Inc.
+ *
+ * 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.cdap.delta.bigquery;
+
+import com.google.gson.stream.JsonWriter;
+import io.cdap.cdap.api.data.schema.Schema;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import java.io.IOException;
+
+@RunWith(MockitoJUnitRunner.class)
+public class StructuredRecordToJsonTest {
+
+  private static final String UPDATED_COLUMN = "updated";
+  public static final Schema DATETIME_SCHEMA = Schema.of(Schema.LogicalType.DATETIME);
+
+  @Mock
+  private JsonWriter jsonWriter;
+
+  @Test
+  public void testWriteDateTimeNanoPrecision() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19 14:31:12.123456789", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19 14:31:12.123456");
+  }
+
+  @Test
+  public void testWriteDateTimeNanoPrecisionAlternateFormat() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19T14:31:12.123456789", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19T14:31:12.123456");
+  }
+
+  @Test
+  public void testWriteDateTimeHundredNanoPrecision() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19 14:31:12.1234567", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19 14:31:12.123456");
+  }
+
+  @Test
+  public void testWriteDateTimeMicrosecondPrecision() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19 14:31:12.123456", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19 14:31:12.123456");
+  }
+
+  @Test
+  public void testWriteDateTimeMillisPrecision() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19 14:31:12.123", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19 14:31:12.123");
+  }
+
+  @Test
+  public void testWriteDateTimeSecondPrecision() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19 14:31:12", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19 14:31:12");
+  }
+
+  @Test
+  public void testWriteDateTimeSecondPrecisionSingleDigit() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19 4:1:2", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19 4:1:2");
+  }
+
+  @Test
+  public void testWriteDateTimeWithoutTime() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-04-19", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-04-19");
+  }
+
+  @Test
+  public void testWriteDateTimeSingleDigitDate() throws IOException {
+    StructuredRecordToJson.write(jsonWriter, UPDATED_COLUMN, "2022-4-1", DATETIME_SCHEMA);
+    Mockito.verify(jsonWriter, Mockito.times(1)).value("2022-4-1");
+  }
+}