Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
com.facebook.presto.jdbc.internal.client.StatementClient Maven / Gradle / Ivy
/*
* 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.internal.client;
import com.facebook.presto.jdbc.internal.client.OkHttpUtil.NullCallback;
import com.facebook.presto.jdbc.internal.spi.type.TimeZoneKey;
import com.facebook.presto.jdbc.internal.guava.base.Joiner;
import com.facebook.presto.jdbc.internal.guava.base.Splitter;
import com.facebook.presto.jdbc.internal.guava.collect.ImmutableMap;
import com.facebook.presto.jdbc.internal.guava.collect.ImmutableSet;
import com.facebook.presto.jdbc.internal.guava.collect.Sets;
import com.facebook.presto.jdbc.internal.airlift.json.JsonCodec;
import com.facebook.presto.jdbc.internal.okhttp3.Headers;
import com.facebook.presto.jdbc.internal.okhttp3.HttpUrl;
import com.facebook.presto.jdbc.internal.okhttp3.MediaType;
import com.facebook.presto.jdbc.internal.okhttp3.OkHttpClient;
import com.facebook.presto.jdbc.internal.okhttp3.Request;
import com.facebook.presto.jdbc.internal.okhttp3.RequestBody;
import javax.annotation.Nullable;
import javax.annotation.concurrent.ThreadSafe;
import java.io.Closeable;
import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.net.URLDecoder;
import java.net.URLEncoder;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_ADDED_PREPARE;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_CATALOG;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_CLEAR_SESSION;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_CLEAR_TRANSACTION_ID;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_CLIENT_INFO;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_CLIENT_TAGS;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_DEALLOCATED_PREPARE;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_LANGUAGE;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_PREPARED_STATEMENT;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_SCHEMA;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_SESSION;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_SET_CATALOG;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_SET_SCHEMA;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_SET_SESSION;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_SOURCE;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_STARTED_TRANSACTION_ID;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_TIME_ZONE;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_TRANSACTION_ID;
import static com.facebook.presto.jdbc.internal.client.PrestoHeaders.PRESTO_USER;
import static com.facebook.presto.jdbc.internal.guava.base.MoreObjects.firstNonNull;
import static com.facebook.presto.jdbc.internal.guava.base.Preconditions.checkState;
import static com.facebook.presto.jdbc.internal.guava.net.HttpHeaders.USER_AGENT;
import static com.facebook.presto.jdbc.internal.airlift.json.JsonCodec.jsonCodec;
import static java.lang.String.format;
import static java.net.HttpURLConnection.HTTP_OK;
import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED;
import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.NANOSECONDS;
@ThreadSafe
public class StatementClient
implements Closeable
{
private static final MediaType MEDIA_TYPE_JSON = MediaType.parse("application/json; charset=utf-8");
private static final JsonCodec QUERY_RESULTS_CODEC = jsonCodec(QueryResults.class);
private static final Splitter SESSION_HEADER_SPLITTER = Splitter.on('=').limit(2).trimResults();
private static final String USER_AGENT_VALUE = StatementClient.class.getSimpleName() +
"/" +
firstNonNull(StatementClient.class.getPackage().getImplementationVersion(), "unknown");
private final OkHttpClient httpClient;
private final boolean debug;
private final String query;
private final AtomicReference currentResults = new AtomicReference<>();
private final AtomicReference setCatalog = new AtomicReference<>();
private final AtomicReference setSchema = new AtomicReference<>();
private final Map setSessionProperties = new ConcurrentHashMap<>();
private final Set resetSessionProperties = Sets.newConcurrentHashSet();
private final Map addedPreparedStatements = new ConcurrentHashMap<>();
private final Set deallocatedPreparedStatements = Sets.newConcurrentHashSet();
private final AtomicReference startedTransactionId = new AtomicReference<>();
private final AtomicBoolean clearTransactionId = new AtomicBoolean();
private final AtomicBoolean closed = new AtomicBoolean();
private final AtomicBoolean gone = new AtomicBoolean();
private final AtomicBoolean valid = new AtomicBoolean(true);
private final TimeZoneKey timeZone;
private final long requestTimeoutNanos;
private final String user;
public StatementClient(OkHttpClient httpClient, ClientSession session, String query)
{
requireNonNull(httpClient, "httpClient is null");
requireNonNull(session, "session is null");
requireNonNull(query, "query is null");
this.httpClient = httpClient;
this.debug = session.isDebug();
this.timeZone = session.getTimeZone();
this.query = query;
this.requestTimeoutNanos = session.getClientRequestTimeout().roundTo(NANOSECONDS);
this.user = session.getUser();
Request request = buildQueryRequest(session, query);
JsonResponse response = JsonResponse.execute(QUERY_RESULTS_CODEC, httpClient, request);
if ((response.getStatusCode() != HTTP_OK) || !response.hasValue()) {
throw requestFailedException("starting query", request, response);
}
processResponse(response.getHeaders(), response.getValue());
}
private Request buildQueryRequest(ClientSession session, String query)
{
HttpUrl url = HttpUrl.get(session.getServer());
if (url == null) {
throw new ClientException("Invalid server URL: " + session.getServer());
}
url = url.newBuilder().encodedPath("/v1/statement").build();
Request.Builder builder = prepareRequest(url)
.post(RequestBody.create(MEDIA_TYPE_JSON, query));
if (session.getSource() != null) {
builder.addHeader(PRESTO_SOURCE, session.getSource());
}
if (session.getClientTags() != null && !session.getClientTags().isEmpty()) {
builder.addHeader(PRESTO_CLIENT_TAGS, Joiner.on(",").join(session.getClientTags()));
}
if (session.getClientInfo() != null) {
builder.addHeader(PRESTO_CLIENT_INFO, session.getClientInfo());
}
if (session.getCatalog() != null) {
builder.addHeader(PRESTO_CATALOG, session.getCatalog());
}
if (session.getSchema() != null) {
builder.addHeader(PRESTO_SCHEMA, session.getSchema());
}
builder.addHeader(PRESTO_TIME_ZONE, session.getTimeZone().getId());
if (session.getLocale() != null) {
builder.addHeader(PRESTO_LANGUAGE, session.getLocale().toLanguageTag());
}
Map property = session.getProperties();
for (Entry entry : property.entrySet()) {
builder.addHeader(PRESTO_SESSION, entry.getKey() + "=" + entry.getValue());
}
Map statements = session.getPreparedStatements();
for (Entry entry : statements.entrySet()) {
builder.addHeader(PRESTO_PREPARED_STATEMENT, urlEncode(entry.getKey()) + "=" + urlEncode(entry.getValue()));
}
builder.addHeader(PRESTO_TRANSACTION_ID, session.getTransactionId() == null ? "NONE" : session.getTransactionId());
return builder.build();
}
public String getQuery()
{
return query;
}
public TimeZoneKey getTimeZone()
{
return timeZone;
}
public boolean isDebug()
{
return debug;
}
public boolean isClosed()
{
return closed.get();
}
public boolean isGone()
{
return gone.get();
}
public boolean isFailed()
{
return currentResults.get().getError() != null;
}
public StatementStats getStats()
{
return currentResults.get().getStats();
}
public QueryStatusInfo currentStatusInfo()
{
checkState(isValid(), "current position is not valid (cursor past end)");
return currentResults.get();
}
public QueryData currentData()
{
checkState(isValid(), "current position is not valid (cursor past end)");
return currentResults.get();
}
public QueryStatusInfo finalStatusInfo()
{
checkState((!isValid()) || isFailed(), "current position is still valid");
return currentResults.get();
}
public Optional getSetCatalog()
{
return Optional.ofNullable(setCatalog.get());
}
public Optional getSetSchema()
{
return Optional.ofNullable(setSchema.get());
}
public Map getSetSessionProperties()
{
return ImmutableMap.copyOf(setSessionProperties);
}
public Set getResetSessionProperties()
{
return ImmutableSet.copyOf(resetSessionProperties);
}
public Map getAddedPreparedStatements()
{
return ImmutableMap.copyOf(addedPreparedStatements);
}
public Set getDeallocatedPreparedStatements()
{
return ImmutableSet.copyOf(deallocatedPreparedStatements);
}
@Nullable
public String getStartedTransactionId()
{
return startedTransactionId.get();
}
public boolean isClearTransactionId()
{
return clearTransactionId.get();
}
public boolean isValid()
{
return valid.get() && (!isGone()) && (!isClosed());
}
private Request.Builder prepareRequest(HttpUrl url)
{
return new Request.Builder()
.addHeader(PRESTO_USER, user)
.addHeader(USER_AGENT, USER_AGENT_VALUE)
.url(url);
}
public boolean advance()
{
URI nextUri = currentStatusInfo().getNextUri();
if (isClosed() || (nextUri == null)) {
valid.set(false);
return false;
}
Request request = prepareRequest(HttpUrl.get(nextUri)).build();
Exception cause = null;
long start = System.nanoTime();
long attempts = 0;
do {
// back-off on retry
if (attempts > 0) {
try {
MILLISECONDS.sleep(attempts * 100);
}
catch (InterruptedException e) {
try {
close();
}
finally {
Thread.currentThread().interrupt();
}
throw new RuntimeException("StatementClient thread was interrupted");
}
}
attempts++;
JsonResponse response;
try {
response = JsonResponse.execute(QUERY_RESULTS_CODEC, httpClient, request);
}
catch (RuntimeException e) {
cause = e;
continue;
}
if ((response.getStatusCode() == HTTP_OK) && response.hasValue()) {
processResponse(response.getHeaders(), response.getValue());
return true;
}
if (response.getStatusCode() != HTTP_UNAVAILABLE) {
throw requestFailedException("fetching next", request, response);
}
}
while (((System.nanoTime() - start) < requestTimeoutNanos) && !isClosed());
gone.set(true);
throw new RuntimeException("Error fetching next", cause);
}
private void processResponse(Headers headers, QueryResults results)
{
setCatalog.set(headers.get(PRESTO_SET_CATALOG));
setSchema.set(headers.get(PRESTO_SET_SCHEMA));
for (String setSession : headers.values(PRESTO_SET_SESSION)) {
List keyValue = SESSION_HEADER_SPLITTER.splitToList(setSession);
if (keyValue.size() != 2) {
continue;
}
setSessionProperties.put(keyValue.get(0), keyValue.get(1));
}
resetSessionProperties.addAll(headers.values(PRESTO_CLEAR_SESSION));
for (String entry : headers.values(PRESTO_ADDED_PREPARE)) {
List keyValue = SESSION_HEADER_SPLITTER.splitToList(entry);
if (keyValue.size() != 2) {
continue;
}
addedPreparedStatements.put(urlDecode(keyValue.get(0)), urlDecode(keyValue.get(1)));
}
for (String entry : headers.values(PRESTO_DEALLOCATED_PREPARE)) {
deallocatedPreparedStatements.add(urlDecode(entry));
}
String startedTransactionId = headers.get(PRESTO_STARTED_TRANSACTION_ID);
if (startedTransactionId != null) {
this.startedTransactionId.set(startedTransactionId);
}
if (headers.get(PRESTO_CLEAR_TRANSACTION_ID) != null) {
clearTransactionId.set(true);
}
currentResults.set(results);
}
private RuntimeException requestFailedException(String task, Request request, JsonResponse response)
{
gone.set(true);
if (!response.hasValue()) {
if (response.getStatusCode() == HTTP_UNAUTHORIZED) {
return new ClientException("Authentication failed" +
Optional.ofNullable(response.getStatusMessage())
.map(message -> ": " + message)
.orElse(""));
}
return new RuntimeException(
format("Error %s at %s returned an invalid response: %s [Error: %s]", task, request.url(), response, response.getResponseBody()),
response.getException());
}
return new RuntimeException(format("Error %s at %s returned HTTP %s", task, request.url(), response.getStatusCode()));
}
public void cancelLeafStage()
{
checkState(!isClosed(), "client is closed");
URI uri = currentStatusInfo().getPartialCancelUri();
if (uri != null) {
httpDelete(uri);
}
}
@Override
public void close()
{
if (!closed.getAndSet(true)) {
URI uri = currentResults.get().getNextUri();
if (uri != null) {
httpDelete(uri);
}
}
}
private void httpDelete(URI uri)
{
Request request = prepareRequest(HttpUrl.get(uri))
.delete()
.build();
httpClient.newCall(request).enqueue(new NullCallback());
}
private static String urlEncode(String value)
{
try {
return URLEncoder.encode(value, "UTF-8");
}
catch (UnsupportedEncodingException e) {
throw new AssertionError(e);
}
}
private static String urlDecode(String value)
{
try {
return URLDecoder.decode(value, "UTF-8");
}
catch (UnsupportedEncodingException e) {
throw new AssertionError(e);
}
}
}
