Examples with SelectionCriteria com.facebook.presto.spi.resourceGroups.SelectionCriteria used on opensource projects

Example 1 with SelectionCriteria

use of com.facebook.presto.spi.resourceGroups.SelectionCriteria in project presto by prestodb.

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(), template);
    SelectionCriteria context = new SelectionCriteria(true, "scheduler.important.testpipeline[5]", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
    assertEquals(selector.match(context).map(SelectionContext::getResourceGroupId), Optional.of(expected));
}

Example 2 with SelectionCriteria

use of com.facebook.presto.spi.resourceGroups.SelectionCriteria in project presto by prestodb.

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.of(sourcePattern), Optional.empty(), Optional.empty(), Optional.empty(), template);
    SelectionCriteria context = new SelectionCriteria(true, "user", Optional.of("scheduler.important.testpipeline[5]"), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
    selector.match(context);
}

Example 3 with SelectionCriteria

use of com.facebook.presto.spi.resourceGroups.SelectionCriteria in project presto by prestodb.

the class TestResourceGroupIdTemplate method testUnresolvedVariableRunTime.

@Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "unresolved variable 'pipeline' in resource group '\\$\\{pipeline\\}', available.*")
public void testUnresolvedVariableRunTime() {
    ResourceGroupIdTemplate template = new ResourceGroupIdTemplate("test.pipeline.${pipeline}.${USER}");
    Pattern sourcePattern = Pattern.compile("scheduler.important.(testpipeline\\[|(?<pipeline>[^\\[]*)).*");
    StaticSelector selector = new StaticSelector(Optional.empty(), Optional.of(sourcePattern), Optional.empty(), Optional.empty(), Optional.empty(), template);
    SelectionCriteria context = new SelectionCriteria(true, "user", Optional.of("scheduler.important.testpipeline[5]"), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty());
    selector.match(context);
}

Example 4 with SelectionCriteria

use of com.facebook.presto.spi.resourceGroups.SelectionCriteria in project presto by prestodb.

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, String slug, int retryCount, SessionContext sessionContext, String query, ResourceGroupManager<C> resourceGroupManager) {
    Session session = null;
    PreparedQuery preparedQuery;
    try {
        if (query.length() > maxQueryLength) {
            int queryLength = query.length();
            query = query.substring(0, maxQueryLength);
            throw new PrestoException(QUERY_TEXT_TOO_LARGE, format("Query text length (%s) exceeds the maximum length (%s)", queryLength, maxQueryLength));
        }
        // decode session
        session = sessionSupplier.createSession(queryId, sessionContext);
        // prepare query
        WarningCollector warningCollector = warningCollectorFactory.create(getWarningHandlingLevel(session));
        preparedQuery = queryPreparer.prepareQuery(session, query, warningCollector);
        query = preparedQuery.getFormattedQuery().orElse(query);
        // select resource group
        Optional<QueryType> queryType = getQueryType(preparedQuery.getStatement().getClass());
        SelectionContext<C> selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria(sessionContext.getIdentity().getPrincipal().isPresent(), sessionContext.getIdentity().getUser(), Optional.ofNullable(sessionContext.getSource()), sessionContext.getClientTags(), sessionContext.getResourceEstimates(), queryType.map(Enum::name)));
        // apply system default session properties (does not override user set properties)
        session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType.map(Enum::name), Optional.of(selectionContext.getResourceGroupId()));
        // mark existing transaction as active
        transactionManager.activateTransaction(session, isTransactionControlStatement(preparedQuery.getStatement()), accessControl);
        DispatchQuery dispatchQuery = dispatchQueryFactory.createDispatchQuery(session, query, preparedQuery, slug, retryCount, selectionContext.getResourceGroupId(), queryType, warningCollector, (dq) -> resourceGroupManager.submit(preparedQuery.getStatement(), dq, selectionContext, queryExecutor));
        boolean queryAdded = queryCreated(dispatchQuery);
        if (queryAdded && !dispatchQuery.isDone()) {
            try {
                clusterStatusSender.registerQuery(dispatchQuery);
                dispatchQuery.startWaitingForPrerequisites();
            } 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(new SessionPropertyManager()).setQueryId(queryId).setIdentity(sessionContext.getIdentity()).setSource(sessionContext.getSource()).build();
        }
        DispatchQuery failedDispatchQuery = failedDispatchQueryFactory.createFailedDispatchQuery(session, query, Optional.empty(), throwable);
        queryCreated(failedDispatchQuery);
    }
}

Example 5 with SelectionCriteria

use of com.facebook.presto.spi.resourceGroups.SelectionCriteria in project presto by prestodb.

the class TestFileResourceGroupConfigurationManager method testQueryTypeConfiguration.

@Test
public void testQueryTypeConfiguration() {
    FileResourceGroupConfigurationManager manager = parse("resource_groups_config_query_type.json");
    List<ResourceGroupSelector> selectors = manager.getSelectors();
    assertMatch(selectors, new SelectionCriteria(true, "test_user", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of("select")), "global.select");
    assertMatch(selectors, new SelectionCriteria(true, "test_user", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of("explain")), "global.explain");
    assertMatch(selectors, new SelectionCriteria(true, "test_user", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of("insert")), "global.insert");
    assertMatch(selectors, new SelectionCriteria(true, "test_user", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of("delete")), "global.delete");
    assertMatch(selectors, new SelectionCriteria(true, "test_user", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of("describe")), "global.describe");
    assertMatch(selectors, new SelectionCriteria(true, "test_user", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of("data_definition")), "global.data_definition");
    assertMatch(selectors, new SelectionCriteria(true, "test_user", Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.of("sth_else")), "global.other");
}