Examples with QueryStateMachine io.trino.execution.QueryStateMachine used on opensource projects

Example 1 with QueryStateMachine

use of io.trino.execution.QueryStateMachine in project trino by trinodb.

the class LocalDispatchQueryFactory method createDispatchQuery.

@Override
public DispatchQuery createDispatchQuery(Session session, Optional<TransactionId> existingTransactionId, String query, PreparedQuery preparedQuery, Slug slug, ResourceGroupId resourceGroup) {
    WarningCollector warningCollector = warningCollectorFactory.create();
    QueryStateMachine stateMachine = QueryStateMachine.begin(existingTransactionId, query, preparedQuery.getPrepareSql(), session, locationFactory.createQueryLocation(session.getQueryId()), resourceGroup, isTransactionControlStatement(preparedQuery.getStatement()), transactionManager, accessControl, executor, metadata, warningCollector, getQueryType(preparedQuery.getStatement()));
    // It is important that `queryCreatedEvent` is called here. Moving it past the `executor.submit` below
    // can result in delivering query-created event after query analysis has already started.
    // That can result in misbehaviour of plugins called during analysis phase (e.g. access control auditing)
    // which depend on the contract that event was already delivered.
    // 
    // Note that for immediate and in-order delivery of query events we depend on synchronous nature of
    // QueryMonitor and EventListenerManager.
    queryMonitor.queryCreatedEvent(stateMachine.getBasicQueryInfo(Optional.empty()));
    ListenableFuture<QueryExecution> queryExecutionFuture = executor.submit(() -> {
        QueryExecutionFactory<?> queryExecutionFactory = executionFactories.get(preparedQuery.getStatement().getClass());
        if (queryExecutionFactory == null) {
            throw new TrinoException(NOT_SUPPORTED, "Unsupported statement type: " + preparedQuery.getStatement().getClass().getSimpleName());
        }
        try {
            return queryExecutionFactory.createQueryExecution(preparedQuery, stateMachine, slug, warningCollector);
        } catch (Throwable e) {
            if (e instanceof Error) {
                if (e instanceof StackOverflowError) {
                    log.error(e, "Unhandled StackOverFlowError; should be handled earlier; to investigate full stacktrace you may need to enable -XX:MaxJavaStackTraceDepth=0 JVM flag");
                } else {
                    log.error(e, "Unhandled Error");
                }
                // wrapping as RuntimeException to guard us from problem that code downstream which investigates queryExecutionFuture may not necessarily handle
                // Error subclass of Throwable well.
                RuntimeException wrappedError = new RuntimeException(e);
                stateMachine.transitionToFailed(wrappedError);
                throw wrappedError;
            }
            stateMachine.transitionToFailed(e);
            throw e;
        }
    });
    return new LocalDispatchQuery(stateMachine, queryExecutionFuture, queryMonitor, clusterSizeMonitor, executor, queryManager::createQuery);
}