Search in sources :

Example 41 with GetTableLayoutRequest

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

the class GetTableLayoutRequestSerDeTest method deserialize.

@Test
public void deserialize() throws IOException {
    logger.info("deserialize: enter");
    InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes());
    GetTableLayoutRequest actual = (GetTableLayoutRequest) mapper.readValue(input, FederationRequest.class);
    logger.info("deserialize: deserialized[{}]", actual);
    assertEquals(expected, actual);
    logger.info("deserialize: exit");
}
Also used : ByteArrayInputStream(java.io.ByteArrayInputStream) ByteArrayInputStream(java.io.ByteArrayInputStream) InputStream(java.io.InputStream) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest) TypedSerDeTest(com.amazonaws.athena.connector.lambda.serde.TypedSerDeTest) Test(org.junit.Test)

Example 42 with GetTableLayoutRequest

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

the class GetTableLayoutRequestSerDeTest method beforeTest.

@Before
public void beforeTest() throws IOException {
    Schema schema = SchemaBuilder.newBuilder().addField("year", new ArrowType.Int(32, true)).addField("month", new ArrowType.Int(32, true)).addField("day", new ArrowType.Int(32, true)).addField("col2", new ArrowType.Utf8()).addField("col3", Types.MinorType.FLOAT8.getType()).addField("col4", Types.MinorType.FLOAT8.getType()).addField("col5", Types.MinorType.FLOAT8.getType()).build();
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col3", SortedRangeSet.copyOf(Types.MinorType.FLOAT8.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.FLOAT8.getType(), -10000D)), false));
    constraintsMap.put("col4", EquatableValueSet.newBuilder(allocator, Types.MinorType.FLOAT8.getType(), false, true).add(1.1D).build());
    constraintsMap.put("col5", new AllOrNoneValueSet(Types.MinorType.FLOAT8.getType(), false, true));
    Constraints constraints = new Constraints(constraintsMap);
    expected = new GetTableLayoutRequest(federatedIdentity, "test-query-id", "test-catalog", new TableName("test-schema", "test-table"), constraints, schema, ImmutableSet.of("year", "month", "day"));
    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "GetTableLayoutRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}
Also used : HashMap(java.util.HashMap) AllOrNoneValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.AllOrNoneValueSet) Schema(org.apache.arrow.vector.types.pojo.Schema) ArrowType(org.apache.arrow.vector.types.pojo.ArrowType) TableName(com.amazonaws.athena.connector.lambda.domain.TableName) Constraints(com.amazonaws.athena.connector.lambda.domain.predicate.Constraints) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) EquatableValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.EquatableValueSet) ValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet) AllOrNoneValueSet(com.amazonaws.athena.connector.lambda.domain.predicate.AllOrNoneValueSet) Before(org.junit.Before)

Example 43 with GetTableLayoutRequest

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

the class HiveMuxMetadataHandlerTest method getPartitions.

@Test
public void getPartitions() throws Exception {
    GetTableLayoutRequest getTableLayoutRequest = Mockito.mock(GetTableLayoutRequest.class);
    Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn("metaHive");
    this.jdbcMetadataHandler.getPartitions(Mockito.mock(BlockWriter.class), getTableLayoutRequest, queryStatusChecker);
    Mockito.verify(this.hiveMetadataHandler, Mockito.times(1)).getPartitions(Mockito.any(BlockWriter.class), Mockito.eq(getTableLayoutRequest), Mockito.eq(queryStatusChecker));
}
Also used : GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) BlockWriter(com.amazonaws.athena.connector.lambda.data.BlockWriter) Test(org.junit.Test)

Example 44 with GetTableLayoutRequest

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

the class HiveMuxMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    GetTableLayoutRequest getTableLayoutRequest = Mockito.mock(GetTableLayoutRequest.class);
    Mockito.when(getTableLayoutRequest.getTableName()).thenReturn(new TableName("testSchema", "testTable"));
    Mockito.when(getTableLayoutRequest.getCatalogName()).thenReturn("metaHive");
    this.jdbcMetadataHandler.doGetTableLayout(this.allocator, getTableLayoutRequest);
    Mockito.verify(this.hiveMetadataHandler, Mockito.times(1)).doGetTableLayout(Mockito.eq(this.allocator), Mockito.eq(getTableLayoutRequest));
}
Also used : TableName(com.amazonaws.athena.connector.lambda.domain.TableName) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) Test(org.junit.Test)

Example 45 with GetTableLayoutRequest

use of com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest 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);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) GetSplitsRequest(com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) Schema(org.apache.arrow.vector.types.pojo.Schema) Connection(java.sql.Connection) DatabaseConnectionConfig(com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) QueryStatusChecker(com.amazonaws.athena.connector.lambda.QueryStatusChecker) FederatedIdentity(com.amazonaws.athena.connector.lambda.security.FederatedIdentity) GetSplitsResponse(com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) BlockWriter(com.amazonaws.athena.connector.lambda.data.BlockWriter) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Before(org.junit.Before)

Aggregations

GetTableLayoutRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest)76 Test (org.junit.Test)71 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)54 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)47 GetTableLayoutResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse)42 Schema (org.apache.arrow.vector.types.pojo.Schema)39 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)33 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)27 ResultSet (java.sql.ResultSet)22 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)22 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)17 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)17 PreparedStatement (java.sql.PreparedStatement)17 HashMap (java.util.HashMap)16 HashSet (java.util.HashSet)16 ArrayList (java.util.ArrayList)15 BlockWriter (com.amazonaws.athena.connector.lambda.data.BlockWriter)14 Map (java.util.Map)14 SchemaBuilder (com.amazonaws.athena.connector.lambda.data.SchemaBuilder)12 ValueSet (com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet)11