Resolve multiple environment variables in configuration values

prestodb · Apr 19, 2024 · 3a9f1c8 · 3a9f1c8
1 parent a09d951
commit 3a9f1c8
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 33 deletions.
diff --git a/bootstrap/pom.xml b/bootstrap/pom.xml
@@ -65,5 +65,20 @@
             <scope>test</scope>
         </dependency>
     </dependencies>
-</project>
 
+    <build>
+        <plugins>
+            <!-- TODO Remove after upgrade to Java 9+ -->
+            <!-- Added to resolve 'Prefer java.lang.StringBuilder' violation -->
+            <plugin>
+                <groupId>org.gaul</groupId>
+                <artifactId>modernizer-maven-plugin</artifactId>
+                <configuration>
+                    <exclusions>
+                        <exclusion>java/lang/StringBuffer."&lt;init&gt;":()V</exclusion>
+                    </exclusions>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
+</project>
diff --git a/bootstrap/src/main/java/com/facebook/airlift/bootstrap/Bootstrap.java b/bootstrap/src/main/java/com/facebook/airlift/bootstrap/Bootstrap.java
@@ -30,6 +30,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableList.Builder;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.inject.Binder;
 import com.google.inject.Guice;
@@ -326,24 +327,27 @@ static SortedMap<String, String> replaceWithEnvironmentVariables(
     {
         SortedMap<String, String> results = new TreeMap<>();
         properties.forEach((key, value) -> {
-            Matcher result = ENV_PATTERN.matcher(value);
-            if (result.matches()) {
-                String variableName = result.group(1);
-                String envValue = environment.get(variableName);
-
-                if (envValue == null) {
-                    String errorMessage = String.format(
-                            "Configuration property `%s` references `%s`, an undefiled environment variable.",
-                            key,
-                            variableName);
-                    onError.accept(key, errorMessage);
+            // TODO: Replace with Matcher.replaceAll(Function) or replace StringBuffer with StringBuilder after upgrade to Java 9+
+            StringBuffer result = new StringBuffer();
+            ImmutableSet.Builder<String> undefinedBuilder = ImmutableSet.builderWithExpectedSize(0);
+            Matcher matcher = ENV_PATTERN.matcher(value);
+            while (matcher.find()) {
+                String variable = matcher.group(1);
+                String variableValue = environment.get(variable);
+                if (variableValue == null) {
+                    undefinedBuilder.add(variable);
                 }
                 else {
-                    results.put(key, envValue);
+                    matcher.appendReplacement(result, variableValue);
                 }
             }
+            ImmutableSet<String> undefined = undefinedBuilder.build();
+            matcher.appendTail(result);
+            if (undefined.isEmpty()) {
+                results.put(key, result.toString());
+            }
             else {
-                results.put(key, value);
+                onError.accept(key, String.format("Configuration property `%s` references undefined environment variable(s): %s", key, undefined));
             }
         });
         return results;

diff --git a/bootstrap/src/test/java/com/facebook/airlift/bootstrap/TestBootstrap.java b/bootstrap/src/test/java/com/facebook/airlift/bootstrap/TestBootstrap.java
@@ -100,29 +100,40 @@ public void testNoStrictConfig()
     @Test
     public void testEnvironmentVariableReplacement()
     {
-        Map<String, String> properties = new HashMap<>();
-        Map<String, String> environment = new HashMap<>();
+        Map<String, String> properties = new ImmutableMap.Builder<String, String>()
+                // Literal value
+                .put("no-change", "no-change")
+                // Correct values with defined environment variables
+                .put("variable", "${ENV:VARIABLE}")
+                .put("mixed-variable", "mixed-${ENV:VARIABLE}")
+                .put("mixed-additional-variable", "mixed-${ENV:ADDITIONAL}-${ENV:VARIABLE}")
+                // Broken values with undefined environment variables
+                .put("missing-variable", "${ENV:MISSING}")
+                .put("mixed-missing-variable", "mixed-${ENV:MISSING}")
+                .put("mixed-missing-additional-variable", "mixed-${ENV:MISSING}-${ENV:VARIABLE}")
+                .put("mixed-missing-repeated-variable", "mixed-${ENV:MISSING}-${ENV:MISSING}-${ENV:VARIABLE}")
+                .put("mixed-missing-many-variables", "mixed-${ENV:MISSING_1}-${ENV:MISSING_2}-${ENV:VARIABLE}")
+                .build();
+        Map<String, String> environment = new ImmutableMap.Builder<String, String>()
+                .put("VARIABLE", "variable")
+                .put("ADDITIONAL", "additional")
+                .build();
         Map<String, String> errors = new HashMap<>();
 
-        properties.put("expected", "${ENV:EXPECTED}");
-        properties.put("invalid", "${ENV:INVALID}");
-        properties.put("no-replacement", "no-replacement");
-        properties.put("mixed-no-change", "${ENV:INVALID}!");
-
-        environment.put("EXPECTED", "expected-replacement");
-
         Map<String, String> results = Bootstrap.replaceWithEnvironmentVariables(properties, environment, errors::put);
 
-        assertEquals(results.get("expected"), "expected-replacement");
-        assertEquals(results.get("no-replacement"), "no-replacement");
-        assertEquals(results.get("mixed-no-change"), "${ENV:INVALID}!");
-
-        assertEquals(3, results.size());
-        assertEquals(1, errors.size());
-
-        assertEquals(
-                errors.get("invalid"),
-                "Configuration property `invalid` references `INVALID`, an undefiled environment variable.");
+        assertEquals(results.size(), 4);
+        assertEquals(results.get("no-change"), "no-change");
+        assertEquals(results.get("variable"), "variable");
+        assertEquals(results.get("mixed-variable"), "mixed-variable");
+        assertEquals(results.get("mixed-additional-variable"), "mixed-additional-variable");
+
+        assertEquals(errors.size(), 5);
+        assertEquals(errors.get("missing-variable"), "Configuration property `missing-variable` references undefined environment variable(s): [MISSING]");
+        assertEquals(errors.get("mixed-missing-variable"), "Configuration property `mixed-missing-variable` references undefined environment variable(s): [MISSING]");
+        assertEquals(errors.get("mixed-missing-additional-variable"), "Configuration property `mixed-missing-additional-variable` references undefined environment variable(s): [MISSING]");
+        assertEquals(errors.get("mixed-missing-repeated-variable"), "Configuration property `mixed-missing-repeated-variable` references undefined environment variable(s): [MISSING]");
+        assertEquals(errors.get("mixed-missing-many-variables"), "Configuration property `mixed-missing-many-variables` references undefined environment variable(s): [MISSING_1, MISSING_2]");
     }
 
     public static class Instance {}