private QueryStateMachine createQueryStateMachineWithTicker(Ticker ticker) { Metadata metadata = MetadataManager.createTestMetadataManager(); TransactionManager transactionManager = createTestTransactionManager(); AccessControl accessControl = new AccessControlManager(transactionManager); QueryStateMachine stateMachine = QueryStateMachine.beginWithTicker( QUERY, TEST_SESSION, LOCATION, new ResourceGroupId("test"), false, transactionManager, accessControl, executor, ticker, metadata, WarningCollector.NOOP); stateMachine.setInputs(INPUTS); stateMachine.setOutput(OUTPUT); stateMachine.setColumns(OUTPUT_FIELD_NAMES, OUTPUT_FIELD_TYPES); stateMachine.setUpdateType(UPDATE_TYPE); stateMachine.setMemoryPool(MEMORY_POOL); for (Entry<String, String> entry : SET_SESSION_PROPERTIES.entrySet()) { stateMachine.addSetSessionProperties(entry.getKey(), entry.getValue()); } RESET_SESSION_PROPERTIES.forEach(stateMachine::addResetSessionProperties); return stateMachine; }
private <T extends Statement> DataDefinitionExecution<T> createDataDefinitionExecution( String query, Session session, ResourceGroupId resourceGroup, T statement, List<Expression> parameters, WarningCollector warningCollector) { @SuppressWarnings("unchecked") DataDefinitionTask<T> task = (DataDefinitionTask<T>) tasks.get(statement.getClass()); checkArgument(task != null, "no task for statement: %s", statement.getClass().getSimpleName()); QueryStateMachine stateMachine = QueryStateMachine.begin( query, session, locationFactory.createQueryLocation(session.getQueryId()), resourceGroup, task.isTransactionControl(), transactionManager, accessControl, executor, metadata, warningCollector); stateMachine.setUpdateType(task.getName()); return new DataDefinitionExecution<>(task, statement, transactionManager, metadata, accessControl, stateMachine, parameters); } }
@Override public VersionedMemoryPoolId getMemoryPool() { return stateMachine.getMemoryPool(); }
@Override public void addStateChangeListener(StateChangeListener<QueryState> stateChangeListener) { try (SetThreadName ignored = new SetThreadName("Query-%s", stateMachine.getQueryId())) { stateMachine.addStateChangeListener(stateChangeListener); } }
@Override public QueryInfo getQueryInfo() { try (SetThreadName ignored = new SetThreadName("Query-%s", stateMachine.getQueryId())) { // acquire reference to scheduler before checking finalQueryInfo, because // state change listener sets finalQueryInfo and then clears scheduler when // the query finishes. SqlQueryScheduler scheduler = queryScheduler.get(); return stateMachine.getFinalQueryInfo().orElseGet(() -> buildQueryInfo(scheduler)); } }
@Override public ListenableFuture<?> execute(Rollback statement, TransactionManager transactionManager, Metadata metadata, AccessControl accessControl, QueryStateMachine stateMachine, List<Expression> parameters) { Session session = stateMachine.getSession(); if (!session.getTransactionId().isPresent()) { throw new PrestoException(NOT_IN_TRANSACTION, "No transaction in progress"); } TransactionId transactionId = session.getTransactionId().get(); stateMachine.clearTransactionId(); transactionManager.asyncAbort(transactionId); return immediateFuture(null); }
private static void assertState(QueryStateMachine stateMachine, QueryState expectedState, Exception expectedException) assertEquals(stateMachine.getQueryId(), TEST_SESSION.getQueryId()); assertEqualSessionsWithoutTransactionId(stateMachine.getSession(), TEST_SESSION); assertSame(stateMachine.getMemoryPool(), MEMORY_POOL); assertEquals(stateMachine.getSetSessionProperties(), SET_SESSION_PROPERTIES); assertEquals(stateMachine.getResetSessionProperties(), RESET_SESSION_PROPERTIES); QueryInfo queryInfo = stateMachine.getQueryInfo(Optional.empty()); assertEquals(queryInfo.getQueryId(), TEST_SESSION.getQueryId()); assertEquals(queryInfo.getSelf(), LOCATION); assertEquals(stateMachine.getQueryState(), expectedState); assertEquals(queryInfo.getState(), expectedState); assertEquals(stateMachine.isDone(), expectedState.isDone());
rootStage.addStateChangeListener(state -> { if (state == FINISHED) { queryStateMachine.transitionToFinishing(); queryStateMachine.transitionToCanceled(); if (queryStateMachine.isDone()) { return; queryStateMachine.transitionToFailed(stage.getStageInfo().getFailureCause().toException()); queryStateMachine.transitionToFailed(new PrestoException(GENERIC_INTERNAL_ERROR, "Query stage was aborted")); else if (queryStateMachine.getQueryState() == QueryState.STARTING) { queryStateMachine.transitionToRunning(); queryStateMachine.addStateChangeListener(newState -> { if (newState.isDone()) { queryStateMachine.updateQueryInfo(Optional.ofNullable(getStageInfo())); stage.addFinalStageInfoListener(status -> queryStateMachine.updateQueryInfo(Optional.ofNullable(getStageInfo())));
assertTrue(stateMachine.transitionToWaitingForResources()); assertState(stateMachine, WAITING_FOR_RESOURCES); assertTrue(stateMachine.transitionToPlanning()); assertState(stateMachine, PLANNING); assertTrue(stateMachine.transitionToStarting()); assertState(stateMachine, STARTING); assertTrue(stateMachine.transitionToRunning()); assertState(stateMachine, RUNNING); assertTrue(stateMachine.transitionToFinishing()); tryGetFutureValue(stateMachine.getStateChange(FINISHING), 2, SECONDS); assertState(stateMachine, FINISHED); QueryStats queryStats = stateMachine.getQueryInfo(Optional.empty()).getQueryStats(); assertEquals(queryStats.getElapsedTime().toMillis(), 1050); assertEquals(queryStats.getQueuedTime().toMillis(), 50);
private void startExecution() try (SetThreadName ignored = new SetThreadName("Query-%s", stateMachine.getQueryId())) { try { if (!stateMachine.transitionToPlanning()) { if (!stateMachine.transitionToStarting()) { if (!stateMachine.isDone()) { scheduler.start();
@Test public void testStartTransaction() throws Exception { Session session = sessionBuilder() .setClientTransactionSupport() .build(); TransactionManager transactionManager = createTestTransactionManager(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "START TRANSACTION", session, URI.create("fake://uri"), true, transactionManager, executor); assertFalse(stateMachine.getSession().getTransactionId().isPresent()); new StartTransactionTask().execute(new StartTransaction(ImmutableList.of()), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); assertFalse(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); assertTrue(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertEquals(transactionManager.getAllTransactionInfos().size(), 1); TransactionInfo transactionInfo = transactionManager.getTransactionInfo(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().get()); assertFalse(transactionInfo.isAutoCommitContext()); }
private QueryStateMachine createQueryStateMachine() { TransactionManager transactionManager = createTestTransactionManager(); QueryStateMachine stateMachine = QueryStateMachine.begin(QUERY_ID, QUERY, TEST_SESSION, LOCATION, false, transactionManager, executor); stateMachine.setInputs(INPUTS); stateMachine.setOutputFieldNames(OUTPUT_FIELD_NAMES); stateMachine.setUpdateType(UPDATE_TYPE); stateMachine.setMemoryPool(MEMORY_POOL); for (Entry<String, String> entry : SET_SESSION_PROPERTIES.entrySet()) { stateMachine.addSetSessionProperties(entry.getKey(), entry.getValue()); } RESET_SESSION_PROPERTIES.forEach(stateMachine::addResetSessionProperties); return stateMachine; }
@Test public void testStartTransaction() { Session session = sessionBuilder() .setClientTransactionSupport() .build(); TransactionManager transactionManager = createTestTransactionManager(); QueryStateMachine stateMachine = createQueryStateMachine("START TRANSACTION", session, transactionManager); assertFalse(stateMachine.getSession().getTransactionId().isPresent()); getFutureValue(new StartTransactionTask().execute(new StartTransaction(ImmutableList.of()), transactionManager, metadata, new AllowAllAccessControl(), stateMachine, emptyList())); assertFalse(stateMachine.getQueryInfo(Optional.empty()).isClearTransactionId()); assertTrue(stateMachine.getQueryInfo(Optional.empty()).getStartedTransactionId().isPresent()); assertEquals(transactionManager.getAllTransactionInfos().size(), 1); TransactionInfo transactionInfo = transactionManager.getTransactionInfo(stateMachine.getQueryInfo(Optional.empty()).getStartedTransactionId().get()); assertFalse(transactionInfo.isAutoCommitContext()); }
queryStateMachine.transitionToFinishing(); queryStateMachine.transitionToFailed(new PrestoException(USER_CANCELED, "Query was canceled")); if (queryStateMachine.isDone()) { return; queryStateMachine.transitionToFailed(stage.getStageInfo().getFailureCause().toException()); queryStateMachine.transitionToFailed(new PrestoException(INTERNAL_ERROR, "Query stage was aborted")); else if (queryStateMachine.getQueryState() == QueryState.STARTING) { queryStateMachine.transitionToRunning();
private void planDistribution(PlanRoot plan) stateMachine.beginDistributedPlanning(); StageExecutionPlan outputStageExecutionPlan = distributedPlanner.plan(plan.getRoot(), stateMachine.getSession()); stateMachine.endDistributedPlanning(); stateMachine.addStateChangeListener(state -> { if (state.isDone()) { closeSplitSources(outputStageExecutionPlan); if (stateMachine.isDone()) { return; stateMachine.setColumns(outputStageExecutionPlan.getFieldNames(), outputStageExecutionPlan.getFragment().getTypes()); nodeScheduler, remoteTaskFactory, stateMachine.getSession(), plan.isSummarizeTaskInfos(), scheduleSplitBatchSize, if (stateMachine.isDone()) { scheduler.abort(); queryScheduler.set(null);
private PlanRoot doAnalyzeQuery() { // time analysis phase stateMachine.beginAnalysis(); // plan query PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator(); LogicalPlanner logicalPlanner = new LogicalPlanner(stateMachine.getSession(), planOptimizers, idAllocator, metadata, sqlParser, statsCalculator, costCalculator, stateMachine.getWarningCollector()); Plan plan = logicalPlanner.plan(analysis); queryPlan.set(plan); // extract inputs List<Input> inputs = new InputExtractor(metadata, stateMachine.getSession()).extractInputs(plan.getRoot()); stateMachine.setInputs(inputs); // extract output Optional<Output> output = new OutputExtractor().extractOutput(plan.getRoot()); stateMachine.setOutput(output); // fragment the plan SubPlan fragmentedPlan = planFragmenter.createSubPlans(stateMachine.getSession(), plan, false); // record analysis time stateMachine.endAnalysis(); boolean explainAnalyze = analysis.getStatement() instanceof Explain && ((Explain) analysis.getStatement()).isAnalyze(); return new PlanRoot(fragmentedPlan, !explainAnalyze, extractConnectors(analysis)); }
requireNonNull(session, "session is null"); requireNonNull(self, "self is null"); this.stateMachine = QueryStateMachine.begin( query, session, stateMachine.getSession(), metadata, sqlParser, this.analysis = analyzer.analyze(preparedQuery.getStatement()); stateMachine.setUpdateType(analysis.getUpdateType()); stateMachine.addStateChangeListener(state -> { if (!state.isDone()) { return;
@Override public QueryId getQueryId() { return stateMachine.getQueryId(); }
private SubPlan doAnalyzeQuery() { // time analysis phase long analysisStart = System.nanoTime(); // analyze query Analyzer analyzer = new Analyzer(stateMachine.getSession(), metadata, sqlParser, accessControl, Optional.of(queryExplainer), experimentalSyntaxEnabled); Analysis analysis = analyzer.analyze(statement); stateMachine.setUpdateType(analysis.getUpdateType()); // plan query PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator(); LogicalPlanner logicalPlanner = new LogicalPlanner(stateMachine.getSession(), planOptimizers, idAllocator, metadata); Plan plan = logicalPlanner.plan(analysis); // extract inputs List<Input> inputs = new InputExtractor(metadata, stateMachine.getSession()).extract(plan.getRoot()); stateMachine.setInputs(inputs); // fragment the plan SubPlan subplan = new PlanFragmenter().createSubPlans(plan); // record analysis time stateMachine.recordAnalysisTime(analysisStart); return subplan; }
@Test public void testUnknownTransactionRollback() throws Exception { TransactionManager transactionManager = createTestTransactionManager(); Session session = sessionBuilder() .setTransactionId(TransactionId.create()) // Use a random transaction ID that is unknown to the system .build(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "ROLLBACK", session, URI.create("fake://uri"), true, transactionManager, executor); new RollbackTask().execute(new Rollback(), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); assertTrue(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); // Still issue clear signal assertFalse(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertTrue(transactionManager.getAllTransactionInfos().isEmpty()); }