use of com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest in project aws-athena-query-federation by awslabs.
the class GetSplitsRequestSerDeTest method deserialize.
@Test
public void deserialize() throws IOException {
logger.info("deserialize: enter");
InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes());
GetSplitsRequest actual = (GetSplitsRequest) mapper.readValue(input, FederationRequest.class);
logger.info("deserialize: deserialized[{}]", actual);
assertEquals(expected, actual);
logger.info("deserialize: exit");
}
use of com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest in project aws-athena-query-federation by awslabs.
the class JdbcMetadataHandlerTest method setup.
@Before
public void setup() {
this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class);
this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
this.secretsManager = Mockito.mock(AWSSecretsManager.class);
this.athena = Mockito.mock(AmazonAthena.class);
Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", "fakedatabase", "fakedatabase://jdbc:fakedatabase://hostname/${testSecret}", "testSecret");
this.jdbcMetadataHandler = new JdbcMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, jdbcConnectionFactory) {
@Override
public Schema getPartitionSchema(final String catalogName) {
return PARTITION_SCHEMA;
}
@Override
public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutRequest getTableLayoutRequest, QueryStatusChecker queryStatusChecker) {
}
@Override
public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest getSplitsRequest) {
return null;
}
};
this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
this.blockAllocator = Mockito.mock(BlockAllocator.class);
}
use of com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest in project aws-athena-query-federation by awslabs.
the class MultiplexingJdbcMetadataHandlerTest method doGetSplits.
@Test
public void doGetSplits() {
GetSplitsRequest getSplitsRequest = Mockito.mock(GetSplitsRequest.class);
Mockito.when(getSplitsRequest.getCatalogName()).thenReturn("fakedatabase");
this.jdbcMetadataHandler.doGetSplits(this.allocator, getSplitsRequest);
Mockito.verify(this.fakeDatabaseHandler, Mockito.times(1)).doGetSplits(Mockito.eq(this.allocator), Mockito.eq(getSplitsRequest));
}
use of com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest in project aws-athena-query-federation by awslabs.
the class PostGreSqlMetadataHandlerTest method doGetSplits.
@Test
public void doGetSplits() 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));
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, null);
GetSplitsResponse getSplitsResponse = this.postGreSqlMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
Set<Map<String, String>> expectedSplits = new HashSet<>();
expectedSplits.add(ImmutableMap.of("partition_schema_name", "s0", "partition_name", "p0"));
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.metadata.GetSplitsRequest in project aws-athena-query-federation by awslabs.
the class RedshiftMetadataHandlerTest method doGetSplits.
@Test
public void doGetSplits() throws Exception {
BlockAllocator blockAllocator = new BlockAllocatorImpl();
Constraints constraints = Mockito.mock(Constraints.class);
TableName tableName = new TableName("testSchema", "testTable");
Schema partitionSchema = this.redshiftMetadataHandler.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));
Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
Mockito.when(this.connection.getMetaData().getSearchStringEscape()).thenReturn(null);
GetTableLayoutResponse getTableLayoutResponse = this.redshiftMetadataHandler.doGetTableLayout(blockAllocator, getTableLayoutRequest);
BlockAllocator splitBlockAllocator = new BlockAllocatorImpl();
GetSplitsRequest getSplitsRequest = new GetSplitsRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, getTableLayoutResponse.getPartitions(), new ArrayList<>(partitionCols), constraints, null);
GetSplitsResponse getSplitsResponse = this.redshiftMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
Set<Map<String, String>> expectedSplits = new HashSet<>();
expectedSplits.add(ImmutableMap.of("partition_schema_name", "s0", "partition_name", "p0"));
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);
}
Aggregations