Examples with QueryPrerequisitesContext com.facebook.presto.spi.prerequisites.QueryPrerequisitesContext used on opensource projects

Example 1 with QueryPrerequisitesContext

use of com.facebook.presto.spi.prerequisites.QueryPrerequisitesContext in project presto by prestodb.

the class TestLocalDispatchQuery method testPrerequisitesQueryFinishedCalled.

@Test
public void testPrerequisitesQueryFinishedCalled() {
    QueryStateMachine stateMachine = createStateMachine();
    CountingEventListener eventListener = new CountingEventListener();
    CompletableFuture<?> prequisitesFuture = new CompletableFuture<>();
    AtomicBoolean queryFinishedCalled = new AtomicBoolean();
    LocalDispatchQuery query = new LocalDispatchQuery(stateMachine, createQueryMonitor(eventListener), immediateFuture(null), createClusterSizeMonitor(0), directExecutor(), dispatchQuery -> {
    }, execution -> {
    }, false, new QueryPrerequisites() {

        @Override
        public CompletableFuture<?> waitForPrerequisites(QueryId queryId, QueryPrerequisitesContext context, WarningCollector warningCollector) {
            return prequisitesFuture;
        }

        @Override
        public void queryFinished(QueryId queryId) {
            queryFinishedCalled.set(true);
        }
    });
    assertEquals(query.getBasicQueryInfo().getState(), WAITING_FOR_PREREQUISITES);
    assertFalse(eventListener.getQueryCompletedEvent().isPresent());
    query.startWaitingForPrerequisites();
    prequisitesFuture.complete(null);
    query.fail(new PrestoException(ABANDONED_QUERY, "foo"));
    assertTrue(queryFinishedCalled.get());
}

Example 2 with QueryPrerequisitesContext

use of com.facebook.presto.spi.prerequisites.QueryPrerequisitesContext in project presto by prestodb.

the class LocalDispatchQuery method startWaitingForPrerequisites.

@Override
public void startWaitingForPrerequisites() {
    // start waiting for prerequisites
    if (isDone()) {
        return;
    }
    try {
        Session session = stateMachine.getSession();
        CompletableFuture<?> prerequisitesFuture = queryPrerequisites.waitForPrerequisites(stateMachine.getQueryId(), new QueryPrerequisitesContext(session.getCatalog(), session.getSchema(), stateMachine.getBasicQueryInfo(Optional.empty()).getQuery(), session.getSystemProperties(), session.getConnectorProperties()), warningCollector);
        addStateChangeListener(state -> {
            if (state.isDone()) {
                queryPrerequisites.queryFinished(stateMachine.getQueryId());
                if (!prerequisitesFuture.isDone()) {
                    prerequisitesFuture.cancel(true);
                }
            }
        });
        prerequisitesFuture.whenCompleteAsync((result, throwable) -> {
            if (throwable != null) {
                fail(throwable);
                return;
            }
            queueQuery();
        }, queryExecutor);
    } catch (Throwable t) {
        fail(t);
        throw t;
    }
}