use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class HiveMetadataHandlerTest method decodeContinuationToken.
@Test
public void decodeContinuationToken() throws Exception {
TableName tableName = new TableName("testSchema", "testTable");
Constraints constraints = Mockito.mock(Constraints.class);
Schema partitionSchema = this.hiveMetadataHandler.getPartitionSchema("testCatalogName");
Set<String> partitionCols = partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
BlockAllocator blockAllocator = new BlockAllocatorImpl();
GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);
GetTableLayoutResponse getTableLayoutResponse = this.hiveMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, "1");
Integer splitRequestToken = 0;
if (getSplitsRequest.hasContinuationToken()) {
splitRequestToken = Integer.valueOf(getSplitsRequest.getContinuationToken());
}
Assert.assertNotNull(splitRequestToken.toString());
}
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 JdbcRecordHandlerTest method readWithConstraint.
@Test
public void readWithConstraint() throws SQLException {
ConstraintEvaluator constraintEvaluator = Mockito.mock(ConstraintEvaluator.class);
Mockito.when(constraintEvaluator.apply(Mockito.anyString(), Mockito.any())).thenReturn(true);
TableName inputTableName = new TableName("testSchema", "testTable");
SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol1", org.apache.arrow.vector.types.Types.MinorType.INT.getType()).build());
expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testCol2", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
expectedSchemaBuilder.addField(FieldBuilder.newBuilder("testPartitionCol", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
Schema fieldSchema = expectedSchemaBuilder.build();
BlockAllocator allocator = new BlockAllocatorImpl();
S3SpillLocation s3SpillLocation = S3SpillLocation.newBuilder().withIsDirectory(true).build();
Split.Builder splitBuilder = Split.newBuilder(s3SpillLocation, null).add("testPartitionCol", String.valueOf("testPartitionValue"));
Constraints constraints = Mockito.mock(Constraints.class, Mockito.RETURNS_DEEP_STUBS);
String[] schema = { "testCol1", "testCol2" };
int[] columnTypes = { Types.INTEGER, Types.VARCHAR };
Object[][] values = { { 1, "testVal1" }, { 2, "testVal2" } };
AtomicInteger rowNumber = new AtomicInteger(-1);
ResultSet resultSet = mockResultSet(schema, columnTypes, values, rowNumber);
Mockito.when(this.preparedStatement.executeQuery()).thenReturn(resultSet);
SpillConfig spillConfig = Mockito.mock(SpillConfig.class);
Mockito.when(spillConfig.getSpillLocation()).thenReturn(s3SpillLocation);
BlockSpiller s3Spiller = new S3BlockSpiller(this.amazonS3, spillConfig, allocator, fieldSchema, constraintEvaluator);
ReadRecordsRequest readRecordsRequest = new ReadRecordsRequest(this.federatedIdentity, "testCatalog", "testQueryId", inputTableName, fieldSchema, splitBuilder.build(), constraints, 1024, 1024);
Mockito.when(amazonS3.putObject(Mockito.anyString(), Mockito.anyString(), Mockito.any(), Mockito.any())).thenAnswer((Answer<PutObjectResult>) invocation -> {
ByteArrayInputStream byteArrayInputStream = (ByteArrayInputStream) invocation.getArguments()[2];
int n = byteArrayInputStream.available();
byte[] bytes = new byte[n];
byteArrayInputStream.read(bytes, 0, n);
String data = new String(bytes, StandardCharsets.UTF_8);
Assert.assertTrue(data.contains("testVal1") || data.contains("testVal2") || data.contains("testPartitionValue"));
return new PutObjectResult();
});
this.jdbcRecordHandler.readWithConstraint(s3Spiller, readRecordsRequest, queryStatusChecker);
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class PostGreSqlMetadataHandlerTest method doGetSplitsContinuation.
@Test
public void doGetSplitsContinuation() throws Exception {
BlockAllocator blockAllocator = new BlockAllocatorImpl();
Constraints constraints = Mockito.mock(Constraints.class);
TableName tableName = new TableName("testSchema", "testTable");
Schema partitionSchema = this.postGreSqlMetadataHandler.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);
PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(PostGreSqlMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);
String[] columns = { "child_schema", "child" };
int[] types = { Types.VARCHAR, Types.VARCHAR };
Object[][] values = { { "s0", "p0" }, { "s1", "p1" } };
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
final String expectedQuery = String.format(PostGreSqlMetadataHandler.GET_PARTITIONS_QUERY, tableName.getTableName(), tableName.getSchemaName());
Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
Mockito.when(this.connection.getMetaData().getSearchStringEscape()).thenReturn(null);
GetTableLayoutResponse getTableLayoutResponse = this.postGreSqlMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, "1");
GetSplitsResponse getSplitsResponse = this.postGreSqlMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
Set<Map<String, String>> expectedSplits = new HashSet<>();
expectedSplits.add(ImmutableMap.of("partition_schema_name", "s1", "partition_name", "p1"));
Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());
Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
Assert.assertEquals(expectedSplits, actualSplits);
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Constraints in project aws-athena-query-federation by awslabs.
the class PostGreSqlMetadataHandlerTest 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.postGreSqlMetadataHandler.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());
PostGreSqlMetadataHandler postGreSqlMetadataHandler = new PostGreSqlMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory);
postGreSqlMetadataHandler.doGetTableLayout(Mockito.mock(BlockAllocator.class), getTableLayoutRequest);
}
Aggregations