Examples with PinotException io.trino.plugin.pinot.PinotException used on opensource projects

Example 1 with PinotException

use of io.trino.plugin.pinot.PinotException in project trino by trinodb.

the class PinotClient method doHttpActionWithHeadersJson.

protected <T> T doHttpActionWithHeadersJson(Request.Builder requestBuilder, Optional<String> requestBody, JsonCodec<T> codec, Multimap<String, String> additionalHeaders) {
    requestBuilder.addHeaders(additionalHeaders);
    requestBuilder.setHeader(ACCEPT, APPLICATION_JSON);
    if (requestBody.isPresent()) {
        requestBuilder.setHeader(CONTENT_TYPE, APPLICATION_JSON);
        requestBuilder.setBodyGenerator(StaticBodyGenerator.createStaticBodyGenerator(requestBody.get(), StandardCharsets.UTF_8));
    }
    Request request = requestBuilder.build();
    JsonResponseHandler<T> responseHandler = createJsonResponseHandler(codec);
    T response = null;
    try {
        response = httpClient.execute(request, responseHandler);
    } catch (UnexpectedResponseException e) {
        throw new PinotException(PinotErrorCode.PINOT_HTTP_ERROR, Optional.empty(), format("Unexpected response status: %d for request %s to url %s, with headers %s, full response %s", e.getStatusCode(), requestBody.orElse(""), request.getUri(), request.getHeaders(), response));
    }
    return response;
}

Example 2 with PinotException

use of io.trino.plugin.pinot.PinotException in project trino by trinodb.

the class PinotClient method submitBrokerQueryJson.

private BrokerResponseNative submitBrokerQueryJson(ConnectorSession session, PinotQueryInfo query) {
    String queryRequest = QUERY_REQUEST_JSON_CODEC.toJson(new QueryRequest(query.getQuery()));
    return doWithRetries(PinotSessionProperties.getPinotRetryCount(session), retryNumber -> {
        String queryHost = getBrokerHost(query.getTable());
        LOG.info("Query '%s' on broker host '%s'", queryHost, query.getQuery());
        Request.Builder builder = Request.Builder.preparePost().setUri(URI.create(format(QUERY_URL_TEMPLATE, queryHost)));
        ImmutableMultimap.Builder<String, String> additionalHeadersBuilder = ImmutableMultimap.builder();
        brokerAuthenticationProvider.getAuthenticationToken().ifPresent(token -> additionalHeadersBuilder.put(AUTHORIZATION, token));
        BrokerResponseNative response = doHttpActionWithHeadersJson(builder, Optional.of(queryRequest), brokerResponseCodec, additionalHeadersBuilder.build());
        if (response.getExceptionsSize() > 0 && response.getProcessingExceptions() != null && !response.getProcessingExceptions().isEmpty()) {
            // so we treat any exception as an error
            throw new PinotException(PINOT_EXCEPTION, Optional.of(query.getQuery()), format("Query %s encountered exception %s", query.getQuery(), response.getProcessingExceptions().get(0)));
        }
        if (response.getNumServersQueried() == 0 || response.getNumServersResponded() == 0 || response.getNumServersQueried() > response.getNumServersResponded()) {
            throw new PinotInsufficientServerResponseException(query, response.getNumServersResponded(), response.getNumServersQueried());
        }
        return response;
    });
}

Example 3 with PinotException

use of io.trino.plugin.pinot.PinotException in project trino by trinodb.

the class PinotQueryClient method queryPinotServerForDataTable.

public Map<ServerInstance, DataTable> queryPinotServerForDataTable(String query, String serverHost, List<String> segments, long connectionTimeoutInMillis, int pinotRetryCount) {
    // TODO: separate into offline and realtime methods
    BrokerRequest brokerRequest;
    try {
        brokerRequest = REQUEST_COMPILER.compileToBrokerRequest(query);
    } catch (SqlCompilationException e) {
        throw new PinotException(PINOT_INVALID_PQL_GENERATED, Optional.of(query), format("Parsing error with on %s, Error = %s", serverHost, e.getMessage()), e);
    }
    ServerInstance serverInstance = pinotHostMapper.getServerInstance(serverHost);
    Map<ServerInstance, List<String>> routingTable = new HashMap<>();
    routingTable.put(serverInstance, new ArrayList<>(segments));
    String tableName = brokerRequest.getQuerySource().getTableName();
    String rawTableName = TableNameBuilder.extractRawTableName(tableName);
    Map<ServerInstance, List<String>> offlineRoutingTable = TableNameBuilder.isOfflineTableResource(tableName) ? routingTable : null;
    Map<ServerInstance, List<String>> realtimeRoutingTable = TableNameBuilder.isRealtimeTableResource(tableName) ? routingTable : null;
    BrokerRequest offlineBrokerRequest = TableNameBuilder.isOfflineTableResource(tableName) ? brokerRequest : null;
    BrokerRequest realtimeBrokerRequest = TableNameBuilder.isRealtimeTableResource(tableName) ? brokerRequest : null;
    AsyncQueryResponse asyncQueryResponse = doWithRetries(pinotRetryCount, requestId -> queryRouter.submitQuery(requestId, rawTableName, offlineBrokerRequest, offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, connectionTimeoutInMillis));
    try {
        Map<ServerRoutingInstance, ServerResponse> response = asyncQueryResponse.getResponse();
        Map<ServerInstance, DataTable> dataTableMap = new HashMap<>();
        for (Map.Entry<ServerRoutingInstance, ServerResponse> entry : response.entrySet()) {
            ServerResponse serverResponse = entry.getValue();
            DataTable dataTable = serverResponse.getDataTable();
            dataTableMap.put(toServerInstance(entry.getKey()), dataTable);
        }
        return dataTableMap;
    } catch (InterruptedException e) {
        throw new PinotException(PINOT_EXCEPTION, Optional.of(query), "Pinot query execution was interrupted", e);
    }
}

Example 4 with PinotException

use of io.trino.plugin.pinot.PinotException in project trino by trinodb.

the class PinotClient method getAllBrokersForTable.

@VisibleForTesting
public List<String> getAllBrokersForTable(String table) {
    ArrayList<String> brokers = sendHttpGetToControllerJson(format(TABLE_INSTANCES_API_TEMPLATE, table), brokersForTableJsonCodec).getBrokers().stream().flatMap(broker -> broker.getInstances().stream()).distinct().map(brokerToParse -> {
        Matcher matcher = BROKER_PATTERN.matcher(brokerToParse);
        if (matcher.matches() && matcher.groupCount() == 2) {
            return pinotHostMapper.getBrokerHost(matcher.group(1), matcher.group(2));
        } else {
            throw new PinotException(PINOT_UNABLE_TO_FIND_BROKER, Optional.empty(), format("Cannot parse %s in the broker instance", brokerToParse));
        }
    }).collect(Collectors.toCollection(ArrayList::new));
    Collections.shuffle(brokers);
    return ImmutableList.copyOf(brokers);
}

Example 5 with PinotException

use of io.trino.plugin.pinot.PinotException in project trino by trinodb.

the class PinotClient method doWithRetries.

public static <T> T doWithRetries(int retries, Function<Integer, T> caller) {
    PinotException firstError = null;
    checkState(retries > 0, "Invalid num of retries %s", retries);
    for (int i = 0; i < retries; ++i) {
        try {
            return caller.apply(i);
        } catch (PinotException e) {
            if (firstError == null) {
                firstError = e;
            }
            if (!e.isRetryable()) {
                throw e;
            }
        }
    }
    throw firstError;
}