Search in sources :

Example 6 with FederationRequest

use of com.amazonaws.athena.connector.lambda.request.FederationRequest 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);
        }
    }
}
Also used : PingRequest(com.amazonaws.athena.connector.lambda.request.PingRequest) UserDefinedFunctionRequest(com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) PingResponse(com.amazonaws.athena.connector.lambda.request.PingResponse) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) InvocationTargetException(java.lang.reflect.InvocationTargetException) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest)

Example 7 with FederationRequest

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

the class UserDefinedFunctionHandlerTest method testRequestTypeValidation.

@Test
public void testRequestTypeValidation() throws Exception {
    FederationRequest federationRequest = new ListSchemasRequest(null, "dummy_catalog", "dummy_qid");
    ObjectMapper objectMapper = VersionedObjectMapperFactory.create(allocator);
    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
    objectMapper.writeValue(byteArrayOutputStream, federationRequest);
    byte[] inputData = byteArrayOutputStream.toByteArray();
    ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(inputData);
    ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    try {
        handler.handleRequest(byteArrayInputStream, outputStream, null);
        fail();
    } catch (Exception e) {
        assertTrue(e.getMessage().contains("Expected a UserDefinedFunctionRequest but found"));
    }
}
Also used : ListSchemasRequest(com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest) ByteArrayInputStream(java.io.ByteArrayInputStream) ByteArrayOutputStream(java.io.ByteArrayOutputStream) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest) Test(org.junit.Test)

Aggregations

FederationRequest (com.amazonaws.athena.connector.lambda.request.FederationRequest)7 ObjectMapper (com.fasterxml.jackson.databind.ObjectMapper)7 BlockAllocator (com.amazonaws.athena.connector.lambda.data.BlockAllocator)6 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)6 IOException (java.io.IOException)5 ByteArrayInputStream (java.io.ByteArrayInputStream)4 ByteArrayOutputStream (java.io.ByteArrayOutputStream)4 FederationResponse (com.amazonaws.athena.connector.lambda.request.FederationResponse)3 PingRequest (com.amazonaws.athena.connector.lambda.request.PingRequest)3 PingResponse (com.amazonaws.athena.connector.lambda.request.PingResponse)3 ListSchemasRequest (com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest)1 MetadataRequest (com.amazonaws.athena.connector.lambda.metadata.MetadataRequest)1 RecordRequest (com.amazonaws.athena.connector.lambda.records.RecordRequest)1 UserDefinedFunctionRequest (com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest)1 InvocationTargetException (java.lang.reflect.InvocationTargetException)1 Test (org.junit.Test)1