Examples with QueryState org.apache.drill.exec.proto.UserBitShared.QueryResult.QueryState used on opensource projects

Example 1 with QueryState

use of org.apache.drill.exec.proto.UserBitShared.QueryResult.QueryState in project drill by apache.

the class TestDrillbitResilience method assertDrillbitsOk.

/**
   * Check that all the drillbits are ok.
   * <p/>
   * <p>The current implementation does this by counting the number of drillbits using a query.
   */
private static void assertDrillbitsOk() {
    final SingleRowListener listener = new SingleRowListener() {

        private final BufferAllocator bufferAllocator = RootAllocatorFactory.newRoot(zkHelper.getConfig());

        private final RecordBatchLoader loader = new RecordBatchLoader(bufferAllocator);

        @Override
        public void rowArrived(final QueryDataBatch queryResultBatch) {
            // load the single record
            final QueryData queryData = queryResultBatch.getHeader();
            try {
                loader.load(queryData.getDef(), queryResultBatch.getData());
            // TODO:  Clean:  DRILL-2933:  That load(...) no longer throws
            // SchemaChangeException, so check/clean catch clause below.
            } catch (final SchemaChangeException e) {
                fail(e.toString());
            }
            assertEquals(1, loader.getRecordCount());
            // there should only be one column
            final BatchSchema batchSchema = loader.getSchema();
            assertEquals(1, batchSchema.getFieldCount());
            // the column should be an integer
            final MaterializedField countField = batchSchema.getColumn(0);
            final MinorType fieldType = countField.getType().getMinorType();
            assertEquals(MinorType.BIGINT, fieldType);
            // get the column value
            final VectorWrapper<?> vw = loader.iterator().next();
            final Object obj = vw.getValueVector().getAccessor().getObject(0);
            assertTrue(obj instanceof Long);
            final Long countValue = (Long) obj;
            // assume this means all the drillbits are still ok
            assertEquals(drillbits.size(), countValue.intValue());
            loader.clear();
        }

        @Override
        public void cleanup() {
            DrillAutoCloseables.closeNoChecked(bufferAllocator);
        }
    };
    try {
        QueryTestUtil.testWithListener(drillClient, QueryType.SQL, "select count(*) from sys.memory", listener);
        listener.waitForCompletion();
        final QueryState state = listener.getQueryState();
        assertTrue(String.format("QueryState should be COMPLETED (and not %s).", state), state == QueryState.COMPLETED);
    } catch (final Exception e) {
        throw new RuntimeException("Couldn't query active drillbits", e);
    }
    final List<DrillPBError> errorList = listener.getErrorList();
    assertTrue("There should not be any errors when checking if Drillbits are OK.", errorList.isEmpty());
}

Example 2 with QueryState

use of org.apache.drill.exec.proto.UserBitShared.QueryResult.QueryState in project drill by apache.

the class TestDrillbitResilience method assertFailsWithException.

/**
   * Given a set of controls, this method ensures TEST_QUERY fails with the given class and desc.
   */
private static void assertFailsWithException(final String controls, final Class<? extends Throwable> exceptionClass, final String exceptionDesc, final String query) {
    setControls(controls);
    final WaitUntilCompleteListener listener = new WaitUntilCompleteListener();
    QueryTestUtil.testWithListener(drillClient, QueryType.SQL, query, listener);
    final Pair<QueryState, Exception> result = listener.waitForCompletion();
    final QueryState state = result.getFirst();
    assertTrue(String.format("Query state should be FAILED (and not %s).", state), state == QueryState.FAILED);
    assertExceptionMessage(result.getSecond(), exceptionClass, exceptionDesc);
}

Example 3 with QueryState

use of org.apache.drill.exec.proto.UserBitShared.QueryResult.QueryState in project drill by apache.

the class TestDrillbitResilience method assertStateCompleted.

/**
   * Given the result of {@link WaitUntilCompleteListener#waitForCompletion}, this method fails if the completed state
   * is not as expected, or if an exception is thrown. The completed state could be COMPLETED or CANCELED. This state
   * is set when {@link WaitUntilCompleteListener#queryCompleted} is called.
   */
private static void assertStateCompleted(final Pair<QueryState, Exception> result, final QueryState expectedState) {
    final QueryState actualState = result.getFirst();
    final Exception exception = result.getSecond();
    if (actualState != expectedState || exception != null) {
        fail(String.format("Query state is incorrect (expected: %s, actual: %s) AND/OR \nException thrown: %s", expectedState, actualState, exception == null ? "none." : exception));
    }
}

Example 4 with QueryState

use of org.apache.drill.exec.proto.UserBitShared.QueryResult.QueryState in project drill by apache.

the class AbstractDisposableUserClientConnection method sendResult.

@Override
public void sendResult(RpcOutcomeListener<Ack> listener, QueryResult result) {
    Preconditions.checkState(result.hasQueryState());
    // Release the wait latch if the query is terminated.
    final QueryState state = result.getQueryState();
    final QueryId queryId = result.getQueryId();
    if (logger.isDebugEnabled()) {
        logger.debug("Result arrived for QueryId: {} with QueryState: {}", QueryIdHelper.getQueryId(queryId), state);
    }
    switch(state) {
        case FAILED:
            error = result.getError(0);
            exception = new UserRemoteException(error);
            latch.countDown();
            break;
        case CANCELED:
        case COMPLETED:
            Preconditions.checkState(result.getErrorCount() == 0);
            latch.countDown();
            break;
        default:
            logger.error("Query with QueryId: {} is in unexpected state: {}", queryId, state);
    }
    // Notify the listener with ACK
    listener.success(Acks.OK, null);
}

Example 5 with QueryState

use of org.apache.drill.exec.proto.UserBitShared.QueryResult.QueryState in project drill by apache.

the class QueryResultHandler method resultArrived.

/**
   * Maps internal low-level API protocol to {@link UserResultsListener}-level API protocol.
   * handles data result messages
   */
public void resultArrived(ByteBuf pBody) throws RpcException {
    final QueryResult queryResult = RpcBus.get(pBody, QueryResult.PARSER);
    final QueryId queryId = queryResult.getQueryId();
    final QueryState queryState = queryResult.getQueryState();
    if (logger.isDebugEnabled()) {
        logger.debug("resultArrived: queryState: {}, queryId = {}", queryState, QueryIdHelper.getQueryId(queryId));
    }
    assert queryResult.hasQueryState() : "received query result without QueryState";
    final boolean isFailureResult = QueryState.FAILED == queryState;
    // CANCELED queries are handled the same way as COMPLETED
    final boolean isTerminalResult;
    switch(queryState) {
        case FAILED:
        case CANCELED:
        case COMPLETED:
            isTerminalResult = true;
            break;
        default:
            logger.error("Unexpected/unhandled QueryState " + queryState + " (for query " + queryId + ")");
            isTerminalResult = false;
            break;
    }
    assert isFailureResult || queryResult.getErrorCount() == 0 : "Error count for the query batch is non-zero but QueryState != FAILED";
    UserResultsListener resultsListener = newUserResultsListener(queryId);
    try {
        if (isFailureResult) {
            // Failure case--pass on via submissionFailed(...).
            resultsListener.submissionFailed(new UserRemoteException(queryResult.getError(0)));
        // Note: Listener is removed in finally below.
        } else if (isTerminalResult) {
            try {
                resultsListener.queryCompleted(queryState);
            } catch (Exception e) {
                resultsListener.submissionFailed(UserException.systemError(e).build(logger));
            }
        } else {
            logger.warn("queryState {} was ignored", queryState);
        }
    } finally {
        if (isTerminalResult) {
            // for it?
            if ((!(resultsListener instanceof BufferingResultsListener) || ((BufferingResultsListener) resultsListener).output != null)) {
                queryIdToResultsListenersMap.remove(queryId, resultsListener);
            }
        }
    }
}