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

Example 1 with ResourceGroupId

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

the class InternalResourceGroupManager method createGroupIfNecessary.

private synchronized void createGroupIfNecessary(SelectionContext<C> context, Executor executor) {
    ResourceGroupId id = context.getResourceGroupId();
    if (!groups.containsKey(id)) {
        InternalResourceGroup group;
        if (id.getParent().isPresent()) {
            createGroupIfNecessary(configurationManager.get().parentGroupContext(context), executor);
            InternalResourceGroup parent = groups.get(id.getParent().get());
            requireNonNull(parent, "parent is null");
            group = parent.getOrCreateSubGroup(id.getLastSegment());
        } else {
            InternalResourceGroup root = new InternalResourceGroup(id.getSegments().get(0), this::exportGroup, executor);
            group = root;
            rootGroups.add(root);
        }
        configurationManager.get().configure(group, context);
        checkState(groups.put(id, group) == null, "Unexpected existing resource group");
    }
}

Example 2 with ResourceGroupId

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

the class TestResetSessionTask method test.

@Test
public void test() {
    Session session = testSessionBuilder(sessionPropertyManager).setSystemProperty("foo", "bar").setCatalogSessionProperty(CATALOG_NAME, "baz", "blah").build();
    QueryStateMachine stateMachine = QueryStateMachine.begin(Optional.empty(), "reset foo", Optional.empty(), session, URI.create("fake://uri"), new ResourceGroupId("test"), false, transactionManager, accessControl, executor, metadata, WarningCollector.NOOP, Optional.empty());
    getFutureValue(new ResetSessionTask(metadata, sessionPropertyManager).execute(new ResetSession(QualifiedName.of(CATALOG_NAME, "baz")), stateMachine, emptyList(), WarningCollector.NOOP));
    Set<String> sessionProperties = stateMachine.getResetSessionProperties();
    assertEquals(sessionProperties, ImmutableSet.of(CATALOG_NAME + ".baz"));
}

Example 3 with ResourceGroupId

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

the class TestSetRoleTask method executeSetRole.

private QueryStateMachine executeSetRole(String statement) {
    SetRole setRole = (SetRole) parser.createStatement(statement, new ParsingOptions());
    QueryStateMachine stateMachine = QueryStateMachine.begin(Optional.empty(), statement, Optional.empty(), testSessionBuilder().setIdentity(Identity.ofUser(USER_NAME)).build(), URI.create("fake://uri"), new ResourceGroupId("test"), false, transactionManager, accessControl, executor, metadata, WarningCollector.NOOP, Optional.empty());
    new SetRoleTask(metadata, accessControl).execute(setRole, stateMachine, ImmutableList.of(), WarningCollector.NOOP);
    return stateMachine;
}

Example 4 with ResourceGroupId

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

the class TestQueryStateMachine method createQueryStateMachineWithTicker.

private QueryStateMachine createQueryStateMachineWithTicker(Ticker ticker) {
    Metadata metadata = createTestMetadataManager();
    TransactionManager transactionManager = createTestTransactionManager();
    AccessControlManager accessControl = new AccessControlManager(transactionManager, emptyEventListenerManager(), new AccessControlConfig(), DefaultSystemAccessControl.NAME);
    accessControl.setSystemAccessControls(List.of(AllowAllSystemAccessControl.INSTANCE));
    QueryStateMachine stateMachine = QueryStateMachine.beginWithTicker(Optional.empty(), QUERY, Optional.empty(), TEST_SESSION, LOCATION, new ResourceGroupId("test"), false, transactionManager, accessControl, executor, ticker, metadata, WarningCollector.NOOP, QUERY_TYPE);
    stateMachine.setInputs(INPUTS);
    stateMachine.setOutput(OUTPUT);
    stateMachine.setColumns(OUTPUT_FIELD_NAMES, OUTPUT_FIELD_TYPES);
    stateMachine.setUpdateType(UPDATE_TYPE);
    for (Entry<String, String> entry : SET_SESSION_PROPERTIES.entrySet()) {
        stateMachine.addSetSessionProperties(entry.getKey(), entry.getValue());
    }
    RESET_SESSION_PROPERTIES.forEach(stateMachine::addResetSessionProperties);
    return stateMachine;
}

Example 5 with ResourceGroupId

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

the class TestPrepareTask method executePrepare.

private Map<String, String> executePrepare(String statementName, Statement statement, String sqlString, Session session) {
    TransactionManager transactionManager = createTestTransactionManager();
    AccessControlManager accessControl = new AccessControlManager(transactionManager, emptyEventListenerManager(), new AccessControlConfig(), DefaultSystemAccessControl.NAME);
    accessControl.setSystemAccessControls(List.of(AllowAllSystemAccessControl.INSTANCE));
    QueryStateMachine stateMachine = QueryStateMachine.begin(Optional.empty(), sqlString, Optional.empty(), session, URI.create("fake://uri"), new ResourceGroupId("test"), false, transactionManager, accessControl, executor, metadata, WarningCollector.NOOP, Optional.empty());
    Prepare prepare = new Prepare(identifier(statementName), statement);
    new PrepareTask(new SqlParser()).execute(prepare, stateMachine, emptyList(), WarningCollector.NOOP);
    return stateMachine.getAddedPreparedStatements();
}