Skip to content

Commit 6f69dbb

Browse files
committed
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
1 parent 6bde5d5 commit 6f69dbb

File tree

3 files changed

+54
-26
lines changed

3 files changed

+54
-26
lines changed

core/src/main/java/io/kestra/core/models/Label.java

+16
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package io.kestra.core.models;
22

3+
import io.kestra.core.utils.MapUtils;
34
import jakarta.validation.constraints.NotNull;
45

56
import java.util.List;
67
import java.util.Map;
8+
import java.util.stream.Collectors;
79

810
public record Label(@NotNull String key, @NotNull String value) {
911
public static final String SYSTEM_PREFIX = "system.";
@@ -14,6 +16,20 @@ public record Label(@NotNull String key, @NotNull String value) {
1416
public static final String APP = SYSTEM_PREFIX + "app";
1517
public static final String READ_ONLY = SYSTEM_PREFIX + "readOnly";
1618

19+
/**
20+
* Static helper method for converting a list of labels to a nested map.
21+
*
22+
* @param labels The list of {@link Label} to be converted.
23+
* @return the nested {@link Map}.
24+
*/
25+
public static Map<String, Object> toNestedMap(List<Label> labels) {
26+
Map<String, Object> asMap = labels.stream()
27+
.filter(label -> label.value() != null && label.key() != null)
28+
// using an accumulator in case labels with the same key exists: the first is kept
29+
.collect(Collectors.toMap(Label::key, Label::value, (first, second) -> first));
30+
return MapUtils.flattenToNestedMap(asMap);
31+
}
32+
1733
/**
1834
* Static helper method for converting a map to a list of labels.
1935
*

core/src/main/java/io/kestra/core/runners/RunVariables.java

+1-26
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,11 @@
1515
import lombok.With;
1616

1717
import java.security.GeneralSecurityException;
18-
import java.util.AbstractMap;
1918
import java.util.HashMap;
2019
import java.util.List;
2120
import java.util.Map;
2221
import java.util.Optional;
2322
import java.util.function.Consumer;
24-
import java.util.stream.Collectors;
2523

2624
/**
2725
* Class for building {@link RunContext} variables.
@@ -294,13 +292,7 @@ public Map<String, Object> build(final RunContextLogger logger) {
294292
}
295293

296294
if (execution.getLabels() != null) {
297-
builder.put("labels", execution.getLabels()
298-
.stream()
299-
.filter(label -> label.value() != null && label.key() != null)
300-
.map(label -> mapLabel(label))
301-
// using an accumulator in case labels with the same key exists: the first is kept
302-
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, (first, second) -> first))
303-
);
295+
builder.put("labels", Label.toNestedMap(execution.getLabels()));
304296
}
305297

306298
if (execution.getVariables() != null) {
@@ -331,22 +323,5 @@ public Map<String, Object> build(final RunContextLogger logger) {
331323
}
332324
}
333325

334-
private static Map.Entry<String, Object> mapLabel(Label label) {
335-
if (label.key().startsWith(Label.SYSTEM_PREFIX)) {
336-
return Map.entry(
337-
label.key().substring(0, Label.SYSTEM_PREFIX.length() - 1),
338-
Map.entry(
339-
label.key().substring(Label.SYSTEM_PREFIX.length()),
340-
label.value()
341-
)
342-
);
343-
} else {
344-
return Map.entry(
345-
label.key(),
346-
label.value()
347-
);
348-
}
349-
}
350-
351326
private RunVariables(){}
352327
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
package io.kestra.core.models;
2+
3+
import org.junit.jupiter.api.Assertions;
4+
import org.junit.jupiter.api.Test;
5+
6+
import java.util.List;
7+
import java.util.Map;
8+
9+
class LabelTest {
10+
11+
@Test
12+
void shouldGetNestedMapGivenDistinctLabels() {
13+
Map<String, Object> result = Label.toNestedMap(List.of(
14+
new Label(Label.USERNAME, "test"),
15+
new Label(Label.CORRELATION_ID, "id"))
16+
);
17+
18+
Assertions.assertEquals(
19+
Map.of("system", Map.of("username", "test", "correlationId", "id")),
20+
result
21+
);
22+
}
23+
24+
@Test
25+
void shouldGetNestedMapGivenDuplicateLabels() {
26+
Map<String, Object> result = Label.toNestedMap(List.of(
27+
new Label(Label.USERNAME, "test1"),
28+
new Label(Label.USERNAME, "test2"),
29+
new Label(Label.CORRELATION_ID, "id"))
30+
);
31+
32+
Assertions.assertEquals(
33+
Map.of("system", Map.of("username", "test1", "correlationId", "id")),
34+
result
35+
);
36+
}
37+
}

0 commit comments

Comments
 (0)