From 8048492501d939c5b8d5676248e69028f6cbfe65 Mon Sep 17 00:00:00 2001 From: AbdulRehman Faraj Date: Fri, 15 Mar 2024 17:44:10 +0000 Subject: [PATCH] Extended QPT to athena-cloudera-hive --- .../athena/connectors/cloudera/HiveMetadataHandler.java | 6 ++++++ .../athena/connectors/cloudera/HiveRecordHandler.java | 9 ++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandler.java b/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandler.java index ac7e6674cc..6f6b38b4bc 100644 --- a/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandler.java +++ b/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveMetadataHandler.java @@ -224,6 +224,11 @@ public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsReq { LOGGER.info("{}: Catalog {}, table {}", getSplitsRequest.getQueryId(), getSplitsRequest.getTableName().getSchemaName(), getSplitsRequest.getTableName().getTableName()); + if (getSplitsRequest.getConstraints().isQueryPassThrough()) { + LOGGER.info("QPT Split Requested"); + return setupQueryPassthroughSplit(getSplitsRequest); + } + int partitionContd = decodeContinuationToken(getSplitsRequest); Set splits = new HashSet<>(); Block partitions = getSplitsRequest.getPartitions(); @@ -271,6 +276,7 @@ public GetDataSourceCapabilitiesResponse doGetDataSourceCapabilities(BlockAlloca LimitPushdownSubType.INTEGER_CONSTANT )); + jdbcQueryPassthrough.addQueryPassthroughCapabilityIfEnabled(capabilities, configOptions); return new GetDataSourceCapabilitiesResponse(request.getCatalogName(), capabilities.build()); } diff --git a/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandler.java b/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandler.java index 9f50bdcaa9..ed5af5284d 100644 --- a/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandler.java +++ b/athena-cloudera-hive/src/main/java/com/amazonaws/athena/connectors/cloudera/HiveRecordHandler.java @@ -74,7 +74,14 @@ public HiveRecordHandler(DatabaseConnectionConfig databaseConnectionConfig, Jdbc @Override public PreparedStatement buildSplitSql(Connection jdbcConnection, String catalogName, TableName tableName, Schema schema, Constraints constraints, Split split) throws SQLException { - PreparedStatement preparedStatement = jdbcSplitQueryBuilder.buildSql(jdbcConnection, null, tableName.getSchemaName(), tableName.getTableName(), schema, constraints, split); + PreparedStatement preparedStatement; + + if (constraints.isQueryPassThrough()) { + preparedStatement = buildQueryPassthroughSql(jdbcConnection, constraints); + } + else { + preparedStatement = jdbcSplitQueryBuilder.buildSql(jdbcConnection, null, tableName.getSchemaName(), tableName.getTableName(), schema, constraints, split); + } preparedStatement.setFetchSize(FETCH_SIZE); return preparedStatement; }