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

Example 6 with PingRequest

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

the class UserDefinedFunctionHandler method handleRequest.

@Override
public final void handleRequest(InputStream inputStream, OutputStream outputStream, Context context) {
    try (BlockAllocator allocator = new BlockAllocatorImpl()) {
        ObjectMapper objectMapper = VersionedObjectMapperFactory.create(allocator);
        try (FederationRequest rawRequest = objectMapper.readValue(inputStream, FederationRequest.class)) {
            if (rawRequest instanceof PingRequest) {
                try (PingResponse response = doPing((PingRequest) rawRequest)) {
                    assertNotNull(response);
                    objectMapper.writeValue(outputStream, response);
                }
                return;
            }
            if (!(rawRequest instanceof UserDefinedFunctionRequest)) {
                throw new RuntimeException("Expected a UserDefinedFunctionRequest but found " + rawRequest.getClass());
            }
            doHandleRequest(allocator, objectMapper, (UserDefinedFunctionRequest) rawRequest, outputStream);
        } catch (Exception ex) {
            throw (ex instanceof RuntimeException) ? (RuntimeException) ex : new RuntimeException(ex);
        }
    }
}

Example 7 with PingRequest

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

the class CompositeHandlerTest method doPing.

@Test
public void doPing() throws Exception {
    PingRequest req = mock(PingRequest.class);
    when(req.getCatalogName()).thenReturn("catalog");
    when(req.getQueryId()).thenReturn("queryId");
    compositeHandler.handleRequest(allocator, req, new ByteArrayOutputStream(), objectMapper);
    verify(mockMetadataHandler, times(1)).doPing(any(PingRequest.class));
}

Example 8 with PingRequest

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

the class PingRequestSerDeTest method testForwardsCompatibility.

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

Example 9 with PingRequest

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

the class PingRequestSerDeTest method beforeTest.

@Before
public void beforeTest() throws IOException {
    expected = new PingRequest(federatedIdentity, "test-catalog", "test-query-id");
    String expectedSerDeFile = utils.getResourceOrFail("serde", "PingRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}