Skip to content

Commit

Permalink
fix(core): fix some labels are lost when having same prefix key
Browse files Browse the repository at this point in the history
This commit fix an issue where only a single system label was injected in the run context

fix: kestra-io/kestra-ee#2697
  • Loading branch information
fhussonnois committed Jan 20, 2025
1 parent 6bde5d5 commit 6f69dbb
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 26 deletions.
16 changes: 16 additions & 0 deletions core/src/main/java/io/kestra/core/models/Label.java
Original file line number Diff line number Diff line change
@@ -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.";
Expand All @@ -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.
*
Expand Down
27 changes: 1 addition & 26 deletions core/src/main/java/io/kestra/core/runners/RunVariables.java
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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(){}
}
37 changes: 37 additions & 0 deletions core/src/test/java/io/kestra/core/models/LabelTest.java
Original file line number Diff line number Diff line change
@@ -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
);
}
}

0 comments on commit 6f69dbb

Please sign in to comment.