feat(lambda): fetch CloudWatch logs after Lambda invocation (#677)

* feat(lambda): fetch CloudWatch logs after Lambda invocation

Fetches CloudWatch Logs for AWS Lambda invocations using
FilterLogEvents paginator and streams log messages into
the task logger.

Includes unit tests covering log retrieval and logging behavior.

* fix(lambda): improve CloudWatch log polling and stabilize unit test

* chore: trigger CI re-run for flaky test

* Refactor fetchAndLogLambdaLogs to use RetryUtils with proper retry policy

* Bump kestraVersion to 1.2.0 for RetryUtils static method

* Fix annotation formatting in Invoke plugin

---------

Co-authored-by: François Delbrayelle <fdelbrayelle@gmail.com>

Author: shivamwayal37

gradle.properties

@@ -1,2 +1,2 @@
 version=1.4.2-SNAPSHOT
-kestraVersion=1.1.0
+kestraVersion=1.2.0

src/main/java/io/kestra/plugin/aws/cloudwatch/CloudWatchLogs.java

@@ -0,0 +1,34 @@
+package io.kestra.plugin.aws.cloudwatch;
+
+import io.kestra.core.exceptions.IllegalVariableEvaluationException;
+import io.kestra.core.runners.RunContext;
+import io.kestra.plugin.aws.AbstractConnection;
+import io.kestra.plugin.aws.ConnectionUtils;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.ToString;
+import lombok.experimental.SuperBuilder;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
+
+@SuperBuilder
+@ToString
+@EqualsAndHashCode
+@Getter
+@NoArgsConstructor
+public class CloudWatchLogs extends AbstractConnection {
+
+    /**
+     * Create a CloudWatchLogsClient using the standard Kestra AWS configuration
+     * (credentials, region, endpoint overrides, etc.).
+     */
+    public CloudWatchLogsClient logsClient(final RunContext runContext)
+        throws IllegalVariableEvaluationException {
+
+        final AwsClientConfig clientConfig = awsClientConfig(runContext);
+
+        return ConnectionUtils
+            .configureSyncClient(clientConfig, CloudWatchLogsClient.builder())
+            .build();
+    }
+}

src/main/java/io/kestra/plugin/aws/lambda/Invoke.java

@@ -1,8 +1,22 @@
 package io.kestra.plugin.aws.lambda;
 
+import java.io.File;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.net.URI;
+import java.nio.file.Files;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Map;
+import java.util.Optional;
+
+import org.apache.http.HttpHeaders;
+import org.apache.http.entity.ContentType;
+
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
+
 import io.kestra.core.exceptions.IllegalVariableEvaluationException;
 import io.kestra.core.models.annotations.Example;
 import io.kestra.core.models.annotations.Metric;
@@ -11,10 +25,13 @@
 import io.kestra.core.models.executions.metrics.Timer;
 import io.kestra.core.models.property.Property;
 import io.kestra.core.models.tasks.RunnableTask;
+import io.kestra.core.models.tasks.retrys.AbstractRetry;
 import io.kestra.core.runners.RunContext;
 import io.kestra.core.serializers.JacksonMapper;
+import io.kestra.core.utils.RetryUtils;
 import io.kestra.plugin.aws.AbstractConnection;
 import io.kestra.plugin.aws.ConnectionUtils;
+import io.kestra.plugin.aws.cloudwatch.CloudWatchLogs;
 import io.kestra.plugin.aws.lambda.Invoke.Output;
 import io.kestra.plugin.aws.s3.ObjectOutput;
 import io.swagger.v3.oas.annotations.media.Schema;
@@ -25,22 +42,17 @@
 import lombok.ToString;
 import lombok.experimental.SuperBuilder;
 import lombok.extern.slf4j.Slf4j;
-import org.apache.http.HttpHeaders;
-import org.apache.http.entity.ContentType;
 import software.amazon.awssdk.core.SdkBytes;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
+import software.amazon.awssdk.services.cloudwatchlogs.model.FilterLogEventsRequest;
+import software.amazon.awssdk.services.cloudwatchlogs.model.FilteredLogEvent;
 import software.amazon.awssdk.services.lambda.LambdaClient;
 import software.amazon.awssdk.services.lambda.model.InvokeRequest;
 import software.amazon.awssdk.services.lambda.model.InvokeResponse;
 import software.amazon.awssdk.services.lambda.model.LambdaException;
+import io.kestra.core.models.tasks.retrys.Exponential;
 
-import java.io.File;
-import java.io.IOException;
-import java.io.UncheckedIOException;
-import java.net.URI;
-import java.nio.file.Files;
-import java.time.Duration;
-import java.util.Map;
-import java.util.Optional;
+import java.util.List;
 
 @SuperBuilder
 @ToString
@@ -94,6 +106,7 @@
         @Metric(name = "duration", type = Timer.TYPE)
     }
 )
+
 @Slf4j
 public class Invoke extends AbstractConnection implements RunnableTask<Output> {
 
@@ -103,32 +116,28 @@ public class Invoke extends AbstractConnection implements RunnableTask<Output> {
     @NotNull
     private Property<String> functionArn;
 
-    @Schema(
-        title = "Function request payload.",
-        description = "Request payload. It's a map of string -> object."
-    )
+    @Schema(title = "Function request payload.", description = "Request payload. It's a map of string -> object.")
     private Property<Map<String, Object>> functionPayload;
 
     @Override
     public Output run(RunContext runContext) throws Exception {
         final long start = System.nanoTime();
+        final Instant invocationStart = Instant.now().minusSeconds(5);
         var functionArn = runContext.render(this.functionArn).as(String.class).orElseThrow();
-        var requestPayload = runContext.render(this.functionPayload).asMap(String.class, Object.class).isEmpty() ?
-            null :
-            runContext.render(this.functionPayload).asMap(String.class, Object.class);
+        var requestPayload = runContext.render(this.functionPayload).asMap(String.class, Object.class).isEmpty() ? null
+                : runContext.render(this.functionPayload).asMap(String.class, Object.class);
         var logger = runContext.logger();
 
         try (var lambda = client(runContext)) {
             var builder = InvokeRequest.builder().functionName(functionArn);
             if (requestPayload != null && requestPayload.size() > 0) {
-                var payload = SdkBytes.fromUtf8String(OBJECT_MAPPER.writeValueAsString(requestPayload)) ;
+                var payload = SdkBytes.fromUtf8String(OBJECT_MAPPER.writeValueAsString(requestPayload));
                 builder.payload(payload);
             }
             InvokeRequest request = builder.build();
             // TODO take care about long-running functions: your client might disconnect during
             // synchronous invocation while it waits for a response. Configure your HTTP client,
             // SDK, firewall, proxy, or operating system to allow for long connections with timeout
-            // or keep-alive settings.
             InvokeResponse res = lambda.invoke(request);
             Optional<String> contentTypeHeader = res.sdkHttpResponse().firstMatchingHeader(HttpHeaders.CONTENT_TYPE);
             ContentType contentType = parseContentType(contentTypeHeader);
@@ -141,13 +150,19 @@ public Output run(RunContext runContext) throws Exception {
                 logger.debug("Lambda {} invoked successfully", functionArn);
             }
             Output out = handleContent(runContext, functionArn, contentType, res.payload());
+            fetchAndLogLambdaLogs(runContext, functionArn, invocationStart);
             runContext.metric(Timer.of("duration", Duration.ofNanos(System.nanoTime() - start)));
             return out;
         } catch (LambdaException e) {
             throw new LambdaInvokeException("Lambda Invoke task execution failed for function: " + functionArn, e);
         }
     }
 
+    @VisibleForTesting
+    CloudWatchLogsClient getCloudWatchLogsClient(RunContext runContext) throws IllegalVariableEvaluationException {
+        return new CloudWatchLogs().logsClient(runContext);
+    }
+
     @VisibleForTesting
     LambdaClient client(final RunContext runContext) throws IllegalVariableEvaluationException {
         final AwsClientConfig clientConfig = awsClientConfig(runContext);
@@ -165,7 +180,7 @@ ContentType parseContentType(Optional<String> contentType) {
                     // Apply charset only if it was provided originally
                     return ContentType.create(known.getMimeType(), parsed.getCharset());
                 }
-            } catch(Exception cte) {
+            } catch (Exception cte) {
                 log.warn("Unable to parse Lambda response content type {}: {}", contentType.get(), cte.getMessage());
             }
         }
@@ -195,6 +210,53 @@ Optional<String> readError(String payload) {
         return Optional.empty();
     }
 
+    @VisibleForTesting
+    void fetchAndLogLambdaLogs(RunContext runContext, String functionArn, Instant startTime) {
+        var logger = runContext.logger();
+        String functionName = extractFunctionName(functionArn);
+        String logGroupName = "/aws/lambda/" + functionName;
+
+        // Explicit retry policy: 5 attempts, 3s interval, maxInterval 3s
+        AbstractRetry retryPolicy = Exponential.builder()
+                .interval(Duration.ofSeconds(3))
+                .maxAttempts(5)
+                .maxInterval(Duration.ofSeconds(10))
+                .build();
+
+        try (CloudWatchLogsClient logsClient = getCloudWatchLogsClient(runContext)) {
+            // Explicitly specify generic type for RetryUtils
+            List<FilteredLogEvent> events = RetryUtils.<List<FilteredLogEvent>, Exception>of(retryPolicy, logger)
+                    .run(
+                            result -> result == null || result.isEmpty(),
+                            () -> {
+                                FilterLogEventsRequest request = FilterLogEventsRequest.builder()
+                                        .logGroupName(logGroupName)
+                                        .startTime(startTime.toEpochMilli())
+                                        .build();
+
+                                var response = logsClient.filterLogEvents(request);
+                                return response.events();
+                            });
+
+            if (events != null) {
+                events.forEach(event -> logger.info("[lambda] {}", event.message().trim()));
+            }
+
+        } catch (Exception e) {
+            logger.warn("Failed to fetch CloudWatch logs for Lambda {}: {}", functionArn, e.getMessage());
+        }
+    }
+
+    @VisibleForTesting
+    private String extractFunctionName(String functionArnOrName) {
+        if (functionArnOrName.contains(":function:")) {
+            // Handle Full ARN
+            return functionArnOrName.split(":function:")[1].split(":")[0];
+        }
+        // Handle just the name
+        return functionArnOrName;
+    }
+
     @VisibleForTesting
     void handleError(String functionArn, ContentType contentType, SdkBytes payload) {
         String errorPayload;
@@ -251,20 +313,14 @@ Output handleContent(RunContext runContext, String functionArn, ContentType cont
     @Getter
     public static class Output extends ObjectOutput implements io.kestra.core.models.tasks.Output {
 
-        @Schema(
-            title = "Response file URI."
-        )
+        @Schema(title = "Response file URI.")
         private final URI uri;
 
-        @Schema(
-            title = "Size of the response content in bytes."
-        )
+        @Schema(title = "Size of the response content in bytes.")
 
         private final Long contentLength;
 
-        @Schema(
-            title = "A standard MIME type describing the format of the content."
-        )
+        @Schema(title = "A standard MIME type describing the format of the content.")
         private final String contentType;
     }
 }

src/test/java/io/kestra/plugin/aws/lambda/InvokeUnitTest.java

@@ -1,45 +1,60 @@
 package io.kestra.plugin.aws.lambda;
 
-import io.kestra.core.exceptions.IllegalVariableEvaluationException;
-import io.kestra.core.models.property.Property;
-import io.kestra.core.runners.RunContext;
-import io.kestra.core.runners.RunContextProperty;
-import io.kestra.core.runners.WorkingDir;
-import io.kestra.core.storages.Storage;
-import io.kestra.plugin.aws.lambda.Invoke.Output;
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.nio.file.Files;
+import java.time.Instant;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import org.apache.http.entity.ContentType;
 import org.junit.jupiter.api.AfterEach;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.BDDMockito.given;
 import org.mockito.Mock;
 import org.mockito.Mock.Strictness;
+
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.mockito.stubbing.Answer;
 import org.slf4j.Logger;
+
+import io.kestra.core.exceptions.IllegalVariableEvaluationException;
+import io.kestra.core.models.property.Property;
+import io.kestra.core.runners.RunContext;
+import io.kestra.core.runners.RunContextProperty;
+import io.kestra.core.runners.WorkingDir;
+import io.kestra.core.storages.Storage;
+import io.kestra.plugin.aws.lambda.Invoke.Output;
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.http.SdkHttpResponse;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
+import software.amazon.awssdk.services.cloudwatchlogs.model.FilterLogEventsRequest;
+import software.amazon.awssdk.services.cloudwatchlogs.model.FilterLogEventsResponse;
+import software.amazon.awssdk.services.cloudwatchlogs.model.FilteredLogEvent;
+import software.amazon.awssdk.services.cloudwatchlogs.paginators.FilterLogEventsIterable;
 import software.amazon.awssdk.services.lambda.LambdaClient;
 import software.amazon.awssdk.services.lambda.model.InvokeRequest;
 import software.amazon.awssdk.services.lambda.model.InvokeResponse;
 
-import java.io.ByteArrayInputStream;
-import java.io.File;
-import java.io.IOException;
-import java.net.URI;
-import java.nio.file.Files;
-import java.util.Collections;
-import java.util.Map;
-import java.util.Optional;
-
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.BDDMockito.given;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.spy;
-
 @ExtendWith(MockitoExtension.class)
 public class InvokeUnitTest {
 
@@ -48,6 +63,9 @@ public class InvokeUnitTest {
     @Mock(strictness = Strictness.LENIENT)
     private RunContext context;
 
+    @Mock
+    private CloudWatchLogsClient logsClient;
+
     @Mock(strictness = Strictness.LENIENT)
     private RunContextProperty runContextProperty;
 
@@ -203,4 +221,37 @@ void givenFunctionArnNoParams_whenInvokeLambda_thenOutputWithFile(
         // Then
         checkOutput(data, res);
     }
+
+    @Test
+    void givenLambdaInvocation_whenLogsAreFetched_thenLogsAreLoggedCorrectly() throws Throwable {
+        // Arrange
+        Instant startTime = Instant.parse("2026-01-15T10:00:00Z");
+
+        FilteredLogEvent logEvent = FilteredLogEvent.builder()
+                .message("Hello from CloudWatch Logs")
+                .timestamp(startTime.plusSeconds(1).toEpochMilli())
+                .build();
+
+        FilterLogEventsResponse response = FilterLogEventsResponse.builder()
+                .events(List.of(logEvent))
+                .build();
+
+        given(logsClient.filterLogEvents(any(FilterLogEventsRequest.class)))
+                .willReturn(response);
+
+        given(context.logger()).willReturn(logger);
+
+        Invoke spyInvoke = spy(invoke);
+        doReturn(logsClient).when(spyInvoke).getCloudWatchLogsClient(any());
+
+        // Act
+        spyInvoke.fetchAndLogLambdaLogs(
+                context,
+                "arn:aws:lambda:ap-south-1:123456789012:function:test-function",
+                startTime);
+
+        // Assert (THIS is what really matters)
+        verify(logger).info("[lambda] {}", "Hello from CloudWatch Logs");
+    }
 }
+