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

Example 1 with BlockAllocatorImpl

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

Example 2 with BlockAllocatorImpl

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

the class ExampleUserDefinedFunctionHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("setUpBefore - enter");
    this.exampleUserDefinedFunctionHandler = new ExampleUserDefinedFunctionHandler();
    this.allocator = new BlockAllocatorImpl();
    this.mapper = VersionedObjectMapperFactory.create(allocator);
}

Example 3 with BlockAllocatorImpl

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

the class CompositeHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    logger.info("{}: enter", testName.getMethodName());
    allocator = new BlockAllocatorImpl();
    objectMapper = ObjectMapperFactory.create(allocator);
    mockMetadataHandler = mock(MetadataHandler.class);
    mockRecordHandler = mock(RecordHandler.class);
    schemaForRead = SchemaBuilder.newBuilder().addField("col1", new ArrowType.Int(32, true)).build();
    when(mockMetadataHandler.doGetTableLayout(any(BlockAllocatorImpl.class), any(GetTableLayoutRequest.class))).thenReturn(new GetTableLayoutResponse("catalog", new TableName("schema", "table"), BlockUtils.newBlock(allocator, "col1", Types.MinorType.BIGINT.getType(), 1L)));
    when(mockMetadataHandler.doListTables(any(BlockAllocatorImpl.class), any(ListTablesRequest.class))).thenReturn(new ListTablesResponse("catalog", Collections.singletonList(new TableName("schema", "table")), null));
    when(mockMetadataHandler.doGetTable(any(BlockAllocatorImpl.class), any(GetTableRequest.class))).thenReturn(new GetTableResponse("catalog", new TableName("schema", "table"), SchemaBuilder.newBuilder().addStringField("col1").build()));
    when(mockMetadataHandler.doListSchemaNames(any(BlockAllocatorImpl.class), any(ListSchemasRequest.class))).thenReturn(new ListSchemasResponse("catalog", Collections.singleton("schema1")));
    when(mockMetadataHandler.doGetSplits(any(BlockAllocatorImpl.class), any(GetSplitsRequest.class))).thenReturn(new GetSplitsResponse("catalog", Split.newBuilder(null, null).build()));
    when(mockMetadataHandler.doPing(any(PingRequest.class))).thenReturn(new PingResponse("catalog", "queryId", "type", 23, 2));
    when(mockRecordHandler.doReadRecords(any(BlockAllocatorImpl.class), any(ReadRecordsRequest.class))).thenReturn(new ReadRecordsResponse("catalog", BlockUtils.newEmptyBlock(allocator, "col", new ArrowType.Int(32, true))));
    compositeHandler = new CompositeHandler(mockMetadataHandler, mockRecordHandler);
}

Example 4 with BlockAllocatorImpl

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

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