Search in sources :

Example 1 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller 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);
}
Also used : Schema(org.apache.arrow.vector.types.pojo.Schema) Connection(java.sql.Connection) S3BlockSpiller(com.amazonaws.athena.connector.lambda.data.S3BlockSpiller) com.amazonaws.athena.connector.lambda.data.writers.extractors(com.amazonaws.athena.connector.lambda.data.writers.extractors) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) ConstraintEvaluator(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintEvaluator) Answer(org.mockito.stubbing.Answer) BlockSpiller(com.amazonaws.athena.connector.lambda.data.BlockSpiller) SQLException(java.sql.SQLException) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) ByteArrayInputStream(java.io.ByteArrayInputStream) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) SchemaBuilder(com.amazonaws.athena.connector.lambda.data.SchemaBuilder) ResultSet(java.sql.ResultSet) Map(java.util.Map) TestBase(com.amazonaws.athena.connectors.jdbc.TestBase) AmazonS3(com.amazonaws.services.s3.AmazonS3) QueryStatusChecker(com.amazonaws.athena.connector.lambda.QueryStatusChecker) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) FederatedIdentity(com.amazonaws.athena.connector.lambda.security.FederatedIdentity) PutObjectResult(com.amazonaws.services.s3.model.PutObjectResult) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) Before(org.junit.Before) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Split(com.amazonaws.athena.connector.lambda.domain.Split) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) Test(org.junit.Test) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) PreparedStatement(java.sql.PreparedStatement) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) StandardCharsets(java.nio.charset.StandardCharsets) FieldBuilder(com.amazonaws.athena.connector.lambda.data.FieldBuilder) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) Mockito(org.mockito.Mockito) S3SpillLocation(com.amazonaws.athena.connector.lambda.domain.spill.S3SpillLocation) SpillConfig(com.amazonaws.athena.connector.lambda.data.SpillConfig) Assert(org.junit.Assert) JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) Collections(java.util.Collections) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) Types(java.sql.Types) PutObjectResult(com.amazonaws.services.s3.model.PutObjectResult) Schema(org.apache.arrow.vector.types.pojo.Schema) ConstraintEvaluator(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintEvaluator) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) SpillConfig(com.amazonaws.athena.connector.lambda.data.SpillConfig) ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) ByteArrayInputStream(java.io.ByteArrayInputStream) S3SpillLocation(com.amazonaws.athena.connector.lambda.domain.spill.S3SpillLocation) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) SchemaBuilder(com.amazonaws.athena.connector.lambda.data.SchemaBuilder) ResultSet(java.sql.ResultSet) S3BlockSpiller(com.amazonaws.athena.connector.lambda.data.S3BlockSpiller) Split(com.amazonaws.athena.connector.lambda.domain.Split) S3BlockSpiller(com.amazonaws.athena.connector.lambda.data.S3BlockSpiller) BlockSpiller(com.amazonaws.athena.connector.lambda.data.BlockSpiller) Test(org.junit.Test)

Example 2 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class MultiplexingJdbcRecordHandlerTest method readWithConstraint.

@Test
public void readWithConstraint() {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("fakedatabase");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
    Mockito.verify(this.fakeJdbcRecordHandler, Mockito.times(1)).readWithConstraint(Mockito.eq(blockSpiller), Mockito.eq(readRecordsRequest), Mockito.eq(queryStatusChecker));
}
Also used : ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) BlockSpiller(com.amazonaws.athena.connector.lambda.data.BlockSpiller) Test(org.junit.Test)

Example 3 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class MultiplexingJdbcRecordHandlerTest method readWithConstraintWithUnsupportedCatalog.

@Test(expected = RuntimeException.class)
public void readWithConstraintWithUnsupportedCatalog() {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("unsupportedCatalog");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
}
Also used : ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) BlockSpiller(com.amazonaws.athena.connector.lambda.data.BlockSpiller) Test(org.junit.Test)

Example 4 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class MySqlMuxJdbcRecordHandlerTest method readWithConstraint.

@Test
public void readWithConstraint() {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("mysql");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
    Mockito.verify(this.mySqlRecordHandler, Mockito.times(1)).readWithConstraint(Mockito.eq(blockSpiller), Mockito.eq(readRecordsRequest), Mockito.eq(queryStatusChecker));
}
Also used : ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) BlockSpiller(com.amazonaws.athena.connector.lambda.data.BlockSpiller) Test(org.junit.Test)

Example 5 with BlockSpiller

use of com.amazonaws.athena.connector.lambda.data.BlockSpiller in project aws-athena-query-federation by awslabs.

the class RedshiftMuxJdbcRecordHandlerTest method readWithConstraint.

@Test
public void readWithConstraint() {
    BlockSpiller blockSpiller = Mockito.mock(BlockSpiller.class);
    ReadRecordsRequest readRecordsRequest = Mockito.mock(ReadRecordsRequest.class);
    Mockito.when(readRecordsRequest.getCatalogName()).thenReturn("redshift");
    this.jdbcRecordHandler.readWithConstraint(blockSpiller, readRecordsRequest, queryStatusChecker);
    Mockito.verify(this.redshiftRecordHandler, Mockito.times(1)).readWithConstraint(Mockito.eq(blockSpiller), Mockito.eq(readRecordsRequest), Mockito.eq(queryStatusChecker));
}
Also used : ReadRecordsRequest(com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest) BlockSpiller(com.amazonaws.athena.connector.lambda.data.BlockSpiller) Test(org.junit.Test)

Aggregations

BlockSpiller (com.amazonaws.athena.connector.lambda.data.BlockSpiller)32 Test (org.junit.Test)31 ReadRecordsRequest (com.amazonaws.athena.connector.lambda.records.ReadRecordsRequest)30 S3BlockSpiller (com.amazonaws.athena.connector.lambda.data.S3BlockSpiller)2 SpillConfig (com.amazonaws.athena.connector.lambda.data.SpillConfig)2 ConstraintEvaluator (com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintEvaluator)2 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)2 S3SpillLocation (com.amazonaws.athena.connector.lambda.domain.spill.S3SpillLocation)2 QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)1 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)1 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)1 FieldBuilder (com.amazonaws.athena.connector.lambda.data.FieldBuilder)1 SchemaBuilder (com.amazonaws.athena.connector.lambda.data.SchemaBuilder)1 com.amazonaws.athena.connector.lambda.data.writers.extractors (com.amazonaws.athena.connector.lambda.data.writers.extractors)1 Split (com.amazonaws.athena.connector.lambda.domain.Split)1 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)1 EquatableValueSet (com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet)1 ValueSet (com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet)1 FederationThrottleException (com.amazonaws.athena.connector.lambda.exceptions.FederationThrottleException)1 GetTableRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableRequest)1