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

Example 21 with ResourceGroupId

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

the class DbSourceExactMatchSelector method match.

@Override
public Optional<SelectionContext<ResourceGroupIdTemplate>> match(SelectionCriteria criteria) {
    if (criteria.getSource().isEmpty()) {
        return Optional.empty();
    }
    try {
        String resourceGroupId = dao.getExactMatchResourceGroup(environment, criteria.getSource().get(), criteria.getQueryType().orElse(""));
        Long start = daoOfflineStart.get();
        if (start != null && daoOfflineStart.compareAndSet(start, null)) {
            log.info("Successfully fetched exact match selectors after %s", nanosSince(start));
        }
        if (resourceGroupId == null) {
            return Optional.empty();
        }
        ResourceGroupId groupId;
        try {
            groupId = resourceGroupIdCodec.fromJson(resourceGroupId);
        } catch (IllegalArgumentException e) {
            log.warn("Failed to decode resource group from DB: %s", resourceGroupId);
            return Optional.empty();
        }
        return Optional.of(new SelectionContext<>(groupId, toTemplate(groupId)));
    } catch (JdbiException e) {
        if (daoOfflineStart.compareAndSet(null, System.nanoTime())) {
            log.warn(e, "Failed to fetch exact match resource group selectors");
        }
        return Optional.empty();
    }
}

Example 22 with ResourceGroupId

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

the class TestResourceGroupIdTemplate method testExpansion.

@Test
public void testExpansion() {
    ResourceGroupIdTemplate template = new ResourceGroupIdTemplate("test.${USER}.${SOURCE}");
    ResourceGroupId expected = new ResourceGroupId(new ResourceGroupId(new ResourceGroupId("test"), "u"), "s");
    assertEquals(template.expandTemplate(new VariableMap(ImmutableMap.of("USER", "u", "SOURCE", "s"))), expected);
    template = new ResourceGroupIdTemplate("test.${USER}");
    assertEquals(template.expandTemplate(new VariableMap(ImmutableMap.of("USER", "alice.smith", "SOURCE", "s"))), new ResourceGroupId(new ResourceGroupId("test"), "alice.smith"));
}

Example 23 with ResourceGroupId

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

the class TestResourceGroupIdTemplate method testExtraction.

@Test
public void testExtraction() {
    ResourceGroupIdTemplate template = new ResourceGroupIdTemplate("test.pipeline.job_${pipeline}_user:${USER}.${USER}");
    ResourceGroupId expected = new ResourceGroupId(new ResourceGroupId(new ResourceGroupId(new ResourceGroupId("test"), "pipeline"), "job_testpipeline_user:user"), "user");
    Pattern sourcePattern = Pattern.compile("scheduler.important.(?<pipeline>[^\\[]*).*");
    StaticSelector selector = new StaticSelector(Optional.empty(), Optional.empty(), Optional.of(sourcePattern), Optional.empty(), Optional.empty(), Optional.empty(), template);
    SelectionCriteria context = new SelectionCriteria(true, "user", ImmutableSet.of(), Optional.of("scheduler.important.testpipeline[5]"), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
    assertEquals(selector.match(context).map(SelectionContext::getResourceGroupId), Optional.of(expected));
}

Example 24 with ResourceGroupId

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

the class TestResourceGroupIdTemplate method testNoSource.

@Test
public void testNoSource() {
    ResourceGroupIdTemplate template = new ResourceGroupIdTemplate("test.pipeline.${pipeline}.${SOURCE}_s");
    ResourceGroupId expected = new ResourceGroupId(new ResourceGroupId(new ResourceGroupId(new ResourceGroupId("test"), "pipeline"), "testpipeline"), "_s");
    Pattern userPattern = Pattern.compile("scheduler.important.(?<pipeline>[^\\[]*).*");
    StaticSelector selector = new StaticSelector(Optional.of(userPattern), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), template);
    SelectionCriteria context = new SelectionCriteria(true, "scheduler.important.testpipeline[5]", ImmutableSet.of(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
    assertEquals(selector.match(context).map(SelectionContext::getResourceGroupId), Optional.of(expected));
}

Example 25 with ResourceGroupId

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

the class TestStaticSelector method testClientTags.

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