From 6f69dbbad454d1cf1a59ef596dd80b0e23150d71 Mon Sep 17 00:00:00 2001
From: Florian Hussonnois <florian.hussonnois@gmail.com>
Date: Fri, 17 Jan 2025 14:34:15 +0100
Subject: [PATCH] fix(core): fix some labels are lost when having same prefix
 key

This commit fix an issue where only a single system label was injected in the run context

fix: kestra-io/kestra-ee#2697
---
 .../java/io/kestra/core/models/Label.java     | 16 ++++++++
 .../io/kestra/core/runners/RunVariables.java  | 27 +-------------
 .../java/io/kestra/core/models/LabelTest.java | 37 +++++++++++++++++++
 3 files changed, 54 insertions(+), 26 deletions(-)
 create mode 100644 core/src/test/java/io/kestra/core/models/LabelTest.java

diff --git a/core/src/main/java/io/kestra/core/models/Label.java b/core/src/main/java/io/kestra/core/models/Label.java
index 9e8fa87f092..1db12b6488e 100644
--- a/core/src/main/java/io/kestra/core/models/Label.java
+++ b/core/src/main/java/io/kestra/core/models/Label.java
@@ -1,9 +1,11 @@
 package io.kestra.core.models;
 
+import io.kestra.core.utils.MapUtils;
 import jakarta.validation.constraints.NotNull;
 
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 public record Label(@NotNull String key, @NotNull String value) {
     public static final String SYSTEM_PREFIX = "system.";
@@ -14,6 +16,20 @@ public record Label(@NotNull String key, @NotNull String value) {
     public static final String APP = SYSTEM_PREFIX + "app";
     public static final String READ_ONLY = SYSTEM_PREFIX + "readOnly";
 
+    /**
+     * Static helper method for converting a list of labels to a nested map.
+     *
+     * @param labels The list of {@link Label} to be converted.
+     * @return the nested {@link Map}.
+     */
+    public static Map<String, Object> toNestedMap(List<Label> labels) {
+        Map<String, Object> asMap = labels.stream()
+            .filter(label -> label.value() != null && label.key() != null)
+            // using an accumulator in case labels with the same key exists: the first is kept
+            .collect(Collectors.toMap(Label::key, Label::value, (first, second) -> first));
+        return MapUtils.flattenToNestedMap(asMap);
+    }
+
     /**
      * Static helper method for converting a map to a list of labels.
      *
diff --git a/core/src/main/java/io/kestra/core/runners/RunVariables.java b/core/src/main/java/io/kestra/core/runners/RunVariables.java
index dfe976f7a45..b62ce4e8c05 100644
--- a/core/src/main/java/io/kestra/core/runners/RunVariables.java
+++ b/core/src/main/java/io/kestra/core/runners/RunVariables.java
@@ -15,13 +15,11 @@
 import lombok.With;
 
 import java.security.GeneralSecurityException;
-import java.util.AbstractMap;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
 
 /**
  * Class for building {@link RunContext} variables.
@@ -294,13 +292,7 @@ public Map<String, Object> build(final RunContextLogger logger) {
                 }
 
                 if (execution.getLabels() != null) {
-                    builder.put("labels", execution.getLabels()
-                        .stream()
-                        .filter(label -> label.value() != null && label.key() != null)
-                        .map(label -> mapLabel(label))
-                        // using an accumulator in case labels with the same key exists: the first is kept
-                        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, (first, second) -> first))
-                    );
+                    builder.put("labels", Label.toNestedMap(execution.getLabels()));
                 }
 
                 if (execution.getVariables() != null) {
@@ -331,22 +323,5 @@ public Map<String, Object> build(final RunContextLogger logger) {
         }
     }
 
-    private static Map.Entry<String, Object> mapLabel(Label label) {
-        if (label.key().startsWith(Label.SYSTEM_PREFIX)) {
-            return Map.entry(
-                label.key().substring(0, Label.SYSTEM_PREFIX.length() - 1),
-                Map.entry(
-                    label.key().substring(Label.SYSTEM_PREFIX.length()),
-                    label.value()
-                )
-            );
-        } else {
-            return Map.entry(
-                label.key(),
-                label.value()
-            );
-        }
-    }
-
     private RunVariables(){}
 }
diff --git a/core/src/test/java/io/kestra/core/models/LabelTest.java b/core/src/test/java/io/kestra/core/models/LabelTest.java
new file mode 100644
index 00000000000..c1c7d8a26ca
--- /dev/null
+++ b/core/src/test/java/io/kestra/core/models/LabelTest.java
@@ -0,0 +1,37 @@
+package io.kestra.core.models;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+
+class LabelTest {
+
+    @Test
+    void shouldGetNestedMapGivenDistinctLabels() {
+        Map<String, Object> result = Label.toNestedMap(List.of(
+            new Label(Label.USERNAME, "test"),
+            new Label(Label.CORRELATION_ID, "id"))
+        );
+
+        Assertions.assertEquals(
+            Map.of("system", Map.of("username", "test", "correlationId", "id")),
+            result
+        );
+    }
+
+    @Test
+    void shouldGetNestedMapGivenDuplicateLabels() {
+        Map<String, Object> result = Label.toNestedMap(List.of(
+            new Label(Label.USERNAME, "test1"),
+            new Label(Label.USERNAME, "test2"),
+            new Label(Label.CORRELATION_ID, "id"))
+        );
+
+        Assertions.assertEquals(
+            Map.of("system", Map.of("username", "test1", "correlationId", "id")),
+            result
+        );
+    }
+}
\ No newline at end of file