Search in sources :

Example 1 with BlockWriter

use of com.amazonaws.athena.connector.lambda.data.BlockWriter 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 2 with BlockWriter

use of com.amazonaws.athena.connector.lambda.data.BlockWriter in project foundry-athena-query-federation-connector by palantir.

the class PartitionFetcher method getAndWritePartitions.

private void getAndWritePartitions(BlockWriter blockWriter, GetTableLayoutRequest request, QueryStatusChecker queryStatusChecker) {
    CatalogLocator locator = FoundryAthenaObjectMapper.objectMapper().convertValue(request.getSchema().getCustomMetadata(), CatalogLocator.class);
    Optional<String> pageToken = Optional.empty();
    while (queryStatusChecker.isQueryRunning()) {
        GetPartitionsResponsePage page = metadataService.getPartitions(authProvider.getAuthHeader(), GetPartitionsRequest.builder().locator(locator).limit(PAGE_SIZE).pageToken(pageToken).build());
        page.getPartitions().forEach(partition -> blockWriter.writeRows((block, rowNum) -> {
            boolean matched = partition.get().entrySet().stream().map(fieldName -> fieldName.getValue().accept(new PartitionValueWriter(block, fieldName.getKey(), rowNum))).reduce(true, Boolean::logicalAnd);
            // if all fields passed then we wrote 1 row
            return matched ? 1 : 0;
        }));
        if (page.getNextPageToken().isPresent()) {
            pageToken = page.getNextPageToken();
        } else {
            return;
        }
    }
}
Also used : CatalogLocator(com.palantir.foundry.athena.api.CatalogLocator) GetPartitionsResponsePage(com.palantir.foundry.athena.api.GetPartitionsResponsePage) BlockWriter(com.amazonaws.athena.connector.lambda.data.BlockWriter) CatalogLocator(com.palantir.foundry.athena.api.CatalogLocator) FoundryAthenaMetadataServiceBlocking(com.palantir.foundry.athena.api.FoundryAthenaMetadataServiceBlocking) GetTableLayoutRequest(com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest) Optional(java.util.Optional) QueryStatusChecker(com.amazonaws.athena.connector.lambda.QueryStatusChecker) VisibleForTesting(com.google.common.annotations.VisibleForTesting) GetPartitionsRequest(com.palantir.foundry.athena.api.GetPartitionsRequest) GetPartitionsResponsePage(com.palantir.foundry.athena.api.GetPartitionsResponsePage)

Example 3 with BlockWriter

use of com.amazonaws.athena.connector.lambda.data.BlockWriter 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

QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)3 BlockWriter (com.amazonaws.athena.connector.lambda.data.BlockWriter)3 GetTableLayoutRequest (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest)3 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)2 GetSplitsRequest (com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest)2 GetSplitsResponse (com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse)2 Before (org.junit.Before)2 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)1 GetTableLayoutResponse (com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse)1 FederatedIdentity (com.amazonaws.athena.connector.lambda.security.FederatedIdentity)1 LocalKeyFactory (com.amazonaws.athena.connector.lambda.security.LocalKeyFactory)1 DatabaseConnectionConfig (com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig)1 JdbcConnectionFactory (com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory)1 JdbcCredentialProvider (com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider)1 AmazonAthena (com.amazonaws.services.athena.AmazonAthena)1 GetTablesRequest (com.amazonaws.services.glue.model.GetTablesRequest)1 GetTablesResult (com.amazonaws.services.glue.model.GetTablesResult)1 Table (com.amazonaws.services.glue.model.Table)1 AWSSecretsManager (com.amazonaws.services.secretsmanager.AWSSecretsManager)1 GetSecretValueRequest (com.amazonaws.services.secretsmanager.model.GetSecretValueRequest)1