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);
}
}
}
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"));
}
}
Aggregations