Examples with MetadataRequest com.amazonaws.athena.connector.lambda.metadata.MetadataRequest used on opensource projects

Example 1 with MetadataRequest

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

the class GlueMetadataHandlerTest method testGetCatalog.

@Test
public void testGetCatalog() {
    // Catalog should be the account from the request
    MetadataRequest req = new GetTableRequest(IdentityUtil.fakeIdentity(), queryId, catalog, new TableName(schema, table));
    String catalog = handler.getCatalog(req);
    assertEquals(IdentityUtil.fakeIdentity().getAccount(), catalog);
    // Catalog should be the account from the lambda context's function arn
    when(mockContext.getInvokedFunctionArn()).thenReturn("arn:aws:lambda:us-east-1:012345678912:function:athena-123");
    req.setContext(mockContext);
    catalog = handler.getCatalog(req);
    assertEquals("012345678912", catalog);
    // Catalog should be the account from the request since function arn is invalid
    when(mockContext.getInvokedFunctionArn()).thenReturn("arn:aws:lambda:us-east-1:012345678912:function:");
    req.setContext(mockContext);
    catalog = handler.getCatalog(req);
    assertEquals(IdentityUtil.fakeIdentity().getAccount(), catalog);
    // Catalog should be the account from the request since function arn is null
    when(mockContext.getInvokedFunctionArn()).thenReturn(null);
    req.setContext(mockContext);
    catalog = handler.getCatalog(req);
    assertEquals(IdentityUtil.fakeIdentity().getAccount(), catalog);
}

Example 2 with MetadataRequest

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

the class MetadataRequestSerializer method doTypedSerialize.

@Override
protected void doTypedSerialize(FederationRequest federationRequest, JsonGenerator jgen, SerializerProvider provider) throws IOException {
    MetadataRequest metadataRequest = (MetadataRequest) federationRequest;
    jgen.writeFieldName(IDENTITY_FIELD);
    identitySerializer.serialize(metadataRequest.getIdentity(), jgen, provider);
    jgen.writeStringField(QUERY_ID_FIELD, metadataRequest.getQueryId());
    jgen.writeStringField(CATALOG_NAME_FIELD, metadataRequest.getCatalogName());
    doRequestSerialize(metadataRequest, jgen, provider);
}

Example 3 with MetadataRequest

use of com.amazonaws.athena.connector.lambda.metadata.MetadataRequest 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 4 with MetadataRequest

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