feat(core): allow subflows to inherit parent flow's labels (#2121)

closes #2114
kestra-io · Sep 22, 2023 · 456b18d · 456b18d
1 parent 444a775
commit 456b18d
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 4 deletions.
diff --git a/core/src/main/java/io/kestra/core/runners/ExecutorService.java b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
@@ -579,7 +579,7 @@ private Executor handleFlowTask(final Executor executor) {
                 );
 
                 try {
-                    Execution execution = flowTask.createExecution(runContext, flowExecutorInterface());
+                    Execution execution = flowTask.createExecution(runContext, flowExecutorInterface(), executor.getExecution());
 
                     WorkerTaskExecution workerTaskExecution = WorkerTaskExecution.builder()
                         .task(flowTask)

diff --git a/core/src/main/java/io/kestra/core/tasks/flows/Flow.java b/core/src/main/java/io/kestra/core/tasks/flows/Flow.java
@@ -95,6 +95,14 @@ public class Flow extends Task implements RunnableTask<Flow.Output> {
     @PluginProperty
     private final Boolean transmitFailed = false;
 
+    @Builder.Default
+    @Schema(
+        title = "Inherit labels from the calling execution",
+        description = "By default, we don't inherit any labels of the calling execution"
+    )
+    @PluginProperty
+    private final Boolean inheritLabels = false;
+
     @Schema(
         title = "Extract outputs from triggered executions.",
         description = "Allow to specify key value (with value rendered), in order to extract any outputs from " +
@@ -117,7 +125,7 @@ public String flowUidWithoutRevision() {
     }
 
     @SuppressWarnings("unchecked")
-    public Execution createExecution(RunContext runContext, FlowExecutorInterface flowExecutorInterface) throws Exception {
+    public Execution createExecution(RunContext runContext, FlowExecutorInterface flowExecutorInterface, Execution currentExecution) throws Exception {
         RunnerUtils runnerUtils = runContext.getApplicationContext().getBean(RunnerUtils.class);
 
         Map<String, String> inputs = new HashMap<>();
@@ -128,6 +136,9 @@ public Execution createExecution(RunContext runContext, FlowExecutorInterface fl
         }
 
         List<Label> labels = new ArrayList<>();
+        if (this.inheritLabels) {
+            labels.addAll(currentExecution.getLabels());
+        }
         if (this.labels != null) {
             for (Map.Entry<String, String> entry: this.labels.entrySet()) {
                 labels.add(new Label(entry.getKey(), runContext.render(entry.getValue())));

diff --git a/core/src/test/java/io/kestra/core/tasks/flows/FlowCaseTest.java b/core/src/test/java/io/kestra/core/tasks/flows/FlowCaseTest.java
@@ -1,17 +1,20 @@
 package io.kestra.core.tasks.flows;
 
 import com.google.common.collect.ImmutableMap;
+import io.kestra.core.models.Label;
 import io.kestra.core.models.executions.Execution;
 import io.kestra.core.models.flows.State;
 import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.queues.QueueInterface;
 import io.kestra.core.runners.RunnerUtils;
 
 import java.time.Duration;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+
 import jakarta.inject.Inject;
 import jakarta.inject.Named;
 import jakarta.inject.Singleton;
@@ -41,7 +44,7 @@ public void invalidOutputs() throws Exception {
     }
 
     @SuppressWarnings({"ResultOfMethodCallIgnored", "unchecked"})
-    void run(String input, State.Type fromState, State.Type triggerState,  int count, String outputs) throws Exception {
+    void run(String input, State.Type fromState, State.Type triggerState, int count, String outputs) throws Exception {
         CountDownLatch countDownLatch = new CountDownLatch(1);
         AtomicReference<Execution> triggered = new AtomicReference<>();
 
@@ -58,7 +61,8 @@ void run(String input, State.Type fromState, State.Type triggerState,  int count
             "task-flow",
             null,
             (f, e) -> ImmutableMap.of("string", input),
-            Duration.ofMinutes(1)
+            Duration.ofMinutes(1),
+            List.of(new Label("mainFlowExecutionLabel", "execFoo"))
         );
 
         countDownLatch.await(1, TimeUnit.MINUTES);
@@ -85,5 +89,12 @@ void run(String input, State.Type fromState, State.Type triggerState,  int count
 
         assertThat(triggered.get().getTaskRunList(), hasSize(count));
         assertThat(triggered.get().getState().getCurrent(), is(triggerState));
+
+        assertThat(triggered.get().getLabels(), hasItems(
+            new Label("mainFlowExecutionLabel", "execFoo"),
+            new Label("mainFlowLabel", "flowFoo"),
+            new Label("launchTaskLabel", "launchFoo"),
+            new Label("switchFlowLabel", "switchFoo")
+        ));
     }
 }
diff --git a/core/src/test/resources/flows/valids/switch.yaml b/core/src/test/resources/flows/valids/switch.yaml
@@ -8,6 +8,9 @@ inputs:
     type: STRING
     defaults: amazing
 
+labels:
+  switchFlowLabel: switchFoo
+
 tasks:
   - id: parent-seq
     type: io.kestra.core.tasks.flows.Switch

diff --git a/core/src/test/resources/flows/valids/task-flow.yaml b/core/src/test/resources/flows/valids/task-flow.yaml
@@ -5,6 +5,9 @@ inputs:
   - name: string
     type: STRING
 
+labels:
+  mainFlowLabel: flowFoo
+
 tasks:
   - id: launch
     type: io.kestra.core.tasks.flows.Flow
@@ -14,5 +17,8 @@ tasks:
       string: "{{ inputs.string }}"
     wait: true
     transmitFailed: true
+    inheritLabels: true
+    labels:
+      launchTaskLabel: launchFoo
     outputs:
       extracted: "{{ outputs.default.value ?? outputs['error-t1'].value }}"