Examples with InvokeRequest com.amazonaws.services.lambda.model.InvokeRequest used on opensource projects

Example 1 with InvokeRequest

use of com.amazonaws.services.lambda.model.InvokeRequest in project nifi by apache.

the class TestPutLambda method testPutLambdaSimple.

@Test
public void testPutLambdaSimple() {
    runner.setProperty(PutLambda.AWS_LAMBDA_FUNCTION_NAME, "test-function");
    runner.enqueue("TestContent");
    InvokeResult invokeResult = new InvokeResult();
    invokeResult.setStatusCode(200);
    invokeResult.setLogResult(Base64.encodeAsString("test-log-result".getBytes()));
    invokeResult.setPayload(ByteBuffer.wrap("test-payload".getBytes()));
    Mockito.when(mockLambdaClient.invoke(Mockito.any(InvokeRequest.class))).thenReturn(invokeResult);
    runner.assertValid();
    runner.run(1);
    ArgumentCaptor<InvokeRequest> captureRequest = ArgumentCaptor.forClass(InvokeRequest.class);
    Mockito.verify(mockLambdaClient, Mockito.times(1)).invoke(captureRequest.capture());
    InvokeRequest request = captureRequest.getValue();
    assertEquals("test-function", request.getFunctionName());
    runner.assertAllFlowFilesTransferred(PutLambda.REL_SUCCESS, 1);
    final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(PutLambda.REL_SUCCESS);
    final MockFlowFile ff0 = flowFiles.get(0);
    ff0.assertAttributeEquals(PutLambda.AWS_LAMBDA_RESULT_STATUS_CODE, "200");
    ff0.assertAttributeEquals(PutLambda.AWS_LAMBDA_RESULT_LOG, "test-log-result");
    ff0.assertAttributeEquals(PutLambda.AWS_LAMBDA_RESULT_PAYLOAD, "test-payload");
}

Example 2 with InvokeRequest

use of com.amazonaws.services.lambda.model.InvokeRequest in project aws-xray-sdk-java by aws.

the class TracingHandlerTest method testLambdaInvokeSubsegmentContainsFunctionName.

@Test
public void testLambdaInvokeSubsegmentContainsFunctionName() {
    // Setup test
    AWSLambda lambda = AWSLambdaClientBuilder.standard().withRequestHandlers(new TracingHandler()).withRegion(Regions.US_EAST_1).withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials("fake", "fake"))).build();
    // Lambda returns "null" on successful fn. with no return value
    mockHttpClient(lambda, "null");
    // Test logic
    Segment segment = AWSXRay.beginSegment("test");
    InvokeRequest request = new InvokeRequest();
    request.setFunctionName("testFunctionName");
    InvokeResult r = lambda.invoke(request);
    System.out.println(r.getStatusCode());
    System.out.println(r);
    Assert.assertEquals(1, segment.getSubsegments().size());
    Assert.assertEquals("Invoke", segment.getSubsegments().get(0).getAws().get("operation"));
    System.out.println(segment.getSubsegments().get(0).getAws());
    Assert.assertEquals("testFunctionName", segment.getSubsegments().get(0).getAws().get("function_name"));
}

Example 3 with InvokeRequest

use of com.amazonaws.services.lambda.model.InvokeRequest in project nifi by apache.

the class PutLambda method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final String functionName = context.getProperty(AWS_LAMBDA_FUNCTION_NAME).getValue();
    final String qualifier = context.getProperty(AWS_LAMBDA_FUNCTION_QUALIFIER).getValue();
    // Max size of message is 6 MB
    if (flowFile.getSize() > MAX_REQUEST_SIZE) {
        getLogger().error("Max size for request body is 6mb but was {} for flow file {} for function {}", new Object[] { flowFile.getSize(), flowFile, functionName });
        session.transfer(flowFile, REL_FAILURE);
        return;
    }
    final AWSLambdaClient client = getClient();
    try {
        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        session.exportTo(flowFile, baos);
        InvokeRequest invokeRequest = new InvokeRequest().withFunctionName(functionName).withLogType(LogType.Tail).withInvocationType(InvocationType.RequestResponse).withPayload(ByteBuffer.wrap(baos.toByteArray())).withQualifier(qualifier);
        long startTime = System.nanoTime();
        InvokeResult result = client.invoke(invokeRequest);
        flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_STATUS_CODE, result.getStatusCode().toString());
        if (!StringUtils.isBlank(result.getLogResult())) {
            flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_LOG, new String(Base64.decode(result.getLogResult()), Charset.defaultCharset()));
        }
        if (result.getPayload() != null) {
            flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_PAYLOAD, new String(result.getPayload().array(), Charset.defaultCharset()));
        }
        if (!StringUtils.isBlank(result.getFunctionError())) {
            flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_FUNCTION_ERROR, result.getFunctionError());
            session.transfer(flowFile, REL_FAILURE);
        } else {
            session.transfer(flowFile, REL_SUCCESS);
            final long totalTimeMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
            session.getProvenanceReporter().send(flowFile, functionName, totalTimeMillis);
        }
    } catch (final InvalidRequestContentException | InvalidParameterValueException | RequestTooLargeException | ResourceNotFoundException | UnsupportedMediaTypeException unrecoverableException) {
        getLogger().error("Failed to invoke lambda {} with unrecoverable exception {} for flow file {}", new Object[] { functionName, unrecoverableException, flowFile });
        flowFile = populateExceptionAttributes(session, flowFile, unrecoverableException);
        session.transfer(flowFile, REL_FAILURE);
    } catch (final TooManyRequestsException retryableServiceException) {
        getLogger().error("Failed to invoke lambda {} with exception {} for flow file {}, therefore penalizing flowfile", new Object[] { functionName, retryableServiceException, flowFile });
        flowFile = populateExceptionAttributes(session, flowFile, retryableServiceException);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
        context.yield();
    } catch (final AmazonServiceException unrecoverableServiceException) {
        getLogger().error("Failed to invoke lambda {} with exception {} for flow file {} sending to fail", new Object[] { functionName, unrecoverableServiceException, flowFile });
        flowFile = populateExceptionAttributes(session, flowFile, unrecoverableServiceException);
        session.transfer(flowFile, REL_FAILURE);
        context.yield();
    } catch (final Exception exception) {
        getLogger().error("Failed to invoke lambda {} with exception {} for flow file {}", new Object[] { functionName, exception, flowFile });
        session.transfer(flowFile, REL_FAILURE);
        context.yield();
    }
}