diff --git a/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/DbResourceGroupConfigurationManager.java b/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/DbResourceGroupConfigurationManager.java index 9e0f8f1a673e..b9cb467ae634 100644 --- a/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/DbResourceGroupConfigurationManager.java +++ b/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/DbResourceGroupConfigurationManager.java @@ -313,7 +313,7 @@ private synchronized Map.Entry new SelectorSpec( selectorRecord.getUserRegex(), - Optional.empty(), + selectorRecord.getUserGroupRegex(), selectorRecord.getSourceRegex(), selectorRecord.getQueryType(), selectorRecord.getClientTags(), diff --git a/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java b/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java index 939c9225869c..ef34fa745159 100644 --- a/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java +++ b/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/ResourceGroupsDao.java @@ -60,7 +60,7 @@ public interface ResourceGroupsDao @UseRowMapper(ResourceGroupSpecBuilder.Mapper.class) List getResourceGroups(@Bind("environment") String environment); - @SqlQuery("SELECT S.resource_group_id, S.priority, S.user_regex, S.source_regex, S.query_type, S.client_tags, S.selector_resource_estimate\n" + + @SqlQuery("SELECT S.resource_group_id, S.priority, S.user_regex, S.source_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" + @@ -72,6 +72,7 @@ public interface ResourceGroupsDao " resource_group_id BIGINT NOT NULL,\n" + " priority BIGINT NOT NULL,\n" + " user_regex VARCHAR(512),\n" + + " user_group_regex VARCHAR(512),\n" + " source_regex VARCHAR(512),\n" + " query_type VARCHAR(512),\n" + " client_tags VARCHAR(512),\n" + diff --git a/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/SelectorRecord.java b/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/SelectorRecord.java index f8049a190780..9fa843df6028 100644 --- a/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/SelectorRecord.java +++ b/plugin/trino-resource-group-managers/src/main/java/io/trino/plugin/resourcegroups/db/SelectorRecord.java @@ -34,6 +34,7 @@ public class SelectorRecord private final long resourceGroupId; private final long priority; private final Optional userRegex; + private final Optional userGroupRegex; private final Optional sourceRegex; private final Optional queryType; private final Optional> clientTags; @@ -43,6 +44,7 @@ public SelectorRecord( long resourceGroupId, long priority, Optional userRegex, + Optional userGroupRegex, Optional sourceRegex, Optional queryType, Optional> clientTags, @@ -51,6 +53,7 @@ public SelectorRecord( this.resourceGroupId = resourceGroupId; this.priority = priority; this.userRegex = requireNonNull(userRegex, "userRegex is null"); + this.userGroupRegex = requireNonNull(userGroupRegex, "userGroupRegex is null"); this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null"); this.queryType = requireNonNull(queryType, "queryType is null"); this.clientTags = requireNonNull(clientTags, "clientTags is null").map(ImmutableList::copyOf); @@ -72,6 +75,11 @@ public Optional getUserRegex() return userRegex; } + public Optional getUserGroupRegex() + { + return userGroupRegex; + } + public Optional getSourceRegex() { return sourceRegex; @@ -106,6 +114,7 @@ public SelectorRecord map(ResultSet resultSet, StatementContext context) resultSet.getLong("resource_group_id"), resultSet.getLong("priority"), Optional.ofNullable(resultSet.getString("user_regex")).map(Pattern::compile), + Optional.ofNullable(resultSet.getString("user_group_regex")).map(Pattern::compile), 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), diff --git a/plugin/trino-resource-group-managers/src/main/resources/db/migration/mysql/V5__add_user_group_to_selectors.sql b/plugin/trino-resource-group-managers/src/main/resources/db/migration/mysql/V5__add_user_group_to_selectors.sql new file mode 100644 index 000000000000..25086c3a89fa --- /dev/null +++ b/plugin/trino-resource-group-managers/src/main/resources/db/migration/mysql/V5__add_user_group_to_selectors.sql @@ -0,0 +1 @@ +ALTER TABLE selectors ADD COLUMN user_group_regex VARCHAR(2048); diff --git a/plugin/trino-resource-group-managers/src/main/resources/db/migration/oracle/V5__add_user_group_to_selectors.sql b/plugin/trino-resource-group-managers/src/main/resources/db/migration/oracle/V5__add_user_group_to_selectors.sql new file mode 100644 index 000000000000..a9de4d420116 --- /dev/null +++ b/plugin/trino-resource-group-managers/src/main/resources/db/migration/oracle/V5__add_user_group_to_selectors.sql @@ -0,0 +1 @@ +ALTER TABLE selectors ADD user_group_regex VARCHAR(2048); diff --git a/plugin/trino-resource-group-managers/src/main/resources/db/migration/postgresql/V5__add_user_group_to_selectors.sql b/plugin/trino-resource-group-managers/src/main/resources/db/migration/postgresql/V5__add_user_group_to_selectors.sql new file mode 100644 index 000000000000..25086c3a89fa --- /dev/null +++ b/plugin/trino-resource-group-managers/src/main/resources/db/migration/postgresql/V5__add_user_group_to_selectors.sql @@ -0,0 +1 @@ +ALTER TABLE selectors ADD COLUMN user_group_regex VARCHAR(2048); diff --git a/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/H2ResourceGroupsDao.java b/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/H2ResourceGroupsDao.java index d36ca8898395..8637576115d8 100644 --- a/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/H2ResourceGroupsDao.java +++ b/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/H2ResourceGroupsDao.java @@ -80,12 +80,13 @@ void updateResourceGroup( void deleteResourceGroup(@Bind("resource_group_id") long resourceGroupId); @SqlUpdate("INSERT INTO selectors\n" + - "(resource_group_id, priority, user_regex, source_regex, query_type, client_tags, selector_resource_estimate)\n" + - "VALUES (:resource_group_id, :priority, :user_regex, :source_regex, :query_type, :client_tags, :selector_resource_estimate)") + "(resource_group_id, priority, user_regex, user_group_regex, source_regex, query_type, client_tags, selector_resource_estimate)\n" + + "VALUES (:resource_group_id, :priority, :user_regex, :user_group_regex, :source_regex, :query_type, :client_tags, :selector_resource_estimate)") void insertSelector( @Bind("resource_group_id") long resourceGroupId, @Bind("priority") long priority, @Bind("user_regex") String userRegex, + @Bind("user_group_regex") String userGroupRegex, @Bind("source_regex") String sourceRegex, @Bind("query_type") String queryType, @Bind("client_tags") String clientTags, diff --git a/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestDbResourceGroupConfigurationManager.java b/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestDbResourceGroupConfigurationManager.java index 42755e136abe..b86fe4ff3f4b 100644 --- a/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestDbResourceGroupConfigurationManager.java +++ b/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestDbResourceGroupConfigurationManager.java @@ -44,6 +44,7 @@ import static io.trino.spi.resourcegroups.SchedulingPolicy.WEIGHTED; import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy; import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; @@ -74,8 +75,8 @@ public void testEnvironments() // two resource groups are the same except the group for the prod environment has a larger softMemoryLimit dao.insertResourceGroup(1, "prod_global", "10MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, prodEnvironment); dao.insertResourceGroup(2, "dev_global", "1MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, devEnvironment); - dao.insertSelector(1, 1, ".*prod_user.*", null, null, null, null); - dao.insertSelector(2, 2, ".*dev_user.*", null, null, null, null); + dao.insertSelector(1, 1, ".*prod_user.*", null, null, null, null, null); + dao.insertSelector(2, 2, ".*dev_user.*", null, null, null, null, null); // check the prod configuration DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager((poolId, listener) -> {}, new DbResourceGroupConfig(), daoProvider.get(), prodEnvironment); @@ -112,7 +113,7 @@ public void testConfiguration() dao.insertResourceGroupsGlobalProperties("cpu_quota_period", "1h"); dao.insertResourceGroup(1, "global", "1MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, ENVIRONMENT); dao.insertResourceGroup(2, "sub", "2MB", 4, 3, 3, null, 5, null, null, null, 1L, ENVIRONMENT); - dao.insertSelector(2, 1, null, null, null, null, null); + dao.insertSelector(2, 1, null, null, null, null, null, null); DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager((poolId, listener) -> {}, new DbResourceGroupConfig(), daoProvider.get(), ENVIRONMENT); AtomicBoolean exported = new AtomicBoolean(); InternalResourceGroup global = new InternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor()); @@ -138,7 +139,7 @@ public void testDuplicateRoots() assertTrue(ex.getCause() instanceof org.h2.jdbc.JdbcException); assertTrue(ex.getCause().getMessage().startsWith("Unique index or primary key violation")); }); - dao.insertSelector(1, 1, null, null, null, null, null); + dao.insertSelector(1, 1, null, null, null, null, null, null); } @Test @@ -156,7 +157,7 @@ public void testDuplicateGroups() assertTrue(ex.getCause() instanceof org.h2.jdbc.JdbcException); assertTrue(ex.getCause().getMessage().startsWith("Unique index or primary key violation")); }); - dao.insertSelector(2, 2, null, null, null, null, null); + dao.insertSelector(2, 2, null, null, null, null, null, null); } @Test @@ -170,7 +171,7 @@ public void testMissing() dao.insertResourceGroup(1, "global", "1MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, ENVIRONMENT); dao.insertResourceGroup(2, "sub", "2MB", 4, 3, 3, null, 5, null, null, null, 1L, ENVIRONMENT); dao.insertResourceGroupsGlobalProperties("cpu_quota_period", "1h"); - dao.insertSelector(2, 1, null, null, null, null, null); + dao.insertSelector(2, 1, null, null, null, null, null, null); DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager((poolId, listener) -> {}, new DbResourceGroupConfig(), daoProvider.get(), ENVIRONMENT); InternalResourceGroup missing = new InternalResourceGroup("missing", (group, export) -> {}, directExecutor()); @@ -190,7 +191,7 @@ public void testReconfig() dao.createSelectorsTable(); dao.insertResourceGroup(1, "global", "1MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, ENVIRONMENT); dao.insertResourceGroup(2, "sub", "2MB", 4, 3, 3, null, 5, null, null, null, 1L, ENVIRONMENT); - dao.insertSelector(2, 1, null, null, null, null, null); + dao.insertSelector(2, 1, null, null, null, null, null, null); dao.insertResourceGroupsGlobalProperties("cpu_quota_period", "1h"); DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager((poolId, listener) -> {}, new DbResourceGroupConfig(), daoProvider.get(), ENVIRONMENT); manager.start(); @@ -228,7 +229,7 @@ public void testExactMatchSelector() dao.createExactMatchSelectorsTable(); dao.insertResourceGroup(1, "global", "1MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, ENVIRONMENT); dao.insertResourceGroup(2, "sub", "2MB", 4, 3, 3, null, 5, null, null, null, 1L, ENVIRONMENT); - dao.insertSelector(2, 1, null, null, null, null, null); + dao.insertSelector(2, 1, null, null, null, null, null, null); dao.insertResourceGroupsGlobalProperties("cpu_quota_period", "1h"); DbResourceGroupConfig config = new DbResourceGroupConfig(); config.setExactMatchSelectorEnabled(true); @@ -267,7 +268,7 @@ public void testSelectorPriority() for (int i = 0; i < numberOfUsers; i++) { int priority = randomPriorities[i]; String user = String.valueOf(priority); - dao.insertSelector(1, priority, user, ".*", null, null, null); + dao.insertSelector(1, priority, user, null, ".*", null, null, null); expectedUsers.add(user); } @@ -335,6 +336,53 @@ public void testRefreshInterval() manager.destroy(); } + @Test + public void testMatchByUserGroups() + { + H2DaoProvider daoProvider = setup("selectors"); + H2ResourceGroupsDao dao = daoProvider.get(); + dao.createResourceGroupsGlobalPropertiesTable(); + dao.createResourceGroupsTable(); + dao.createSelectorsTable(); + dao.insertResourceGroup(1, "group", "100%", 100, 100, 100, null, null, null, null, null, null, ENVIRONMENT); + dao.insertSelector(1, 1, null, "first matching|second matching", null, null, null, null); + + DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager( + (poolId, listener) -> {}, + new DbResourceGroupConfig().setMaxRefreshInterval(new io.airlift.units.Duration(1, MILLISECONDS)), + daoProvider.get(), + ENVIRONMENT); + + assertThat(manager.match(userGroupsSelectionCriteria("not matching"))).isEmpty(); + assertThat(manager.match(userGroupsSelectionCriteria("first matching"))) + .map(SelectionContext::getContext) + .isEqualTo(Optional.of(new ResourceGroupIdTemplate("group"))); + } + + @Test + public void testMatchByUsersAndGroups() + { + H2DaoProvider daoProvider = setup("selectors"); + H2ResourceGroupsDao dao = daoProvider.get(); + dao.createResourceGroupsGlobalPropertiesTable(); + dao.createResourceGroupsTable(); + dao.createSelectorsTable(); + dao.insertResourceGroup(1, "group", "100%", 100, 100, 100, null, null, null, null, null, null, ENVIRONMENT); + dao.insertSelector(1, 1, "Matching user", "Matching group", null, null, null, null); + + DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager( + (poolId, listener) -> {}, + new DbResourceGroupConfig().setMaxRefreshInterval(new io.airlift.units.Duration(1, MILLISECONDS)), + daoProvider.get(), + ENVIRONMENT); + + assertThat(manager.match(userAndUserGroupsSelectionCriteria("Matching user", "Not matching group"))).isEmpty(); + assertThat(manager.match(userAndUserGroupsSelectionCriteria("Not matching user", "Matching group"))).isEmpty(); + assertThat(manager.match(userAndUserGroupsSelectionCriteria("Matching user", "Matching group"))) + .map(SelectionContext::getContext) + .isEqualTo(Optional.of(new ResourceGroupIdTemplate("group"))); + } + private static void assertEqualsResourceGroup( InternalResourceGroup group, String softMemoryLimit, @@ -357,4 +405,23 @@ private static void assertEqualsResourceGroup( assertEquals(group.getSoftCpuLimit(), softCpuLimit); assertEquals(group.getHardCpuLimit(), hardCpuLimit); } + + private static SelectionCriteria userGroupsSelectionCriteria(String... groups) + { + return new SelectionCriteria(true, "test_user", ImmutableSet.copyOf(groups), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty()); + } + + private static SelectionCriteria userAndUserGroupsSelectionCriteria(String user, String group, String... groups) + { + return new SelectionCriteria( + true, + user, + ImmutableSet.builder() + .add(group) + .add(groups).build(), + Optional.empty(), + ImmutableSet.of(), + EMPTY_RESOURCE_ESTIMATES, + Optional.empty()); + } } diff --git a/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestResourceGroupsDao.java b/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestResourceGroupsDao.java index dd3a7d257771..0ed1b03631da 100644 --- a/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestResourceGroupsDao.java +++ b/plugin/trino-resource-group-managers/src/test/java/io/trino/plugin/resourcegroups/db/TestResourceGroupsDao.java @@ -126,6 +126,7 @@ private static void testSelectorInsert(H2ResourceGroupsDao dao, Map records = dao.getSelectors(ENVIRONMENT); compareSelectors(map, records); } @@ -168,6 +171,7 @@ private static void testSelectorUpdate(H2ResourceGroupsDao dao, Map map) { - SelectorRecord updated = new SelectorRecord(2, 3L, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()); + SelectorRecord updated = new SelectorRecord(2, 3L, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()); map.put(2L, updated); dao.updateSelector(2, null, null, null, "ping.*", "ping_source", LIST_STRING_CODEC.toJson(ImmutableList.of("tag1"))); compareSelectors(map, dao.getSelectors(ENVIRONMENT)); @@ -186,6 +190,7 @@ private static void testSelectorUpdateNull(H2ResourceGroupsDao dao, Map map) { dao.updateSelector(3, null, null, null, "admin_user", ".*", LIST_STRING_CODEC.toJson(ImmutableList.of("tag1", "tag2"))); - SelectorRecord nullRegexes = new SelectorRecord(3L, 2L, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()); + SelectorRecord nullRegexes = new SelectorRecord(3L, 2L, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()); map.put(3L, nullRegexes); compareSelectors(map, dao.getSelectors(ENVIRONMENT)); dao.deleteSelector(3, null, null, null); @@ -219,11 +224,12 @@ private static void testSelectorMultiDelete(H2ResourceGroupsDao dao, Map