@Override
public ListenableFuture<TransactionResult> transactionAsync() {
  final ListenableFuture<Response> 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));
  }, MoreExecutors.directExecutor());
}
 

@Override
public ListenableFuture<RollbackResult> rollbackAsync(final ListenableFuture<TransactionResult> txn) {
  final ListenableFuture<Response> 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());
  }, MoreExecutors.directExecutor());
  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));
  }, MoreExecutors.directExecutor());
}
 

@Override
public ListenableFuture<AllocateIdsResult> executeAsync(final AllocateIds statement) {
  final ListenableFuture<Response> 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));
  }, MoreExecutors.directExecutor());
}
 

@Override
public ListenableFuture<QueryResult> executeAsync(final List<KeyQuery> statements, final ListenableFuture<TransactionResult> txn) {
  final ListenableFuture<Response> httpResponse = Futures.transformAsync(txn, result -> {
    final List<com.google.datastore.v1.Key> 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());
  }, MoreExecutors.directExecutor());
  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));
  }, MoreExecutors.directExecutor());
}
 

private ListenableFuture<MutationResult> executeAsyncMutations(final List<Mutation> mutations, final ListenableFuture<TransactionResult> txn) {
  final ListenableFuture<Response> 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());
  }, MoreExecutors.directExecutor());
  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));
  }, MoreExecutors.directExecutor());
}
 

@Override
public ListenableFuture<QueryResult> executeAsync(final Query statement, final ListenableFuture<TransactionResult> txn) {
  final ListenableFuture<Response> 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());
  }, MoreExecutors.directExecutor());
  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));
  }, MoreExecutors.directExecutor());
}
 

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;
}
 

/**
 * @param abstractGoogleProtoClient Google protocol buffer client
 * @param requestMethod HTTP Method
 * @param uriTemplate URI template for the path relative to the base URL. If it starts with a "/"
 *        the base path from the base URL will be stripped out. The URI template can also be a
 *        full URL. URI template expansion is done using
 *        {@link UriTemplate#expand(String, String, Object, boolean)}
 * @param message message to serialize or {@code null} for none
 * @param responseClass response class to parse into
 */
protected AbstractGoogleProtoClientRequest(AbstractGoogleProtoClient abstractGoogleProtoClient,
    String requestMethod, String uriTemplate, MessageLite message, Class<T> responseClass) {
  super(abstractGoogleProtoClient, requestMethod, uriTemplate, message == null
      ? null : new ProtoHttpContent(message), responseClass);
  this.message = message;
}