Skip to content

Commit

Permalink
Allow apply both environment variables and system properties to user …
Browse files Browse the repository at this point in the history 
…and table configs, Environment variables take precedence over system properties (#13011)
  • Loading branch information
@xiangfu0
xiangfu0 committed Apr 27, 2024
1 parent 5fc89ce commit 0be51ca
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 23 deletions.
4 changes: 2 additions & 2 deletions pinot-common/src/main/java/org/apache/pinot/common/metadata/ZKMetadataProvider.java
View file
Expand Up @@ -313,7 +313,7 @@ public static UserConfig getUserConfig(ZkHelixPropertyStore<ZNRecord> propertySt
}
try {
UserConfig userConfig = AccessControlUserConfigUtils.fromZNRecord(znRecord);
return ConfigUtils.applyConfigWithEnvVariables(userConfig);
return ConfigUtils.applyConfigWithEnvVariablesAndSystemProperties(userConfig);
} catch (Exception e) {
LOGGER.error("Caught exception while getting user configuration for user: {}", username, e);
return null;
Expand Down Expand Up @@ -422,7 +422,7 @@ private static TableConfig toTableConfig(@Nullable ZNRecord znRecord) {
}
try {
TableConfig tableConfig = TableConfigUtils.fromZNRecord(znRecord);
return ConfigUtils.applyConfigWithEnvVariables(tableConfig);
return ConfigUtils.applyConfigWithEnvVariablesAndSystemProperties(tableConfig);
} catch (Exception e) {
LOGGER.error("Caught exception while creating table config from ZNRecord: {}", znRecord.getId(), e);
return null;
Expand Down
41 changes: 26 additions & 15 deletions pinot-spi/src/main/java/org/apache/pinot/spi/config/ConfigUtils.java
View file
Expand Up @@ -23,6 +23,7 @@
import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.JsonNodeType;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import org.apache.pinot.spi.utils.JsonUtils;
Expand All @@ -35,45 +36,55 @@ private ConfigUtils() {
private static final Map<String, String> ENVIRONMENT_VARIABLES = System.getenv();

/**
* Apply environment variables to any given BaseJsonConfig.
* Apply system properties and environment variables to any given BaseJsonConfig.
* Environment variables take precedence over system properties.
* Since the System properties are mutable, this method will read it at runtime.
*
* @return Config with environment variable applied.
* @return Config with both system properties and environment variables applied.
*/
public static <T extends BaseJsonConfig> T applyConfigWithEnvVariables(T config) {
return applyConfigWithEnvVariables(ENVIRONMENT_VARIABLES, config);
public static <T extends BaseJsonConfig> T applyConfigWithEnvVariablesAndSystemProperties(T config) {
Map<String, String> combinedMap = new HashMap<>();
// Add all system properties to the map
System.getProperties().forEach((key, value) -> combinedMap.put(String.valueOf(key), String.valueOf(value)));
// Add all environment variables to the map, potentially overwriting system properties
combinedMap.putAll(ENVIRONMENT_VARIABLES);
return applyConfigWithEnvVariablesAndSystemProperties(combinedMap, config);
}

/**
* Apply environment variables to any given BaseJsonConfig.
* Apply a map of config to any given BaseJsonConfig with templates.
*
* @return Config with environment variable applied.
* @return Config with the configs applied.
*/
public static <T extends BaseJsonConfig> T applyConfigWithEnvVariables(Map<String, String> environment, T config) {
public static <T extends BaseJsonConfig> T applyConfigWithEnvVariablesAndSystemProperties(
Map<String, String> configValues, T configTemplate) {
JsonNode jsonNode;
try {
jsonNode = applyConfigWithEnvVariables(environment, config.toJsonNode());
jsonNode = applyConfigWithEnvVariablesAndSystemProperties(configValues, configTemplate.toJsonNode());
} catch (RuntimeException e) {
throw new RuntimeException(String
.format("Unable to apply environment variables on json config class [%s].", config.getClass().getName()), e);
.format("Unable to apply environment variables on json config class [%s].",
configTemplate.getClass().getName()), e);
}
try {
return (T) JsonUtils.jsonNodeToObject(jsonNode, config.getClass());
return (T) JsonUtils.jsonNodeToObject(jsonNode, configTemplate.getClass());
} catch (IOException e) {
throw new RuntimeException(String
.format("Unable to read JsonConfig to class [%s] after applying environment variables, jsonConfig is: '%s'.",
config.getClass().getName(), jsonNode.toString()), e);
configTemplate.getClass().getName(), jsonNode.toString()), e);
}
}

private static JsonNode applyConfigWithEnvVariables(Map<String, String> environment, JsonNode jsonNode) {
private static JsonNode applyConfigWithEnvVariablesAndSystemProperties(Map<String, String> configValues,
JsonNode jsonNode) {
final JsonNodeType nodeType = jsonNode.getNodeType();
switch (nodeType) {
case OBJECT:
if (!jsonNode.isEmpty()) {
Iterator<Map.Entry<String, JsonNode>> iterator = jsonNode.fields();
while (iterator.hasNext()) {
final Map.Entry<String, JsonNode> next = iterator.next();
next.setValue(applyConfigWithEnvVariables(environment, next.getValue()));
next.setValue(applyConfigWithEnvVariablesAndSystemProperties(configValues, next.getValue()));
}
}
break;
Expand All @@ -82,7 +93,7 @@ private static JsonNode applyConfigWithEnvVariables(Map<String, String> environm
ArrayNode arrayNode = (ArrayNode) jsonNode;
for (int i = 0; i < arrayNode.size(); i++) {
JsonNode arrayElement = arrayNode.get(i);
arrayNode.set(i, applyConfigWithEnvVariables(environment, arrayElement));
arrayNode.set(i, applyConfigWithEnvVariablesAndSystemProperties(configValues, arrayElement));
}
}
break;
Expand All @@ -91,7 +102,7 @@ private static JsonNode applyConfigWithEnvVariables(Map<String, String> environm
if (field.startsWith("${") && field.endsWith("}")) {
String[] envVarSplits = field.substring(2, field.length() - 1).split(":", 2);
String envVarKey = envVarSplits[0];
String value = environment.get(envVarKey);
String value = configValues.get(envVarKey);
if (value != null) {
return JsonNodeFactory.instance.textNode(value);
} else if (envVarSplits.length > 1) {
Expand Down
30 changes: 24 additions & 6 deletions pinot-spi/src/test/java/org/apache/pinot/spi/config/ConfigUtilsTest.java
View file
Expand Up @@ -39,6 +39,25 @@ public class ConfigUtilsTest {

@Test
public void testIndexing() {
Map<String, String> environment =
ImmutableMap.of("LOAD_MODE", "MMAP", "AWS_ACCESS_KEY", "default_aws_access_key", "AWS_SECRET_KEY",
"default_aws_secret_key");
testIndexingWithConfig(environment);
}

@Test
public void testIndexingWithSystemProperties() {
// Use default System properties
System.setProperty("LOAD_MODE", "MMAP");
System.setProperty("AWS_ACCESS_KEY", "default_aws_access_key");
System.setProperty("AWS_SECRET_KEY", "default_aws_secret_key");
testIndexingWithConfig(null);
System.clearProperty("LOAD_MODE");
System.clearProperty("AWS_ACCESS_KEY");
System.clearProperty("AWS_SECRET_KEY");
}

private void testIndexingWithConfig(Map<String, String> configOverride) {
IndexingConfig indexingConfig = new IndexingConfig();
indexingConfig.setLoadMode("${LOAD_MODE}");
indexingConfig.setAggregateMetrics(true);
Expand Down Expand Up @@ -80,12 +99,11 @@ public void testIndexing() {
streamConfigMap.put(StreamConfigProperties.constructStreamProperty(streamType, "aws.secretKey"),
"${AWS_SECRET_KEY}");
indexingConfig.setStreamConfigs(streamConfigMap);

Map<String, String> environment =
ImmutableMap.of("LOAD_MODE", "MMAP", "AWS_ACCESS_KEY", "default_aws_access_key", "AWS_SECRET_KEY",
"default_aws_secret_key");

indexingConfig = ConfigUtils.applyConfigWithEnvVariables(environment, indexingConfig);
if (configOverride != null) {
indexingConfig = ConfigUtils.applyConfigWithEnvVariablesAndSystemProperties(configOverride, indexingConfig);
} else {
indexingConfig = ConfigUtils.applyConfigWithEnvVariablesAndSystemProperties(indexingConfig);
}
assertEquals(indexingConfig.getLoadMode(), "MMAP");
assertTrue(indexingConfig.isAggregateMetrics());
assertEquals(indexingConfig.getInvertedIndexColumns(), invertedIndexColumns);
Expand Down

0 comments on commit 0be51ca

Please sign in to comment.