Examples with ResourceGroupId io.trino.spi.resourcegroups.ResourceGroupId used on opensource projects

Example 26 with ResourceGroupId

use of io.trino.spi.resourcegroups.ResourceGroupId in project trino by trinodb.

the class TestQueuesDb method testRunningQuery.

@Test(timeOut = 60_000)
public void testRunningQuery() throws Exception {
    queryRunner.execute("SELECT COUNT(*), clerk FROM orders GROUP BY clerk");
    while (true) {
        ResourceGroupInfo global = queryRunner.getCoordinator().getResourceGroupManager().get().tryGetResourceGroupInfo(new ResourceGroupId(new ResourceGroupId("global"), "bi-user")).orElseThrow(() -> new IllegalStateException("Resource group not found"));
        if (global.getSoftMemoryLimit().toBytes() > 0) {
            break;
        }
        TimeUnit.SECONDS.sleep(2);
    }
}

Example 27 with ResourceGroupId

use of io.trino.spi.resourcegroups.ResourceGroupId in project trino by trinodb.

the class TestQueuesDb method assertResourceGroupWithClientTags.

private void assertResourceGroupWithClientTags(Set<String> clientTags, ResourceGroupId expectedResourceGroup) throws InterruptedException {
    Session session = testSessionBuilder().setCatalog("tpch").setSchema("sf100000").setSource("client_tags").setClientTags(clientTags).build();
    QueryId queryId = createQuery(queryRunner, session, LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED));
    Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId();
    assertTrue(resourceGroupId.isPresent(), "Query should have a resource group");
    assertEquals(resourceGroupId.get(), expectedResourceGroup, format("Expected: '%s' resource group, found: %s", expectedResourceGroup, resourceGroupId.get()));
}

Example 28 with ResourceGroupId

use of io.trino.spi.resourcegroups.ResourceGroupId in project trino by trinodb.

the class TestDbSourceExactMatchSelector method testMatch.

@Test
public void testMatch() {
    ResourceGroupId resourceGroupId1 = new ResourceGroupId(ImmutableList.of("global", "test", "user", "insert"));
    ResourceGroupId resourceGroupId2 = new ResourceGroupId(ImmutableList.of("global", "test", "user", "select"));
    dao.insertExactMatchSelector("test", "@test@test_pipeline", INSERT.name(), CODEC.toJson(resourceGroupId1));
    dao.insertExactMatchSelector("test", "@test@test_pipeline", SELECT.name(), CODEC.toJson(resourceGroupId2));
    DbSourceExactMatchSelector selector = new DbSourceExactMatchSelector("test", dao);
    assertEquals(selector.match(new SelectionCriteria(true, "testuser", ImmutableSet.of(), Optional.of("@test@test_pipeline"), ImmutableSet.of("tag"), EMPTY_RESOURCE_ESTIMATES, Optional.empty())), Optional.empty());
    assertEquals(selector.match(new SelectionCriteria(true, "testuser", ImmutableSet.of(), Optional.of("@test@test_pipeline"), ImmutableSet.of("tag"), EMPTY_RESOURCE_ESTIMATES, Optional.of(INSERT.name()))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId1));
    assertEquals(selector.match(new SelectionCriteria(true, "testuser", ImmutableSet.of(), Optional.of("@test@test_pipeline"), ImmutableSet.of("tag"), EMPTY_RESOURCE_ESTIMATES, Optional.of(SELECT.name()))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId2));
    assertEquals(selector.match(new SelectionCriteria(true, "testuser", ImmutableSet.of(), Optional.of("@test@test_pipeline"), ImmutableSet.of("tag"), EMPTY_RESOURCE_ESTIMATES, Optional.of(DELETE.name()))), Optional.empty());
    assertEquals(selector.match(new SelectionCriteria(true, "testuser", ImmutableSet.of(), Optional.of("@test@test_new"), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of(INSERT.name()))), Optional.empty());
}

Example 29 with ResourceGroupId

use of io.trino.spi.resourcegroups.ResourceGroupId in project trino by trinodb.

the class TestStaticSelector method testSelectorResourceEstimate.

@Test
public void testSelectorResourceEstimate() {
    ResourceGroupId resourceGroupId = new ResourceGroupId(new ResourceGroupId("global"), "foo");
    StaticSelector smallQuerySelector = new StaticSelector(Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.of(new SelectorResourceEstimate(Optional.of(new Range<>(Optional.empty(), Optional.of(new Duration(5, MINUTES)))), Optional.empty(), Optional.of(new Range<>(Optional.empty(), Optional.of(DataSize.valueOf("500MB")))))), Optional.empty(), new ResourceGroupIdTemplate("global.foo"));
    assertEquals(smallQuerySelector.match(newSelectionCriteria("userA", null, ImmutableSet.of("tag1", "tag2"), new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.of(DataSize.of(400, MEGABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
    assertEquals(smallQuerySelector.match(newSelectionCriteria("A.user", "a source b", ImmutableSet.of("tag1"), new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.of(DataSize.of(600, MEGABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.empty());
    assertEquals(smallQuerySelector.match(newSelectionCriteria("userB", "source", ImmutableSet.of(), new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.empty()))).map(SelectionContext::getResourceGroupId), Optional.empty());
    StaticSelector largeQuerySelector = new StaticSelector(Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.of(new SelectorResourceEstimate(Optional.empty(), Optional.empty(), Optional.of(new Range<>(Optional.of(DataSize.valueOf("5TB")), Optional.empty())))), Optional.empty(), new ResourceGroupIdTemplate("global.foo"));
    assertEquals(largeQuerySelector.match(newSelectionCriteria("userA", null, ImmutableSet.of("tag1", "tag2"), new ResourceEstimates(Optional.of(java.time.Duration.ofHours(100)), Optional.empty(), Optional.of(DataSize.of(4, TERABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.empty());
    assertEquals(largeQuerySelector.match(newSelectionCriteria("A.user", "a source b", ImmutableSet.of("tag1"), new ResourceEstimates(Optional.empty(), Optional.empty(), Optional.of(DataSize.of(6, TERABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
    assertEquals(largeQuerySelector.match(newSelectionCriteria("userB", "source", ImmutableSet.of(), new ResourceEstimates(Optional.of(java.time.Duration.ofSeconds(1)), Optional.of(java.time.Duration.ofSeconds(1)), Optional.of(DataSize.of(6, TERABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
}

Example 30 with ResourceGroupId

use of io.trino.spi.resourcegroups.ResourceGroupId in project trino by trinodb.

the class TestStaticSelector method testUserRegex.

@Test
public void testUserRegex() {
    ResourceGroupId resourceGroupId = new ResourceGroupId(new ResourceGroupId("global"), "foo");
    StaticSelector selector = new StaticSelector(Optional.of(Pattern.compile("user.*")), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), new ResourceGroupIdTemplate("global.foo"));
    assertEquals(selector.match(newSelectionCriteria("userA", null, ImmutableSet.of("tag1"), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
    assertEquals(selector.match(newSelectionCriteria("userB", "source", ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
    assertEquals(selector.match(newSelectionCriteria("A.user", null, ImmutableSet.of("tag1"), EMPTY_RESOURCE_ESTIMATES)), Optional.empty());
}