use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class HiveMetadataHandlerTest method doGetSplits.
@Test
public void doGetSplits() throws Exception {
BlockAllocator blockAllocator = new BlockAllocatorImpl();
String[] schema = { "data_type", "col_name" };
Object[][] values = { { "INTEGER", "case_number" }, { "VARCHAR", "case_location" }, { "TIMESTAMP", "case_instance" }, { "DATE", "case_date" } };
AtomicInteger rowNumber = new AtomicInteger(-1);
ResultSet resultSet = mockResultSet(schema, values, rowNumber);
String[] columns3 = { "col" };
int[] types3 = { Types.VARCHAR };
Object[][] values4 = { { "Partitioned:true" } };
ResultSet resultSet2 = mockResultSet(columns3, types3, values4, new AtomicInteger(-1));
Constraints constraints = Mockito.mock(Constraints.class);
TableName tempTableName = new TableName("testSchema", "testTable");
Schema partitionSchema = this.hiveMetadataHandler.getPartitionSchema("testCatalogName");
Set<String> partitionCols = new HashSet<>(Arrays.asList("partition"));
GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tempTableName, constraints, partitionSchema, partitionCols);
String value2 = "case_date=01-01-2000/case_number=0/case_instance=89898989/case_location=__HIVE_DEFAULT_PARTITION__";
String value3 = "case_date=02-01-2000/case_number=1/case_instance=89898990/case_location=Hyderabad";
String[] columns2 = { "Partition" };
int[] types2 = { Types.VARCHAR };
Object[][] values1 = { { value2 }, { value3 } };
Mockito.when(jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(connection);
String tableName = getTableLayoutRequest.getTableName().getTableName().toUpperCase();
PreparedStatement preparestatement1 = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(HiveMetadataHandler.GET_METADATA_QUERY + tableName)).thenReturn(preparestatement1);
final String getPartitionExistsSql = "show table extended like " + tableName;
final String getPartitionDetailsSql = "show partitions " + getTableLayoutRequest.getTableName().getTableName().toUpperCase();
Statement statement1 = Mockito.mock(Statement.class);
Mockito.when(this.connection.createStatement()).thenReturn(statement1);
ResultSet resultSet1 = mockResultSet(columns2, types2, values1, new AtomicInteger(-1));
Mockito.when(preparestatement1.executeQuery()).thenReturn(resultSet);
Mockito.when(statement1.executeQuery(getPartitionDetailsSql)).thenReturn(resultSet1);
Mockito.when(statement1.executeQuery(getPartitionExistsSql)).thenReturn(resultSet2);
Mockito.when(resultSet2.getString(1)).thenReturn("PARTITIONED:true");
GetTableLayoutResponse getTableLayoutResponse = this.hiveMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tempTableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, null);
GetSplitsResponse getSplitsResponse = this.hiveMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
Assert.assertEquals(2, getSplitsResponse.getSplits().size());
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class HiveMetadataHandlerTest method doGetTableLayout.
@Test
public void doGetTableLayout() throws Exception {
BlockAllocator blockAllocator = new BlockAllocatorImpl();
String[] schema = { "data_type", "col_name" };
Object[][] values = { { "INTEGER", "case_number" }, { "VARCHAR", "case_location" }, { "TIMESTAMP", "case_instance" }, { "DATE", "case_date" } };
AtomicInteger rowNumber = new AtomicInteger(-1);
ResultSet resultSet = mockResultSet(schema, values, rowNumber);
Constraints constraints = Mockito.mock(Constraints.class);
TableName tempTableName = new TableName("testSchema", "testTable");
Schema partitionSchema = this.hiveMetadataHandler.getPartitionSchema("testCatalogName");
Set<String> partitionCols = new HashSet<>(Arrays.asList("partition"));
GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tempTableName, constraints, partitionSchema, partitionCols);
String value2 = "case_date=01-01-2000/case_number=0/case_instance=89898989/case_location=__HIVE_DEFAULT_PARTITION__";
String value3 = "case_date=02-01-2000/case_number=1/case_instance=89898990/case_location=Hyderabad";
String[] columns2 = { "Partition" };
int[] types2 = { Types.VARCHAR };
Object[][] values1 = { { value3 }, { value2 } };
String[] columns3 = { "col" };
int[] types3 = { Types.VARCHAR };
Object[][] values4 = { { "PARTITIONED:true" } };
ResultSet resultSet2 = mockResultSet(columns3, types3, values4, new AtomicInteger(-1));
Mockito.when(jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(connection);
String tableName = getTableLayoutRequest.getTableName().getTableName().toUpperCase();
PreparedStatement preparestatement1 = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(HiveMetadataHandler.GET_METADATA_QUERY + tableName)).thenReturn(preparestatement1);
final String getPartitionExistsSql = "show table extended like " + tableName;
final String getPartitionDetailsSql = "show partitions " + getTableLayoutRequest.getTableName().getTableName().toUpperCase();
Statement statement1 = Mockito.mock(Statement.class);
Mockito.when(this.connection.createStatement()).thenReturn(statement1);
ResultSet resultSet1 = mockResultSet(columns2, types2, values1, new AtomicInteger(-1));
Mockito.when(preparestatement1.executeQuery()).thenReturn(resultSet);
Mockito.when(statement1.executeQuery(getPartitionDetailsSql)).thenReturn(resultSet1);
Mockito.when(statement1.executeQuery(getPartitionExistsSql)).thenReturn(resultSet2);
Mockito.when(resultSet2.getString(1)).thenReturn("PARTITIONED:true");
GetTableLayoutResponse getTableLayoutResponse = this.hiveMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
List<String> expectedValues = new ArrayList<>();
for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
expectedValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
}
Assert.assertEquals(expectedValues.get(0), "[partition : case_date=02-01-2000 and case_number=1 and case_instance=89898990 and case_location='Hyderabad']");
SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
expectedSchemaBuilder.addField(FieldBuilder.newBuilder("partition", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
Schema expectedSchema = expectedSchemaBuilder.build();
Assert.assertEquals(expectedSchema, getTableLayoutResponse.getPartitions().getSchema());
Assert.assertEquals(tempTableName, getTableLayoutResponse.getTableName());
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class HiveMetadataHandlerTest method doGetTableLayoutWithSQLException.
@Test(expected = RuntimeException.class)
public void doGetTableLayoutWithSQLException() throws Exception {
Constraints constraints = Mockito.mock(Constraints.class);
TableName tableName = new TableName("testSchema", "testTable");
Schema partitionSchema = this.hiveMetadataHandler.getPartitionSchema("testCatalogName");
Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
Connection connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
JdbcConnectionFactory jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
Mockito.when(jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(connection);
Mockito.when(connection.getMetaData().getSearchStringEscape()).thenThrow(new SQLException());
HiveMetadataHandler implalaMetadataHandler = new HiveMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
implalaMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class HiveMuxRecordHandlerTest method buildSplitSql.
@Test
public void buildSplitSql() throws SQLException {
ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("recordHive");
Connection jdbcConnection = Mockito.mock(Connection.class);
TableName tableName = new TableName("testSchema", "tableName");
Schema schema = Mockito.mock(Schema.class);
Constraints constraints = Mockito.mock(Constraints.class);
Split split = Mockito.mock(Split.class);
this.jdbcRecordHandler.buildSplitSql(jdbcConnection, "recordHive", tableName, schema, constraints, split);
Mockito.verify(this.hiveRecordHandler, Mockito.times(1)).buildSplitSql(Mockito.eq(jdbcConnection), Mockito.eq("recordHive"), Mockito.eq(tableName), Mockito.eq(schema), Mockito.eq(constraints), Mockito.eq(split));
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class ImpalaMuxRecordHandlerTest method buildSplitSql.
@Test
public void buildSplitSql() throws SQLException {
ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("recordImpala");
Connection jdbcConnection = Mockito.mock(Connection.class);
TableName tableName = new TableName("testSchema", "tableName");
Schema schema = Mockito.mock(Schema.class);
Constraints constraints = Mockito.mock(Constraints.class);
Split split = Mockito.mock(Split.class);
this.jdbcRecordHandler.buildSplitSql(jdbcConnection, "recordImpala", tableName, schema, constraints, split);
Mockito.verify(this.impalaRecordHandler, Mockito.times(1)).buildSplitSql(Mockito.eq(jdbcConnection), Mockito.eq("recordImpala"), Mockito.eq(tableName), Mockito.eq(schema), Mockito.eq(constraints), Mockito.eq(split));
}
Aggregations