com.amazonaws.services.lambda.model.InvokeResult Java Examples

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

@Test
public void testInvokeLambdaFunctionAsynchronousError() throws Exception {
    InvokeConfig invokeConfig = new InvokeConfig("function", "{\"key1\": \"value1\"}", false, null);

    when(awsLambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(new InvokeResult().withFunctionError("Handled"));

    try {
        lambdaInvokeService.invokeLambdaFunction(invokeConfig);
        fail("Should fail with LambdaInvokeException");
    } catch (LambdaInvokeException lie){
        assertEquals("Function returned error of type: Handled", lie.getMessage());
    }

    verify(awsLambdaClient, times(1)).invoke(invokeRequestArg.capture());
    InvokeRequest invokeRequest = invokeRequestArg.getValue();
    assertEquals("function", invokeRequest.getFunctionName());
    assertEquals("{\"key1\": \"value1\"}", new String(invokeRequest.getPayload().array(), Charset.forName("UTF-8")));
    assertEquals(InvocationType.Event.toString(), invokeRequest.getInvocationType());
    verify(jenkinsLogger).log(eq("Lambda invoke request:%n%s%nPayload:%n%s%n"), anyVararg());
    verify(jenkinsLogger).log(eq("Lambda invoke response:%n%s%nPayload:%n%s%n"), anyVararg());
}
 

private boolean invokeRule(AWSLambda awsLambdaClient, Rule ruleDetails, String invocationId, List<Map<String, Object>> additionalRuleParams) {
	String ruleParams = ruleDetails.getRuleParams();
	if(invocationId != null) {
		Map<String, Object> ruleParamDetails;
		try {
			ruleParamDetails = mapper.readValue(ruleDetails.getRuleParams(), new TypeReference<Map<String, Object>>(){});
			ruleParamDetails.put("invocationId", invocationId);
			ruleParamDetails.put("additionalParams", mapper.writeValueAsString(additionalRuleParams));
			ruleParams = mapper.writeValueAsString(ruleParamDetails);
		} catch (Exception exception) {
			log.error(UNEXPECTED_ERROR_OCCURRED, exception);
		} 
	}
	String functionName = config.getRule().getLambda().getFunctionName();
	ByteBuffer payload = ByteBuffer.wrap(ruleParams.getBytes());
	InvokeRequest invokeRequest = new InvokeRequest().withFunctionName(functionName).withPayload(payload);
	InvokeResult invokeResult = awsLambdaClient.invoke(invokeRequest);
	if (invokeResult.getStatusCode() == 200) {
		return true;
	} else {
		return false;
	}
}
 

@Before
  public void setUp() throws Exception{
      awsLambdaClient = mock(AWSLambdaClient.class);
      amazonCloudWatchEvents = mock(AmazonCloudWatchEvents.class);
      invokeResult = mock(InvokeResult.class);
      PowerMockito.whenNew(ObjectMapper.class).withNoArguments().thenReturn(mapper);
      invokeRequest = Mockito.spy(new InvokeRequest());
      putRuleRequest = Mockito.spy(new PutRuleRequest());
      putRuleResult = Mockito.spy(new PutRuleResult());
      putTargetsResult = Mockito.spy(new PutTargetsResult());
      putTargetsRequest = Mockito.spy(new PutTargetsRequest());
      AWSCredentials awsCredentialsProperty = buildAWSCredentials();
when(config.getAws()).thenReturn(awsCredentialsProperty);
      PowerMockito.whenNew(AWSLambdaClient.class).withAnyArguments().thenReturn(awsLambdaClient);  
      when(amazonClient.getAWSLambdaClient(anyString())).thenReturn(awsLambdaClient);
when(amazonClient.getAmazonCloudWatchEvents(anyString())).thenReturn(amazonCloudWatchEvents);
  }
 

@Before
  public void setUp() throws Exception{
      awsLambdaClient = mock(AWSLambdaClient.class);
      amazonCloudWatchEvents = mock(AmazonCloudWatchEvents.class);
     // invokeRequest = mock(InvokeRequest.class);
      invokeResult = mock(InvokeResult.class);
      PowerMockito.whenNew(ObjectMapper.class).withNoArguments().thenReturn(mapper);
      invokeRequest = Mockito.spy(new InvokeRequest());
      getPolicyResult = Mockito.spy(new GetPolicyResult());
      putRuleRequest = Mockito.spy(new PutRuleRequest());
      putRuleResult = Mockito.spy(new PutRuleResult());
      putTargetsResult = Mockito.spy(new PutTargetsResult());
      putTargetsRequest = Mockito.spy(new PutTargetsRequest());
      AWSCredentials awsCredentialsProperty = buildAWSCredentials();
when(config.getAws()).thenReturn(awsCredentialsProperty);
      PowerMockito.whenNew(AWSLambdaClient.class).withAnyArguments().thenReturn(awsLambdaClient);  
      when(amazonClient.getAWSLambdaClient(anyString())).thenReturn(awsLambdaClient);
when(amazonClient.getAmazonCloudWatchEvents(anyString())).thenReturn(amazonCloudWatchEvents);
  }
 

@Test
public void testInvokeLambdaFunctionAsynchronous() throws Exception {
    InvokeConfig invokeConfig = new InvokeConfig("function", "{\"key1\": \"value1\"}", false, null);

    when(awsLambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(new InvokeResult());

    String result = lambdaInvokeService.invokeLambdaFunction(invokeConfig);

    verify(awsLambdaClient, times(1)).invoke(invokeRequestArg.capture());
    InvokeRequest invokeRequest = invokeRequestArg.getValue();
    assertEquals("function", invokeRequest.getFunctionName());
    assertEquals("{\"key1\": \"value1\"}", new String(invokeRequest.getPayload().array(), Charset.forName("UTF-8")));
    assertEquals(InvocationType.Event.toString(), invokeRequest.getInvocationType());
    verify(jenkinsLogger).log(eq("Lambda invoke request:%n%s%nPayload:%n%s%n"), anyVararg());
    verify(jenkinsLogger).log(eq("Lambda invoke response:%n%s%nPayload:%n%s%n"), anyVararg());
    assertEquals("", result);
}
 

public String translate(final String testPhrase, final String language) {
    final Map<String, Slot> slots = new HashMap<>();
    slots.put("termA", Slot.builder().withName("termA").withValue(testPhrase).build());
    slots.put("termB", Slot.builder().withName("termB").build());
    slots.put("language", Slot.builder().withName("language").withValue(language).build());
    final SpeechletRequestEnvelope envelope = givenIntentSpeechletRequestEnvelope("Translate", slots);
    final ObjectMapper mapper = new ObjectMapper();
    String response = null;
    try {
        final AWSLambdaClient awsLambda = new AWSLambdaClient();
        final InvokeRequest invokeRequest = new InvokeRequest()
                .withInvocationType(InvocationType.RequestResponse)
                .withFunctionName(lambdaName)
                .withPayload(mapper.writeValueAsString(envelope));
        final InvokeResult invokeResult = awsLambda.invoke(invokeRequest);
        response = new String(invokeResult.getPayload().array());
    } catch (JsonProcessingException e) {
        log.error(e.getMessage());
    }
    return response;
}
 

/**
 * Synchronously or asynchronously invokes an AWS Lambda function.
 * If synchronously invoked, the AWS Lambda log is collected and the response payload is returned
 * @param invokeConfig AWS Lambda invocation configuration
 * @return response payload
 */
public String invokeLambdaFunction(InvokeConfig invokeConfig) throws LambdaInvokeException {
    InvokeRequest invokeRequest = new InvokeRequest()
            .withFunctionName(invokeConfig.getFunctionName())
            .withPayload(invokeConfig.getPayload());

    if(invokeConfig.isSynchronous()){
        invokeRequest
                .withInvocationType(InvocationType.RequestResponse)
                .withLogType(LogType.Tail);
    } else {
        invokeRequest
                .withInvocationType(InvocationType.Event);
    }
    logger.log("Lambda invoke request:%n%s%nPayload:%n%s%n", invokeRequest.toString(), invokeConfig.getPayload());

    InvokeResult invokeResult = client.invoke(invokeRequest);
    String payload = "";
    if(invokeResult.getPayload() != null){
        payload = new String(invokeResult.getPayload().array(), Charset.forName("UTF-8"));
    }
    logger.log("Lambda invoke response:%n%s%nPayload:%n%s%n", invokeResult.toString(), payload);

    if(invokeResult.getLogResult() != null){
        logger.log("Log:%n%s%n", new String(Base64.decode(invokeResult.getLogResult()), Charset.forName("UTF-8")));
    }

    if(StringUtils.isNotEmpty(invokeResult.getFunctionError())){
        throw new LambdaInvokeException("Function returned error of type: " + invokeResult.getFunctionError());
    }

    return payload;
}
 

/**
 * mediate to AWS Lambda
 * @param messageContext - contains the payload
 * @return true
 */
public boolean mediate(MessageContext messageContext) {
    org.apache.axis2.context.MessageContext axis2MessageContext = ((Axis2MessageContext) messageContext)
            .getAxis2MessageContext();

    String payload;
    if (JsonUtil.hasAJsonPayload(axis2MessageContext)) {
        payload = JsonUtil.jsonPayloadToString(axis2MessageContext);
    } else {
        payload = "{}";
    }

    InvokeResult invokeResult = invokeLambda(payload);

    if (invokeResult != null) {
        if (log.isDebugEnabled()) {
            log.debug("AWS Lambda function: " + resourceName + " is invoked successfully.");
        }
        JsonUtil.setJsonStream(axis2MessageContext, new ByteArrayInputStream(invokeResult.getPayload().array()));
        axis2MessageContext.setProperty(APIMgtGatewayConstants.HTTP_SC, invokeResult.getStatusCode());
        axis2MessageContext.setProperty(APIMgtGatewayConstants.REST_MESSAGE_TYPE, APIConstants.APPLICATION_JSON_MEDIA_TYPE);
        axis2MessageContext.setProperty(APIMgtGatewayConstants.REST_CONTENT_TYPE, APIConstants.APPLICATION_JSON_MEDIA_TYPE);
        axis2MessageContext.removeProperty(APIConstants.NO_ENTITY_BODY);
    } else {
        if (log.isDebugEnabled()) {
            log.debug("Failed to invoke AWS Lambda function: " + resourceName);
        }
        axis2MessageContext.setProperty(APIMgtGatewayConstants.HTTP_SC, APIMgtGatewayConstants.HTTP_SC_CODE);
        axis2MessageContext.setProperty(APIConstants.NO_ENTITY_BODY, true);
    }

    return true;
}
 

@Test
public void testInvokeLambdaFunctionSynchronous() throws Exception {
    final String logBase64 = "bGFtYmRh";
    final String log = "lambda";
    final String requestPayload = "{\"key1\": \"value1\"}";
    final String responsePayload = "{\"key2\": \"value2\"}";

    InvokeConfig invokeConfig = new InvokeConfig("function", requestPayload, true, null);

    InvokeResult invokeResult = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload.getBytes()));

    when(awsLambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult);

    String result = lambdaInvokeService.invokeLambdaFunction(invokeConfig);

    verify(awsLambdaClient, times(1)).invoke(invokeRequestArg.capture());
    InvokeRequest invokeRequest = invokeRequestArg.getValue();
    assertEquals("function", invokeRequest.getFunctionName());
    assertEquals(LogType.Tail.toString(), invokeRequest.getLogType());
    assertEquals(requestPayload, new String(invokeRequest.getPayload().array(), Charset.forName("UTF-8")));
    assertEquals(InvocationType.RequestResponse.toString(), invokeRequest.getInvocationType());

    ArgumentCaptor<String> stringArgs = ArgumentCaptor.forClass(String.class);
    verify(jenkinsLogger).log(eq("Lambda invoke request:%n%s%nPayload:%n%s%n"), stringArgs.capture());
    List<String> stringArgValues = stringArgs.getAllValues();
    assertEquals(Arrays.asList(invokeRequest.toString(), requestPayload), stringArgValues);

    verify(jenkinsLogger).log(eq("Log:%n%s%n"), eq(log));

    stringArgs = ArgumentCaptor.forClass(String.class);
    verify(jenkinsLogger).log(eq("Lambda invoke response:%n%s%nPayload:%n%s%n"), stringArgs.capture());
    stringArgValues = stringArgs.getAllValues();
    assertEquals(Arrays.asList(invokeResult.toString(), responsePayload), stringArgValues);

    assertEquals(responsePayload, result);
}
 

@Test
public void testPerformFailure() throws IOException, ExecutionException, InterruptedException {
    List<JsonParameterVariables> jsonParameterVariables = new ArrayList<JsonParameterVariables>();
    jsonParameterVariables.add(new JsonParameterVariables("KEY", "$.key2"));
    LambdaInvokeBuildStepVariables clone = new LambdaInvokeBuildStepVariables(false, "accessKeyId", Secret.fromString("secretKey"), "eu-west-1", "function", "payload", true, jsonParameterVariables);
    LambdaInvokeBuildStepVariables spy = Mockito.spy(clone);

    when(original.getClone()).thenReturn(spy);
    when(spy.getLambdaClientConfig()).thenReturn(clientConfig);
    when(clientConfig.getClient()).thenReturn(lambdaClient);
    final String logBase64 = "bGFtYmRh";
    final String responsePayload = "{\"errorMessage\":\"event_fail\"}";

    InvokeResult invokeResult = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload.getBytes()))
            .withFunctionError("Unhandled");

    when(lambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult);

    FreeStyleProject p = j.createFreeStyleProject();
    p.getBuildersList().add(new LambdaInvokeBuildStep(original));
    FreeStyleBuild build = p.scheduleBuild2(0).get();
    EnvVars environment = build.getEnvironment(new LogTaskListener(LOGGER, Level.INFO));

    assertEquals(null, environment.get("KEY"));
    assertEquals(Result.FAILURE, build.getResult());
}
 

@Test
public void testPerform() throws IOException, ExecutionException, InterruptedException {
    List<JsonParameterVariables> jsonParameterVariables = new ArrayList<JsonParameterVariables>();
    jsonParameterVariables.add(new JsonParameterVariables("KEY", "$.key2"));
    LambdaInvokeBuildStepVariables clone = new LambdaInvokeBuildStepVariables(false, "accessKeyId", Secret.fromString("secretKey"), "eu-west-1", "function", "payload", true, jsonParameterVariables);
    LambdaInvokeBuildStepVariables spy = Mockito.spy(clone);

    when(original.getClone()).thenReturn(spy);
    when(spy.getLambdaClientConfig()).thenReturn(clientConfig);
    when(clientConfig.getClient()).thenReturn(lambdaClient);
    final String logBase64 = "bGFtYmRh";
    final String responsePayload = "{\"key2\": \"value2\"}";

    InvokeResult invokeResult = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload.getBytes()));

    when(lambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult);

    FreeStyleProject p = j.createFreeStyleProject();
    p.getBuildersList().add(new LambdaInvokeBuildStep(original));
    FreeStyleBuild build = p.scheduleBuild2(0).get();
    EnvVars environment = build.getEnvironment(new LogTaskListener(LOGGER, Level.INFO));

    assertEquals("value2", environment.get("KEY"));
    assertEquals(Result.SUCCESS, build.getResult());
}
 

@Test
public void testPerformInstanceRole() throws IOException, ExecutionException, InterruptedException {
    List<JsonParameterVariables> jsonParameterVariables = new ArrayList<>();
    jsonParameterVariables.add(new JsonParameterVariables("KEY", "$.key2"));
    LambdaInvokeVariables clone = new LambdaInvokeVariables(true, "accessKeyId", Secret.fromString("secretKey"), "eu-west-1", "function", "payload", true, true, jsonParameterVariables);
    LambdaInvokeVariables spy = Mockito.spy(clone);

    when(original.getClone()).thenReturn(spy);
    when(spy.getLambdaClientConfig()).thenReturn(clientConfig);
    when(clientConfig.getClient()).thenReturn(lambdaClient);
    final String logBase64 = "bGFtYmRh";
    final String responsePayload = "{\"key2\": \"value2\"}";

    InvokeResult invokeResult = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload.getBytes()));

    when(lambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult);

    FreeStyleProject p = j.createFreeStyleProject();
    p.getPublishersList().add(new LambdaInvokePublisher(Arrays.asList(original, original)));
    FreeStyleBuild build = p.scheduleBuild2(0).get();
    EnvVars environment = build.getEnvironment(new LogTaskListener(LOGGER, Level.INFO));

    assertEquals("value2", environment.get("KEY"));
    assertEquals(Result.SUCCESS, build.getResult());
}
 

@Test
public void testPerform() throws IOException, ExecutionException, InterruptedException {
    List<JsonParameterVariables> jsonParameterVariables = new ArrayList<>();
    jsonParameterVariables.add(new JsonParameterVariables("KEY", "$.key2"));
    LambdaInvokeVariables clone = new LambdaInvokeVariables(false, "accessKeyId", Secret.fromString("secretKey"), "eu-west-1", "function", "payload", true, true, jsonParameterVariables);
    LambdaInvokeVariables spy = Mockito.spy(clone);

    when(original.getClone()).thenReturn(spy);
    when(spy.getLambdaClientConfig()).thenReturn(clientConfig);
    when(clientConfig.getClient()).thenReturn(lambdaClient);
    final String logBase64 = "bGFtYmRh";
    final String responsePayload = "{\"key2\": \"value2\"}";

    InvokeResult invokeResult = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload.getBytes()));

    when(lambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult);

    FreeStyleProject p = j.createFreeStyleProject();
    p.getPublishersList().add(new LambdaInvokePublisher(Arrays.asList(original, original)));
    FreeStyleBuild build = p.scheduleBuild2(0).get();
    EnvVars environment = build.getEnvironment(new LogTaskListener(LOGGER, Level.INFO));

    assertEquals("value2", environment.get("KEY"));
    assertEquals(Result.SUCCESS, build.getResult());
}
 

private Object decodeResult(FunctionDefinition definition, Argument<O> outputType, InvokeResult invokeResult) {
    Integer statusCode = invokeResult.getStatusCode();
    if (statusCode >= STATUS_CODE_ERROR) {
        throw new FunctionExecutionException("Error executing AWS Lambda [" + definition.getName() + "]: " + invokeResult.getFunctionError());
    }
    io.micronaut.core.io.buffer.ByteBuffer byteBuffer = byteBufferFactory.copiedBuffer(invokeResult.getPayload());

    return jsonMediaTypeCodec.decode(outputType, byteBuffer);
}
 

public Optional<AlexaResponse> fire(AlexaRequest request, String payload) {
    final InvocationType invocationType = request.expectsResponse() ? InvocationType.RequestResponse : InvocationType.Event;
    final InvokeRequest invokeRequest = new InvokeRequest()
            .withInvocationType(invocationType)
            .withFunctionName(lambdaFunctionName)
            .withPayload(payload);
    log.info(String.format("->[INFO] Invoke lambda function '%s'.", lambdaFunctionName));
    log.debug(String.format("->[INFO] with request payload '%s'.", payload));
    final InvokeResult invokeResult = lambdaClient.invoke(invokeRequest);
    return invocationType.equals(InvocationType.RequestResponse) ?
            Optional.of(new AlexaResponse(request, payload, new String(invokeResult.getPayload().array()))) : Optional.empty();
}
 

@Override
protected Object run() throws Exception {
	TaskListener listener = this.getContext().get(TaskListener.class);
	AWSLambda client = AWSClientFactory.create(AWSLambdaClientBuilder.standard(), this.getContext());

	String functionName = this.step.getFunctionName();

	listener.getLogger().format("Invoke Lambda function %s%n", functionName);

	InvokeRequest request = new InvokeRequest();
	request.withFunctionName(functionName);
	request.withPayload(this.step.getPayloadAsString());
	request.withLogType(LogType.Tail);

	InvokeResult result = client.invoke(request);

	listener.getLogger().append(this.getLogResult(result));
	String functionError = result.getFunctionError();
	if (functionError != null) {
		throw new RuntimeException("Invoke lambda failed! " + this.getPayloadAsString(result));
	}
	if (this.step.isReturnValueAsString()) {
		return this.getPayloadAsString(result);
	} else {
		return JsonUtils.fromString(this.getPayloadAsString(result));
	}
}
 

private void invokeRule(AWSLambda awsLambdaClient, String lambdaFunctionName, String params) {
	InvokeRequest invokeRequest = new InvokeRequest().withFunctionName(lambdaFunctionName).withPayload(ByteBuffer.wrap(params.getBytes()));
	InvokeResult invokeResult = awsLambdaClient.invoke(invokeRequest);
	if (invokeResult.getStatusCode() == 200) {
		invokeResult.getPayload();
	} else {
		log.error("Received a non-OK response from AWS: "+invokeResult.getStatusCode());
	}
}
 

/**
 * Invoke rule.
 *
 * @param ruleId the rule id
 * @param runTimeParams the run time params
 * @return true, if successful
 */
@SuppressWarnings("unchecked")
private boolean invokeRule(final String ruleId,
        Map<String, String> runTimeParams) {
    RuleInstance ruleInstance = ruleInstanceService
            .getRuleInstanceByRuleId(ruleId);
    String ruleParams = ruleInstance.getRuleParams();
    Map<String, Object> ruleParamDetails = (Map<String, Object>) CommonUtil
            .deSerializeToObject(ruleParams);
    if (runTimeParams != null) {
        ruleParamDetails.put(additionalParams,
                formatAdditionalParameters(runTimeParams));
    }
    ruleParams = CommonUtil.serializeToString(ruleParamDetails);
    AWSLambdaAsyncClient awsLambdaClient = getAWSLambdaAsyncClient();
    InvokeRequest invokeRequest = new InvokeRequest().withFunctionName(
            ruleLambdaFunctionName).withPayload(
            ByteBuffer.wrap(ruleParams.getBytes()));
    InvokeResult invokeResult = awsLambdaClient.invoke(invokeRequest);
    if (invokeResult.getStatusCode() == TWO_HUNDRED) {
        ByteBuffer responsePayload = invokeResult.getPayload();
        log.error("Return Value :" + new String(responsePayload.array()));
        return true;
    } else {
        log.error("Received a non-OK response from AWS: "
                + invokeResult.getStatusCode());
        return false;
    }
}
 

/**
 * Invokes the remote lambda function and returns the decompressed response as bytes.
 */
@Override
protected void invoke(final Marker marker, byte[] bytes, boolean fireAndForget, final Handler<AsyncResult<byte[]>> callback) {
  final RemoteFunctionConfig remoteFunction = getConnectorConfig().remoteFunction;
  logger.debug(marker, "Invoking remote lambda function with id '{}' Event size is: {}", remoteFunction.id, bytes.length);

  InvokeRequest invokeReq = new InvokeRequest()
      .withFunctionName(((AWSLambda) remoteFunction).lambdaARN)
      .withPayload(ByteBuffer.wrap(bytes))
      .withInvocationType(fireAndForget ? InvocationType.Event : InvocationType.RequestResponse);

  asyncClient.invokeAsync(invokeReq, new AsyncHandler<InvokeRequest, InvokeResult>() {
    @Override
    public void onError(Exception exception) {
      if (callback == null) {
        logger.error(marker, "Error sending event to remote lambda function", exception);
      }
      else {
        callback.handle(Future.failedFuture(getWHttpException(marker, exception)));
      }
    }

    @Override
    public void onSuccess(InvokeRequest request, InvokeResult result) {
      byte[] responseBytes = new byte[result.getPayload().remaining()];
      result.getPayload().get(responseBytes);
      callback.handle(Future.succeededFuture(responseBytes));
    }
  });
}
 

public static void main(String[] args) {

         /*
        Function names appear as arn:aws:lambda:us-west-2:335556330391:function:HelloFunction
        you can retrieve the value by looking at the function in the AWS Console
         */
        if (args.length < 1) {
            System.out.println("Please specify a function name");
            System.exit(1);
        }

        // snippet-start:[lambda.java1.invoke.main]
        String functionName = args[0];

        InvokeRequest invokeRequest = new InvokeRequest()
                .withFunctionName(functionName)
                .withPayload("{\n" +
                        " \"Hello \": \"Paris\",\n" +
                        " \"countryCode\": \"FR\"\n" +
                        "}");
        InvokeResult invokeResult = null;

        try {
            AWSLambda awsLambda = AWSLambdaClientBuilder.standard()
                    .withCredentials(new ProfileCredentialsProvider())
                    .withRegion(Regions.US_WEST_2).build();

            invokeResult = awsLambda.invoke(invokeRequest);

            String ans = new String(invokeResult.getPayload().array(), StandardCharsets.UTF_8);

            //write out the return value
            System.out.println(ans);

        } catch (ServiceException e) {
            System.out.println(e);
        }

        System.out.println(invokeResult.getStatusCode());
        // snippet-end:[lambda.java1.invoke.main]
    }
 

private String getLogResult(InvokeResult result) {
	return new String(Base64.getDecoder().decode(result.getLogResult()), StandardCharsets.UTF_8);
}
 

@Test
public void testPerformFailure() throws IOException, ExecutionException, InterruptedException {
    List<JsonParameterVariables> jsonParameterVariables = new ArrayList<>();
    jsonParameterVariables.add(new JsonParameterVariables("KEY", "$.key2"));
    LambdaInvokeVariables clone = new LambdaInvokeVariables(false, "accessKeyId", Secret.fromString("secretKey"), "eu-west-1", "function", "payload", true, true, jsonParameterVariables);
    LambdaInvokeVariables spy = Mockito.spy(clone);

    when(original.getClone()).thenReturn(spy);
    when(spy.getLambdaClientConfig()).thenReturn(clientConfig);
    when(clientConfig.getClient()).thenReturn(lambdaClient);
    final String logBase64 = "bGFtYmRh";
    final String responsePayload = "{\"key2\": \"value2\"}";

    InvokeResult invokeResult = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload.getBytes()));

    when(lambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult);


    LambdaInvokeVariables spy2 = Mockito.spy(clone);
    when(original2.getClone()).thenReturn(spy2);
    when(spy2.getLambdaClientConfig()).thenReturn(clientConfig2);
    when(clientConfig2.getClient()).thenReturn(lambdaClient2);
    final String responsePayload2 = "{\"errorMessage\":\"event_fail\"}";

    InvokeResult invokeResult2 = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload2.getBytes()))
            .withFunctionError("Unhandled");

    when(lambdaClient2.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult2);


    FreeStyleProject p = j.createFreeStyleProject();
    p.getPublishersList().add(new LambdaInvokePublisher(Arrays.asList(original, original2)));
    FreeStyleBuild build = p.scheduleBuild2(0).get();
    EnvVars environment = build.getEnvironment(new LogTaskListener(LOGGER, Level.INFO));

    assertEquals("value2", environment.get("KEY"));
    assertEquals(Result.FAILURE, build.getResult());
}
 

private String getPayloadAsString(InvokeResult result) {
	return new String(result.getPayload().array(), StandardCharsets.UTF_8);
}
 

@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();
    }
}
 

@Test
public void testInvokeLambdaFunctionSynchronousError() throws Exception {
    final String logBase64 = "bGFtYmRh";
    final String log = "lambda";
    final String requestPayload = "{\"key1\": \"value1\"}";
    final String responsePayload = "{\"errorMessage\":\"event_fail\"}";

    InvokeConfig invokeConfig = new InvokeConfig("function", requestPayload, true, null);

    InvokeResult invokeResult = new InvokeResult()
            .withLogResult(logBase64)
            .withPayload(ByteBuffer.wrap(responsePayload.getBytes()))
            .withFunctionError("Unhandled");

    when(awsLambdaClient.invoke(any(InvokeRequest.class)))
            .thenReturn(invokeResult);

    try {
        lambdaInvokeService.invokeLambdaFunction(invokeConfig);
        fail("Should fail with LambdaInvokeException");
    } catch (LambdaInvokeException lie){
        assertEquals("Function returned error of type: Unhandled", lie.getMessage());
    }

    verify(awsLambdaClient, times(1)).invoke(invokeRequestArg.capture());
    InvokeRequest invokeRequest = invokeRequestArg.getValue();
    assertEquals("function", invokeRequest.getFunctionName());
    assertEquals(LogType.Tail.toString(), invokeRequest.getLogType());
    assertEquals(requestPayload, new String(invokeRequest.getPayload().array(), Charset.forName("UTF-8")));
    assertEquals(InvocationType.RequestResponse.toString(), invokeRequest.getInvocationType());

    ArgumentCaptor<String> stringArgs = ArgumentCaptor.forClass(String.class);
    verify(jenkinsLogger).log(eq("Lambda invoke request:%n%s%nPayload:%n%s%n"), stringArgs.capture());
    List<String> stringArgValues = stringArgs.getAllValues();
    assertEquals(Arrays.asList(invokeRequest.toString(), requestPayload), stringArgValues);

    verify(jenkinsLogger).log(eq("Log:%n%s%n"), eq(log));

    stringArgs = ArgumentCaptor.forClass(String.class);
    verify(jenkinsLogger).log(eq("Lambda invoke response:%n%s%nPayload:%n%s%n"), stringArgs.capture());
    stringArgValues = stringArgs.getAllValues();
    assertEquals(Arrays.asList(invokeResult.toString(), responsePayload), stringArgValues);
}
 

/**
 * Executes a lambda function and returns the result of the execution.
 */
@Override
public cfData execute( cfSession _session, cfArgStructData argStruct ) throws cfmRunTimeException {

	AmazonKey amazonKey = getAmazonKey( _session, argStruct );

	// Arguments to extract
	String payload = getNamedStringParam( argStruct, "payload", null );
	String functionName = getNamedStringParam( argStruct, "function", null );
	String qualifier = getNamedStringParam( argStruct, "qualifier", null );

	try {

		// Construct the Lambda Client
		InvokeRequest invokeRequest = new InvokeRequest();
		invokeRequest.setInvocationType( InvocationType.RequestResponse );
		invokeRequest.setLogType( LogType.Tail );
		invokeRequest.setFunctionName( functionName );
		invokeRequest.setPayload( payload );
		if ( qualifier != null ) {
			invokeRequest.setQualifier( qualifier );
		}

		// Lambda client must be created with credentials
		BasicAWSCredentials awsCreds = new BasicAWSCredentials( amazonKey.getKey(), amazonKey.getSecret() );
		AWSLambda awsLambda = AWSLambdaClientBuilder.standard()
				.withRegion( amazonKey.getAmazonRegion().toAWSRegion().getName() )
				.withCredentials( new AWSStaticCredentialsProvider( awsCreds ) ).build();

		// Execute and process the results
		InvokeResult result = awsLambda.invoke( invokeRequest );

		// Convert the returned result
		ByteBuffer resultPayload = result.getPayload();
		String resultJson = new String( resultPayload.array(), "UTF-8" );
		Map<String, Object> resultMap = Jackson.fromJsonString( resultJson, Map.class );

		return tagUtils.convertToCfData( resultMap );

	} catch ( Exception e ) {
		throwException( _session, "AmazonLambdaExecute: " + e.getMessage() );
		return cfBooleanData.FALSE;
	}

}
 

@Test
public void test() {
  Map<String, String> props = new HashMap<String, String>() {{
    put(FUNCTION_NAME_CONFIG, FUNCTION_NAME);
    put(REGION_CONFIG, REGION);
  }};

  Set<TopicPartition> assignment = new HashSet<>();
  assignment.add(TOPIC_PARTITION);
  assignment.add(TOPIC_PARTITION2);
  MockSinkTaskContext context = new MockSinkTaskContext(assignment);

  AwsLambdaSinkTask task = new AwsLambdaSinkTask();


  Collection<SinkRecord> records = new ArrayList<>();
  int partition = 1;
  String key = "key1";

  Schema valueSchema = SchemaBuilder.struct()
    .name("com.example.Person")
    .field("name", STRING_SCHEMA)
    .field("age", Schema.INT32_SCHEMA)
    .build();

  String bobbyMcGee = "Bobby McGee";
  int value21 = 21;

  Struct value = new Struct(valueSchema)
    .put("name", bobbyMcGee)
    .put("age", value21);

  long offset = 100;
  long timestamp = 200L;
  SinkRecord sinkRecord = new SinkRecord(
    TOPIC,
    partition,
    STRING_SCHEMA,
    key,
    valueSchema,
    value,
    offset,
    timestamp,
    TimestampType.CREATE_TIME);
  records.add(sinkRecord);

  String payload = "{\"schema\":" +
    "{\"type\":\"struct\"," +
    "\"fields\":[" +
    "{\"type\":\"string\",\"optional\":false,\"field\":\"name\"}," +
    "{\"type\":\"int32\",\"optional\":false,\"field\":\"age\"}" +
    "]," +
    "\"optional\":false," +
    "\"name\":\"com.example.Person\"}," +
    "\"payload\":{\"name\":\"Bobby McGee\",\"age\":21}}";

  task.setClient(new AbstractAWSLambda() {
    @Override
    public InvokeResult invoke(final InvokeRequest request) {
      assertEquals(FUNCTION_NAME, request.getFunctionName());
      assertEquals(payload, new String(request.getPayload().array()));
      return null;
    }
  });

  task.initialize(context);
  task.start(props);
  task.put(records);
}
 

/**
 * Invokes an AWS Lambda Function in sync mode using AWS Java SDK
 *
 * @param identity          Access key associated with your Amazon AWS or IAM account.
 *                          Example: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
 * @param credential        Secret access key ID associated with your Amazon AWS or IAM account.
 * @param proxyHost         Optional - proxy server used to connect to Amazon API. If empty no proxy will be used.
 *                          Default: ""
 * @param proxyPort         Optional - proxy server port. You must either specify values for both proxyHost and
 *                          proxyPort inputs or leave them both empty.
 *                          Default: ""
 * @param proxyUsername     Optional - proxy server user name.
 *                          Default: ""
 * @param proxyPassword     Optional - proxy server password associated with the proxyUsername input value.
 *                          Default: ""
 * @param function          Lambda function name to call
 *                          Example: "helloWord"
 * @param qualifier         Optional - Lambda function version or alias
 *                          Example: ":1"
 *                          Default: "$LATEST"
 * @param payload           Optional - Lambda function payload in JSON format
 *                          Example: "{"key1":"value1", "key1":"value2"}"
 *                          Default: null
 * @return                  A map with strings as keys and strings as values that contains: outcome of the action, returnCode of the
 *                          operation, or failure message and the exception if there is one
 */
@Action(name = "Invoke AWS Lambda Function",
        outputs = {
                @Output(Outputs.RETURN_CODE),
                @Output(Outputs.RETURN_RESULT),
                @Output(Outputs.EXCEPTION)
        },
        responses = {
                @Response(text = Outputs.SUCCESS, field = Outputs.RETURN_CODE, value = Outputs.SUCCESS_RETURN_CODE,
                        matchType = MatchType.COMPARE_EQUAL, responseType = ResponseType.RESOLVED),
                @Response(text = Outputs.FAILURE, field = Outputs.RETURN_CODE, value = Outputs.FAILURE_RETURN_CODE,
                        matchType = MatchType.COMPARE_EQUAL, responseType = ResponseType.ERROR)
        }
)
public Map<String, String> execute(
        @Param(value = IDENTITY,   required = true)  String identity,
        @Param(value = CREDENTIAL, required = true, encrypted = true)  String credential,
        @Param(value = REGION,     required = true)  String region,
        @Param(value = PROXY_HOST)                   String proxyHost,
        @Param(value = PROXY_PORT)                   String proxyPort,
        @Param(value = PROXY_USERNAME)               String proxyUsername,
        @Param(value = PROXY_PASSWORD)               String proxyPassword,
        @Param(value = FUNCTION_NAME, required = true) String function,
        @Param(value = FUNCTION_QUALIFIER)           String qualifier,
        @Param(value = FUNCTION_PAYLOAD)             String payload,
        @Param(value = CONNECT_TIMEOUT)              String connectTimeoutMs,
        @Param(value = EXECUTION_TIMEOUT)            String execTimeoutMs) {

    proxyPort = defaultIfEmpty(proxyPort, DefaultValues.PROXY_PORT);
    connectTimeoutMs = defaultIfEmpty(connectTimeoutMs, DefaultValues.CONNECT_TIMEOUT);
    execTimeoutMs = defaultIfBlank(execTimeoutMs, DefaultValues.EXEC_TIMEOUT);
    qualifier = defaultIfBlank(qualifier, DefaultValues.DEFAULT_FUNCTION_QUALIFIER);

    ClientConfiguration lambdaClientConf = AmazonWebServiceClientUtil.getClientConfiguration(proxyHost, proxyPort, proxyUsername, proxyPassword, connectTimeoutMs, execTimeoutMs);

    AWSLambdaAsyncClient client = (AWSLambdaAsyncClient) AWSLambdaAsyncClientBuilder.standard()
            .withClientConfiguration(lambdaClientConf)
            .withCredentials(new AWSStaticCredentialsProvider(new BasicAWSCredentials(identity, credential)))
            .withRegion(region)
            .build();

    InvokeRequest invokeRequest = new InvokeRequest()
            .withFunctionName(function)
            .withQualifier(qualifier)
            .withPayload(payload)
            .withSdkClientExecutionTimeout(Integer.parseInt(execTimeoutMs));

    try {
        InvokeResult invokeResult = client.invoke(invokeRequest);
        return OutputUtilities.getSuccessResultsMap(new String(invokeResult.getPayload().array()));
    } catch (Exception e) {
        return OutputUtilities.getFailureResultsMap(e);
    }
}
 

@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();
    }
}