com.google.cloud.datastore.TransactionImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of google-cloud-datastore Show documentation
Java idiomatic client for Google Cloud Datastore.
There is a newer version: 2.24.2
Show newest version
/*
 * 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 static com.google.cloud.datastore.ReadOption.transactionId;

import com.google.api.core.BetaApi;
import com.google.cloud.datastore.models.ExplainOptions;
import com.google.cloud.datastore.telemetry.TraceUtil;
import com.google.common.collect.ImmutableList;
import com.google.datastore.v1.ReadOptions;
import com.google.datastore.v1.TransactionOptions;
import com.google.protobuf.ByteString;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
import javax.annotation.Nonnull;

final class TransactionImpl extends BaseDatastoreBatchWriter implements Transaction {

  private final DatastoreImpl datastore;
  private final ByteString transactionId;
  private boolean rolledback;

  private final ReadOptionProtoPreparer readOptionProtoPreparer;

  @Nonnull private final TraceUtil traceUtil;

  static class ResponseImpl implements Transaction.Response {

    private final com.google.datastore.v1.CommitResponse response;
    private final int numAutoAllocatedIds;

    ResponseImpl(com.google.datastore.v1.CommitResponse response, int numAutoAllocatedIds) {
      this.response = response;
      this.numAutoAllocatedIds = numAutoAllocatedIds;
    }

    @Override
    public List getGeneratedKeys() {
      Iterator results =
          response.getMutationResultsList().iterator();
      List generated = new ArrayList<>(numAutoAllocatedIds);
      for (int i = 0; i < numAutoAllocatedIds; i++) {
        generated.add(Key.fromPb(results.next().getKey()));
      }
      return generated;
    }
  }

  TransactionImpl(DatastoreImpl datastore) {
    this(datastore, null);
  }

  TransactionImpl(DatastoreImpl datastore, TransactionOptions options) {
    super("transaction");
    this.datastore = datastore;
    com.google.datastore.v1.BeginTransactionRequest.Builder requestPb =
        com.google.datastore.v1.BeginTransactionRequest.newBuilder();

    requestPb.setProjectId(this.datastore.getOptions().getProjectId());
    requestPb.setDatabaseId(this.datastore.getOptions().getDatabaseId());

    if (options != null) {
      requestPb.setTransactionOptions(options);
    }

    transactionId = datastore.requestTransactionId(requestPb);
    this.readOptionProtoPreparer = new ReadOptionProtoPreparer();
    this.traceUtil = datastore.getOptions().getTraceUtil();
  }

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

  @Override
  public Iterator get(Key... keys) {
    validateActive();
    Optional readOptions =
        this.readOptionProtoPreparer.prepare(ImmutableList.of(transactionId(transactionId)));
    return datastore.get(readOptions, keys);
  }

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

  @Override
  public  QueryResults run(Query query) {
    validateActive();
    Optional readOptions =
        this.readOptionProtoPreparer.prepare(ImmutableList.of(transactionId(transactionId)));
    return datastore.run(readOptions, query, null);
  }

  @Override
  @BetaApi
  public  QueryResults run(Query query, ExplainOptions explainOptions) {
    validateActive();
    Optional readOptions =
        this.readOptionProtoPreparer.prepare(ImmutableList.of(transactionId(transactionId)));
    return datastore.run(readOptions, query, explainOptions.toPb());
  }

  @Override
  public AggregationResults runAggregation(AggregationQuery query) {
    return datastore.runAggregation(query, transactionId(transactionId));
  }

  @Override
  @BetaApi
  public AggregationResults runAggregation(AggregationQuery query, ExplainOptions explainOptions) {
    return datastore.runAggregation(query, explainOptions, transactionId(transactionId));
  }

  @Override
  public Transaction.Response commit() {
    validateActive();
    List mutationsPb = toMutationPbList();
    com.google.datastore.v1.CommitRequest.Builder requestPb =
        com.google.datastore.v1.CommitRequest.newBuilder();
    requestPb.setMode(com.google.datastore.v1.CommitRequest.Mode.TRANSACTIONAL);
    requestPb.setTransaction(transactionId);
    requestPb.addAllMutations(mutationsPb);
    requestPb.setProjectId(datastore.getOptions().getProjectId());
    requestPb.setDatabaseId(datastore.getOptions().getDatabaseId());
    com.google.datastore.v1.CommitResponse responsePb = datastore.commit(requestPb.build());
    deactivate();
    return new ResponseImpl(responsePb, toAddAutoId().size());
  }

  @Override
  public void rollback() {
    if (rolledback) {
      return;
    }
    validateActive();
    datastore.rollbackTransaction(transactionId);
    deactivate();
    rolledback = true;
  }

  @Override
  public Datastore getDatastore() {
    return datastore;
  }

  @Override
  public ByteString getTransactionId() {
    return transactionId;
  }
}