Examples with QueryId io.trino.spi.QueryId used on opensource projects

Example 31 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestFileBasedSystemAccessControl method testSchemaOperationsReadOnly.

@Test
public void testSchemaOperationsReadOnly() {
    TransactionManager transactionManager = createTestTransactionManager();
    AccessControlManager accessControlManager = newAccessControlManager(transactionManager, "catalog_read_only.json");
    transaction(transactionManager, accessControlManager).execute(transactionId -> {
        Set<String> aliceSchemas = ImmutableSet.of("schema");
        assertEquals(accessControlManager.filterSchemas(new SecurityContext(transactionId, alice, queryId), "alice-catalog", aliceSchemas), aliceSchemas);
        assertEquals(accessControlManager.filterSchemas(new SecurityContext(transactionId, bob, queryId), "alice-catalog", aliceSchemas), ImmutableSet.of());
        accessControlManager.checkCanShowSchemas(new SecurityContext(transactionId, alice, queryId), "alice-catalog");
    });
    assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanCreateSchema(new SecurityContext(transactionId, alice, queryId), aliceSchema);
    })).isInstanceOf(AccessDeniedException.class).hasMessage("Access Denied: Cannot create schema alice-catalog.schema");
    assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanDropSchema(new SecurityContext(transactionId, alice, queryId), aliceSchema);
    })).isInstanceOf(AccessDeniedException.class).hasMessage("Access Denied: Cannot drop schema alice-catalog.schema");
    assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanRenameSchema(new SecurityContext(transactionId, alice, queryId), aliceSchema, "new-schema");
    })).isInstanceOf(AccessDeniedException.class).hasMessage("Access Denied: Cannot rename schema from alice-catalog.schema to new-schema");
    assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanCreateSchema(new SecurityContext(transactionId, bob, queryId), aliceSchema);
    })).isInstanceOf(AccessDeniedException.class).hasMessage("Access Denied: Cannot access catalog alice-catalog");
}

Example 32 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestFileBasedSystemAccessControl method testRefreshing.

@Test
public void testRefreshing() throws Exception {
    TransactionManager transactionManager = createTestTransactionManager();
    AccessControlManager accessControlManager = new AccessControlManager(transactionManager, emptyEventListenerManager(), new AccessControlConfig(), DefaultSystemAccessControl.NAME);
    File configFile = newTemporaryFile();
    configFile.deleteOnExit();
    copy(new File(getResourcePath("catalog.json")), configFile);
    accessControlManager.setSystemAccessControl(FileBasedSystemAccessControl.NAME, ImmutableMap.of(SECURITY_CONFIG_FILE, configFile.getAbsolutePath(), SECURITY_REFRESH_PERIOD, "1ms"));
    transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanCreateView(new SecurityContext(transactionId, alice, queryId), aliceView);
        accessControlManager.checkCanCreateView(new SecurityContext(transactionId, alice, queryId), aliceView);
        accessControlManager.checkCanCreateView(new SecurityContext(transactionId, alice, queryId), aliceView);
    });
    copy(new File(getResourcePath("security-config-file-with-unknown-rules.json")), configFile);
    sleep(2);
    assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanCreateView(new SecurityContext(transactionId, alice, queryId), aliceView);
    })).isInstanceOf(IllegalArgumentException.class).hasMessageStartingWith("Invalid JSON file");
    // test if file based cached control was not cached somewhere
    assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanCreateView(new SecurityContext(transactionId, alice, queryId), aliceView);
    })).isInstanceOf(IllegalArgumentException.class).hasMessageStartingWith("Invalid JSON file");
    copy(new File(getResourcePath("catalog.json")), configFile);
    sleep(2);
    transaction(transactionManager, accessControlManager).execute(transactionId -> {
        accessControlManager.checkCanCreateView(new SecurityContext(transactionId, alice, queryId), aliceView);
    });
}

Example 33 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class BaseJdbcConnectorTest method testCancellation.

@Test(timeOut = 60_000)
public void testCancellation() throws Exception {
    if (!hasBehavior(SUPPORTS_CANCELLATION)) {
        throw new SkipException("Cancellation is not supported by given connector");
    }
    try (TestView sleepingView = createSleepingView(new Duration(1, MINUTES))) {
        String query = "SELECT * FROM " + sleepingView.getName();
        Future<?> future = executor.submit(() -> assertQueryFails(query, "Query killed. Message: Killed by test"));
        QueryId queryId = getQueryId(query);
        assertEventually(() -> assertRemoteQueryStatus(sleepingView.getName(), RUNNING));
        assertUpdate(format("CALL system.runtime.kill_query(query_id => '%s', message => '%s')", queryId, "Killed by test"));
        future.get();
        assertEventually(() -> assertRemoteQueryStatus(sleepingView.getName(), CANCELLED));
    }
}

Example 34 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class BaseJdbcConnectorTest method getQueryId.

private QueryId getQueryId(String query) throws Exception {
    for (int i = 0; i < 100; i++) {
        MaterializedResult queriesResult = getQueryRunner().execute(format("SELECT query_id FROM system.runtime.queries WHERE query = '%s' AND query NOT LIKE '%%system.runtime.queries%%'", query));
        int rowCount = queriesResult.getRowCount();
        if (rowCount == 0) {
            Thread.sleep(100);
            continue;
        }
        checkState(rowCount == 1, "Too many (%s) query ids were found for: %s", rowCount, query);
        return new QueryId((String) queriesResult.getOnlyValue());
    }
    throw new IllegalStateException("Query id not found for: " + query);
}

Example 35 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestTaskExecutor method testMinDriversPerTaskWhenTargetConcurrencyIncreases.

@Test
public void testMinDriversPerTaskWhenTargetConcurrencyIncreases() {
    MultilevelSplitQueue splitQueue = new MultilevelSplitQueue(2);
    TestingTicker ticker = new TestingTicker();
    // create a task executor with min/max drivers per task to be 2
    TaskExecutor taskExecutor = new TaskExecutor(4, 1, 2, 2, splitQueue, ticker);
    taskExecutor.start();
    try {
        TaskHandle testTaskHandle = taskExecutor.addTask(new TaskId(new StageId(new QueryId("test"), 0), 0, 0), // make sure buffer is underutilized
        () -> 0, 1, new Duration(1, MILLISECONDS), OptionalInt.of(2));
        // create 3 splits
        int batchCount = 3;
        TestingJob[] splits = new TestingJob[3];
        Phaser[] phasers = new Phaser[batchCount];
        for (int batch = 0; batch < batchCount; batch++) {
            phasers[batch] = new Phaser();
            phasers[batch].register();
            TestingJob split = new TestingJob(ticker, new Phaser(), new Phaser(), phasers[batch], 1, 0);
            splits[batch] = split;
        }
        taskExecutor.enqueueSplits(testTaskHandle, false, ImmutableList.copyOf(splits));
        // wait until first split starts
        waitUntilSplitsStart(ImmutableList.of(splits[0]));
        // remaining splits shouldn't start because initial split concurrency is 1
        assertSplitStates(0, splits);
        // complete first split (SplitConcurrencyController for TaskHandle should increase concurrency since buffer is underutilized)
        phasers[0].arriveAndDeregister();
        // 2 remaining splits should be started
        waitUntilSplitsStart(ImmutableList.of(splits[1], splits[2]));
    } finally {
        taskExecutor.stop();
    }
}