Added queryText regex pattern to match in resource group selectors #27129
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Reviewer's GuideIntroduce an optional queryText regex parameter to the resource group selection pipeline by extending the core SPI, dispatch manager, selector specs and records, DAOs, static and DB configuration managers, and updating matching logic, tests, and documentation to prioritize queries based on SQL text patterns. Entity relationship diagram for updated selectors tableerDiagram
SELECTORS {
VARCHAR(512) resource_group_id
INTEGER priority
VARCHAR(512) user_regex
VARCHAR(512) source_regex
VARCHAR(512) original_user_regex
VARCHAR(512) authenticated_user_regex
VARCHAR(512) query_text
VARCHAR(512) query_type
VARCHAR(512) client_tags
VARCHAR(1024) selector_resource_estimate
VARCHAR(512) user_group_regex
}
RESOURCE_GROUPS {
VARCHAR(512) resource_group_id
VARCHAR(512) environment
}
SELECTORS ||--o{ RESOURCE_GROUPS : "resource_group_id"
Class diagram for SelectorSpec, SelectorRecord, and SelectionCriteria changesclassDiagram
class SelectorSpec {
+Optional<Pattern> originalUserRegex
+Optional<Pattern> authenticatedUserRegex
+Optional<Pattern> sourceRegex
+Optional<Pattern> queryText
+Optional<String> queryType
+Optional<List<String>> clientTags
+Optional<SelectorResourceEstimate> selectorResourceEstimate
+ResourceGroupIdTemplate group
+getQueryText()
}
class SelectorRecord {
+Optional<Pattern> originalUserRegex
+Optional<Pattern> authenticatedUserRegex
+Optional<Pattern> sourceRegex
+Optional<Pattern> queryText
+Optional<String> queryType
+Optional<List<String>> clientTags
+Optional<SelectorResourceEstimate> selectorResourceEstimate
+getQueryText()
}
class SelectionCriteria {
+boolean authenticated
+String user
+Set<String> userGroups
+Optional<String> source
+Set<String> clientTags
+ResourceEstimates resourceEstimates
+String queryText
+Optional<String> queryType
+getQueryText()
}
Flow diagram for query dispatch with queryText-based resource group selectionflowchart TD
Q["Incoming SQL Query"] --> S["DispatchManager creates SelectionCriteria (includes queryText)"]
S --> M["ResourceGroupSelector matches against queryText regex"]
M --> G["Assigns to ResourceGroupIdTemplate (priority group if matched)"]
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey there - I've reviewed your changes and they look great!
Prompt for AI Agents
Please address the comments from this code review:
## Individual Comments
### Comment 1
<location> `plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/StaticSelector.java:67` </location>
<code_context>
requireNonNull(sourceRegex, "sourceRegex is null");
requireNonNull(clientTags, "clientTags is null");
requireNonNull(selectorResourceEstimate, "selectorResourceEstimate is null");
+ requireNonNull(queryType, "queryText is null");
requireNonNull(queryType, "queryType is null");
this.group = requireNonNull(group, "group is null");
</code_context>
<issue_to_address>
**issue (bug_risk):** Incorrect requireNonNull parameter and error message for queryText.
The requireNonNull check for queryText should use queryText as the parameter and the error message should reference queryText, not queryType.
</issue_to_address>
### Comment 2
<location> `core/trino-spi/src/main/java/io/trino/spi/resourcegroups/SelectionCriteria.java:123` </location>
<code_context>
.add("source=" + source)
.add("clientTags=" + clientTags)
.add("resourceEstimates=" + resourceEstimates)
+ .add("queryText=queryText")
.add("queryType=" + queryType)
.toString();
</code_context>
<issue_to_address>
**issue (bug_risk):** Incorrect toString output for queryText field.
Replace .add("queryText=queryText") with .add("queryText=" + queryText) to show the actual field value.
</issue_to_address>
### Comment 3
<location> `plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/TestStaticSelector.java:332-341` </location>
<code_context>
private static final long MEMORY_POOL_SIZE = 31415926535900L; // arbitrary uneven value for testing
+ private static final String QUERY = "select 1";
@Test
public void testInvalid()
</code_context>
<issue_to_address>
**suggestion (testing):** Missing edge case tests for queryText regex matching.
Add tests for null, empty, very long, partially matching, and special character queryText values to ensure robust regex handling.
</issue_to_address>
### Comment 4
<location> `plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestDbSourceExactMatchSelector.java:57-62` </location>
<code_context>
+ assertThat(selector.match(new SelectionCriteria(true, "testuser", ImmutableSet.of(), "testuser", Optional.empty(), Optional.of("@test@test_pipeline"), ImmutableSet.of("tag"), EMPTY_RESOURCE_ESTIMATES, QUERY, Optional.empty()))).isEqualTo(Optional.empty());
</code_context>
<issue_to_address>
**suggestion (testing):** No tests for queryText matching in DB selectors.
Please add tests to ensure DB selectors handle queryText regex matching correctly, including cases with null or empty queryText.
</issue_to_address>
### Comment 5
<location> `plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestDbResourceGroupConfigurationManager.java:102-105` </location>
<code_context>
assertThat(manager.getSelectors()).hasSize(1);
ResourceGroupSelector prodSelector = manager.getSelectors().get(0);
- ResourceGroupId prodResourceGroupId = prodSelector.match(new SelectionCriteria(true, "prod_user", ImmutableSet.of(), "prod_user", Optional.empty(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty())).get().getResourceGroupId();
+ ResourceGroupId prodResourceGroupId = prodSelector.match(new SelectionCriteria(true, "prod_user", ImmutableSet.of(), "prod_user", Optional.empty(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, QUERY, Optional.empty())).get().getResourceGroupId();
assertThat(prodResourceGroupId.toString()).isEqualTo("prod_global");
</code_context>
<issue_to_address>
**suggestion (testing):** No tests for queryText-based selector configuration in DB resource group manager.
Add tests to cover both matching and non-matching scenarios for selectors using queryText regex in the DB resource group manager.
</issue_to_address>
### Comment 6
<location> `plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestResourceGroupsDao.java:131-135` </location>
<code_context>
Optional.of(Pattern.compile(".*")),
Optional.empty(),
Optional.empty(),
+ Optional.empty(),
</code_context>
<issue_to_address>
**suggestion (testing):** No test for queryText field persistence in selector DAO.
Add tests to ensure queryText is properly handled in the DAO, covering insert, update, delete, and edge cases like null and empty strings.
```suggestion
Optional.of(Pattern.compile(".*")),
Optional.empty(),
Optional.empty(),
Optional.empty(),
Optional.empty()));
@Test
public void testInsertSelectorWithQueryText() {
SelectorRecord record = new SelectorRecord(
Optional.of(Pattern.compile("user1")),
Optional.of(Pattern.compile("role1")),
Optional.of(Pattern.compile("source1")),
Optional.of("SELECT * FROM test_table"),
Optional.of(EXPLAIN.name())
);
long id = dao.insertSelector(record);
SelectorRecord loaded = dao.getSelector(id);
assertEquals(loaded.getQueryText(), Optional.of("SELECT * FROM test_table"));
}
@Test
public void testUpdateSelectorQueryText() {
SelectorRecord record = new SelectorRecord(
Optional.of(Pattern.compile("user2")),
Optional.of(Pattern.compile("role2")),
Optional.of(Pattern.compile("source2")),
Optional.of("SELECT 1"),
Optional.of(EXPLAIN.name())
);
long id = dao.insertSelector(record);
SelectorRecord updated = new SelectorRecord(
record.getUserRegex(),
record.getRoleRegex(),
record.getSourceRegex(),
Optional.of("SELECT 2"),
record.getClientTags()
);
dao.updateSelector(id, updated);
SelectorRecord loaded = dao.getSelector(id);
assertEquals(loaded.getQueryText(), Optional.of("SELECT 2"));
}
@Test
public void testDeleteSelectorWithQueryText() {
SelectorRecord record = new SelectorRecord(
Optional.of(Pattern.compile("user3")),
Optional.of(Pattern.compile("role3")),
Optional.of(Pattern.compile("source3")),
Optional.of("DELETE FROM test_table"),
Optional.of(EXPLAIN.name())
);
long id = dao.insertSelector(record);
dao.deleteSelector(id);
assertNull(dao.getSelector(id));
}
@Test
public void testSelectorWithNullQueryText() {
SelectorRecord record = new SelectorRecord(
Optional.of(Pattern.compile("user4")),
Optional.of(Pattern.compile("role4")),
Optional.of(Pattern.compile("source4")),
Optional.empty(),
Optional.of(EXPLAIN.name())
);
long id = dao.insertSelector(record);
SelectorRecord loaded = dao.getSelector(id);
assertEquals(loaded.getQueryText(), Optional.empty());
}
@Test
public void testSelectorWithEmptyQueryText() {
SelectorRecord record = new SelectorRecord(
Optional.of(Pattern.compile("user5")),
Optional.of(Pattern.compile("role5")),
Optional.of(Pattern.compile("source5")),
Optional.of(""),
Optional.of(EXPLAIN.name())
);
long id = dao.insertSelector(record);
SelectorRecord loaded = dao.getSelector(id);
assertEquals(loaded.getQueryText(), Optional.of(""));
}
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
...ino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/StaticSelector.java
Outdated
Show resolved
Hide resolved
core/trino-spi/src/main/java/io/trino/spi/resourcegroups/SelectionCriteria.java
Outdated
Show resolved
Hide resolved
gertjanal
force-pushed
the
resource-selector-query-text
branch
from
44baab8 to
3df6a9b
Compare
xkrogen
reviewed
Oct 28, 2025
core/trino-spi/src/main/java/io/trino/spi/resourcegroups/SelectionCriteria.java
Show resolved
Hide resolved
...trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/SelectorSpec.java
Outdated
Show resolved
Hide resolved
...source-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java
Outdated
Show resolved
Hide resolved
...source-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
This doesn't have any tests that actually leverage your new functionality; it only updates existing tests. We should update H2ResourceGroupsDao to support the new regex field, and add corresponding test cases here as well as TestResourceGroupsDao
There was a problem hiding this comment.
I can't really figure out how to add a test to TestResourceGroupsDao, there is no test for any other pattern in that class. Can you tell me how to fix this?
There was a problem hiding this comment.
@gertjanal you can see an example in #24662:
- Changes to H2ResourceGroupsDao: https://github.com/trinodb/trino/pull/24662/files#diff-8e773d414c808d8bcd3d376576a2b514315ca6debb43d1eebd5ab27ecd653a6a
- Changes to TestResourcesGroupsDao: https://github.com/trinodb/trino/pull/24662/files#diff-fb183a6bc30d5c35512673713ffbda4961c6919b10ec8bbec745b6f6c4f6aa79
...resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/TestStaticSelector.java
Show resolved
Hide resolved
.../src/test/java/io/trino/plugin/resourcegroups/TestFileResourceGroupConfigurationManager.java
Show resolved
Hide resolved
gertjanal
force-pushed
the
resource-selector-query-text
branch
from
3df6a9b to
7b80187
Compare
gertjanal
force-pushed
the
resource-selector-query-text
branch
from
7b80187 to
4fcae56
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Some queries are quick and simple, like
SELECT 1and a query to list all tables. These queries should not have to wait in the queue to be executed, but should return quickly for a responsive UI.With this PR, a resource selector
queryTextregex pattern can be added to match these simple types of queries and the query can be passed on to a priority resource group.Additional context and related issues
This feature is also available in the Starburst version: https://docs.starburst.io/latest/admin/resource-groups.html#selector-rules
Release notes
( ) This is not user-visible or is docs only, and no release notes are required.
( ) Release notes are required. Please propose a release note for me.
( X ) Release notes are required, with the following suggested text:
Summary by Sourcery
Introduce
queryTextregex support for resource group selectors to route queries based on their SQL content, updating SPI, plugin, database mapping, dispatch logic, tests, and documentation.New Features:
queryTextregex in resource group selectors to match SQL query stringsEnhancements:
Documentation:
queryTextselector property in resource groups documentationTests: