Examples with BlockAllocator com.amazonaws.athena.connector.lambda.data.BlockAllocator used on opensource projects

Example 1 with BlockAllocator

use of com.amazonaws.athena.connector.lambda.data.BlockAllocator 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;
    });
}

Example 2 with BlockAllocator

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

the class ObjectMapperUtil method assertBaseSerialization.

private static <T> void assertBaseSerialization(Object object) {
    Class<?> clazz = object.getClass();
    if (object instanceof FederationRequest)
        clazz = FederationRequest.class;
    else if (object instanceof FederationResponse) {
        clazz = FederationResponse.class;
    }
    try (BlockAllocator allocator = new BlockAllocatorImpl()) {
        // check SerDe write, SerDe read
        ByteArrayOutputStream serDeOut = new ByteArrayOutputStream();
        ObjectMapper serDe = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_TWO);
        serDe.writeValue(serDeOut, object);
        byte[] serDeOutput = serDeOut.toByteArray();
        assertEquals(object, serDe.readValue(new ByteArrayInputStream(serDeOutput), clazz));
    } catch (IOException | AssertionError ex) {
        throw new RuntimeException(ex);
    }
}

Example 3 with BlockAllocator

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

the class ObjectMapperUtil method assertSerializationBackwardsCompatibleV3toV2.

private static <T> void assertSerializationBackwardsCompatibleV3toV2(Object object) {
    Class<?> clazz = object.getClass();
    if (object instanceof FederationRequest)
        clazz = FederationRequest.class;
    else if (object instanceof FederationResponse) {
        clazz = FederationResponse.class;
    }
    try (BlockAllocator allocator = new BlockAllocatorImpl()) {
        // check SerDe write, SerDe read
        ByteArrayOutputStream serDeOut = new ByteArrayOutputStream();
        ObjectMapper serDeV3 = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_THREE);
        ObjectMapper serDeV2 = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_TWO);
        serDeV3.writeValue(serDeOut, object);
        byte[] serDeOutput = serDeOut.toByteArray();
        assertEquals(object, serDeV2.readValue(new ByteArrayInputStream(serDeOutput), clazz));
    } catch (IOException | AssertionError ex) {
        throw new RuntimeException(ex);
    }
}

Example 4 with BlockAllocator

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

the class ObjectMapperUtil method assertSerializationForwardsCompatibleV2toV3.

private static <T> void assertSerializationForwardsCompatibleV2toV3(Object object) {
    Class<?> clazz = object.getClass();
    if (object instanceof FederationRequest)
        clazz = FederationRequest.class;
    else if (object instanceof FederationResponse) {
        clazz = FederationResponse.class;
    }
    try (BlockAllocator allocator = new BlockAllocatorImpl()) {
        // check SerDe write, SerDe read
        ByteArrayOutputStream serDeOut = new ByteArrayOutputStream();
        ObjectMapper serDeV3 = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_THREE);
        ObjectMapper serDeV2 = VersionedObjectMapperFactory.create(allocator, SERDE_VERSION_TWO);
        serDeV2.writeValue(serDeOut, object);
        byte[] serDeOutput = serDeOut.toByteArray();
        assertEquals(object, serDeV3.readValue(new ByteArrayInputStream(serDeOutput), clazz));
    } catch (IOException | AssertionError ex) {
        throw new RuntimeException(ex);
    }
}

Example 5 with BlockAllocator

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

the class BigQueryMetadataHandlerTest method testDoGetSplits.

@Test
public void testDoGetSplits() throws Exception {
    BlockAllocator blockAllocator = new BlockAllocatorImpl();
    PowerMockito.mockStatic(BigQueryUtils.class);
    when(BigQueryUtils.getBigQueryClient()).thenReturn(bigQuery);
    when(BigQueryUtils.getEnvVar("concurrencyLimit")).thenReturn("10");
    GetSplitsRequest request = new GetSplitsRequest(federatedIdentity, QUERY_ID, CATALOG, TABLE_NAME, mock(Block.class), Collections.<String>emptyList(), new Constraints(new HashMap<>()), null);
    // added schema with integer column countCol
    List<Field> testSchemaFields = Arrays.asList(Field.of("countCol", LegacySQLTypeName.INTEGER));
    com.google.cloud.bigquery.Schema tableSchema = Schema.of(testSchemaFields);
    // mocked table row count as 15
    List<FieldValue> bigQueryRowValue = Arrays.asList(FieldValue.of(FieldValue.Attribute.PRIMITIVE, "15"));
    FieldValueList fieldValueList = FieldValueList.of(bigQueryRowValue, FieldList.of(testSchemaFields));
    List<FieldValueList> tableRows = Arrays.asList(fieldValueList);
    when(job.isDone()).thenReturn(false).thenReturn(true);
    Page<FieldValueList> pageNoSchema = new BigQueryPage<>(tableRows);
    TableResult result = new TableResult(tableSchema, tableRows.size(), pageNoSchema);
    when(job.getQueryResults()).thenReturn(result);
    QueryStatusChecker queryStatusChecker = mock(QueryStatusChecker.class);
    when(queryStatusChecker.isQueryRunning()).thenReturn(true);
    GetSplitsResponse response = bigQueryMetadataHandler.doGetSplits(blockAllocator, request);
    assertNotNull(response);
}