Java tutorial
/* * 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.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.ReadOption.EventualConsistency; 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.ReadOptions.ReadConsistency; import com.google.datastore.v1.TransactionOptions; import com.google.protobuf.ByteString; 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.Set; import java.util.concurrent.Callable; final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datastore { private final DatastoreRpc datastoreRpc; private final RetrySettings retrySettings; private static final ExceptionHandler TRANSACTION_EXCEPTION_HANDLER = TransactionExceptionHandler.build(); DatastoreImpl(DatastoreOptions options) { super(options); this.datastoreRpc = options.getDatastoreRpcV1(); retrySettings = MoreObjects.firstNonNull(options.getRetrySettings(), ServiceOptions.getNoRetrySettings()); } @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<T> implements Callable<T> { private final Datastore datastore; private final TransactionCallable<T> callable; private volatile TransactionOptions options; private volatile Transaction transaction; ReadWriteTransactionCallable(Datastore datastore, TransactionCallable<T> 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> T runInTransaction(final TransactionCallable<T> callable) { final DatastoreImpl self = this; try { return RetryHelper.runWithRetries(new ReadWriteTransactionCallable<T>(self, callable, null), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } @Override public <T> T runInTransaction(final TransactionCallable<T> callable, TransactionOptions transactionOptions) { final DatastoreImpl self = this; try { return RetryHelper.runWithRetries( new ReadWriteTransactionCallable<T>(self, callable, transactionOptions), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } @Override public <T> QueryResults<T> run(Query<T> query) { return run(null, query); } @Override public <T> QueryResults<T> run(Query<T> query, ReadOption... options) { return run(toReadOptionsPb(options), query); } <T> QueryResults<T> run(com.google.datastore.v1.ReadOptions readOptionsPb, Query<T> query) { return new QueryResultsImpl<>(this, readOptionsPb, query); } com.google.datastore.v1.RunQueryResponse runQuery(final com.google.datastore.v1.RunQueryRequest requestPb) { try { return RetryHelper.runWithRetries(new Callable<com.google.datastore.v1.RunQueryResponse>() { @Override public com.google.datastore.v1.RunQueryResponse call() throws DatastoreException { return datastoreRpc.runQuery(requestPb); } }, retrySettings, EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } @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<Key> 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()); } com.google.datastore.v1.AllocateIdsResponse responsePb = allocateIds(requestPb.build()); ImmutableList.Builder<Key> 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) { try { return RetryHelper.runWithRetries(new Callable<com.google.datastore.v1.AllocateIdsResponse>() { @Override public com.google.datastore.v1.AllocateIdsResponse call() throws DatastoreException { return datastoreRpc.allocateIds(requestPb); } }, retrySettings, EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } 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<Entity> add(FullEntity<?>... entities) { if (entities.length == 0) { return Collections.emptyList(); } List<com.google.datastore.v1.Mutation> mutationsPb = new ArrayList<>(); Map<Key, Entity> completeEntities = new LinkedHashMap<>(); for (FullEntity<?> entity : entities) { Entity completeEntity = null; if (entity.getKey() instanceof Key) { completeEntity = Entity.convert((FullEntity<Key>) 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<com.google.datastore.v1.MutationResult> mutationResults = commitResponse.getMutationResultsList() .iterator(); ImmutableList.Builder<Entity> 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<Entity> get(Key... keys) { return get(null, keys); } @Override public Iterator<Entity> get(Iterable<Key> keys, ReadOption... options) { return get(toReadOptionsPb(options), Iterables.toArray(keys, Key.class)); } private static com.google.datastore.v1.ReadOptions toReadOptionsPb(ReadOption... options) { com.google.datastore.v1.ReadOptions readOptionsPb = null; if (options != null && ReadOption.asImmutableMap(options).containsKey(EventualConsistency.class)) { readOptionsPb = com.google.datastore.v1.ReadOptions.newBuilder() .setReadConsistency(ReadConsistency.EVENTUAL).build(); } return readOptionsPb; } @Override public List<Entity> fetch(Key... keys) { return DatastoreHelper.fetch(this, keys); } @Override public List<Entity> fetch(Iterable<Key> keys, ReadOption... options) { return DatastoreHelper.fetch(this, Iterables.toArray(keys, Key.class), options); } Iterator<Entity> get(com.google.datastore.v1.ReadOptions readOptionsPb, final Key... keys) { if (keys.length == 0) { return Collections.emptyIterator(); } com.google.datastore.v1.LookupRequest.Builder requestPb = com.google.datastore.v1.LookupRequest .newBuilder(); if (readOptionsPb != null) { requestPb.setReadOptions(readOptionsPb); } for (Key k : Sets.newLinkedHashSet(Arrays.asList(keys))) { requestPb.addKeys(k.toPb()); } return new ResultsIterator(requestPb); } final class ResultsIterator extends AbstractIterator<Entity> { private final com.google.datastore.v1.LookupRequest.Builder requestPb; Iterator<com.google.datastore.v1.EntityResult> 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) { try { return RetryHelper.runWithRetries(new Callable<com.google.datastore.v1.LookupResponse>() { @Override public com.google.datastore.v1.LookupResponse call() throws DatastoreException { return datastoreRpc.lookup(requestPb); } }, retrySettings, EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } @Override public void update(Entity... entities) { if (entities.length > 0) { List<com.google.datastore.v1.Mutation> mutationsPb = new ArrayList<>(); Map<Key, Entity> 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<Entity> put(FullEntity<?>... entities) { if (entities.length == 0) { return Collections.emptyList(); } List<com.google.datastore.v1.Mutation> mutationsPb = new ArrayList<>(); Map<Key, Entity> 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<Key>) 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<com.google.datastore.v1.MutationResult> mutationResults = commitResponse.getMutationResultsList() .iterator(); ImmutableList.Builder<Entity> 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<com.google.datastore.v1.Mutation> mutationsPb = new ArrayList<>(); Set<Key> 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<com.google.datastore.v1.Mutation> 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.addAllMutations(mutationsPb); return commit(requestPb.build()); } com.google.datastore.v1.CommitResponse commit(final com.google.datastore.v1.CommitRequest requestPb) { try { return RetryHelper.runWithRetries(new Callable<com.google.datastore.v1.CommitResponse>() { @Override public com.google.datastore.v1.CommitResponse call() throws DatastoreException { return datastoreRpc.commit(requestPb); } }, retrySettings, EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } 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) { try { return RetryHelper.runWithRetries(new Callable<com.google.datastore.v1.BeginTransactionResponse>() { @Override public com.google.datastore.v1.BeginTransactionResponse call() throws DatastoreException { return datastoreRpc.beginTransaction(requestPb); } }, retrySettings, EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } void rollbackTransaction(ByteString transaction) { com.google.datastore.v1.RollbackRequest.Builder requestPb = com.google.datastore.v1.RollbackRequest .newBuilder(); requestPb.setTransaction(transaction); rollback(requestPb.build()); } void rollback(final com.google.datastore.v1.RollbackRequest requestPb) { try { RetryHelper.runWithRetries(new Callable<Void>() { @Override public Void call() throws DatastoreException { datastoreRpc.rollback(requestPb); return null; } }, retrySettings, EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); } } }