Skip to content

Added isExecuteImmediate selector to resource groups #27079

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Added isExecuteImmediate selector to resource groups #27079

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion core/trino-main/src/main/java/io/trino/dispatcher/DispatchManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,8 @@ private <C> void createQueryInternal(QueryId queryId, Span querySpan, Slug slug,
sessionContext.getSource(),
sessionContext.getClientTags(),
sessionContext.getResourceEstimates(),
queryType));
queryType,
preparedQuery.isExecuteImmediate()));

// apply system default session properties (does not override user set properties)
session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType, selectionContext.getResourceGroupId());
Expand Down
16 changes: 15 additions & 1 deletion core/trino-main/src/main/java/io/trino/execution/QueryPreparer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,11 +60,13 @@ public PreparedQuery prepareQuery(Session session, Statement wrappedStatement)
{
Statement statement = wrappedStatement;
Optional<String> prepareSql = Optional.empty();
boolean isExecuteImmediate = false;
if (statement instanceof Execute executeStatement) {
prepareSql = Optional.of(session.getPreparedStatementFromExecute(executeStatement));
statement = sqlParser.createStatement(prepareSql.get());
}
else if (statement instanceof ExecuteImmediate executeImmediateStatement) {
isExecuteImmediate = true;
statement = sqlParser.createStatement(
executeImmediateStatement.getStatement().getValue(),
executeImmediateStatement.getStatement().getLocation().orElseThrow());
Expand All @@ -86,7 +88,7 @@ else if (wrappedStatement instanceof ExecuteImmediate executeImmediateStatement)
}
validateParameters(statement, parameters);

return new PreparedQuery(statement, parameters, prepareSql);
return new PreparedQuery(statement, parameters, prepareSql, isExecuteImmediate);
}

private static void validateParameters(Statement node, List<Expression> parameterValues)
Expand All @@ -105,12 +107,19 @@ public static class PreparedQuery
private final Statement statement;
private final List<Expression> parameters;
private final Optional<String> prepareSql;
private final boolean isExecuteImmediate;

public PreparedQuery(Statement statement, List<Expression> parameters, Optional<String> prepareSql)
{
this(statement, parameters, prepareSql, false);
}

public PreparedQuery(Statement statement, List<Expression> parameters, Optional<String> prepareSql, boolean isExecuteImmediate)
{
this.statement = requireNonNull(statement, "statement is null");
this.parameters = ImmutableList.copyOf(requireNonNull(parameters, "parameters is null"));
this.prepareSql = requireNonNull(prepareSql, "prepareSql is null");
this.isExecuteImmediate = isExecuteImmediate;
}

public Statement getStatement()
Expand All @@ -127,5 +136,10 @@ public Optional<String> getPrepareSql()
{
return prepareSql;
}

public boolean isExecuteImmediate()
{
return isExecuteImmediate;
}
}
}
11 changes: 10 additions & 1 deletion core/trino-spi/src/main/java/io/trino/spi/resourcegroups/SelectionCriteria.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ public final class SelectionCriteria
private final Set<String> clientTags;
private final ResourceEstimates resourceEstimates;
private final Optional<String> queryType;
private final boolean isExecuteImmediate;

public SelectionCriteria(
boolean authenticated,
Expand All @@ -42,7 +43,8 @@ public SelectionCriteria(
Optional<String> source,
Set<String> clientTags,
ResourceEstimates resourceEstimates,
Optional<String> queryType)
Optional<String> queryType,
boolean isExecuteImmediate)
{
this.authenticated = authenticated;
this.user = requireNonNull(user, "user is null");
Expand All @@ -53,6 +55,7 @@ public SelectionCriteria(
this.clientTags = Set.copyOf(requireNonNull(clientTags, "clientTags is null"));
this.resourceEstimates = requireNonNull(resourceEstimates, "resourceEstimates is null");
this.queryType = requireNonNull(queryType, "queryType is null");
this.isExecuteImmediate = isExecuteImmediate;
}

public boolean isAuthenticated()
Expand Down Expand Up @@ -100,6 +103,11 @@ public Optional<String> getQueryType()
return queryType;
}

public boolean isExecuteImmediate()
{
return isExecuteImmediate;
}

@Override
public String toString()
{
Expand All @@ -113,6 +121,7 @@ public String toString()
.add("clientTags=" + clientTags)
.add("resourceEstimates=" + resourceEstimates)
.add("queryType=" + queryType)
.add("isExecuteImmediate=" + isExecuteImmediate)
.toString();
}
}
5 changes: 5 additions & 0 deletions docs/src/main/sphinx/admin/resource-groups.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,11 @@ documentation](https://docs.oracle.com/en/java/javase/24/docs/api/java.base/java
- `clientTags` (optional): list of tags. To match, every tag in this list must be in the list of
client-provided tags associated with the query.

- `isExecuteImmediate` (optional): Some clients use `EXECUTE IMMEDIATE ...`
to run quick queries directly. Use `true` when this selector should match a
`EXECUTE IMMEDIATE '...'` query. Use `false` to match all other
queries, omit to match any.
Comment on lines +209 to +212
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion (typo): Change 'a EXECUTE IMMEDIATE' to 'an EXECUTE IMMEDIATE' for grammatical correctness.

Use 'an EXECUTE IMMEDIATE ... query' for correct grammar.

Suggested change
- `isExecuteImmediate` (optional): Some clients use `EXECUTE IMMEDIATE ...`
to run quick queries directly. Use `true` when this selector should match a
`EXECUTE IMMEDIATE '...'` query. Use `false` to match all other
queries, omit to match any.
- `isExecuteImmediate` (optional): Some clients use `EXECUTE IMMEDIATE ...`
to run quick queries directly. Use `true` when this selector should match an
`EXECUTE IMMEDIATE ...` query. Use `false` to match all other
queries, omit to match any.


- `group` (required): the group these queries will run in.

All rules within a single selector are combined using a logical `AND`. Therefore
Expand Down
1 change: 1 addition & 0 deletions ...rs/src/main/java/io/trino/plugin/resourcegroups/AbstractResourceConfigurationManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ protected List<ResourceGroupSelector> buildSelectors(ManagerSpec managerSpec)
spec.getClientTags(),
spec.getResourceEstimate(),
spec.getQueryType(),
spec.isExecuteImmediate(),
spec.getGroup()));
}
return selectors.build();
Expand Down
15 changes: 13 additions & 2 deletions ...no-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/SelectorSpec.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ public class SelectorSpec
private final Optional<String> queryType;
private final Optional<List<String>> clientTags;
private final Optional<SelectorResourceEstimate> selectorResourceEstimate;
private final Optional<Boolean> isExecuteImmediate;
private final ResourceGroupIdTemplate group;

@JsonCreator
Expand All @@ -46,6 +47,7 @@ public SelectorSpec(
@JsonProperty("queryType") Optional<String> queryType,
@JsonProperty("clientTags") Optional<List<String>> clientTags,
@JsonProperty("selectorResourceEstimate") Optional<SelectorResourceEstimate> selectorResourceEstimate,
@JsonProperty("isExecuteImmediate") Optional<Boolean> isExecuteImmediate,
@JsonProperty("group") ResourceGroupIdTemplate group)
{
this.userRegex = requireNonNull(userRegex, "userRegex is null");
Expand All @@ -56,6 +58,7 @@ public SelectorSpec(
this.queryType = requireNonNull(queryType, "queryType is null");
this.clientTags = requireNonNull(clientTags, "clientTags is null");
this.selectorResourceEstimate = requireNonNull(selectorResourceEstimate, "selectorResourceEstimate is null");
this.isExecuteImmediate = requireNonNull(isExecuteImmediate, "isExecuteImmediate is null");
this.group = requireNonNull(group, "group is null");
}

Expand Down Expand Up @@ -99,6 +102,11 @@ public Optional<SelectorResourceEstimate> getResourceEstimate()
return selectorResourceEstimate;
}

public Optional<Boolean> isExecuteImmediate()
{
return isExecuteImmediate;
}

public ResourceGroupIdTemplate getGroup()
{
return group;
Expand All @@ -125,7 +133,8 @@ public boolean equals(Object other)
sourceRegex.map(Pattern::pattern).equals(that.sourceRegex.map(Pattern::pattern))) &&
sourceRegex.map(Pattern::flags).equals(that.sourceRegex.map(Pattern::flags)) &&
queryType.equals(that.queryType) &&
clientTags.equals(that.clientTags);
clientTags.equals(that.clientTags) &&
isExecuteImmediate.equals(that.isExecuteImmediate);
}

@Override
Expand All @@ -144,7 +153,8 @@ public int hashCode()
sourceRegex.map(Pattern::pattern),
sourceRegex.map(Pattern::flags),
queryType,
clientTags);
clientTags,
isExecuteImmediate);
}

@Override
Expand All @@ -164,6 +174,7 @@ public String toString()
.add("sourceFlags", sourceRegex.map(Pattern::flags))
.add("queryType", queryType)
.add("clientTags", clientTags)
.add("isExecuteImmediate", isExecuteImmediate)
.toString();
}
}
4 changes: 4 additions & 0 deletions ...-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/StaticSelector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ public StaticSelector(
Optional<List<String>> clientTags,
Optional<SelectorResourceEstimate> selectorResourceEstimate,
Optional<String> queryType,
Optional<Boolean> isExecuteImmediate,
ResourceGroupIdTemplate group)
{
this.userRegex = requireNonNull(userRegex, "userRegex is null");
Expand All @@ -64,6 +65,7 @@ public StaticSelector(
requireNonNull(clientTags, "clientTags is null");
requireNonNull(selectorResourceEstimate, "selectorResourceEstimate is null");
requireNonNull(queryType, "queryType is null");
requireNonNull(isExecuteImmediate, "isExecuteImmediate is null");
this.group = requireNonNull(group, "group is null");

HashSet<String> variableNames = new HashSet<>(ImmutableList.of(USER_VARIABLE, SOURCE_VARIABLE));
Expand Down Expand Up @@ -92,6 +94,8 @@ public StaticSelector(
new BasicMatcher(criteria -> selectorResourceEstimateValue.match(criteria.getResourceEstimates()))))
.add(clientTags.map(clientTagsValue ->
new BasicMatcher(criteria -> criteria.getTags().containsAll(clientTagsValue))))
.add(isExecuteImmediate.map(isExecuteImmediateValue ->
new BasicMatcher(criteria -> isExecuteImmediateValue == criteria.isExecuteImmediate())))
.build()
.stream()
.flatMap(Optional::stream) // remove any empty optionals
Expand Down
1 change: 1 addition & 0 deletions .../src/main/java/io/trino/plugin/resourcegroups/db/DbResourceGroupConfigurationManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,7 @@ private synchronized Map.Entry<ManagerSpec, Map<ResourceGroupIdTemplate, Resourc
selectorRecord.getQueryType(),
selectorRecord.getClientTags(),
selectorRecord.getSelectorResourceEstimate(),
selectorRecord.isExecuteImmediate(),
resourceGroupIdTemplateMap.get(selectorRecord.getResourceGroupId()))
).collect(Collectors.toList());

Expand Down
3 changes: 2 additions & 1 deletion ...rce-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_type, S.client_tags, S.selector_resource_estimate, S.user_group_regex, S.is_execute_immediate\n" +
"FROM selectors S\n" +
"JOIN resource_groups R ON (S.resource_group_id = R.resource_group_id)\n" +
"WHERE R.environment = :environment\n" +
Expand All @@ -81,6 +81,7 @@ public interface ResourceGroupsDao
" query_type VARCHAR(512),\n" +
" client_tags VARCHAR(512),\n" +
" selector_resource_estimate VARCHAR(1024),\n" +
" is_execute_immediate VARCHAR(6),\n" +
" FOREIGN KEY (resource_group_id) REFERENCES resource_groups (resource_group_id)\n" +
")")
void createSelectorsTable();
Expand Down
13 changes: 11 additions & 2 deletions ...source-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/SelectorRecord.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ public class SelectorRecord
private final Optional<String> queryType;
private final Optional<List<String>> clientTags;
private final Optional<SelectorResourceEstimate> selectorResourceEstimate;
private final Optional<Boolean> isExecuteImmediate;

public SelectorRecord(
long resourceGroupId,
Expand All @@ -52,7 +53,8 @@ public SelectorRecord(
Optional<Pattern> sourceRegex,
Optional<String> queryType,
Optional<List<String>> clientTags,
Optional<SelectorResourceEstimate> selectorResourceEstimate)
Optional<SelectorResourceEstimate> selectorResourceEstimate,
Optional<Boolean> isExecuteImmediate)
{
this.resourceGroupId = resourceGroupId;
this.priority = priority;
Expand All @@ -64,6 +66,7 @@ public SelectorRecord(
this.queryType = requireNonNull(queryType, "queryType is null");
this.clientTags = clientTags.map(ImmutableList::copyOf);
this.selectorResourceEstimate = requireNonNull(selectorResourceEstimate, "selectorResourceEstimate is null");
this.isExecuteImmediate = requireNonNull(isExecuteImmediate, "isExecuteImmediate is null");
}

public long getResourceGroupId()
Expand Down Expand Up @@ -116,6 +119,11 @@ public Optional<SelectorResourceEstimate> getSelectorResourceEstimate()
return selectorResourceEstimate;
}

public Optional<Boolean> isExecuteImmediate()
{
return isExecuteImmediate;
}

public static class Mapper
implements RowMapper<SelectorRecord>
{
Expand All @@ -136,7 +144,8 @@ public SelectorRecord map(ResultSet resultSet, StatementContext context)
Optional.ofNullable(resultSet.getString("source_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));
Optional.ofNullable(resultSet.getString("selector_resource_estimate")).map(SELECTOR_RESOURCE_ESTIMATE_JSON_CODEC::fromJson),
Optional.ofNullable(resultSet.getString("is_execute_immediate")).map(Boolean::parseBoolean));
}
}
}
17 changes: 10 additions & 7 deletions ...c/test/java/io/trino/plugin/resourcegroups/TestFileResourceGroupConfigurationManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,7 +264,8 @@ public void testDocsExample()
Optional.of("jdbc#powerfulbi"),
ImmutableSet.of("hipri"),
EMPTY_RESOURCE_ESTIMATES,
Optional.of("select")));
Optional.of("select"),
false));
assertThat(selectionContext.getResourceGroupId().toString()).isEqualTo("global.adhoc.bi-powerfulbi.Alice");
TestingResourceGroup resourceGroup = new TestingResourceGroup(selectionContext.getResourceGroupId());
manager.configure(resourceGroup, selectionContext);
Expand All @@ -281,7 +282,8 @@ public void testDocsExample()
Optional.empty(),
ImmutableSet.of(),
EMPTY_RESOURCE_ESTIMATES,
Optional.empty()));
Optional.empty(),
false));
assertThat(selectionContext.getResourceGroupId().toString()).isEqualTo("global.adhoc.other.Amanda");
resourceGroup = new TestingResourceGroup(selectionContext.getResourceGroupId());
manager.configure(resourceGroup, selectionContext);
Expand Down Expand Up @@ -339,7 +341,7 @@ private static ManagerSpec parseManagerSpec(String fileName)

private static SelectionCriteria userAndSourceSelectionCriteria(String user, String source)
{
return new SelectionCriteria(true, user, ImmutableSet.of(), user, Optional.empty(), Optional.of(source), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
return new SelectionCriteria(true, user, ImmutableSet.of(), user, Optional.empty(), Optional.of(source), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty(), false);
}

private static SelectionCriteria userSelectionCriteria(String user)
Expand All @@ -349,17 +351,17 @@ private static SelectionCriteria userSelectionCriteria(String user)

private static SelectionCriteria identitySelectionCriteria(String user, String originalUser, Optional<String> authenticatedUser)
{
return new SelectionCriteria(true, user, ImmutableSet.of(), originalUser, authenticatedUser, Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
return new SelectionCriteria(true, user, ImmutableSet.of(), originalUser, authenticatedUser, Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty(), false);
}

private static SelectionCriteria queryTypeSelectionCriteria(String queryType)
{
return new SelectionCriteria(true, "test_user", ImmutableSet.of(), "test_user", Optional.empty(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of(queryType));
return new SelectionCriteria(true, "test_user", ImmutableSet.of(), "test_user", Optional.empty(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of(queryType), false);
}

private static SelectionCriteria userGroupsSelectionCriteria(String... groups)
{
return new SelectionCriteria(true, "test_user", ImmutableSet.copyOf(groups), "test_user", Optional.empty(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
return new SelectionCriteria(true, "test_user", ImmutableSet.copyOf(groups), "test_user", Optional.empty(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty(), false);
}

private static SelectionCriteria userAndUserGroupsSelectionCriteria(String user, String group, String... groups)
Expand All @@ -375,6 +377,7 @@ private static SelectionCriteria userAndUserGroupsSelectionCriteria(String user,
Optional.empty(),
ImmutableSet.of(),
EMPTY_RESOURCE_ESTIMATES,
Optional.empty());
Optional.empty(),
false);
}
}
Loading
Loading