From 9ebc772b7cb0a843641309eecde68b2faf8b5955 Mon Sep 17 00:00:00 2001 From: Trianz-Akshay <108925344+Trianz-Akshay@users.noreply.github.com> Date: Wed, 6 Nov 2024 20:17:17 +0530 Subject: [PATCH] Issue 2326 - Timestamp constraint in Cloudera Hive connector (#2353) --- .../cloudera/HiveQueryStringBuilder.java | 20 ++++++++++ .../cloudera/HiveRecordHandlerTest.java | 37 +++++++++++++++++++ .../jdbc/manager/JdbcSplitQueryBuilder.java | 4 +- .../saphana/SaphanaQueryStringBuilder.java | 2 +- 4 files changed, 60 insertions(+), 3 deletions(-) diff --git a/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveQueryStringBuilder.java b/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveQueryStringBuilder.java index 96d4e0f360..a2a10d985a 100644 --- a/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveQueryStringBuilder.java +++ b/athena-cloudera-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 getPartitionWhereClauses(Split split) } return Collections.emptyList(); } + + protected String toPredicate(String columnName, String operator, Object value, ArrowType type, + List 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/athena-cloudera-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandlerTest.java b/athena-cloudera-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandlerTest.java index 108474f096..cd1dba8155 100644 --- a/athena-cloudera-hive/src/test/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandlerTest.java +++ b/athena-cloudera-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() + .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/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcSplitQueryBuilder.java b/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcSplitQueryBuilder.java index a532d02aa2..ad71e20ad1 100644 --- a/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcSplitQueryBuilder.java +++ b/athena-jdbc/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 accumulator) + protected String toPredicate(String columnName, String operator, Object value, ArrowType type, + List accumulator) { accumulator.add(new TypeAndValue(type, value)); return quote(columnName) + " " + operator + " ?"; diff --git a/athena-saphana/src/main/java/com/amazonaws/athena/connectors/saphana/SaphanaQueryStringBuilder.java b/athena-saphana/src/main/java/com/amazonaws/athena/connectors/saphana/SaphanaQueryStringBuilder.java index 16ec948679..bdb70c367e 100644 --- a/athena-saphana/src/main/java/com/amazonaws/athena/connectors/saphana/SaphanaQueryStringBuilder.java +++ b/athena-saphana/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 accumulator) { accumulator.add(new TypeAndValue(type, value));