Fixing Schema variable support in Resource Group Template

Due to a validation check in StaticSelector, resource group was not supporting ${SCHEMA} in resource group name. Fixing it as part of the PR.
jkhaliqi · Jun 2, 2023 · cc1cec1 · cc1cec1
1 parent 179854d
commit cc1cec1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/...ource-group-managers/src/main/java/com/facebook/presto/resourceGroups/StaticSelector.java b/...ource-group-managers/src/main/java/com/facebook/presto/resourceGroups/StaticSelector.java
@@ -75,7 +75,7 @@ public StaticSelector(
         this.schema = requireNonNull(schema, "schema is null");
         this.group = requireNonNull(group, "group is null");
 
-        HashSet<String> variableNames = new HashSet<>(ImmutableList.of(USER_VARIABLE, SOURCE_VARIABLE));
+        HashSet<String> variableNames = new HashSet<>(ImmutableList.of(USER_VARIABLE, SOURCE_VARIABLE, SCHEMA_VARIABLE));
         userRegex.ifPresent(u -> addNamedGroups(u, variableNames));
         sourceRegex.ifPresent(s -> addNamedGroups(s, variableNames));
         this.variableNames = ImmutableSet.copyOf(variableNames);

diff --git a/...e-group-managers/src/test/java/com/facebook/presto/resourceGroups/TestStaticSelector.java b/...e-group-managers/src/test/java/com/facebook/presto/resourceGroups/TestStaticSelector.java
@@ -102,7 +102,7 @@ public void testSchema()
                 Optional.empty(),
                 Optional.empty(),
                 Optional.of("schema1"),
-                new ResourceGroupIdTemplate("global.schema1"));
+                new ResourceGroupIdTemplate("global.${SCHEMA}"));
         assertEquals(selector.match(newSelectionCriteria("userA", null, "schema1", ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
         assertEquals(selector.match(newSelectionCriteria("userB", "source", "schema2", ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)), Optional.empty());
         assertEquals(selector.match(newSelectionCriteria("userA", null, "schema1", ImmutableSet.of("tag1"), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));