Repository: kestra-io/kestra (java)

Issue:
Add a dedicated task allowing to set custom Execution Labels from a flow
### Feature description

Possible syntax:

```yaml
id: labels_test
namespace: dev

inputs:
  - name: myinput
    type: STRING

tasks:
  - id: labels
    type: io.kestra.core.tasks.executions.Labels
    labels:
      mykey: "{{ inputs.myinput }}"
```

This will also address this request in a straightforward and explicit manner https://github.com/kestra-io/kestra/issues/1760 as it's otherwise difficult to sometimes add labels from inputs (nested JSON etc)

Additionally, this task allows setting execution labels based on specific task outputs, flow progress, etc. - many use cases enabled by adding this.

Pull request:
feat/update-execution-labels-at-runtime
closes #2632

Patch:
diff --git a/core/src/main/java/io/kestra/core/models/tasks/ExecutionUpdatableTask.java b/core/src/main/java/io/kestra/core/models/tasks/ExecutionUpdatableTask.java
new file mode 100644
index 000000000..ac2f2b3a8
--- /dev/null
+++ b/core/src/main/java/io/kestra/core/models/tasks/ExecutionUpdatableTask.java
@@ -0,0 +1,11 @@
+package io.kestra.core.models.tasks;
+
+import io.kestra.core.models.executions.Execution;
+import io.kestra.core.runners.RunContext;
+
+/**
+ * Interface for tasks that modify the execution at runtime.
+ */
+public interface ExecutionUpdatableTask {
+    Execution update(Execution execution, RunContext runContext) throws Exception;
+}
diff --git a/core/src/main/java/io/kestra/core/runners/ExecutorService.java b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
index d816db8ff..ee0ba4d3d 100644
--- a/core/src/main/java/io/kestra/core/runners/ExecutorService.java
+++ b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
@@ -10,11 +10,7 @@ import io.kestra.core.models.executions.TaskRun;
 import io.kestra.core.models.executions.TaskRunAttempt;
 import io.kestra.core.models.flows.Flow;
 import io.kestra.core.models.flows.State;
-import io.kestra.core.models.tasks.ExecutableTask;
-import io.kestra.core.models.tasks.FlowableTask;
-import io.kestra.core.models.tasks.Output;
-import io.kestra.core.models.tasks.ResolvedTask;
-import io.kestra.core.models.tasks.Task;
+import io.kestra.core.models.tasks.*;
 import io.kestra.core.services.ConditionService;
 import io.kestra.core.tasks.flows.Pause;
 import io.kestra.core.tasks.flows.WorkingDirectory;
@@ -62,7 +58,7 @@ public class ExecutorService {
 
     public Executor checkConcurrencyLimit(Executor executor, Flow flow, Execution execution, long count) {
         if (count >= flow.getConcurrency().getLimit()) {
-            return switch(flow.getConcurrency().getBehavior()) {
+            return switch (flow.getConcurrency().getBehavior()) {
                 case QUEUE -> {
                     var newExecution = execution.withState(State.Type.QUEUED);
 
@@ -125,6 +121,9 @@ public class ExecutorService {
             // search for worker task result
             executor = this.handleChildWorkerTaskResult(executor);
 
+            // search for execution updating tasks
+            executor = this.handleExecutionUpdatingTask(executor);
+
             // search for flow task
             executor = this.handleExecutableTask(executor);
         } catch (Exception e) {
@@ -192,8 +191,7 @@ public class ExecutorService {
             Optional<State.Type> state;
             try {
                 state = flowableParent.resolveState(runContext, execution, parentTaskRun);
-            }
-            catch (Exception e) {
+            } catch (Exception e) {
                 // This will lead to the next task being still executed but at least Kestra will not crash.
                 // This is the best we can do, Flowable task should not fail, so it's a kind of panic mode.
                 runContext.logger().error("Unable to resolve state from the Flowable task: " + e.getMessage(), e);
@@ -455,7 +453,7 @@ public class ExecutorService {
             .filter(Objects::nonNull)
             .collect(Collectors.toList());
 
-        if(executor.getExecution().getState().getCurrent() != State.Type.PAUSED) {
+        if (executor.getExecution().getState().getCurrent() != State.Type.PAUSED) {
             return executor
                 .withExecution(executor.getExecution().withState(State.Type.PAUSED), "handlePausedDelay")
                 .withWorkerTaskDelays(list, "handlePausedDelay");
@@ -623,8 +621,7 @@ public class ExecutorService {
                                 .withTaskRun(executableTaskRun.withState(State.Type.SUCCESS)),
                             "handleExecutableTaskRunning.noExecution"
                         );
-                    }
-                    else {
+                    } else {
                         executions.addAll(workerTaskExecutions);
                         if (!executableTask.waitForExecution()) {
                             // send immediately all workerTaskResult to ends the executable task
@@ -639,8 +636,7 @@ public class ExecutorService {
                             }
                         }
                     }
-                }
-                catch (Exception e) {
+                } catch (Exception e) {
                     workerTaskResults.add(WorkerTaskResult.builder()
                         .taskRun(workerTask.getTaskRun().withState(State.Type.FAILED)
                             .withAttempts(Collections.singletonList(
@@ -667,6 +663,47 @@ public class ExecutorService {
         return resultExecutor;
     }
 
+    private Executor handleExecutionUpdatingTask(final Executor executor) {
+        List<WorkerTaskResult> workerTaskResults = new ArrayList<>();
+
+        executor.getWorkerTasks()
+            .removeIf(workerTask -> {
+                if (!(workerTask.getTask() instanceof ExecutionUpdatableTask)) {
+                    return false;
+                }
+
+                var executionUpdatingTask = (ExecutionUpdatableTask) workerTask.getTask();
+
+                try {
+                    executor.withExecution(
+                        executionUpdatingTask.update(executor.getExecution(), workerTask.getRunContext())
+                            .withTaskRun(workerTask.getTaskRun().withState(State.Type.RUNNING)),
+                        "handleExecutionUpdatingTask.updateExecution"
+                    );
+
+                    workerTaskResults.add(
+                        ExecutableUtils.workerTaskResult(workerTask.getTaskRun().withState(State.Type.SUCCESS))
+                    );
+                } catch (Exception e) {
+                    workerTaskResults.add(WorkerTaskResult.builder()
+                        .taskRun(workerTask.getTaskRun().withState(State.Type.FAILED)
+                            .withAttempts(Collections.singletonList(
+                                TaskRunAttempt.builder().state(new State().withState(State.Type.FAILED)).build()
+                            ))
+                        )
+                        .build());
+                    executor.withException(e, "handleExecutionUpdatingTask");
+                }
+                return true;
+            });
+
+        if (!workerTaskResults.isEmpty()) {
+            return executor.withWorkerTaskResults(workerTaskResults, "handleExecutionUpdatingTask.workerTaskResults");
+        }
+
+        return executor;
+    }
+
     public Execution addDynamicTaskRun(Execution execution, Flow flow, WorkerTaskResult workerTaskResult) throws InternalException {
         ArrayList<TaskRun> taskRuns = new ArrayList<>(execution.getTaskRunList());
 
diff --git a/core/src/main/java/io/kestra/core/tasks/executions/Labels.java b/core/src/main/java/io/kestra/core/tasks/executions/Labels.java
new file mode 100644
index 000000000..fa2c5b4db
--- /dev/null
+++ b/core/src/main/java/io/kestra/core/tasks/executions/Labels.java
@@ -0,0 +1,145 @@
+package io.kestra.core.tasks.executions;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import io.kestra.core.exceptions.IllegalVariableEvaluationException;
+import io.kestra.core.models.Label;
+import io.kestra.core.models.annotations.Example;
+import io.kestra.core.models.annotations.Plugin;
+import io.kestra.core.models.annotations.PluginProperty;
+import io.kestra.core.models.executions.Execution;
+import io.kestra.core.models.tasks.ExecutionUpdatableTask;
+import io.kestra.core.models.tasks.Task;
+import io.kestra.core.runners.RunContext;
+import io.kestra.core.serializers.JacksonMapper;
+import io.swagger.v3.oas.annotations.media.Schema;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.ToString;
+import lombok.experimental.SuperBuilder;
+import org.slf4j.Logger;
+
+import javax.validation.constraints.NotNull;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+import static io.kestra.core.utils.Rethrow.throwBiConsumer;
+import static io.kestra.core.utils.Rethrow.throwFunction;
+
+@SuperBuilder
+@ToString
+@EqualsAndHashCode
+@Getter
+@NoArgsConstructor
+@Schema(
+    title = "Allow to add or overwrite labels for the current execution at runtime."
+)
+@Plugin(
+    examples = {
+        @Example(
+            title = "Add labels based on a webhook payload",
+            full = true,
+            code = {
+                "id: webhook-based-labels",
+                "namespace: com.company",
+                "tasks:",
+                "  - id: update-labels-with-map",
+                "    type: io.kestra.core.tasks.executions.Labels",
+                "    labels:",
+                "      customerId: \"{{trigger.body.customerId}}\"",
+                "  - id: by-list",
+                "    type: io.kestra.core.tasks.executions.Labels",
+                "    labels:",
+                "      - key: orderId",
+                "        value: \"{{trigger.body.orderId}}\"",
+                "      - key: orderType",
+                "        value: \"{{trigger.body.orderType}}\"",
+                "triggers:",
+                "  - id: webhook",
+                "    key: order-webhook",
+                "    type: io.kestra.core.models.triggers.types.Webhook",
+                "    conditions:",
+                "      - type: io.kestra.core.models.conditions.types.VariableCondition",
+                "        expression: \"{{ trigger.body.customerId is defined and trigger.body.orderId is defined and trigger.body.orderType is defined }}\""
+            }
+        )
+    }
+)
+public class Labels extends Task implements ExecutionUpdatableTask {
+    @Schema(
+        title = "Labels to add to the current execution.",
+        description = "The value should result in a list of labels or a labelKey:labelValue map",
+        anyOf = {
+            String.class,
+            Label[].class,
+            Map.class
+        }
+    )
+    @PluginProperty(dynamic = true, additionalProperties = String.class)
+    @NotNull
+    private Object labels;
+
+    @Override
+    public Execution update(Execution execution, RunContext runContext) throws Exception {
+        Logger logger = runContext.logger();
+
+        Object labelsObject = labels;
+        if (labels instanceof String labelStr) {
+            try {
+                labelsObject = JacksonMapper.toObject(runContext.render(labelStr));
+            } catch (JsonProcessingException e) {
+                throw new IllegalVariableEvaluationException(e);
+            }
+        }
+
+        if (labelsObject instanceof List<?> labelsList) {
+            labelsObject = labelsList.stream()
+                .map(throwFunction(label -> {
+                        if (label instanceof Map<?, ?> labelMap) {
+                            return Map.entry(
+                                (String) labelMap.get("key"),
+                                (String) labelMap.get("value")
+                            );
+                        } else {
+                            throw new IllegalVariableEvaluationException("Unknown value type: " + label.getClass());
+                        }
+                    })
+                ).collect(Collectors.toMap(
+                    Map.Entry::getKey,
+                    Map.Entry::getValue
+                ));
+        }
+
+        if (labelsObject instanceof Map<?, ?> labelsMap) {
+            Map<String, String> newLabels = Optional.ofNullable(execution.getLabels()).orElse(Collections.emptyList()).stream()
+                .collect(Collectors.toMap(
+                    Label::key,
+                    Label::value,
+                    (labelValue1, labelValue2) -> labelValue2
+                ));
+            labelsMap.forEach(throwBiConsumer((key, value) -> {
+                String renderedKey = runContext.render((String) key);
+                String renderedValue = runContext.render((String) value);
+                logger.info("Adding label {} with value {}", renderedKey, renderedValue);
+                newLabels.put(
+                    renderedKey,
+                    renderedValue
+                );
+            }));
+
+            return execution.toBuilder()
+                .labels(newLabels.entrySet().stream()
+                    .map(throwFunction(entry -> new Label(
+                        entry.getKey(),
+                        entry.getValue()
+                    )))
+                    .toList())
+                .build();
+        } else {
+            throw new IllegalVariableEvaluationException("Unknown value type: " + labels.getClass());
+        }
+    }
+}
diff --git a/core/src/main/java/io/kestra/core/tasks/flows/EachParallel.java b/core/src/main/java/io/kestra/core/tasks/flows/EachParallel.java
index f410da35d..53cf81741 100644
--- a/core/src/main/java/io/kestra/core/tasks/flows/EachParallel.java
+++ b/core/src/main/java/io/kestra/core/tasks/flows/EachParallel.java
@@ -103,7 +103,7 @@ public class EachParallel extends Parallel implements FlowableTask<VoidOutput> {
     @PluginProperty(dynamic = true)
     @Schema(
         title = "The list of values for this task",
-        description = "The value car be passed as a String, a list of String, or a list of objects",
+        description = "The value can be passed as a String, a list of String, or a list of objects",
         anyOf = {String.class, Object[].class}
     )
     private Object value;


Test patch:
diff --git a/core/src/test/java/io/kestra/core/tasks/flows/RuntimeLabelsTest.java b/core/src/test/java/io/kestra/core/tasks/flows/RuntimeLabelsTest.java
new file mode 100644
index 000000000..665d33a57
--- /dev/null
+++ b/core/src/test/java/io/kestra/core/tasks/flows/RuntimeLabelsTest.java
@@ -0,0 +1,68 @@
+package io.kestra.core.tasks.flows;
+
+import io.kestra.core.models.Label;
+import io.kestra.core.models.executions.Execution;
+import io.kestra.core.models.flows.State;
+import io.kestra.core.runners.AbstractMemoryRunnerTest;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeoutException;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.*;
+
+class RuntimeLabelsTest extends AbstractMemoryRunnerTest {
+    @Test
+    void update() throws TimeoutException {
+        Execution execution = runnerUtils.runOne(
+            null,
+            "io.kestra.tests",
+            "labels-update-task",
+            null,
+            (flow, createdExecution) -> Map.of(
+                "labelsJson", "{\"keyFromJson\": \"valueFromJson\"}",
+                "labelsMapKey", "keyFromMap",
+                "labelsMapValue", "valueFromMap",
+                "labelsListKey", "keyFromList",
+                "labelsListValue", "valueFromList"
+            ),
+            null,
+            List.of(
+                new Label("keyFromExecution", "valueFromExecution"),
+                new Label("overriddenExecutionLabelKey", "executionValueThatWillGetOverridden")
+            )
+        );
+
+        assertThat(execution.getTaskRunList().size(), is(4));
+        assertThat(execution.getState().getCurrent(), is(State.Type.SUCCESS));
+
+        String labelsOverriderTaskRunId = execution.findTaskRunsByTaskId("override-labels").get(0).getId();
+        assertThat(execution.getLabels(), containsInAnyOrder(
+            is(new Label("flowLabelKey", "flowLabelValue")),
+            is(new Label("overriddenFlowLabelKey", "io.kestra.tests.labels-update-task")),
+            is(new Label("keyFromJson", "valueFromJson")),
+            is(new Label("keyFromMap", "valueFromMap")),
+            is(new Label("keyFromList", "valueFromList")),
+            is(new Label("keyFromExecution", "valueFromExecution")),
+            is(new Label("overriddenExecutionLabelKey", labelsOverriderTaskRunId))
+        ));
+    }
+
+
+    @Test
+    void noNpeOnNullPreviousExecutionLabels() throws TimeoutException {
+        Execution execution = runnerUtils.runOne(
+            null,
+            "io.kestra.tests",
+            "npe-labels-update-task"
+        );
+
+        assertThat(execution.getTaskRunList().size(), is(1));
+        assertThat(execution.getState().getCurrent(), is(State.Type.SUCCESS));
+
+        String labelsTaskRunId = execution.findTaskRunsByTaskId("labels").get(0).getId();
+        assertThat(execution.getLabels(), contains(is(new Label("someLabel", labelsTaskRunId))));
+    }
+}
diff --git a/core/src/test/resources/flows/valids/labels-update-task.yml b/core/src/test/resources/flows/valids/labels-update-task.yml
new file mode 100644
index 000000000..36d7ccd85
--- /dev/null
+++ b/core/src/test/resources/flows/valids/labels-update-task.yml
@@ -0,0 +1,34 @@
+id: labels-update-task
+namespace: io.kestra.tests
+labels:
+  flowLabelKey: flowLabelValue
+  overriddenFlowLabelKey: flowValueThatWillGetOverridden
+inputs:
+  - name: labelsJson
+    type: JSON
+  - name: labelsMapKey
+    type: STRING
+  - name: labelsMapValue
+    type: STRING
+  - name: labelsListKey
+    type: STRING
+  - name: labelsListValue
+    type: STRING
+tasks:
+  - id: from-render
+    type: io.kestra.core.tasks.executions.Labels
+    labels: "{{ inputs.labelsJson }}"
+  - id: from-map
+    type: io.kestra.core.tasks.executions.Labels
+    labels:
+      "{{ inputs.labelsMapKey }}": "{{ inputs.labelsMapValue }}"
+  - id: from-list
+    type: io.kestra.core.tasks.executions.Labels
+    labels:
+      - key: "{{ inputs.labelsListKey }}"
+        value: "{{ inputs.labelsListValue }}"
+  - id: override-labels
+    type: io.kestra.core.tasks.executions.Labels
+    labels:
+      overriddenFlowLabelKey: "{{ flow.namespace ~ '.' ~ flow.id }}"
+      overriddenExecutionLabelKey: "{{ taskrun.id }}"
\ No newline at end of file
diff --git a/core/src/test/resources/flows/valids/npe-labels-update-task.yml b/core/src/test/resources/flows/valids/npe-labels-update-task.yml
new file mode 100644
index 000000000..e0a1928cb
--- /dev/null
+++ b/core/src/test/resources/flows/valids/npe-labels-update-task.yml
@@ -0,0 +1,7 @@
+id: npe-labels-update-task
+namespace: io.kestra.tests
+tasks:
+  - id: labels
+    type: io.kestra.core.tasks.executions.Labels
+    labels:
+      someLabel: "{{ taskrun.id }}"
\ No newline at end of file
diff --git a/jdbc/src/test/java/io/kestra/jdbc/repository/AbstractJdbcFlowRepositoryTest.java b/jdbc/src/test/java/io/kestra/jdbc/repository/AbstractJdbcFlowRepositoryTest.java
index 152722d28..21f171c1c 100644
--- a/jdbc/src/test/java/io/kestra/jdbc/repository/AbstractJdbcFlowRepositoryTest.java
+++ b/jdbc/src/test/java/io/kestra/jdbc/repository/AbstractJdbcFlowRepositoryTest.java
@@ -37,7 +37,7 @@ public abstract class AbstractJdbcFlowRepositoryTest extends io.kestra.core.repo
 
     @Test
     protected void find() {
-        List<Flow> save = flowRepository.find(Pageable.from(1, 100, Sort.of(Sort.Order.asc("id"))), null, null, null, null);
+        List<Flow> save = flowRepository.find(Pageable.unpaged(), null, null, null, null);
         assertThat((long) save.size(), is(Helpers.FLOWS_COUNT));
 
         save = flowRepository.find(Pageable.from(1, 10, Sort.UNSORTED), "trigger-multiplecondition", null, null, null);