Skip to content

Commit 8decfab

Browse files
bHacklvbHacklv
committed
fix: ENV VAR kv parsing to handle json values
- JOB_KUBE_ANNOTATIONS can have JSON as value - This fixes the parsin of such values
1 parent 70f3fb3 commit 8decfab

File tree

4 files changed

+129
-29
lines changed

4 files changed

+129
-29
lines changed

‎airbyte-commons-with-dependencies/build.gradle.kts

+1
Original file line numberDiff line numberDiff line change
@@ -18,4 +18,5 @@ dependencies {
1818

1919
testImplementation(libs.mockito.core)
2020
testImplementation(libs.bundles.micronaut.test)
21+
testImplementation(libs.assertj.core)
2122
}

‎airbyte-commons-with-dependencies/src/main/java/io/airbyte/commons/workers/config/EnvUtils.java

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package io.airbyte.commons.workers.config;
2+
3+
import lombok.AccessLevel;
4+
import lombok.NoArgsConstructor;
5+
import lombok.extern.slf4j.Slf4j;
6+
7+
import java.util.Collections;
8+
import java.util.HashMap;
9+
import java.util.Map;
10+
import java.util.regex.Matcher;
11+
import java.util.regex.Pattern;
12+
13+
@Slf4j
14+
@NoArgsConstructor(access = AccessLevel.PRIVATE)
15+
class EnvUtils {
16+
17+
private static final Pattern KEY_JSON_VALUE_PATTERN = Pattern.compile("(?<key>[^=]+)\\s*?=\\s*?(?<value>\\{[^=]+})\\s*,?\\s*");
18+
private static final Pattern KEY_VALUE_PATTERN = Pattern.compile("(?<key>[^=]+)=(?<value>[^=,]+),?");
19+
private static final String KEY_GROUP_ALIAS = "key";
20+
private static final String VALUE_GROUP_ALIAS = "value";
21+
22+
/**
23+
* Splits key value pairs from the input string into a map. Each kv-pair is separated by a ','.
24+
* </br>
25+
* The key and the value are separated by '='.
26+
* <p>
27+
* For example:- The following represents two map entries
28+
* </p>
29+
* - key1=value1,key2=value2
30+
* </br>
31+
* - key1={key11: value1},key2={key22: value2}
32+
* </br>
33+
* - key1={key11: value11, key12: value12},key2={key21: value21, key22: value22}
34+
*
35+
* @param input string
36+
* @return map containing kv pairs
37+
*/
38+
public static Map<String, String> splitKVPairsFromEnvString(final String input) {
39+
if (input == null || input.isBlank()) {
40+
return Map.of();
41+
}
42+
final Map<String, String> jsonValuesMatchResult = match(input, KEY_JSON_VALUE_PATTERN);
43+
return jsonValuesMatchResult.isEmpty()
44+
? getKVPairsMatchedWithSimplePattern(input)
45+
: jsonValuesMatchResult;
46+
}
47+
48+
private static Map<String, String> match(final String input, final Pattern pattern) {
49+
final Matcher matcher = pattern.matcher(input);
50+
final Map<String, String> kvResult = new HashMap<>();
51+
while (matcher.find()) {
52+
kvResult.put(matcher.group(KEY_GROUP_ALIAS).trim(), matcher.group(VALUE_GROUP_ALIAS).trim());
53+
}
54+
return kvResult;
55+
}
56+
57+
private static Map<String, String> getKVPairsMatchedWithSimplePattern(final String input) {
58+
final Map<String, String> stringMatchResult = match(input, KEY_VALUE_PATTERN);
59+
if (stringMatchResult.isEmpty()) {
60+
log.warn("No valid key value pairs found in the input string: {}", input);
61+
return Collections.emptyMap();
62+
}
63+
return stringMatchResult;
64+
}
65+
66+
}

‎airbyte-commons-with-dependencies/src/main/java/io/airbyte/commons/workers/config/WorkerConfigsProvider.java

+5-29
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44

55
package io.airbyte.commons.workers.config;
66

7-
import com.google.common.base.Splitter;
87
import com.google.common.base.Strings;
98
import io.airbyte.config.ResourceRequirements;
109
import io.airbyte.config.ResourceRequirementsType;
@@ -24,7 +23,6 @@
2423
import java.util.Optional;
2524
import java.util.regex.Matcher;
2625
import java.util.regex.Pattern;
27-
import java.util.stream.Collectors;
2826
import lombok.extern.slf4j.Slf4j;
2927

3028
/**
@@ -210,25 +208,25 @@ private WorkerConfigs getConfig(final KubeResourceKey key) {
210208
.orElseThrow(() -> new NoSuchElementException(String.format("Unable to find config: {variant:%s, type:%s, subtype:%s}",
211209
key.variant, key.type, key.subType)));
212210

213-
final Map<String, String> isolatedNodeSelectors = splitKVPairsFromEnvString(workerConfigsDefaults.isolatedNodeSelectors);
211+
final Map<String, String> isolatedNodeSelectors = EnvUtils.splitKVPairsFromEnvString(workerConfigsDefaults.isolatedNodeSelectors);
214212
validateIsolatedPoolConfigInitialization(workerConfigsDefaults.useCustomNodeSelector(), isolatedNodeSelectors);
215213

216214
// if annotations are not defined for this specific resource, then fallback to the default
217215
// resource's annotations
218216
final Map<String, String> annotations;
219217
if (Strings.isNullOrEmpty(kubeResourceConfig.getAnnotations())) {
220-
annotations = splitKVPairsFromEnvString(workerConfigsDefaults.defaultKubeResourceConfig.getAnnotations());
218+
annotations = EnvUtils.splitKVPairsFromEnvString(workerConfigsDefaults.defaultKubeResourceConfig.getAnnotations());
221219
} else {
222-
annotations = splitKVPairsFromEnvString(kubeResourceConfig.getAnnotations());
220+
annotations = EnvUtils.splitKVPairsFromEnvString(kubeResourceConfig.getAnnotations());
223221
}
224222

225223
return new WorkerConfigs(
226224
getResourceRequirementsFrom(kubeResourceConfig, workerConfigsDefaults.defaultKubeResourceConfig()),
227225
TolerationPOJO.getJobKubeTolerations(workerConfigsDefaults.jobKubeTolerations()),
228-
splitKVPairsFromEnvString(kubeResourceConfig.getNodeSelectors()),
226+
EnvUtils.splitKVPairsFromEnvString(kubeResourceConfig.getNodeSelectors()),
229227
workerConfigsDefaults.useCustomNodeSelector() ? Optional.of(isolatedNodeSelectors) : Optional.empty(),
230228
annotations,
231-
splitKVPairsFromEnvString(kubeResourceConfig.getLabels()),
229+
EnvUtils.splitKVPairsFromEnvString(kubeResourceConfig.getLabels()),
232230
workerConfigsDefaults.mainContainerImagePullSecret(),
233231
workerConfigsDefaults.mainContainerImagePullPolicy());
234232
}
@@ -323,28 +321,6 @@ private Optional<KubeResourceKey> parseKubeResourceKey(final String value) {
323321
return Optional.empty();
324322
}
325323

326-
/**
327-
* Splits key value pairs from the input string into a map. Each kv-pair is separated by a ','. The
328-
* key and the value are separated by '='.
329-
* <p>
330-
* For example:- The following represents two map entries
331-
* </p>
332-
* key1=value1,key2=value2
333-
*
334-
* @param input string
335-
* @return map containing kv pairs
336-
*/
337-
private Map<String, String> splitKVPairsFromEnvString(final String input) {
338-
if (input == null || input.isBlank()) {
339-
return Map.of();
340-
}
341-
return Splitter.on(",")
342-
.splitToStream(input)
343-
.filter(s -> !Strings.isNullOrEmpty(s) && s.contains("="))
344-
.map(s -> s.split("="))
345-
.collect(Collectors.toMap(s -> s[0].trim(), s -> s[1].trim()));
346-
}
347-
348324
private ResourceRequirements getResourceRequirementsFrom(final KubeResourceConfig kubeResourceConfig, final KubeResourceConfig defaultConfig) {
349325
return new ResourceRequirements()
350326
.withCpuLimit(useDefaultIfEmpty(kubeResourceConfig.getCpuLimit(), defaultConfig.getCpuLimit()))

‎airbyte-commons-with-dependencies/src/test/java/io/airbyte/commons/workers/config/EnvUtilsTest.java

+57
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package io.airbyte.commons.workers.config;
2+
3+
import org.junit.jupiter.params.ParameterizedTest;
4+
import org.junit.jupiter.params.provider.Arguments;
5+
import org.junit.jupiter.params.provider.MethodSource;
6+
7+
import java.util.Collections;
8+
import java.util.Map;
9+
import java.util.stream.Stream;
10+
11+
import static org.assertj.core.api.Assertions.assertThat;
12+
13+
class EnvUtilsTest {
14+
15+
@ParameterizedTest
16+
@MethodSource("splitKVPairsFromEnvString")
17+
void splitKVPairsFromEnvString(String input, Map<String, String> expected) {
18+
final Map<String, String> result = EnvUtils.splitKVPairsFromEnvString(input);
19+
assertThat(result).isEqualTo(expected);
20+
}
21+
22+
private static Stream<Arguments> splitKVPairsFromEnvString() {
23+
return Stream.of(
24+
// unmatched
25+
Arguments.of("key1", Collections.emptyMap()),
26+
Arguments.of("key1,value", Collections.emptyMap()),
27+
Arguments.of("key1-value", Collections.emptyMap()),
28+
Arguments.of("key1:value", Collections.emptyMap()),
29+
// matched k:v pairs
30+
Arguments.of("key1=value1", Map.of("key1", "value1")),
31+
Arguments.of("key1 = value1", Map.of("key1", "value1")),
32+
Arguments.of("key1=value1,key2=value2", Map.of("key1", "value1", "key2", "value2")),
33+
Arguments.of("key1 = value1, key2 = value2", Map.of("key1", "value1", "key2", "value2")),
34+
// matched k:jsonV pairs
35+
Arguments.of("key1={value1}", Map.of("key1", "{value1}")),
36+
Arguments.of("key1={ value1 }", Map.of("key1", "{ value1 }")),
37+
Arguments.of("key1 = { value1 }", Map.of("key1", "{ value1 }")),
38+
Arguments.of("key1={value1},key2={value2}", Map.of("key1", "{value1}", "key2", "{value2}")),
39+
Arguments.of("key1= {value1} , key2={value2}", Map.of("key1", "{value1}", "key2", "{value2}")),
40+
Arguments.of("key1= {value1 } , key2= { value2}", Map.of("key1", "{value1 }", "key2", "{ value2}")),
41+
Arguments.of("key1={key11: value11},key2={key22: value22}", Map.of(
42+
"key1", "{key11: value11}",
43+
"key2", "{key22: value22}"
44+
)),
45+
Arguments.of("key1={key11: value11, key12: value12},key2={key21: value21, key22: value22}", Map.of(
46+
"key1", "{key11: value11, key12: value12}",
47+
"key2", "{key21: value21, key22: value22}"
48+
)),
49+
Arguments.of("key1={key11: value11, key12: value12, key13: {key131: value131}},"
50+
+ "key2={key21: value21, key22: value22, key23: {key231: value231}}", Map.of(
51+
"key1", "{key11: value11, key12: value12, key13: {key131: value131}}",
52+
"key2", "{key21: value21, key22: value22, key23: {key231: value231}}"
53+
))
54+
);
55+
}
56+
57+
}

0 commit comments

Comments
 (0)