Examples with BasicQueryInfo io.prestosql.server.BasicQueryInfo used on opensource projects

Example 6 with BasicQueryInfo

use of io.prestosql.server.BasicQueryInfo in project hetu-core by openlookeng.

the class TestStateUpdater method createBasicQueryInfo.

private BasicQueryInfo createBasicQueryInfo() {
    QueryInfo queryInfo = Mockito.mock(QueryInfo.class);
    when(queryInfo.getQueryStats()).then(new Returns(Mockito.mock(QueryStats.class)));
    Duration mockInterval = new Duration(MINIMUM_UPDATE_INTERVAL, MILLISECONDS);
    when(queryInfo.getQueryStats().getQueuedTime()).then(new Returns(mockInterval));
    when(queryInfo.getQueryStats().getElapsedTime()).then(new Returns(mockInterval));
    when(queryInfo.getQueryStats().getExecutionTime()).then(new Returns(mockInterval));
    when(queryInfo.getQueryStats().getRawInputDataSize()).then(new Returns(Mockito.mock(DataSize.class)));
    String mockQueryId = MOCK_QUERY_ID;
    QueryId queryId = new QueryId(mockQueryId);
    when(queryInfo.getQueryId()).then(new Returns(queryId));
    SessionRepresentation sessionRepresentation = TEST_SESSION.toSessionRepresentation();
    when(queryInfo.getSession()).then(new Returns(sessionRepresentation));
    ResourceGroupId resourceGroupId = new ResourceGroupId(GLOBAL_RESOURCE_ID);
    Optional<ResourceGroupId> optionalResourceGroupId = Optional.of(resourceGroupId);
    when(queryInfo.getResourceGroupId()).then(new Returns(optionalResourceGroupId));
    when(queryInfo.getState()).then(new Returns(QueryState.FINISHED));
    URI mockURI = URI.create(URI_LOCALHOST);
    when(queryInfo.getSelf()).then(new Returns(mockURI));
    String mockQuery = QUERY_STRING;
    when(queryInfo.getQuery()).then(new Returns(mockQuery));
    Optional<String> preparedQuery = Optional.of(PREPARED_QUERY_MOCK_DATA);
    when(queryInfo.getPreparedQuery()).then(new Returns(preparedQuery));
    ErrorCode errorCode = new ErrorCode(ERROR_CODE_VALUE_INDEX_TIME_NO_INVOCATION, ERROR_CODE_TEST, USER_ERROR);
    when(queryInfo.getErrorCode()).then(new Returns(errorCode));
    Set<BlockedReason> setBlockedReason = new HashSet<>();
    setBlockedReason.add(BlockedReason.WAITING_FOR_MEMORY);
    when(queryInfo.getQueryStats().getBlockedReasons()).then(new Returns(setBlockedReason));
    when(queryInfo.getQueryStats().getProgressPercentage()).then(new Returns(OptionalDouble.empty()));
    BasicQueryInfo basicQueryInfo = new BasicQueryInfo(queryInfo);
    return basicQueryInfo;
}

Example 7 with BasicQueryInfo

use of io.prestosql.server.BasicQueryInfo in project hetu-core by openlookeng.

the class TestStateUpdater method mockDispatchQueryData.

private DispatchQuery mockDispatchQueryData(boolean userError) {
    DispatchQuery dispatchQuery = Mockito.mock(LocalDispatchQuery.class);
    BasicQueryInfo basicQueryInfo = createBasicQueryInfo();
    when(dispatchQuery.getBasicQueryInfo()).then(new Returns(basicQueryInfo));
    when(dispatchQuery.getSession()).then(new Returns(TEST_SESSION));
    ErrorCode errorCode;
    if (!userError) {
        errorCode = CLUSTER_OUT_OF_MEMORY.toErrorCode();
    } else {
        errorCode = new ErrorCode(ERROR_CODE_VALUE_INDEX_TIME_NO_INVOCATION, ERROR_CODE_TEST, USER_ERROR);
    }
    Optional<ErrorCode> optionalErrorCode = Optional.of(errorCode);
    when(dispatchQuery.getErrorCode()).then(new Returns(optionalErrorCode));
    DataSize userDataSize = new DataSize(USER_DATA_SIZE, DataSize.Unit.BYTE);
    DataSize totalDataSize = new DataSize(TOTAL_DATA_SIZE, DataSize.Unit.BYTE);
    when(dispatchQuery.getUserMemoryReservation()).then(new Returns(userDataSize));
    when(dispatchQuery.getTotalMemoryReservation()).then(new Returns(totalDataSize));
    when(dispatchQuery.getTotalCpuTime()).then(new Returns(new Duration(ERROR_CODE_VALUE_INDEX_TIME_NO_INVOCATION, MILLISECONDS)));
    when(dispatchQuery.getExecutionStartTime()).then(new Returns(Optional.of(new DateTime(DateTimeZone.UTC))));
    return dispatchQuery;
}

Example 8 with BasicQueryInfo

use of io.prestosql.server.BasicQueryInfo in project hetu-core by openlookeng.

the class TestQueuesDb method testSelectorPriority.

@Test(timeOut = 60_000)
public void testSelectorPriority() throws Exception {
    InternalResourceGroupManager<?> manager = queryRunner.getCoordinator().getResourceGroupManager().get();
    QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
    DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager();
    QueryId firstQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, firstQuery, RUNNING);
    Optional<ResourceGroupId> resourceGroup = queryManager.getFullQueryInfo(firstQuery).getResourceGroupId();
    assertTrue(resourceGroup.isPresent());
    assertEquals(resourceGroup.get().toString(), "global.user-user.dashboard-user");
    // create a new resource group that rejects all queries submitted to it
    // Hetu: add parameters softReservedMemory and hardReservedConcurrency
    dao.insertResourceGroup(8, "reject-all-queries", "1MB", "1MB", 0, 0, 0, 0, null, null, null, null, null, "RECENT_QUERIES", 3L, TEST_ENVIRONMENT);
    // add a new selector that has a higher priority than the existing dashboard selector and that routes queries to the "reject-all-queries" resource group
    dao.insertSelector(8, 200, "user.*", "(?i).*dashboard.*", null, null, null);
    // reload the configuration
    dbConfigurationManager.load();
    QueryId secondQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, secondQuery, FAILED);
    DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
    BasicQueryInfo basicQueryInfo = dispatchManager.getQueryInfo(secondQuery);
    assertEquals(basicQueryInfo.getErrorCode(), QUERY_QUEUE_FULL.toErrorCode());
}

Example 9 with BasicQueryInfo

use of io.prestosql.server.BasicQueryInfo in project hetu-core by openlookeng.

the class StateFetcher method createExpiredState.

private SharedQueryState createExpiredState(BasicQueryInfo oldQueryInfo, SharedQueryState oldState) {
    BasicQueryInfo newQueryInfo = new BasicQueryInfo(oldQueryInfo.getQueryId(), oldQueryInfo.getSession(), oldQueryInfo.getResourceGroupId(), QueryState.FAILED, oldQueryInfo.getMemoryPool(), oldQueryInfo.isScheduled(), oldQueryInfo.getSelf(), oldQueryInfo.getQuery(), oldQueryInfo.getPreparedQuery(), oldQueryInfo.getQueryStats(), ErrorType.INTERNAL_ERROR, QUERY_EXPIRE.toErrorCode());
    SharedQueryState newState = new SharedQueryState(newQueryInfo, Optional.of(QUERY_EXPIRE.toErrorCode()), oldState.getUserMemoryReservation(), oldState.getTotalMemoryReservation(), oldState.getTotalCpuTime(), oldState.getStateUpdateTime(), oldState.getExecutionStartTime());
    return newState;
}

Example 10 with BasicQueryInfo

use of io.prestosql.server.BasicQueryInfo in project hetu-core by openlookeng.

the class StateFetcher method handleExpiredQueryState.

private void handleExpiredQueryState(SharedQueryState state) {
    // State store hasn't been loaded yet
    final StateStore stateStore = stateStoreProvider.getStateStore();
    if (stateStore == null) {
        return;
    }
    Lock lock = null;
    boolean locked = false;
    try {
        lock = stateStore.getLock(HANDLE_EXPIRED_QUERY_LOCK_NAME);
        locked = lock.tryLock(DEFAULT_ACQUIRED_LOCK_TIME_MS, TimeUnit.MILLISECONDS);
        if (locked) {
            LOG.debug(String.format("EXPIRED!!! REMOVING... Id: %s, state: %s, uri: %s, query: %s", state.getBasicQueryInfo().getQueryId().getId(), state.getBasicQueryInfo().getState().toString(), state.getBasicQueryInfo().getSelf().toString(), state.getBasicQueryInfo().getQuery()));
            // remove expired query from oom
            StateCollection stateCollection = stateStore.getStateCollection(OOM_QUERY_STATE_COLLECTION_NAME);
            removeState(stateCollection, Optional.of(state.getBasicQueryInfo().getQueryId()), LOG);
            // update query to failed in QUERY_STATE_COLLECTION_NAME if exists
            stateCollection = stateStore.getStateCollection(QUERY_STATE_COLLECTION_NAME);
            StateCollection finishStateCollection = stateStore.getStateCollection(FINISHED_QUERY_STATE_COLLECTION_NAME);
            if (stateCollection != null && stateCollection.getType().equals(StateCollection.Type.MAP)) {
                String queryState = ((StateMap<String, String>) stateCollection).get(state.getBasicQueryInfo().getQueryId().getId());
                if (queryState != null) {
                    BasicQueryInfo oldQueryInfo = state.getBasicQueryInfo();
                    SharedQueryState newState = createExpiredState(oldQueryInfo, state);
                    String stateJson = MAPPER.writeValueAsString(newState);
                    ((StateMap) finishStateCollection).put(newState.getBasicQueryInfo().getQueryId().getId(), stateJson);
                    removeState(stateCollection, Optional.of(state.getBasicQueryInfo().getQueryId()), LOG);
                }
            }
        }
    } catch (Exception e) {
        LOG.error("Error handleExpiredQueryState: " + e.getMessage());
    } finally {
        if (locked) {
            lock.unlock();
        }
    }
}