com.spotify.asyncdatastoreclient.DatastoreImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of async-datastore-client Show documentation
There is a newer version: 3.1.0
Show newest version
/*
 * Copyright (c) 2011-2015 Spotify AB
 *
 * 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.spotify.asyncdatastoreclient;

import com.google.api.client.auth.oauth2.Credential;
import com.google.api.client.http.protobuf.ProtoHttpContent;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
import com.google.datastore.v1.BeginTransactionResponse;
import com.google.datastore.v1.CommitRequest;
import com.google.datastore.v1.CommitResponse;
import com.google.datastore.v1.LookupRequest;
import com.google.datastore.v1.LookupResponse;
import com.google.datastore.v1.Mutation;
import com.google.datastore.v1.PartitionId;
import com.google.datastore.v1.ReadOptions;
import com.google.datastore.v1.RollbackRequest;
import com.google.datastore.v1.RollbackResponse;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
import com.google.protobuf.ByteString;
import com.ning.http.client.AsyncHttpClient;
import com.ning.http.client.AsyncHttpClientConfig;
import com.ning.http.client.Response;
import com.ning.http.client.extra.ListenableFutureAdapter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.zip.GZIPInputStream;

/**
 * The Datastore implementation.
 */
final class DatastoreImpl implements Datastore {

  private static final Logger log = LoggerFactory.getLogger(Datastore.class);

  private static final String VERSION = "1.0.0";
  private static final String USER_AGENT = "Datastore-Java-Client/" + VERSION + " (gzip)";

  private final DatastoreConfig config;
  private final AsyncHttpClient client;
  private final String prefixUri;

  private final ScheduledExecutorService executor;
  private volatile String accessToken;

  DatastoreImpl(final DatastoreConfig config) {
    this.config = config;
    final AsyncHttpClientConfig httpConfig = new AsyncHttpClientConfig.Builder()
        .setConnectTimeout(config.getConnectTimeout())
        .setRequestTimeout(config.getRequestTimeout())
        .setMaxConnections(config.getMaxConnections())
        .setMaxRequestRetry(config.getRequestRetry())
        .setCompressionEnforced(true)
        .build();

    client = new AsyncHttpClient(httpConfig);
    prefixUri = String.format("%s/%s/projects/%s:", config.getHost(), config.getVersion(), config.getProject());

    executor = Executors.newSingleThreadScheduledExecutor();

    if (config.getCredential() != null) {
      // block while retrieving an access token for the first time
      refreshAccessToken();

      // wake up every 10 seconds to check if access token has expired
      executor.scheduleAtFixedRate(this::refreshAccessToken, 10, 10, TimeUnit.SECONDS);
    }
  }

  @Override
  public void close() {
    executor.shutdown();
    client.close();
  }

  private void refreshAccessToken() {
    final Credential credential = config.getCredential();
    final Long expiresIn = credential.getExpiresInSeconds();

    // trigger refresh if token null or is about to expire
    if (this.accessToken == null
        || credential.getAccessToken() == null
        || expiresIn != null && expiresIn <= 60) {
      try {
        credential.refreshToken();
        final String accessTokenLocal = credential.getAccessToken();
        if (accessTokenLocal != null) {
          this.accessToken = accessTokenLocal;
        }
      } catch (final IOException e) {
        log.error("Storage exception", Throwables.getRootCause(e));
      }
    }
  }

  private static boolean isSuccessful(final int statusCode) {
    return statusCode >= 200 && statusCode < 300;
  }

  private AsyncHttpClient.BoundRequestBuilder prepareRequest(final String method, final ProtoHttpContent payload) throws IOException {
    final AsyncHttpClient.BoundRequestBuilder builder = client.preparePost(prefixUri + method);
    builder.addHeader("Authorization", "Bearer " + accessToken);
    builder.addHeader("Content-Type", "application/x-protobuf");
    builder.addHeader("User-Agent", USER_AGENT);
    builder.addHeader("Accept-Encoding", "gzip");
    builder.setContentLength((int) payload.getLength());
    builder.setBody(payload.getMessage().toByteArray());
    return builder;
  }

  private InputStream streamResponse(final Response response) throws IOException {
    final InputStream input = response.getResponseBodyAsStream();
    final boolean compressed = "gzip".equals(response.getHeader("Content-Encoding"));
    return compressed ? new GZIPInputStream(input) : input;
  }

  @Override
  public TransactionResult transaction() throws DatastoreException {
    return Futures.getChecked(transactionAsync(), DatastoreException.class);
  }

  @Override
  public ListenableFuture transactionAsync() {
    final ListenableFuture httpResponse;
    try {
      final BeginTransactionRequest.Builder request = BeginTransactionRequest.newBuilder();
      final ProtoHttpContent payload = new ProtoHttpContent(request.build());
      httpResponse = ListenableFutureAdapter.asGuavaFuture(prepareRequest("beginTransaction", payload).execute());
    } catch (final Exception e) {
      return Futures.immediateFailedFuture(new DatastoreException(e));
    }
    return Futures.transformAsync(httpResponse, response -> {
      if (!isSuccessful(response.getStatusCode())) {
        throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
      }
      final BeginTransactionResponse transaction = BeginTransactionResponse.parseFrom(streamResponse(response));
      return Futures.immediateFuture(TransactionResult.build(transaction));
    });
  }

  @Override
  public RollbackResult rollback(final TransactionResult txn) throws DatastoreException {
    return Futures.getChecked(rollbackAsync(Futures.immediateFuture(txn)), DatastoreException.class);
  }

  @Override
  public ListenableFuture rollbackAsync(final ListenableFuture txn) {
    final ListenableFuture httpResponse = Futures.transformAsync(txn, result -> {
      final ByteString transaction = result.getTransaction();
      if (transaction == null) {
        throw new DatastoreException("Invalid transaction.");
      }
      final RollbackRequest.Builder request = RollbackRequest.newBuilder();
      final ProtoHttpContent payload = new ProtoHttpContent(request.build());
      return ListenableFutureAdapter.asGuavaFuture(prepareRequest("rollback", payload).execute());
    });
    return Futures.transformAsync(httpResponse, response -> {
      if (!isSuccessful(response.getStatusCode())) {
        throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
      }
      final RollbackResponse rollback = RollbackResponse.parseFrom(streamResponse(response));
      return Futures.immediateFuture(RollbackResult.build(rollback));
    });
  }

  @Override
  public MutationResult commit(final TransactionResult txn) throws DatastoreException {
    return Futures.getChecked(executeAsync((MutationStatement) null, Futures.immediateFuture(txn)), DatastoreException.class);
  }

  @Override
  public ListenableFuture commitAsync(final ListenableFuture txn) {
    return executeAsync((MutationStatement) null, txn);
  }

  @Override
  public AllocateIdsResult execute(final AllocateIds statement) throws DatastoreException {
    return Futures.getChecked(executeAsync(statement), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final AllocateIds statement) {
    final ListenableFuture httpResponse;
    try {
      final AllocateIdsRequest.Builder request = AllocateIdsRequest.newBuilder()
          .addAllKeys(statement.getPb(config.getNamespace()));
      final ProtoHttpContent payload = new ProtoHttpContent(request.build());
      httpResponse = ListenableFutureAdapter.asGuavaFuture(prepareRequest("allocateIds", payload).execute());
    } catch (final Exception e) {
      return Futures.immediateFailedFuture(new DatastoreException(e));
    }
    return Futures.transformAsync(httpResponse, response -> {
      if (!isSuccessful(response.getStatusCode())) {
        throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
      }
      final AllocateIdsResponse allocate = AllocateIdsResponse.parseFrom(streamResponse(response));
      return Futures.immediateFuture(AllocateIdsResult.build(allocate));
    });
  }

  @Override
  public QueryResult execute(final KeyQuery statement) throws DatastoreException {
    return Futures.getChecked(executeAsync(statement), DatastoreException.class);
  }

  @Override
  public QueryResult execute(final List statements) throws DatastoreException {
    return Futures.getChecked(executeAsync(statements), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final KeyQuery statement) {
    return executeAsync(statement, Futures.immediateFuture(TransactionResult.build()));
  }

  @Override
  public ListenableFuture executeAsync(final List statements) {
    return executeAsync(statements, Futures.immediateFuture(TransactionResult.build()));
  }

  @Override
  public QueryResult execute(final KeyQuery statement, final TransactionResult txn) throws DatastoreException {
    return Futures.getChecked(executeAsync(statement, Futures.immediateFuture(txn)), DatastoreException.class);
  }

  @Override
  public QueryResult execute(final List statements, final TransactionResult txn) throws DatastoreException {
    return Futures.getChecked(executeAsync(statements, Futures.immediateFuture(txn)), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final KeyQuery statement, final ListenableFuture txn) {
    return executeAsync(ImmutableList.of(statement), txn);
  }

  @Override
  public ListenableFuture executeAsync(final List statements, final ListenableFuture txn) {
    final ListenableFuture httpResponse = Futures.transformAsync(txn, result -> {
      final List keys = statements
        .stream().map(s -> s.getKey().getPb(config.getNamespace())).collect(Collectors.toList());
      final LookupRequest.Builder request = LookupRequest.newBuilder().addAllKeys(keys);
      final ByteString transaction = result.getTransaction();
      if (transaction != null) {
        request.setReadOptions(ReadOptions.newBuilder().setTransaction(transaction));
      }
      final ProtoHttpContent payload = new ProtoHttpContent(request.build());
      return ListenableFutureAdapter.asGuavaFuture(prepareRequest("lookup", payload).execute());
    });
    return Futures.transformAsync(httpResponse, response -> {
      if (!isSuccessful(response.getStatusCode())) {
        throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
      }
      final LookupResponse query = LookupResponse.parseFrom(streamResponse(response));
      return Futures.immediateFuture(QueryResult.build(query));
    });
  }

  @Override
  public MutationResult execute(final MutationStatement statement) throws DatastoreException {
    return Futures.getChecked(executeAsync(statement), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final MutationStatement statement) {
    return executeAsync(statement, Futures.immediateFuture(TransactionResult.build()));
  }

  @Override
  public MutationResult execute(final MutationStatement statement, final TransactionResult txn) throws DatastoreException {
    return Futures.getChecked(executeAsync(statement, Futures.immediateFuture(txn)), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final MutationStatement statement, final ListenableFuture txn) {
    final List mutations = Optional
      .ofNullable(statement)
      .flatMap(s -> Optional.of(ImmutableList.of(s.getPb(config.getNamespace()))))
      .orElse(ImmutableList.of());

    return executeAsyncMutations(mutations, txn);
  }

  @Override
  public MutationResult execute(final Batch batch) throws DatastoreException {
    return Futures.getChecked(executeAsync(batch), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final Batch batch) {
    return executeAsync(batch, Futures.immediateFuture(TransactionResult.build()));
  }

  @Override
  public MutationResult execute(final Batch batch, final TransactionResult txn) throws DatastoreException {
    return Futures.getChecked(executeAsync(batch, Futures.immediateFuture(txn)), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final Batch batch, final ListenableFuture txn) {
    return executeAsyncMutations(batch.getPb(config.getNamespace()), txn);
  }

  private ListenableFuture executeAsyncMutations(final List mutations, final ListenableFuture txn) {
    final ListenableFuture httpResponse = Futures.transformAsync(txn, result -> {
      final CommitRequest.Builder request = CommitRequest.newBuilder();
      if (mutations != null) {
        request.addAllMutations(mutations);
      }

      final ByteString transaction = result.getTransaction();
      if (transaction != null) {
        request.setTransaction(transaction);
      } else {
        request.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
      }
      final ProtoHttpContent payload = new ProtoHttpContent(request.build());
      return ListenableFutureAdapter.asGuavaFuture(prepareRequest("commit", payload).execute());
    });
    return Futures.transformAsync(httpResponse, response -> {
      if (!isSuccessful(response.getStatusCode())) {
        throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
      }
      final CommitResponse commit = CommitResponse.parseFrom(streamResponse(response));
      return Futures.immediateFuture(MutationResult.build(commit));
    });
  }

  @Override
  public QueryResult execute(final Query statement) throws DatastoreException {
    return Futures.getChecked(executeAsync(statement), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final Query statement) {
    return executeAsync(statement, Futures.immediateFuture(TransactionResult.build()));
  }

  @Override
  public QueryResult execute(final Query statement, final TransactionResult txn) throws DatastoreException {
    return Futures.getChecked(executeAsync(statement, Futures.immediateFuture(txn)), DatastoreException.class);
  }

  @Override
  public ListenableFuture executeAsync(final Query statement, final ListenableFuture txn) {
    final ListenableFuture httpResponse = Futures.transformAsync(txn, result -> {
      final String namespace = config.getNamespace();
      final RunQueryRequest.Builder request = RunQueryRequest.newBuilder()
        .setQuery(statement.getPb(namespace != null ? namespace : ""));
      if (namespace != null) {
        request.setPartitionId(PartitionId.newBuilder().setNamespaceId(namespace));
      }
      final ByteString transaction = result.getTransaction();
      if (transaction != null) {
        request.setReadOptions(ReadOptions.newBuilder().setTransaction(transaction));
      }
      final ProtoHttpContent payload = new ProtoHttpContent(request.build());
      return ListenableFutureAdapter.asGuavaFuture(prepareRequest("runQuery", payload).execute());
    });
    return Futures.transformAsync(httpResponse, response -> {
      if (!isSuccessful(response.getStatusCode())) {
        throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
      }
      final RunQueryResponse query = RunQueryResponse.parseFrom(streamResponse(response));
      return Futures.immediateFuture(QueryResult.build(query));
    });
  }
}