Search in sources :

Example 1 with FederationRequest

use of com.amazonaws.athena.connector.lambda.request.FederationRequest 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);
    }
}
Also used : ByteArrayOutputStream(java.io.ByteArrayOutputStream) IOException(java.io.IOException) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) ByteArrayInputStream(java.io.ByteArrayInputStream) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) FederationResponse(com.amazonaws.athena.connector.lambda.request.FederationResponse) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest)

Example 2 with FederationRequest

use of com.amazonaws.athena.connector.lambda.request.FederationRequest 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);
    }
}
Also used : ByteArrayOutputStream(java.io.ByteArrayOutputStream) IOException(java.io.IOException) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) ByteArrayInputStream(java.io.ByteArrayInputStream) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) FederationResponse(com.amazonaws.athena.connector.lambda.request.FederationResponse) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest)

Example 3 with FederationRequest

use of com.amazonaws.athena.connector.lambda.request.FederationRequest 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);
    }
}
Also used : ByteArrayOutputStream(java.io.ByteArrayOutputStream) IOException(java.io.IOException) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) ByteArrayInputStream(java.io.ByteArrayInputStream) BlockAllocator(com.amazonaws.athena.connector.lambda.data.BlockAllocator) FederationResponse(com.amazonaws.athena.connector.lambda.request.FederationResponse) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest)

Example 4 with FederationRequest

use of com.amazonaws.athena.connector.lambda.request.FederationRequest 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);
        }
    }
}
Also used : PingRequest(com.amazonaws.athena.connector.lambda.request.PingRequest) MetadataRequest(com.amazonaws.athena.connector.lambda.metadata.MetadataRequest) 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) IOException(java.io.IOException) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest)

Example 5 with FederationRequest

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

the class RecordHandler 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 RecordRequest)) {
                throw new RuntimeException("Expected a RecordRequest but found " + rawReq.getClass());
            }
            doHandleRequest(allocator, objectMapper, (RecordRequest) rawReq, outputStream);
        } catch (Exception ex) {
            logger.warn("handleRequest: Completed with an exception.", ex);
            throw (ex instanceof RuntimeException) ? (RuntimeException) ex : new RuntimeException(ex);
        }
    }
}
Also used : RecordRequest(com.amazonaws.athena.connector.lambda.records.RecordRequest) PingRequest(com.amazonaws.athena.connector.lambda.request.PingRequest) 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) IOException(java.io.IOException) FederationRequest(com.amazonaws.athena.connector.lambda.request.FederationRequest)

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