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

Example 1 with SelectionCriteria

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

the class DispatchManager method createQueryInternal.

/**
 * Creates and registers a dispatch query with the query tracker.  This method will never fail to register a query with the query
 * tracker.  If an error occurs while creating a dispatch query, a failed dispatch will be created and registered.
 */
private <C> void createQueryInternal(QueryId queryId, Slug slug, SessionContext sessionContext, String query, ResourceGroupManager<C> resourceGroupManager) {
    Session session = null;
    PreparedQuery preparedQuery = null;
    try {
        if (query.length() > maxQueryLength) {
            int queryLength = query.length();
            query = query.substring(0, maxQueryLength);
            throw new TrinoException(QUERY_TEXT_TOO_LARGE, format("Query text length (%s) exceeds the maximum length (%s)", queryLength, maxQueryLength));
        }
        // decode session
        session = sessionSupplier.createSession(queryId, sessionContext);
        // check query execute permissions
        accessControl.checkCanExecuteQuery(sessionContext.getIdentity());
        // prepare query
        preparedQuery = queryPreparer.prepareQuery(session, query);
        // select resource group
        Optional<String> queryType = getQueryType(preparedQuery.getStatement()).map(Enum::name);
        SelectionContext<C> selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria(sessionContext.getIdentity().getPrincipal().isPresent(), sessionContext.getIdentity().getUser(), sessionContext.getIdentity().getGroups(), sessionContext.getSource(), sessionContext.getClientTags(), sessionContext.getResourceEstimates(), queryType));
        // apply system default session properties (does not override user set properties)
        session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType, selectionContext.getResourceGroupId());
        DispatchQuery dispatchQuery = dispatchQueryFactory.createDispatchQuery(session, sessionContext.getTransactionId(), query, preparedQuery, slug, selectionContext.getResourceGroupId());
        boolean queryAdded = queryCreated(dispatchQuery);
        if (queryAdded && !dispatchQuery.isDone()) {
            try {
                resourceGroupManager.submit(dispatchQuery, selectionContext, dispatchExecutor);
            } catch (Throwable e) {
                // dispatch query has already been registered, so just fail it directly
                dispatchQuery.fail(e);
            }
        }
    } catch (Throwable throwable) {
        // creation must never fail, so register a failed query in this case
        if (session == null) {
            session = Session.builder(sessionPropertyManager).setQueryId(queryId).setIdentity(sessionContext.getIdentity()).setSource(sessionContext.getSource().orElse(null)).build();
        }
        Optional<String> preparedSql = Optional.ofNullable(preparedQuery).flatMap(PreparedQuery::getPrepareSql);
        DispatchQuery failedDispatchQuery = failedDispatchQueryFactory.createFailedDispatchQuery(session, query, preparedSql, Optional.empty(), throwable);
        queryCreated(failedDispatchQuery);
    }
}

Example 2 with SelectionCriteria

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

the class TestDbResourceGroupConfigurationManager method testEnvironments.

@Test
public void testEnvironments() {
    H2DaoProvider daoProvider = setup("test_configuration");
    H2ResourceGroupsDao dao = daoProvider.get();
    dao.createResourceGroupsGlobalPropertiesTable();
    dao.createResourceGroupsTable();
    dao.createSelectorsTable();
    String prodEnvironment = "prod";
    String devEnvironment = "dev";
    dao.insertResourceGroupsGlobalProperties("cpu_quota_period", "1h");
    // 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, null);
    dao.insertSelector(2, 2, ".*dev_user.*", null, null, null, null, null);
    // check the prod configuration
    DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager(listener -> {
    }, new DbResourceGroupConfig(), daoProvider.get(), prodEnvironment);
    List<ResourceGroupSpec> groups = manager.getRootGroups();
    assertEquals(groups.size(), 1);
    InternalResourceGroup prodGlobal = new InternalResourceGroup("prod_global", (group, export) -> {
    }, directExecutor());
    manager.configure(prodGlobal, new SelectionContext<>(prodGlobal.getId(), new ResourceGroupIdTemplate("prod_global")));
    assertEqualsResourceGroup(prodGlobal, "10MB", 1000, 100, 100, WEIGHTED, DEFAULT_WEIGHT, true, Duration.ofHours(1), Duration.ofDays(1));
    assertEquals(manager.getSelectors().size(), 1);
    ResourceGroupSelector prodSelector = manager.getSelectors().get(0);
    ResourceGroupId prodResourceGroupId = prodSelector.match(new SelectionCriteria(true, "prod_user", ImmutableSet.of(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty())).get().getResourceGroupId();
    assertEquals(prodResourceGroupId.toString(), "prod_global");
    // check the dev configuration
    manager = new DbResourceGroupConfigurationManager(listener -> {
    }, new DbResourceGroupConfig(), daoProvider.get(), devEnvironment);
    assertEquals(groups.size(), 1);
    InternalResourceGroup devGlobal = new InternalResourceGroup("dev_global", (group, export) -> {
    }, directExecutor());
    manager.configure(devGlobal, new SelectionContext<>(prodGlobal.getId(), new ResourceGroupIdTemplate("dev_global")));
    assertEqualsResourceGroup(devGlobal, "1MB", 1000, 100, 100, WEIGHTED, DEFAULT_WEIGHT, true, Duration.ofHours(1), Duration.ofDays(1));
    assertEquals(manager.getSelectors().size(), 1);
    ResourceGroupSelector devSelector = manager.getSelectors().get(0);
    ResourceGroupId devResourceGroupId = devSelector.match(new SelectionCriteria(true, "dev_user", ImmutableSet.of(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty())).get().getResourceGroupId();
    assertEquals(devResourceGroupId.toString(), "dev_global");
}

Example 3 with SelectionCriteria

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

the class TestResourceGroupIdTemplate method testUnresolvedVariableLoadTime.

@Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "unresolved variables \\[user\\] in resource group ID.*")
public void testUnresolvedVariableLoadTime() {
    ResourceGroupIdTemplate template = new ResourceGroupIdTemplate("test.pipeline.${pipeline}.${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());
    selector.match(context);
}

Example 4 with SelectionCriteria

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

the class TestResourceGroupIdTemplate method testNoMatch.

@Test
public void testNoMatch() {
    ResourceGroupIdTemplate template = new ResourceGroupIdTemplate("test.pipeline.${pipeline}.${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.testpipeline[5]"), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
    assertFalse(selector.match(context).isPresent());
}

Example 5 with SelectionCriteria

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

the class StaticSelector method match.

@Override
public Optional<SelectionContext<ResourceGroupIdTemplate>> match(SelectionCriteria criteria) {
    Map<String, String> variables = new HashMap<>();
    if (userRegex.isPresent()) {
        Matcher userMatcher = userRegex.get().matcher(criteria.getUser());
        if (!userMatcher.matches()) {
            return Optional.empty();
        }
        addVariableValues(userRegex.get(), criteria.getUser(), variables);
    }
    if (userGroupRegex.isPresent() && criteria.getUserGroups().stream().noneMatch(group -> userGroupRegex.get().matcher(group).matches())) {
        return Optional.empty();
    }
    if (sourceRegex.isPresent()) {
        String source = criteria.getSource().orElse("");
        if (!sourceRegex.get().matcher(source).matches()) {
            return Optional.empty();
        }
        addVariableValues(sourceRegex.get(), source, variables);
    }
    if (!clientTags.isEmpty() && !criteria.getTags().containsAll(clientTags)) {
        return Optional.empty();
    }
    if (selectorResourceEstimate.isPresent() && !selectorResourceEstimate.get().match(criteria.getResourceEstimates())) {
        return Optional.empty();
    }
    if (queryType.isPresent()) {
        String contextQueryType = criteria.getQueryType().orElse("");
        if (!queryType.get().equalsIgnoreCase(contextQueryType)) {
            return Optional.empty();
        }
    }
    variables.putIfAbsent(USER_VARIABLE, criteria.getUser());
    // Special handling for source, which is an optional field that is part of the standard variables
    variables.putIfAbsent(SOURCE_VARIABLE, criteria.getSource().orElse(""));
    ResourceGroupId id = group.expandTemplate(new VariableMap(variables));
    return Optional.of(new SelectionContext<>(id, group));
}