Issue 2326 - Timestamp constraint in Cloudera Hive connector (#2353)

awslabs · Nov 6, 2024 · 9ebc772 · 9ebc772
1 parent d35068b
commit 9ebc772
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 3 deletions.
diff --git a/...a-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveQueryStringBuilder.java b/...a-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveQueryStringBuilder.java
@@ -22,7 +22,10 @@
 import com.amazonaws.athena.connector.lambda.domain.Split;
 import com.amazonaws.athena.connectors.jdbc.manager.FederationExpressionParser;
 import com.amazonaws.athena.connectors.jdbc.manager.JdbcSplitQueryBuilder;
+import com.amazonaws.athena.connectors.jdbc.manager.TypeAndValue;
 import com.google.common.base.Strings;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.pojo.ArrowType;
 
 import java.util.Collections;
 import java.util.List;
@@ -56,4 +59,21 @@ protected List<String> getPartitionWhereClauses(Split split)
         }
         return Collections.emptyList();
     }
+
+    protected String toPredicate(String columnName, String operator, Object value, ArrowType type,
+                                 List<TypeAndValue> accumulator)
+    {
+        Types.MinorType minorType = Types.getMinorTypeForArrowType(type);
+        if (minorType.equals(Types.MinorType.DATEMILLI) && isOperatorEquals(operator)) {
+            accumulator.add(new TypeAndValue(type, value));
+            return quote(columnName) + " " + operator + " cast(? as timestamp)";
+        }
+        // Default to parent's behavior
+        return super.toPredicate(columnName, operator, value, type, accumulator);
+    }
+
+    private boolean isOperatorEquals(String operator)
+    {
+        return operator.equals("=");
+    }
 }
diff --git a/...ra-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandlerTest.java b/...ra-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandlerTest.java
@@ -51,6 +51,10 @@
 import java.sql.Date;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
+import java.sql.Timestamp;
+import java.time.LocalDateTime;
+import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
 import java.util.Collections;
 import java.util.concurrent.TimeUnit;
 
@@ -136,4 +140,37 @@ public void buildSplitSql()
         Mockito.verify(preparedStatement, Mockito.times(1))
                 .setDate(1, expectedDate);
     }
+
+    @Test
+    public void buildSplitSqlTimestamp()
+            throws SQLException
+    {
+        TableName tableName = new TableName("testSchema", "testTable");
+
+        SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
+        schemaBuilder.addField(FieldBuilder.newBuilder("testCol1", Types.MinorType.DATEMILLI.getType()).build());
+        schemaBuilder.addField(FieldBuilder.newBuilder("partition", Types.MinorType.VARCHAR.getType()).build());
+        Schema schema = schemaBuilder.build();
+
+        Split split = Mockito.mock(Split.class);
+        Mockito.when(split.getProperties()).thenReturn(Collections.singletonMap("partition", "p0"));
+        Mockito.when(split.getProperty(Mockito.eq("partition"))).thenReturn("p0");
+
+        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
+        LocalDateTime timestamp = LocalDateTime.parse("2024-10-03 12:34:56", formatter);
+        ValueSet valueSet2 = getSingleValueSet(timestamp);
+        Constraints constraints = Mockito.mock(Constraints.class);
+        Mockito.when(constraints.getSummary()).thenReturn(new ImmutableMap.Builder<String, ValueSet>()
+                .put("testCol1", valueSet2)
+                .build());
+        PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
+        Mockito.when(this.connection.prepareStatement(nullable(String.class))).thenReturn(expectedPreparedStatement);
+        PreparedStatement preparedStatement = this.hiveRecordHandler.buildSplitSql(this.connection, "testCatalogName", tableName, schema, constraints, split);
+        Assert.assertEquals(expectedPreparedStatement, preparedStatement);
+        LocalDateTime timestampExp = LocalDateTime.parse("2024-10-03 12:34:56", formatter);
+        Timestamp expectedTimestamp = new Timestamp(timestamp.toInstant(ZoneOffset.UTC).toEpochMilli());
+        Assert.assertEquals(expectedPreparedStatement, preparedStatement);
+        Mockito.verify(preparedStatement, Mockito.times(1))
+                .setTimestamp(1, expectedTimestamp);
+    }
 }
diff --git a/...dbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcSplitQueryBuilder.java b/...dbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcSplitQueryBuilder.java
@@ -336,8 +336,8 @@ else if (singleValues.size() > 1) {
         return "(" + Joiner.on(" OR ").join(disjuncts) + ")";
     }
 
-    private String toPredicate(String columnName, String operator, Object value, ArrowType type,
-            List<TypeAndValue> accumulator)
+    protected String toPredicate(String columnName, String operator, Object value, ArrowType type,
+                                 List<TypeAndValue> accumulator)
     {
         accumulator.add(new TypeAndValue(type, value));
         return quote(columnName) + " " + operator + " ?";

diff --git a/...hana/src/main/java/com/amazonaws/athena/connectors/saphana/SaphanaQueryStringBuilder.java b/...hana/src/main/java/com/amazonaws/athena/connectors/saphana/SaphanaQueryStringBuilder.java
@@ -350,7 +350,7 @@ else if (singleValues.size() > 1) {
         return "(" + Joiner.on(" OR ").join(disjuncts) + ")";
     }
 
-    private String toPredicate(String columnName, String operator, Object value, ArrowType type,
+    protected String toPredicate(String columnName, String operator, Object value, ArrowType type,
                                List<TypeAndValue> accumulator)
     {
         accumulator.add(new TypeAndValue(type, value));