feat(cloudwatch): display lambda logs (#681)

* 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

* chore: sync with main and trigger CI for #677

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

---------

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.1.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,74 @@ Optional<String> readError(String payload) {
         return Optional.empty();
     }
 
+    @VisibleForTesting
+    void fetchAndLogLambdaLogs(RunContext runContext, String functionArn, Instant startTime) {
+        var logger = runContext.logger();
+        String functionName;
+
+        try {
+            functionName = extractFunctionName(functionArn);
+        } catch (Exception e) {
+            logger.warn("Unable to determine Lambda function name from ARN: {}", functionArn);
+            return;
+        }
+
+        String logGroupName = "/aws/lambda/" + functionName;
+
+        // Polling logic: 5 attempts, waiting 3 seconds between each
+        int maxAttempts = 5;
+        long sleepMillis = 3000;
+        boolean logsFound = false;
+
+        try (CloudWatchLogsClient logsClient = getCloudWatchLogsClient(runContext)) {
+            for (int i = 0; i < maxAttempts; i++) {
+                FilterLogEventsRequest request = FilterLogEventsRequest.builder()
+                        .logGroupName(logGroupName)
+                        .startTime(startTime.toEpochMilli())
+                        .build();
+
+                var response = logsClient.filterLogEvents(request);
+                var events = response.events();
+
+                if (events != null && !events.isEmpty()) {
+                    events.stream()
+                            .limit(1000)
+                            .map(FilteredLogEvent::message)
+                            .filter(message -> message != null && !message.isBlank())
+                            .forEach(message -> logger.info("[lambda] {}", message.trim()));
+
+                    logsFound = true;
+                    break; // Exit early if found logs
+                }
+
+                // Wait before next retry, but not after the last attempt
+                if (i < maxAttempts - 1) {
+                    Thread.sleep(sleepMillis);
+                }
+            }
+
+            if (!logsFound) {
+                logger.debug("No CloudWatch logs found for {} after {} attempts.", functionName, maxAttempts);
+            }
+
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            logger.warn("Log fetching interrupted for Lambda {}", functionArn);
+        } 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 +334,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;
     }
 }