Examples with PingResponse com.amazonaws.athena.connector.lambda.request.PingResponse used on opensource projects

Example 1 with PingResponse

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

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

the class PingResponseSerDeTest method testForwardsCompatibility.

@Test
public void testForwardsCompatibility() throws IOException {
    logger.info("testForwardsCompatibility: enter");
    String expectedSerDeFile = utils.getResourceOrFail("serde", "PingResponseForwardsCompatible.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
    InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes());
    PingResponse actual = (PingResponse) mapper.readValue(input, FederationResponse.class);
    logger.info("testForwardsCompatibility: deserialized[{}]", actual);
    assertEquals(expected, actual);
    logger.info("testForwardsCompatibility: exit");
}

Example 3 with PingResponse

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

the class FederationServiceProvider method getService.

public static FederationService getService(String lambdaFunction, FederatedIdentity identity, String catalog) {
    FederationService service = serviceCache.get(lambdaFunction);
    if (service != null) {
        return service;
    }
    service = LambdaInvokerFactory.builder().lambdaClient(AWSLambdaClientBuilder.defaultClient()).objectMapper(VersionedObjectMapperFactory.create(BLOCK_ALLOCATOR)).lambdaFunctionNameResolver(new Mapper(lambdaFunction)).build(FederationService.class);
    PingRequest pingRequest = new PingRequest(identity, catalog, generateQueryId());
    PingResponse pingResponse = (PingResponse) service.call(pingRequest);
    int actualSerDeVersion = pingResponse.getSerDeVersion();
    log.info("SerDe version for function {}, catalog {} is {}", lambdaFunction, catalog, actualSerDeVersion);
    if (actualSerDeVersion != SERDE_VERSION) {
        service = LambdaInvokerFactory.builder().lambdaClient(AWSLambdaClientBuilder.defaultClient()).objectMapper(VersionedObjectMapperFactory.create(BLOCK_ALLOCATOR, actualSerDeVersion)).lambdaFunctionNameResolver(new Mapper(lambdaFunction)).build(FederationService.class);
    }
    serviceCache.put(lambdaFunction, service);
    return service;
}

Example 4 with PingResponse

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

the class CompositeHandler method handleRequest.

/**
 * Handles routing the request to the appropriate Handler, either MetadataHandler or RecordHandler.
 *
 * @param allocator The BlockAllocator to use for Apache Arrow Resources.
 * @param rawReq The request object itself.
 * @param outputStream The OutputStream to which all responses should be written.
 * @param objectMapper The ObjectMapper that can be used for serializing responses.
 * @throws Exception
 * @note that PingRequests are routed to the MetadataHandler even though both MetadataHandler and RecordHandler
 * implemented PingRequest handling.
 */
public final void handleRequest(BlockAllocator allocator, FederationRequest rawReq, OutputStream outputStream, ObjectMapper objectMapper) throws Exception {
    if (rawReq instanceof PingRequest) {
        try (PingResponse response = metadataHandler.doPing((PingRequest) rawReq)) {
            assertNotNull(response);
            objectMapper.writeValue(outputStream, response);
        }
        return;
    }
    if (rawReq instanceof MetadataRequest) {
        metadataHandler.doHandleRequest(allocator, objectMapper, (MetadataRequest) rawReq, outputStream);
    } else if (rawReq instanceof RecordRequest) {
        recordHandler.doHandleRequest(allocator, objectMapper, (RecordRequest) rawReq, outputStream);
    } else if (udfhandler != null && rawReq instanceof UserDefinedFunctionRequest) {
        udfhandler.doHandleRequest(allocator, objectMapper, (UserDefinedFunctionRequest) rawReq, outputStream);
    } else {
        throw new IllegalArgumentException("Unknown request class " + rawReq.getClass());
    }
}

Example 5 with PingResponse

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

the class MetadataHandler method handleRequest.

public final void handleRequest(InputStream inputStream, OutputStream outputStream, final Context context) throws IOException {
    try (BlockAllocator allocator = new BlockAllocatorImpl()) {
        ObjectMapper objectMapper = VersionedObjectMapperFactory.create(allocator);
        try (FederationRequest rawReq = objectMapper.readValue(inputStream, FederationRequest.class)) {
            if (rawReq instanceof PingRequest) {
                try (PingResponse response = doPing((PingRequest) rawReq)) {
                    assertNotNull(response);
                    objectMapper.writeValue(outputStream, response);
                }
                return;
            }
            if (!(rawReq instanceof MetadataRequest)) {
                throw new RuntimeException("Expected a MetadataRequest but found " + rawReq.getClass());
            }
            ((MetadataRequest) rawReq).setContext(context);
            doHandleRequest(allocator, objectMapper, (MetadataRequest) rawReq, outputStream);
        } catch (Exception ex) {
            logger.warn("handleRequest: Completed with an exception.", ex);
            throw (ex instanceof RuntimeException) ? (RuntimeException) ex : new RuntimeException(ex);
        }
    }
}