• Home
• com.google.cloud
• google-cloud-datastore
• 2.21.2
• source code
• DatastoreImpl.java

×

Project download

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.google.cloud.datastore.DatastoreImpl Maven / Gradle / Ivy

/*
 * Copyright 2015 Google LLC
 *
 * 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.google.cloud.datastore;

import com.google.api.core.BetaApi;
import com.google.api.gax.retrying.RetrySettings;
import com.google.cloud.BaseService;
import com.google.cloud.ExceptionHandler;
import com.google.cloud.RetryHelper;
import com.google.cloud.RetryHelper.RetryHelperException;
import com.google.cloud.ServiceOptions;
import com.google.cloud.datastore.execution.AggregationQueryExecutor;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;
import com.google.datastore.v1.ExplainOptions;
import com.google.datastore.v1.ReadOptions;
import com.google.datastore.v1.ReserveIdsRequest;
import com.google.datastore.v1.TransactionOptions;
import com.google.protobuf.ByteString;
import io.opencensus.common.Scope;
import io.opencensus.trace.Span;
import io.opencensus.trace.Status;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Callable;

final class DatastoreImpl extends BaseService implements Datastore {

  private final DatastoreRpc datastoreRpc;
  private final RetrySettings retrySettings;
  private static final ExceptionHandler TRANSACTION_EXCEPTION_HANDLER =
      TransactionExceptionHandler.build();
  private static final ExceptionHandler TRANSACTION_OPERATION_EXCEPTION_HANDLER =
      TransactionOperationExceptionHandler.build();
  private final TraceUtil traceUtil = TraceUtil.getInstance();

  private final ReadOptionProtoPreparer readOptionProtoPreparer;
  private final AggregationQueryExecutor aggregationQueryExecutor;

  DatastoreImpl(DatastoreOptions options) {
    super(options);
    this.datastoreRpc = options.getDatastoreRpcV1();
    retrySettings =
        MoreObjects.firstNonNull(options.getRetrySettings(), ServiceOptions.getNoRetrySettings());

    readOptionProtoPreparer = new ReadOptionProtoPreparer();
    aggregationQueryExecutor =
        new AggregationQueryExecutor(
            new RetryAndTraceDatastoreRpcDecorator(datastoreRpc, traceUtil, retrySettings, options),
            options);
  }

  @Override
  public Batch newBatch() {
    return new BatchImpl(this);
  }

  @Override
  public Transaction newTransaction(TransactionOptions transactionOptions) {
    return new TransactionImpl(this, transactionOptions);
  }

  @Override
  public Transaction newTransaction() {
    return new TransactionImpl(this);
  }

  static class ReadWriteTransactionCallable implements Callable {

    private final Datastore datastore;
    private final TransactionCallable callable;
    private volatile TransactionOptions options;
    private volatile Transaction transaction;

    ReadWriteTransactionCallable(
        Datastore datastore, TransactionCallable callable, TransactionOptions options) {
      this.datastore = datastore;
      this.callable = callable;
      this.options = options;
      this.transaction = null;
    }

    Datastore getDatastore() {
      return datastore;
    }

    TransactionOptions getOptions() {
      return options;
    }

    Transaction getTransaction() {
      return transaction;
    }

    void setPrevTransactionId(ByteString transactionId) {
      TransactionOptions.ReadWrite readWrite =
          TransactionOptions.ReadWrite.newBuilder().setPreviousTransaction(transactionId).build();
      options = options.toBuilder().setReadWrite(readWrite).build();
    }

    @Override
    public T call() throws DatastoreException {
      transaction = datastore.newTransaction(options);
      try {
        T value = callable.run(transaction);
        transaction.commit();
        return value;
      } catch (Exception ex) {
        transaction.rollback();
        throw DatastoreException.propagateUserException(ex);
      } finally {
        if (transaction.isActive()) {
          transaction.rollback();
        }
        if (options != null
            && options.getModeCase().equals(TransactionOptions.ModeCase.READ_WRITE)) {
          setPrevTransactionId(transaction.getTransactionId());
        }
      }
    }
  }

  @Override
  public  T runInTransaction(final TransactionCallable callable) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          new ReadWriteTransactionCallable(this, callable, null),
          retrySettings,
          TRANSACTION_EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  @Override
  public  T runInTransaction(
      final TransactionCallable callable, TransactionOptions transactionOptions) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          new ReadWriteTransactionCallable(this, callable, transactionOptions),
          retrySettings,
          TRANSACTION_EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  @Override
  public  QueryResults run(Query query) {
    return run(Optional.empty(), query, null);
  }

  @Override
  public  QueryResults run(Query query, ReadOption... options) {
    return run(toReadOptionsPb(options), query, null);
  }

  @Override
  @BetaApi
  public  QueryResults run(
      Query query,
      com.google.cloud.datastore.models.ExplainOptions explainOptions,
      ReadOption... options) {
    return run(toReadOptionsPb(options), query, explainOptions.toPb());
  }

  @SuppressWarnings("unchecked")
   QueryResults run(
      Optional readOptionsPb, Query query, ExplainOptions explainOptions) {
    return new QueryResultsImpl(
        this, readOptionsPb, (RecordQuery) query, query.getNamespace(), explainOptions);
  }

  @Override
  public AggregationResults runAggregation(AggregationQuery query) {
    return aggregationQueryExecutor.execute(query, null);
  }

  @Override
  public AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
    return aggregationQueryExecutor.execute(query, null, options);
  }

  @Override
  @BetaApi
  public AggregationResults runAggregation(
      AggregationQuery query, com.google.cloud.datastore.models.ExplainOptions explainOptions) {
    return aggregationQueryExecutor.execute(query, explainOptions);
  }

  @Override
  @BetaApi
  public AggregationResults runAggregation(
      AggregationQuery query,
      com.google.cloud.datastore.models.ExplainOptions explainOptions,
      ReadOption... options) {
    return aggregationQueryExecutor.execute(query, explainOptions, options);
  }

  com.google.datastore.v1.RunQueryResponse runQuery(
      final com.google.datastore.v1.RunQueryRequest requestPb) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          () -> datastoreRpc.runQuery(requestPb),
          retrySettings,
          requestPb.getReadOptions().getTransaction().isEmpty()
              ? EXCEPTION_HANDLER
              : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  @Override
  public Key allocateId(IncompleteKey key) {
    return DatastoreHelper.allocateId(this, key);
  }

  private boolean verifyIncompleteKeyType(IncompleteKey... keys) {
    for (IncompleteKey key : keys) {
      if (key instanceof Key) {
        return false;
      }
    }
    return true;
  }

  @Override
  public List allocateId(IncompleteKey... keys) {
    Preconditions.checkArgument(
        verifyIncompleteKeyType(keys), "keys must be IncompleteKey instances");
    if (keys.length == 0) {
      return Collections.emptyList();
    }
    com.google.datastore.v1.AllocateIdsRequest.Builder requestPb =
        com.google.datastore.v1.AllocateIdsRequest.newBuilder();
    for (IncompleteKey key : keys) {
      requestPb.addKeys(trimNameOrId(key).toPb());
    }
    requestPb.setProjectId(getOptions().getProjectId());
    requestPb.setDatabaseId(getOptions().getDatabaseId());
    com.google.datastore.v1.AllocateIdsResponse responsePb = allocateIds(requestPb.build());
    ImmutableList.Builder keyList = ImmutableList.builder();
    for (com.google.datastore.v1.Key keyPb : responsePb.getKeysList()) {
      keyList.add(Key.fromPb(keyPb));
    }
    return keyList.build();
  }

  private com.google.datastore.v1.AllocateIdsResponse allocateIds(
      final com.google.datastore.v1.AllocateIdsRequest requestPb) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ALLOCATEIDS);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          new Callable() {
            @Override
            public com.google.datastore.v1.AllocateIdsResponse call() throws DatastoreException {
              return datastoreRpc.allocateIds(requestPb);
            }
          },
          retrySettings,
          EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  private IncompleteKey trimNameOrId(IncompleteKey key) {
    if (key instanceof Key) {
      return IncompleteKey.newBuilder(key).build();
    }
    return key;
  }

  @Override
  public Entity add(FullEntity entity) {
    return DatastoreHelper.add(this, entity);
  }

  @SuppressWarnings("unchecked")
  @Override
  public List add(FullEntity... entities) {
    if (entities.length == 0) {
      return Collections.emptyList();
    }
    List mutationsPb = new ArrayList<>();
    Map completeEntities = new LinkedHashMap<>();
    for (FullEntity entity : entities) {
      Entity completeEntity = null;
      if (entity.getKey() instanceof Key) {
        completeEntity = Entity.convert((FullEntity) entity);
      }
      if (completeEntity != null) {
        if (completeEntities.put(completeEntity.getKey(), completeEntity) != null) {
          throw DatastoreException.throwInvalidRequest(
              "Duplicate entity with the key %s", entity.getKey());
        }
      } else {
        Preconditions.checkArgument(entity.hasKey(), "Entity %s is missing a key", entity);
      }
      mutationsPb.add(
          com.google.datastore.v1.Mutation.newBuilder().setInsert(entity.toPb()).build());
    }
    com.google.datastore.v1.CommitResponse commitResponse = commitMutation(mutationsPb);
    Iterator mutationResults =
        commitResponse.getMutationResultsList().iterator();
    ImmutableList.Builder responseBuilder = ImmutableList.builder();
    for (FullEntity entity : entities) {
      Entity completeEntity = completeEntities.get(entity.getKey());
      if (completeEntity != null) {
        responseBuilder.add(completeEntity);
        mutationResults.next();
      } else {
        responseBuilder.add(
            Entity.newBuilder(Key.fromPb(mutationResults.next().getKey()), entity).build());
      }
    }
    return responseBuilder.build();
  }

  @Override
  public Entity get(Key key) {
    return DatastoreHelper.get(this, key);
  }

  @Override
  public Entity get(Key key, ReadOption... options) {
    return DatastoreHelper.get(this, key, options);
  }

  @Override
  public Iterator get(Key... keys) {
    return get(Optional.empty(), keys);
  }

  @Override
  public Iterator get(Iterable keys, ReadOption... options) {
    return get(toReadOptionsPb(options), Iterables.toArray(keys, Key.class));
  }

  private Optional toReadOptionsPb(ReadOption... options) {
    if (options == null) {
      return Optional.empty();
    }
    return this.readOptionProtoPreparer.prepare(Arrays.asList(options));
  }

  @Override
  public List fetch(Key... keys) {
    return DatastoreHelper.fetch(this, keys);
  }

  @Override
  public List fetch(Iterable keys, ReadOption... options) {
    return DatastoreHelper.fetch(this, Iterables.toArray(keys, Key.class), options);
  }

  Iterator get(Optional readOptionsPb, final Key... keys) {
    if (keys.length == 0) {
      return Collections.emptyIterator();
    }
    com.google.datastore.v1.LookupRequest.Builder requestPb =
        com.google.datastore.v1.LookupRequest.newBuilder();
    readOptionsPb.ifPresent(requestPb::setReadOptions);
    for (Key k : Sets.newLinkedHashSet(Arrays.asList(keys))) {
      requestPb.addKeys(k.toPb());
    }
    requestPb.setProjectId(getOptions().getProjectId());
    requestPb.setDatabaseId(getOptions().getDatabaseId());
    return new ResultsIterator(requestPb);
  }

  final class ResultsIterator extends AbstractIterator {

    private final com.google.datastore.v1.LookupRequest.Builder requestPb;
    Iterator iter;

    ResultsIterator(com.google.datastore.v1.LookupRequest.Builder requestPb) {
      this.requestPb = requestPb;
      loadResults();
    }

    private void loadResults() {
      com.google.datastore.v1.LookupResponse responsePb = lookup(requestPb.build());
      iter = responsePb.getFoundList().iterator();
      requestPb.clearKeys();
      if (responsePb.getDeferredCount() > 0) {
        requestPb.addAllKeys(responsePb.getDeferredList());
      }
    }

    @SuppressWarnings("unchecked")
    @Override
    protected Entity computeNext() {
      while (!iter.hasNext()) {
        if (requestPb.getKeysCount() == 0) {
          return endOfData();
        }
        loadResults();
      }
      return Entity.fromPb(iter.next().getEntity());
    }
  }

  com.google.datastore.v1.LookupResponse lookup(
      final com.google.datastore.v1.LookupRequest requestPb) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          new Callable() {
            @Override
            public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
              return datastoreRpc.lookup(requestPb);
            }
          },
          retrySettings,
          requestPb.getReadOptions().getTransaction().isEmpty()
              ? EXCEPTION_HANDLER
              : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  @Override
  public List reserveIds(Key... keys) {
    ReserveIdsRequest.Builder requestPb = ReserveIdsRequest.newBuilder();
    for (Key key : keys) {
      requestPb.addKeys(key.toPb());
    }
    requestPb.setProjectId(getOptions().getProjectId());
    requestPb.setDatabaseId(getOptions().getDatabaseId());
    com.google.datastore.v1.ReserveIdsResponse responsePb = reserveIds(requestPb.build());
    ImmutableList.Builder keyList = ImmutableList.builder();
    if (responsePb.isInitialized()) {
      for (Key key : keys) {
        keyList.add(key);
      }
    }
    return keyList.build();
  }

  com.google.datastore.v1.ReserveIdsResponse reserveIds(
      final com.google.datastore.v1.ReserveIdsRequest requestPb) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RESERVEIDS);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          new Callable() {
            @Override
            public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreException {
              return datastoreRpc.reserveIds(requestPb);
            }
          },
          retrySettings,
          EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  @Override
  public void update(Entity... entities) {
    if (entities.length > 0) {
      List mutationsPb = new ArrayList<>();
      Map dedupEntities = new LinkedHashMap<>();
      for (Entity entity : entities) {
        dedupEntities.put(entity.getKey(), entity);
      }
      for (Entity entity : dedupEntities.values()) {
        mutationsPb.add(
            com.google.datastore.v1.Mutation.newBuilder().setUpdate(entity.toPb()).build());
      }
      commitMutation(mutationsPb);
    }
  }

  @Override
  public Entity put(FullEntity entity) {
    return DatastoreHelper.put(this, entity);
  }

  @SuppressWarnings("unchecked")
  @Override
  public List put(FullEntity... entities) {
    if (entities.length == 0) {
      return Collections.emptyList();
    }
    List mutationsPb = new ArrayList<>();
    Map dedupEntities = new LinkedHashMap<>();
    for (FullEntity entity : entities) {
      Preconditions.checkArgument(entity.hasKey(), "Entity %s is missing a key", entity);
      if (entity.getKey() instanceof Key) {
        Entity completeEntity = Entity.convert((FullEntity) entity);
        dedupEntities.put(completeEntity.getKey(), completeEntity);
      } else {
        mutationsPb.add(
            com.google.datastore.v1.Mutation.newBuilder().setUpsert(entity.toPb()).build());
      }
    }
    for (Entity entity : dedupEntities.values()) {
      mutationsPb.add(
          com.google.datastore.v1.Mutation.newBuilder().setUpsert(entity.toPb()).build());
    }
    com.google.datastore.v1.CommitResponse commitResponse = commitMutation(mutationsPb);
    Iterator mutationResults =
        commitResponse.getMutationResultsList().iterator();
    ImmutableList.Builder responseBuilder = ImmutableList.builder();
    for (FullEntity entity : entities) {
      Entity completeEntity = dedupEntities.get(entity.getKey());
      if (completeEntity != null) {
        responseBuilder.add(completeEntity);
      } else {
        responseBuilder.add(
            Entity.newBuilder(Key.fromPb(mutationResults.next().getKey()), entity).build());
      }
    }
    return responseBuilder.build();
  }

  @Override
  public void delete(Key... keys) {
    if (keys.length > 0) {
      List mutationsPb = new ArrayList<>();
      Set dedupKeys = new LinkedHashSet<>(Arrays.asList(keys));
      for (Key key : dedupKeys) {
        mutationsPb.add(
            com.google.datastore.v1.Mutation.newBuilder().setDelete(key.toPb()).build());
      }
      commitMutation(mutationsPb);
    }
  }

  @Override
  public KeyFactory newKeyFactory() {
    return DatastoreHelper.newKeyFactory(getOptions());
  }

  private com.google.datastore.v1.CommitResponse commitMutation(
      List mutationsPb) {
    com.google.datastore.v1.CommitRequest.Builder requestPb =
        com.google.datastore.v1.CommitRequest.newBuilder();
    requestPb.setMode(com.google.datastore.v1.CommitRequest.Mode.NON_TRANSACTIONAL);
    requestPb.setProjectId(getOptions().getProjectId());
    requestPb.setDatabaseId(getOptions().getDatabaseId());
    requestPb.addAllMutations(mutationsPb);
    return commit(requestPb.build());
  }

  com.google.datastore.v1.CommitResponse commit(
      final com.google.datastore.v1.CommitRequest requestPb) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          () -> datastoreRpc.commit(requestPb),
          retrySettings,
          requestPb.getTransaction().isEmpty()
              ? EXCEPTION_HANDLER
              : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  ByteString requestTransactionId(
      com.google.datastore.v1.BeginTransactionRequest.Builder requestPb) {
    return beginTransaction(requestPb.build()).getTransaction();
  }

  com.google.datastore.v1.BeginTransactionResponse beginTransaction(
      final com.google.datastore.v1.BeginTransactionRequest requestPb) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_BEGINTRANSACTION);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      return RetryHelper.runWithRetries(
          new Callable() {
            @Override
            public com.google.datastore.v1.BeginTransactionResponse call()
                throws DatastoreException {
              return datastoreRpc.beginTransaction(requestPb);
            }
          },
          retrySettings,
          EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }

  void rollbackTransaction(ByteString transaction) {
    com.google.datastore.v1.RollbackRequest.Builder requestPb =
        com.google.datastore.v1.RollbackRequest.newBuilder();
    requestPb.setTransaction(transaction);
    requestPb.setProjectId(getOptions().getProjectId());
    requestPb.setDatabaseId(getOptions().getDatabaseId());
    rollback(requestPb.build());
  }

  void rollback(final com.google.datastore.v1.RollbackRequest requestPb) {
    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ROLLBACK);
    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
      RetryHelper.runWithRetries(
          new Callable() {
            @Override
            public Void call() throws DatastoreException {
              datastoreRpc.rollback(requestPb);
              return null;
            }
          },
          retrySettings,
          EXCEPTION_HANDLER,
          getOptions().getClock());
    } catch (RetryHelperException e) {
      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
      throw DatastoreException.translateAndThrow(e);
    } finally {
      span.end(TraceUtil.END_SPAN_OPTIONS);
    }
  }
}