Added queryText regex pattern to match in resource group selector

Author: gertjanal

core/trino-main/src/main/java/io/trino/dispatcher/DispatchManager.java

@@ -235,6 +235,7 @@ private <C> void createQueryInternal(QueryId queryId, Span querySpan, Slug slug,
                     sessionContext.getSource(),
                     sessionContext.getClientTags(),
                     sessionContext.getResourceEstimates(),
+                    query,
                     queryType));
 
             // apply system default session properties (does not override user set properties)

core/trino-spi/pom.xml

@@ -291,6 +291,19 @@
                                     <new>method void io.trino.spi.block.ShortArrayBlockEncoding::&lt;init&gt;(boolean)</new>
                                     <justification>ShortArrayBlockEncoding need to accept a parameter to enable SIMD support</justification>
                                 </item>
+                                <item>
+                                    <ignore>true</ignore>
+                                    <code>java.method.numberOfParametersChanged</code>
+                                    <old>method void io.trino.spi.resourcegroups.SelectionCriteria::&lt;init&gt;(boolean, java.lang.String, java.util.Set&lt;java.lang.String&gt;, java.lang.String, java.util.Optional&lt;java.lang.String&gt;, java.util.Optional&lt;java.lang.String&gt;, java.util.Set&lt;java.lang.String&gt;, io.trino.spi.session.ResourceEstimates, java.util.Optional&lt;java.lang.String&gt;)</old>
+                                    <new>method void io.trino.spi.resourcegroups.SelectionCriteria::&lt;init&gt;(boolean, java.lang.String, java.util.Set&lt;java.lang.String&gt;, java.lang.String, java.util.Optional&lt;java.lang.String&gt;, java.util.Optional&lt;java.lang.String&gt;, java.util.Set&lt;java.lang.String&gt;, io.trino.spi.session.ResourceEstimates, java.lang.String, java.util.Optional&lt;java.lang.String&gt;)</new>
+                                    <justification>Added queryText before queryType parameter</justification>
+                                </item>
+                                <item>
+                                    <ignore>true</ignore>
+                                    <code>java.method.removed</code>
+                                    <old>method void io.trino.spi.resourcegroups.SelectionCriteria::&lt;init&gt;(boolean, java.lang.String, java.util.Set&lt;java.lang.String&gt;, java.lang.String, java.util.Optional&lt;java.lang.String&gt;, java.util.Optional&lt;java.lang.String&gt;, java.util.Set&lt;java.lang.String&gt;, io.trino.spi.session.ResourceEstimates, java.util.Optional&lt;java.lang.String&gt;)</old>
+                                    <justification>Added queryText before queryType parameter</justification>
+                                </item>
                             </differences>
                         </revapi.differences>
                     </analysisConfiguration>

core/trino-spi/src/main/java/io/trino/spi/resourcegroups/SelectionCriteria.java

@@ -31,6 +31,7 @@ public final class SelectionCriteria
     private final Optional<String> source;
     private final Set<String> clientTags;
     private final ResourceEstimates resourceEstimates;
+    private final String queryText;
     private final Optional<String> queryType;
 
     public SelectionCriteria(
@@ -42,6 +43,7 @@ public SelectionCriteria(
             Optional<String> source,
             Set<String> clientTags,
             ResourceEstimates resourceEstimates,
+            String queryText,
             Optional<String> queryType)
     {
         this.authenticated = authenticated;
@@ -52,6 +54,7 @@ public SelectionCriteria(
         this.source = requireNonNull(source, "source is null");
         this.clientTags = Set.copyOf(requireNonNull(clientTags, "clientTags is null"));
         this.resourceEstimates = requireNonNull(resourceEstimates, "resourceEstimates is null");
+        this.queryText = requireNonNull(queryText, "queryText is null");
         this.queryType = requireNonNull(queryType, "queryType is null");
     }
 
@@ -95,6 +98,11 @@ public ResourceEstimates getResourceEstimates()
         return resourceEstimates;
     }
 
+    public String getQueryText()
+    {
+        return queryText;
+    }
+
     public Optional<String> getQueryType()
     {
         return queryType;
@@ -112,6 +120,7 @@ public String toString()
                 .add("source=" + source)
                 .add("clientTags=" + clientTags)
                 .add("resourceEstimates=" + resourceEstimates)
+                .add("queryText=" + queryText)
                 .add("queryType=" + queryType)
                 .toString();
     }

docs/src/main/sphinx/admin/resource-groups.md

@@ -181,6 +181,8 @@ documentation](https://docs.oracle.com/en/java/javase/24/docs/api/java.base/java
 
 - `source` (optional): Java regex to match against source string.
 
+- `queryText` (optional): regex to match against the SQL query string.
+
 - `queryType` (optional): string to match against the type of the query submitted:
 
   - `SELECT`: [SELECT](/sql/select) queries.
@@ -251,7 +253,7 @@ Templates allow administrators to construct resource group trees dynamically. Fo
 the `pipeline_${USER}` group, `${USER}` is expanded to the name of the user that submitted
 the query. `${SOURCE}` is also supported, which is expanded to the source that submitted the
 query. You may also use custom named variables in the regular expressions for `user`, `source`,
-`originalUser`, and `authenticatedUser`.
+`originalUser`, `authenticatedUser` and `queryText`.
 
 There are six selectors, that define which queries run in which resource group:
 

plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/AbstractResourceConfigurationManager.java

@@ -107,6 +107,7 @@ protected List<ResourceGroupSelector> buildSelectors(ManagerSpec managerSpec)
                     spec.getSourceRegex(),
                     spec.getClientTags(),
                     spec.getResourceEstimate(),
+                    spec.getQueryText(),
                     spec.getQueryType(),
                     spec.getGroup()));
         }

plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/SelectorSpec.java

@@ -31,6 +31,7 @@ public class SelectorSpec
     private final Optional<Pattern> originalUserRegex;
     private final Optional<Pattern> authenticatedUserRegex;
     private final Optional<Pattern> sourceRegex;
+    private final Optional<Pattern> queryText;
     private final Optional<String> queryType;
     private final Optional<List<String>> clientTags;
     private final Optional<SelectorResourceEstimate> selectorResourceEstimate;
@@ -43,6 +44,7 @@ public SelectorSpec(
             @JsonProperty("originalUser") Optional<Pattern> originalUserRegex,
             @JsonProperty("authenticatedUser") Optional<Pattern> authenticatedUserRegex,
             @JsonProperty("source") Optional<Pattern> sourceRegex,
+            @JsonProperty("queryText") Optional<Pattern> queryText,
             @JsonProperty("queryType") Optional<String> queryType,
             @JsonProperty("clientTags") Optional<List<String>> clientTags,
             @JsonProperty("selectorResourceEstimate") Optional<SelectorResourceEstimate> selectorResourceEstimate,
@@ -53,6 +55,7 @@ public SelectorSpec(
         this.originalUserRegex = requireNonNull(originalUserRegex, "originalUserRegex is null");
         this.authenticatedUserRegex = requireNonNull(authenticatedUserRegex, "authenticatedUserRegex is null");
         this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null");
+        this.queryText = requireNonNull(queryText, "queryText is null");
         this.queryType = requireNonNull(queryType, "queryType is null");
         this.clientTags = requireNonNull(clientTags, "clientTags is null");
         this.selectorResourceEstimate = requireNonNull(selectorResourceEstimate, "selectorResourceEstimate is null");
@@ -84,6 +87,11 @@ public Optional<Pattern> getSourceRegex()
         return sourceRegex;
     }
 
+    public Optional<Pattern> getQueryText()
+    {
+        return queryText;
+    }
+
     public Optional<String> getQueryType()
     {
         return queryType;
@@ -124,6 +132,8 @@ public boolean equals(Object other)
                 authenticatedUserRegex.map(Pattern::flags).equals(that.authenticatedUserRegex.map(Pattern::flags)) &&
                 sourceRegex.map(Pattern::pattern).equals(that.sourceRegex.map(Pattern::pattern))) &&
                 sourceRegex.map(Pattern::flags).equals(that.sourceRegex.map(Pattern::flags)) &&
+                queryText.map(Pattern::pattern).equals(that.queryText.map(Pattern::pattern)) &&
+                queryText.map(Pattern::flags).equals(that.queryText.map(Pattern::flags)) &&
                 queryType.equals(that.queryType) &&
                 clientTags.equals(that.clientTags);
     }
@@ -143,6 +153,7 @@ public int hashCode()
                 authenticatedUserRegex.map(Pattern::flags),
                 sourceRegex.map(Pattern::pattern),
                 sourceRegex.map(Pattern::flags),
+                queryText,
                 queryType,
                 clientTags);
     }
@@ -162,6 +173,7 @@ public String toString()
                 .add("authenticatedUserFlags", authenticatedUserRegex.map(Pattern::flags))
                 .add("sourceRegex", sourceRegex)
                 .add("sourceFlags", sourceRegex.map(Pattern::flags))
+                .add("queryText", queryText)
                 .add("queryType", queryType)
                 .add("clientTags", clientTags)
                 .toString();

plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/StaticSelector.java

@@ -53,6 +53,7 @@ public StaticSelector(
             Optional<Pattern> sourceRegex,
             Optional<List<String>> clientTags,
             Optional<SelectorResourceEstimate> selectorResourceEstimate,
+            Optional<Pattern> queryTextRegex,
             Optional<String> queryType,
             ResourceGroupIdTemplate group)
     {
@@ -63,6 +64,7 @@ public StaticSelector(
         requireNonNull(sourceRegex, "sourceRegex is null");
         requireNonNull(clientTags, "clientTags is null");
         requireNonNull(selectorResourceEstimate, "selectorResourceEstimate is null");
+        requireNonNull(queryTextRegex, "queryTextRegex is null");
         requireNonNull(queryType, "queryType is null");
         this.group = requireNonNull(group, "group is null");
 
@@ -86,6 +88,7 @@ public StaticSelector(
                 }))
                 .add(userGroupRegex.map(userGroupRegexValue ->
                             new BasicMatcher(criteria -> criteria.getUserGroups().stream().anyMatch(userGroup -> userGroupRegexValue.matcher(userGroup).matches()))))
+                .add(queryTextRegex.map(queryTextRegexValue -> new PatternMatcher(variableNames, queryTextRegexValue, SelectionCriteria::getQueryText)))
                 .add(queryType.map(queryTypeValue ->
                             new BasicMatcher(criteria -> queryTypeValue.equalsIgnoreCase(criteria.getQueryType().orElse("")))))
                 .add(selectorResourceEstimate.map(selectorResourceEstimateValue ->

plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/DbResourceGroupConfigurationManager.java

@@ -388,6 +388,7 @@ private synchronized Map.Entry<ManagerSpec, Map<ResourceGroupIdTemplate, Resourc
                                 selectorRecord.getOriginalUserRegex(),
                                 selectorRecord.getAuthenticatedUserRegex(),
                                 selectorRecord.getSourceRegex(),
+                                selectorRecord.getQueryTextRegex(),
                                 selectorRecord.getQueryType(),
                                 selectorRecord.getClientTags(),
                                 selectorRecord.getSelectorResourceEstimate(),

plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java

@@ -62,7 +62,7 @@ public interface ResourceGroupsDao
     @UseRowMapper(ResourceGroupSpecBuilder.Mapper.class)
     List<ResourceGroupSpecBuilder> getResourceGroups(@Bind("environment") String environment);
 
-    @SqlQuery("SELECT S.resource_group_id, S.priority, S.user_regex, S.source_regex, S.original_user_regex, S.authenticated_user_regex, S.query_type, S.client_tags, S.selector_resource_estimate, S.user_group_regex\n" +
+    @SqlQuery("SELECT S.resource_group_id, S.priority, S.user_regex, S.source_regex, S.original_user_regex, S.authenticated_user_regex, S.query_text_regex, S.query_type, S.client_tags, S.selector_resource_estimate, S.user_group_regex\n" +
             "FROM selectors S\n" +
             "JOIN resource_groups R ON (S.resource_group_id = R.resource_group_id)\n" +
             "WHERE R.environment = :environment\n" +
@@ -78,6 +78,7 @@ public interface ResourceGroupsDao
             "  original_user_regex VARCHAR(512),\n" +
             "  authenticated_user_regex VARCHAR(512),\n" +
             "  source_regex VARCHAR(512),\n" +
+            "  query_text_regex VARCHAR(1024),\n" +
             "  query_type VARCHAR(512),\n" +
             "  client_tags VARCHAR(512),\n" +
             "  selector_resource_estimate VARCHAR(1024),\n" +

plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/SelectorRecord.java

@@ -38,6 +38,7 @@ public class SelectorRecord
     private final Optional<Pattern> originalUserRegex;
     private final Optional<Pattern> authenticatedUserRegex;
     private final Optional<Pattern> sourceRegex;
+    private final Optional<Pattern> queryTextRegex;
     private final Optional<String> queryType;
     private final Optional<List<String>> clientTags;
     private final Optional<SelectorResourceEstimate> selectorResourceEstimate;
@@ -50,6 +51,7 @@ public SelectorRecord(
             Optional<Pattern> originalUserRegex,
             Optional<Pattern> authenticatedUserRegex,
             Optional<Pattern> sourceRegex,
+            Optional<Pattern> queryTextRegex,
             Optional<String> queryType,
             Optional<List<String>> clientTags,
             Optional<SelectorResourceEstimate> selectorResourceEstimate)
@@ -61,6 +63,7 @@ public SelectorRecord(
         this.originalUserRegex = requireNonNull(originalUserRegex, "originalUserRegex is null");
         this.authenticatedUserRegex = requireNonNull(authenticatedUserRegex, "authenticatedUserRegex is null");
         this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null");
+        this.queryTextRegex = requireNonNull(queryTextRegex, "queryTextRegex is null");
         this.queryType = requireNonNull(queryType, "queryType is null");
         this.clientTags = clientTags.map(ImmutableList::copyOf);
         this.selectorResourceEstimate = requireNonNull(selectorResourceEstimate, "selectorResourceEstimate is null");
@@ -101,6 +104,11 @@ public Optional<Pattern> getSourceRegex()
         return sourceRegex;
     }
 
+    public Optional<Pattern> getQueryTextRegex()
+    {
+        return queryTextRegex;
+    }
+
     public Optional<String> getQueryType()
     {
         return queryType;
@@ -134,6 +142,7 @@ public SelectorRecord map(ResultSet resultSet, StatementContext context)
                     Optional.ofNullable(resultSet.getString("original_user_regex")).map(Pattern::compile),
                     Optional.ofNullable(resultSet.getString("authenticated_user_regex")).map(Pattern::compile),
                     Optional.ofNullable(resultSet.getString("source_regex")).map(Pattern::compile),
+                    Optional.ofNullable(resultSet.getString("query_text_regex")).map(Pattern::compile),
                     Optional.ofNullable(resultSet.getString("query_type")),
                     Optional.ofNullable(resultSet.getString("client_tags")).map(LIST_STRING_CODEC::fromJson),
                     Optional.ofNullable(resultSet.getString("selector_resource_estimate")).map(SELECTOR_RESOURCE_ESTIMATE_JSON_CODEC::fromJson));