Search in sources :

Example 1 with LocalKeyFactory

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

the class ExampleMetadataHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("setUpBefore - enter");
    allocator = new BlockAllocatorImpl();
    metadataHandler = new ExampleMetadataHandler(new LocalKeyFactory(), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "spill-bucket", "spill-prefix");
    logger.info("setUpBefore - exit");
}
Also used : BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Example 2 with LocalKeyFactory

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

the class GlueMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    logger.info("{}: enter", testName.getMethodName());
    handler = new GlueMetadataHandler(mockGlue, new LocalKeyFactory(), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "glue-test", "spill-bucket", "spill-prefix") {

        @Override
        public GetTableLayoutResponse doGetTableLayout(BlockAllocator blockAllocator, GetTableLayoutRequest request) {
            throw new UnsupportedOperationException();
        }

        @Override
        public void getPartitions(BlockWriter blockWriter, GetTableLayoutRequest request, QueryStatusChecker queryStatusChecker) throws Exception {
            throw new UnsupportedOperationException();
        }

        @Override
        public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest request) {
            throw new UnsupportedOperationException();
        }
    };
    allocator = new BlockAllocatorImpl();
    // doListTables pagination.
    when(mockGlue.getTables(any(GetTablesRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        GetTablesRequest request = (GetTablesRequest) invocationOnMock.getArguments()[0];
        String nextToken = request.getNextToken();
        int pageSize = request.getMaxResults() == null ? UNLIMITED_PAGE_SIZE_VALUE : request.getMaxResults();
        assertEquals(accountId, request.getCatalogId());
        assertEquals(schema, request.getDatabaseName());
        GetTablesResult mockResult = mock(GetTablesResult.class);
        if (pageSize == UNLIMITED_PAGE_SIZE_VALUE) {
            // Simulate full list of tables returned from Glue.
            when(mockResult.getTableList()).thenReturn(unPaginatedTables);
            when(mockResult.getNextToken()).thenReturn(null);
        } else {
            // Simulate paginated list of tables returned from Glue.
            List<Table> paginatedTables = unPaginatedTables.stream().sorted(Comparator.comparing(Table::getName)).filter(table -> nextToken == null || table.getName().compareTo(nextToken) >= 0).limit(pageSize + 1).collect(Collectors.toList());
            if (paginatedTables.size() > pageSize) {
                when(mockResult.getNextToken()).thenReturn(paginatedTables.get(pageSize).getName());
                when(mockResult.getTableList()).thenReturn(paginatedTables.subList(0, pageSize));
            } else {
                when(mockResult.getNextToken()).thenReturn(null);
                when(mockResult.getTableList()).thenReturn(paginatedTables);
            }
        }
        return mockResult;
    });
}
Also used : Table(com.amazonaws.services.glue.model.Table) GetSplitsRequest(com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest) GetTablesResult(com.amazonaws.services.glue.model.GetTablesResult) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) GetTableLayoutResponse(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse) QueryStatusChecker(com.amazonaws.athena.connector.lambda.QueryStatusChecker) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) GetSplitsResponse(com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse) InvocationOnMock(org.mockito.invocation.InvocationOnMock) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) BlockWriter(com.amazonaws.athena.connector.lambda.data.BlockWriter) GetTablesRequest(com.amazonaws.services.glue.model.GetTablesRequest) Before(org.junit.Before)

Example 3 with LocalKeyFactory

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

the class NeptuneMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    logger.info("setUpBefore - enter");
    allocator = new BlockAllocatorImpl();
    handler = new NeptuneMetadataHandler(glue, neptuneConnection, new LocalKeyFactory(), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "spill-bucket", "spill-prefix");
    logger.info("setUpBefore - exit");
}
Also used : BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Example 4 with LocalKeyFactory

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

the class RedisMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    logger.info("{}: enter", testName.getMethodName());
    when(mockFactory.getOrCreateConn(eq(decodedEndpoint), anyBoolean(), anyBoolean(), anyString())).thenReturn(mockConnection);
    when(mockConnection.sync()).thenReturn(mockSyncCommands);
    handler = new RedisMetadataHandler(mockGlue, new LocalKeyFactory(), mockSecretsManager, mockAthena, mockFactory, "bucket", "prefix");
    allocator = new BlockAllocatorImpl();
    when(mockSecretsManager.getSecretValue(any(GetSecretValueRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
        GetSecretValueRequest request = invocation.getArgumentAt(0, GetSecretValueRequest.class);
        if ("endpoint".equalsIgnoreCase(request.getSecretId())) {
            return new GetSecretValueResult().withSecretString(decodedEndpoint);
        }
        throw new RuntimeException("Unknown secret " + request.getSecretId());
    });
}
Also used : GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) InvocationOnMock(org.mockito.invocation.InvocationOnMock) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Example 5 with LocalKeyFactory

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

the class TimestreamMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    handler = new TimestreamMetadataHandler(mockTsQuery, mockTsMeta, mockGlue, new LocalKeyFactory(), mockSecretsManager, mockAthena, "spillBucket", "spillPrefix");
    allocator = new BlockAllocatorImpl();
}
Also used : BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Aggregations

LocalKeyFactory (com.amazonaws.athena.connector.lambda.security.LocalKeyFactory)19 Before (org.junit.Before)15 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)14 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)5 Test (org.junit.Test)4 InvocationOnMock (org.mockito.invocation.InvocationOnMock)4 QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)2 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)2 AmazonS3 (com.amazonaws.services.s3.AmazonS3)2 GetSecretValueRequest (com.amazonaws.services.secretsmanager.model.GetSecretValueRequest)2 GetSecretValueResult (com.amazonaws.services.secretsmanager.model.GetSecretValueResult)2 Field (org.apache.arrow.vector.types.pojo.Field)2 Schema (org.apache.arrow.vector.types.pojo.Schema)2 Mockito.anyString (org.mockito.Mockito.anyString)2 Block (com.amazonaws.athena.connector.lambda.data.Block)1 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)1 BlockWriter (com.amazonaws.athena.connector.lambda.data.BlockWriter)1 S3BlockSpillReader (com.amazonaws.athena.connector.lambda.data.S3BlockSpillReader)1 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)1 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)1