Examples with RecordRequest com.amazonaws.athena.connector.lambda.records.RecordRequest used on opensource projects

Example 1 with RecordRequest

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

use of com.amazonaws.athena.connector.lambda.records.RecordRequest 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);
        }
    }
}