Fix ValidationException for Timestream connector (#1888)

awslabs · May 16, 2024 · 907f5bb · 907f5bb
1 parent b3d46a9
commit 907f5bb
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/...ream/src/main/java/com/amazonaws/athena/connectors/timestream/query/PredicateBuilder.java b/...ream/src/main/java/com/amazonaws/athena/connectors/timestream/query/PredicateBuilder.java
@@ -30,12 +30,23 @@
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 
+import java.time.LocalDateTime;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
+import java.time.format.DateTimeFormatterBuilder;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
 
 public class PredicateBuilder
 {
+    // We use a specific format to use the full precision provided by Timestream.
+    // Additionally, it prevents generating invalid format like `2024-12-31 00:11:22.`.
+    private static final DateTimeFormatter TIMESTAMP_FORMATTER = new DateTimeFormatterBuilder()
+            .appendPattern("yyyy-MM-dd HH:mm:ss.nnnnnnnnn")
+            .toFormatter()
+            .withZone(ZoneId.of("UTC"));
+
     private PredicateBuilder() {}
 
     public static List<String> buildConjucts(
@@ -157,6 +168,8 @@ private static String quoteValue(Object value, ArrowType type)
         switch (Types.getMinorTypeForArrowType(type)) {
             case VARCHAR:
                 return "\'" + value + "\'";
+            case DATEMILLI:
+                return "\'" + ((LocalDateTime) value).format(TIMESTAMP_FORMATTER) + "\'";
             default:
                 return String.valueOf(value);
         }

diff --git a/...rc/test/java/com/amazonaws/athena/connectors/timestream/query/SelectQueryBuilderTest.java b/...rc/test/java/com/amazonaws/athena/connectors/timestream/query/SelectQueryBuilderTest.java
@@ -36,9 +36,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
+import java.time.LocalDateTime;
+import java.util.*;
 
 import static com.amazonaws.athena.connector.lambda.domain.predicate.Constraints.DEFAULT_NO_LIMIT;
 import static com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.VIEW_METADATA_FIELD;
@@ -104,6 +103,37 @@ public void build()
         logger.info("build: exit");
     }
 
+    @Test
+    public void buildWithTime() {
+        logger.info("build: enter");
+
+        String expected = "SELECT val FROM \"myDatabase\".\"myTable\" WHERE ((\"time1\" > '2024-04-05 09:31:12.000000000')) AND ((\"time0\" > '2024-04-05 09:31:12.142000000'))";
+
+        Map<String, ValueSet> constraintsMap = new HashMap<>();
+        constraintsMap.put("time0", SortedRangeSet.copyOf(Types.MinorType.DATEMILLI.getType(),
+                ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.DATEMILLI.getType(),
+                        LocalDateTime.of(2024, 4, 5, 9, 31, 12, 142000000))), false));
+        constraintsMap.put("time1", SortedRangeSet.copyOf(Types.MinorType.DATEMILLI.getType(),
+                ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.DATEMILLI.getType(),
+                        LocalDateTime.of(2024, 4, 5, 9, 31, 12))), false));
+
+        Schema schema = SchemaBuilder.newBuilder()
+                .addField("val", Types.MinorType.DATEMILLI.getType())
+                .build();
+
+        String actual = queryFactory.createSelectQueryBuilder(VIEW_METADATA_FIELD)
+                .withDatabaseName("myDatabase")
+                .withTableName("myTable")
+                .withProjection(schema)
+                .withConjucts(new Constraints(constraintsMap, Collections.emptyList(), Collections.emptyList(), DEFAULT_NO_LIMIT))
+                .build().replace("\n", "");
+
+        logger.info("build: actual[{}]", actual);
+        assertEquals(expected, actual);
+
+        logger.info("build: exit");
+    }
+
     @Test
     public void buildWithView()
     {