Java tutorial
/* * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.facebook.presto.jdbc; import com.facebook.presto.jdbc.client.ClientSession; import com.facebook.presto.jdbc.client.PrestoHeaders; import com.facebook.presto.jdbc.client.QueryHttpClient; import com.facebook.presto.jdbc.client.QueryResults; import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.http.HttpHeaders; import org.apache.http.HttpResponse; import org.apache.http.HttpStatus; import org.apache.http.client.config.RequestConfig; import org.apache.http.client.methods.HttpDelete; import org.apache.http.client.methods.HttpGet; import org.apache.http.client.methods.HttpPost; import org.apache.http.client.methods.HttpRequestBase; import org.apache.http.entity.StringEntity; import org.apache.http.impl.nio.client.CloseableHttpAsyncClient; import java.io.IOException; import java.net.URI; import java.nio.charset.Charset; import java.util.Map; import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicBoolean; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.net.HttpHeaders.USER_AGENT; import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; import static java.lang.String.format; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.MINUTES; import static org.apache.http.entity.ContentType.APPLICATION_JSON; import static org.apache.http.entity.ContentType.parse; public class ApacheQueryHttpClient implements QueryHttpClient { private static final RequestConfig requestConfig = RequestConfig.custom().setSocketTimeout(10000) .setConnectTimeout(10000).build(); private final CloseableHttpAsyncClient httpAsyncClient; private final ObjectMapper mapper; private final String userAgent; private final AtomicBoolean closed = new AtomicBoolean(); public ApacheQueryHttpClient(CloseableHttpAsyncClient httpAsyncClient, ObjectMapper mapper, String userAgent) { checkNotNull(httpAsyncClient, "httpAsyncClient is null"); checkNotNull(mapper, "mapper is null"); checkNotNull(userAgent, "userAgent is null"); this.httpAsyncClient = httpAsyncClient; this.mapper = mapper; this.userAgent = userAgent; } @Override public QueryResults startQuery(ClientSession session, String query) { HttpPost request = buildQueryRequest(session, query); try { HttpResponse response = httpAsyncClient.execute(request, null).get(); if (response.getStatusLine().getStatusCode() != HttpStatus.SC_OK) { throw requestFailedStatus("starting query", request, response.getStatusLine().getReasonPhrase()); } return parseResult(request, response); } catch (InterruptedException e) { throw requestFailedException("starting query", request, e); } catch (ExecutionException e) { throw requestFailedException("starting query", request, e); } } @Override public void deleteAsync(URI uri) { HttpDelete request = new HttpDelete(uri); request.setHeader(USER_AGENT, userAgent); request.setConfig(requestConfig); httpAsyncClient.execute(request, null); } @Override public boolean delete(URI uri) { throw new UnsupportedOperationException(); } @Override public QueryResults execute(URI uri) throws RuntimeException { HttpGet request = new HttpGet(uri); request.setConfig(requestConfig); request.setHeader(USER_AGENT, userAgent); Exception cause = null; long start = System.nanoTime(); long attempts = 0; do { // back-off on retry if (attempts > 0) { sleepUninterruptibly(attempts * 100, MILLISECONDS); } attempts++; HttpResponse response; Future<HttpResponse> responseFuture = null; try { responseFuture = httpAsyncClient.execute(request, null); response = responseFuture.get(); } catch (InterruptedException e) { try { if (responseFuture != null) { responseFuture.cancel(true); } } finally { Thread.currentThread().interrupt(); } throw new RuntimeException("ApacheQueryHttpClient interrupted"); } catch (ExecutionException e) { cause = e; continue; } if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) { return parseResult(request, response); } if (response.getStatusLine().getStatusCode() != HttpStatus.SC_SERVICE_UNAVAILABLE) { throw requestFailedStatus("fetching next", request, response.getStatusLine().getReasonPhrase()); } } while ((System.nanoTime() - start) < MINUTES.toNanos(2) && !isClosed()); throw new RuntimeException("Error fetching next", cause); } @Override public Map<String, String> getSetSessionProperties() { throw new UnsupportedOperationException(); } @Override public Set<String> getResetSessionProperties() { throw new UnsupportedOperationException(); } @Override public boolean isClosed() { return closed.get(); } @Override public void close() { closed.set(true); } private RuntimeException requestFailedException(String task, HttpRequestBase request, Exception exception) { return new RuntimeException(format("Error " + task + " at %s returned an invalid response: %s", request.getURI(), exception.getMessage()), exception); } private RuntimeException requestFailedStatus(String task, HttpRequestBase request, String reason) { return new RuntimeException( format("Error " + task + " at %s failed with status %s", request.getURI(), reason)); } private HttpPost buildQueryRequest(ClientSession session, String query) { HttpPost post = new HttpPost(session.getServer()); post.setEntity(new StringEntity(query, Charset.forName("UTF-8"))); if (session.getUser() != null) { post.setHeader(PrestoHeaders.PRESTO_USER, session.getUser()); } if (session.getSource() != null) { post.setHeader(PrestoHeaders.PRESTO_SOURCE, session.getSource()); } if (session.getCatalog() != null) { post.setHeader(PrestoHeaders.PRESTO_CATALOG, session.getCatalog()); } if (session.getSchema() != null) { post.setHeader(PrestoHeaders.PRESTO_SCHEMA, session.getSchema()); } post.setHeader(PrestoHeaders.PRESTO_TIME_ZONE, session.getTimeZoneId()); // In Java 6 we don't have toLanguageTag //String localeId = session.getLocale().toLanguageTag()); String localeId = LocaleHelper.buildLanguageTag(session.getLocale()); post.setHeader(PrestoHeaders.PRESTO_LANGUAGE, localeId); post.setHeader(USER_AGENT, userAgent); post.setConfig(requestConfig); Map<String, String> property = session.getProperties(); for (Map.Entry<String, String> entry : property.entrySet()) { post.addHeader(PrestoHeaders.PRESTO_SESSION, entry.getKey() + "=" + entry.getValue()); } return post; } private QueryResults parseResult(HttpRequestBase request, HttpResponse response) { String contentType = response.getFirstHeader(HttpHeaders.CONTENT_TYPE).getValue(); if (!isApplicationJson(contentType)) { throw new RuntimeException(format("Error parsing result. Wrong content type. Got %s and expected %s", contentType, APPLICATION_JSON.toString())); } try { return mapper.readValue(response.getEntity().getContent(), QueryResults.class); } catch (IOException e) { throw requestFailedException("parse result", request, e); } } private boolean isApplicationJson(String contentType) { return parse(contentType).getMimeType().equals(APPLICATION_JSON.getMimeType()); } }