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

Example 11 with QueryId

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

the class TestBroadcastOutputBuffer method createBroadcastBuffer.

private BroadcastOutputBuffer createBroadcastBuffer(OutputBuffers outputBuffers, DataSize dataSize, AggregatedMemoryContext memoryContext, Executor notificationExecutor) {
    BroadcastOutputBuffer buffer = new BroadcastOutputBuffer(TASK_INSTANCE_ID, new OutputBufferStateMachine(new TaskId(new StageId(new QueryId("query"), 0), 0, 0), stateNotificationExecutor), dataSize, () -> memoryContext.newLocalMemoryContext("test"), notificationExecutor, () -> {
    });
    buffer.setOutputBuffers(outputBuffers);
    return buffer;
}

Example 12 with QueryId

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

the class TestQueryIdGenerator method testCreateNextQueryId.

@Test
public void testCreateNextQueryId() {
    TestIdGenerator idGenerator = new TestIdGenerator();
    long millis = epochMillis(2001, 7, 14, 1, 2, 3, 4);
    idGenerator.setNow(millis);
    // generate ids to 99,999
    for (int i = 0; i < 100_000; i++) {
        assertEquals(idGenerator.createNextQueryId(), new QueryId(format("20010714_010203_%05d_%s", i, idGenerator.getCoordinatorId())));
    }
    // next id will cause counter to roll, but we need to add a second to the time or code will block for ever
    millis += 1000;
    idGenerator.setNow(millis);
    for (int i = 0; i < 100_000; i++) {
        assertEquals(idGenerator.createNextQueryId(), new QueryId(format("20010714_010204_%05d_%s", i, idGenerator.getCoordinatorId())));
    }
    // move forward one more second and generate 100 ids
    millis += 1000;
    idGenerator.setNow(millis);
    for (int i = 0; i < 100; i++) {
        assertEquals(idGenerator.createNextQueryId(), new QueryId(format("20010714_010205_%05d_%s", i, idGenerator.getCoordinatorId())));
    }
    // move forward one more second and verify counter not reset
    millis += 1000;
    idGenerator.setNow(millis);
    for (int i = 100; i < 200; i++) {
        assertEquals(idGenerator.createNextQueryId(), new QueryId(format("20010714_010206_%05d_%s", i, idGenerator.getCoordinatorId())));
    }
    // now we move to the start of the next day, and the counter should reset
    millis = epochMillis(2001, 7, 15, 0, 0, 0, 0);
    idGenerator.setNow(millis);
    for (int i = 0; i < 90_123; i++) {
        assertEquals(idGenerator.createNextQueryId(), new QueryId(format("20010715_000000_%05d_%s", i, idGenerator.getCoordinatorId())));
    }
    // moving forward one second with counter close to the limit causes it to roll
    millis += 1000;
    idGenerator.setNow(millis);
    for (int i = 0; i < 100_000; i++) {
        assertEquals(idGenerator.createNextQueryId(), new QueryId(format("20010715_000001_%05d_%s", i, idGenerator.getCoordinatorId())));
    }
}

Example 13 with QueryId

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

the class TestFileBasedSystemAccessControl method testSchemaOperations.

@Test
public void testSchemaOperations() {
    TransactionManager transactionManager = createTestTransactionManager();
    AccessControlManager accessControlManager = newAccessControlManager(transactionManager, "catalog.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.checkCanCreateSchema(new SecurityContext(transactionId, alice, queryId), aliceSchema);
        accessControlManager.checkCanDropSchema(new SecurityContext(transactionId, alice, queryId), aliceSchema);
        accessControlManager.checkCanRenameSchema(new SecurityContext(transactionId, alice, queryId), aliceSchema, "new-schema");
        accessControlManager.checkCanShowSchemas(new SecurityContext(transactionId, alice, queryId), "alice-catalog");
    });
    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 14 with QueryId

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

the class TestMemoryRevokingScheduler method testScheduleMemoryRevoking.

@Test
public void testScheduleMemoryRevoking() throws Exception {
    QueryContext q1 = getOrCreateQueryContext(new QueryId("q1"));
    QueryContext q2 = getOrCreateQueryContext(new QueryId("q2"));
    SqlTask sqlTask1 = newSqlTask(q1.getQueryId());
    SqlTask sqlTask2 = newSqlTask(q2.getQueryId());
    TaskContext taskContext1 = getOrCreateTaskContext(sqlTask1);
    PipelineContext pipelineContext11 = taskContext1.addPipelineContext(0, false, false, false);
    DriverContext driverContext111 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext1 = driverContext111.addOperatorContext(1, new PlanNodeId("na"), "na");
    OperatorContext operatorContext2 = driverContext111.addOperatorContext(2, new PlanNodeId("na"), "na");
    DriverContext driverContext112 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext3 = driverContext112.addOperatorContext(3, new PlanNodeId("na"), "na");
    TaskContext taskContext2 = getOrCreateTaskContext(sqlTask2);
    PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false);
    DriverContext driverContext211 = pipelineContext21.addDriverContext();
    OperatorContext operatorContext4 = driverContext211.addOperatorContext(4, new PlanNodeId("na"), "na");
    OperatorContext operatorContext5 = driverContext211.addOperatorContext(5, new PlanNodeId("na"), "na");
    Collection<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(memoryPool, () -> tasks, executor, 1.0, 1.0);
    allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2, operatorContext3, operatorContext4, operatorContext5);
    assertMemoryRevokingNotRequested();
    requestMemoryRevoking(scheduler);
    assertEquals(10, memoryPool.getFreeBytes());
    assertMemoryRevokingNotRequested();
    LocalMemoryContext revocableMemory1 = operatorContext1.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory3 = operatorContext3.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory4 = operatorContext4.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory5 = operatorContext5.localRevocableMemoryContext();
    revocableMemory1.setBytes(3);
    revocableMemory3.setBytes(6);
    assertEquals(1, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good - no revoking needed
    assertMemoryRevokingNotRequested();
    revocableMemory4.setBytes(7);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we need to revoke 3 and 6
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // yet another revoking request should not change anything
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // lets revoke some bytes
    revocableMemory1.setBytes(0);
    operatorContext1.resetMemoryRevokingRequested();
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext3);
    assertEquals(-3, memoryPool.getFreeBytes());
    // and allocate some more
    revocableMemory5.setBytes(3);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good with just OC3 in process of revoking
    assertMemoryRevokingRequestedFor(operatorContext3);
    // and allocate some more
    revocableMemory5.setBytes(4);
    assertEquals(-7, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // no we have to trigger revoking for OC4
    assertMemoryRevokingRequestedFor(operatorContext3, operatorContext4);
}

Example 15 with QueryId

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

the class TestMemoryRevokingScheduler method testImmediateMemoryRevoking.

/**
 * Test that when a {@link MemoryPool} is over-allocated, revocable memory is revoked without delay (although asynchronously).
 */
@Test
public void testImmediateMemoryRevoking() throws Exception {
    // Given
    SqlTask sqlTask = newSqlTask(new QueryId("query"));
    OperatorContext operatorContext = createContexts(sqlTask);
    allOperatorContexts = ImmutableSet.of(operatorContext);
    List<SqlTask> tasks = ImmutableList.of(sqlTask);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(memoryPool, () -> tasks, executor, 1.0, 1.0);
    // no periodic check initiated
    scheduler.registerPoolListeners();
    // When
    operatorContext.localRevocableMemoryContext().setBytes(12);
    awaitAsynchronousCallbacksRun();
    // Then
    assertMemoryRevokingRequestedFor(operatorContext);
}