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

Example 11 with QueryId

use of org.apache.drill.exec.proto.UserBitShared.QueryId 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);
            }
        }
    }
}

Example 12 with QueryId

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

the class ProfileResources method cancelQuery.

@GET
@Path("/profiles/cancel/{queryid}")
@Produces(MediaType.TEXT_PLAIN)
public String cancelQuery(@PathParam("queryid") String queryId) {
    QueryId id = QueryIdHelper.getQueryIdFromString(queryId);
    // first check local running
    Foreman f = work.getBee().getForemanForQueryId(id);
    if (f != null) {
        checkOrThrowQueryCancelAuthorization(f.getQueryContext().getQueryUserName(), queryId);
        f.cancel();
        return String.format("Cancelled query %s on locally running node.", queryId);
    }
    // then check remote running
    try {
        final TransientStore<QueryInfo> running = work.getContext().getProfileStoreContext().getRunningProfileStore();
        final QueryInfo info = running.get(queryId);
        checkOrThrowQueryCancelAuthorization(info.getUser(), queryId);
        Ack a = work.getContext().getController().getTunnel(info.getForeman()).requestCancelQuery(id).checkedGet(2, TimeUnit.SECONDS);
        if (a.getOk()) {
            return String.format("Query %s canceled on node %s.", queryId, info.getForeman().getAddress());
        } else {
            return String.format("Attempted to cancel query %s on %s but the query is no longer active on that node.", queryId, info.getForeman().getAddress());
        }
    } catch (Exception e) {
        logger.debug("Failure to find query as running profile.", e);
        return String.format("Failure attempting to cancel query %s.  Unable to find information about where query is actively running.", queryId);
    }
}

Example 13 with QueryId

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

the class TestCustomTunnel method ensureRoundTrip.

@Test
public void ensureRoundTrip() throws Exception {
    final DrillbitContext context = getDrillbitContext();
    final TestCustomMessageHandler handler = new TestCustomMessageHandler(context.getEndpoint(), false);
    context.getController().registerCustomHandler(1001, handler, DrillbitEndpoint.PARSER);
    final ControlTunnel loopbackTunnel = context.getController().getTunnel(context.getEndpoint());
    final CustomTunnel<DrillbitEndpoint, QueryId> tunnel = loopbackTunnel.getCustomTunnel(1001, DrillbitEndpoint.class, QueryId.PARSER);
    CustomFuture<QueryId> future = tunnel.send(context.getEndpoint());
    assertEquals(expectedId, future.get());
}

Example 14 with QueryId

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

the class QueryBuilder method produceSummary.

private QuerySummary produceSummary(BufferingQueryEventListener listener) throws Exception {
    long start = System.currentTimeMillis();
    int recordCount = 0;
    int batchCount = 0;
    QueryId queryId = null;
    QueryState state = null;
    loop: for (; ; ) {
        QueryEvent event = listener.get();
        switch(event.type) {
            case BATCH:
                batchCount++;
                recordCount += event.batch.getHeader().getRowCount();
                event.batch.release();
                break;
            case EOF:
                state = event.state;
                break loop;
            case ERROR:
                throw event.error;
            case QUERY_ID:
                queryId = event.queryId;
                break;
            default:
                throw new IllegalStateException("Unexpected event: " + event.type);
        }
    }
    long end = System.currentTimeMillis();
    long elapsed = end - start;
    return new QuerySummary(queryId, recordCount, batchCount, elapsed, state);
}

Example 15 with QueryId

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

the class TestDrillbitResilience method passThrough.

// To test pause and resume. Test hangs and times out if resume did not happen.
@Test
public void passThrough() {
    final long before = countAllocatedMemory();
    final WaitUntilCompleteListener listener = new WaitUntilCompleteListener() {

        @Override
        public void queryIdArrived(final QueryId queryId) {
            super.queryIdArrived(queryId);
            final ExtendedLatch trigger = new ExtendedLatch(1);
            (new ResumingThread(queryId, ex, trigger)).start();
            trigger.countDown();
        }
    };
    final String controls = Controls.newBuilder().addPause(PojoRecordReader.class, "read-next").build();
    setControls(controls);
    QueryTestUtil.testWithListener(drillClient, QueryType.SQL, TEST_QUERY, listener);
    final Pair<QueryState, Exception> result = listener.waitForCompletion();
    assertStateCompleted(result, QueryState.COMPLETED);
    final long after = countAllocatedMemory();
    assertEquals(String.format("We are leaking %d bytes", after - before), before, after);
}