com.google.cloud.datastore.BaseDatastoreBatchWriter 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
Show all versions of google-cloud-datastore Show documentation
Java idiomatic client for Google Cloud Datastore.
/*
* 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.common.base.Preconditions;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
/** Base class for DatastoreBatchWriter. */
public abstract class BaseDatastoreBatchWriter implements DatastoreBatchWriter {
private final String name;
private final Map> toAdd = new LinkedHashMap<>();
private final List> toAddAutoId = new LinkedList<>();
private final Map> toUpdate = new LinkedHashMap<>();
private final Map> toPut = new LinkedHashMap<>();
private final Set toDelete = new LinkedHashSet<>();
private boolean active = true;
BaseDatastoreBatchWriter(String name) {
this.name = name;
}
@SuppressWarnings("unchecked")
@Override
public final void addWithDeferredIdAllocation(FullEntity... entities) {
validateActive();
for (FullEntity entity : entities) {
IncompleteKey key = entity.getKey();
Preconditions.checkArgument(key != null, "Entity must have a key");
if (key instanceof Key) {
addInternal((FullEntity) entity);
} else {
toAddAutoId.add((FullEntity) entity);
}
}
}
private void addInternal(FullEntity entity) {
Key key = entity.getKey();
if (toAdd.containsKey(key) || toUpdate.containsKey(key) || toPut.containsKey(key)) {
throw newInvalidRequest(
"Entity with the key %s was already added or updated in this %s", entity.getKey(), name);
}
if (toDelete.remove(key)) {
toPut.put(key, entity);
} else {
toAdd.put(key, entity);
}
}
@Override
public final Entity add(FullEntity entity) {
return DatastoreHelper.add(this, entity);
}
@SuppressWarnings("unchecked")
@Override
public final List add(FullEntity... entities) {
validateActive();
List incompleteKeys = Lists.newArrayListWithExpectedSize(entities.length);
for (FullEntity entity : entities) {
IncompleteKey key = entity.getKey();
Preconditions.checkArgument(key != null, "Entity must have a key");
if (!(key instanceof Key)) {
incompleteKeys.add(key);
}
}
Iterator allocated;
if (!incompleteKeys.isEmpty()) {
IncompleteKey[] toAllocate = Iterables.toArray(incompleteKeys, IncompleteKey.class);
allocated = getDatastore().allocateId(toAllocate).iterator();
} else {
allocated = Collections.emptyIterator();
}
List answer = Lists.newArrayListWithExpectedSize(entities.length);
for (FullEntity entity : entities) {
if (entity.getKey() instanceof Key) {
addInternal((FullEntity) entity);
answer.add(Entity.convert((FullEntity) entity));
} else {
Entity entityWithAllocatedId = Entity.newBuilder(allocated.next(), entity).build();
addInternal(entityWithAllocatedId);
answer.add(entityWithAllocatedId);
}
}
return answer;
}
@SafeVarargs
@Override
public final void update(Entity... entities) {
validateActive();
for (Entity entity : entities) {
Key key = entity.getKey();
if (toDelete.contains(key)) {
throw newInvalidRequest(
"Entity with the key %s was already deleted in this %s", entity.getKey(), name);
}
if (toAdd.remove(key) != null || toPut.containsKey(key)) {
toPut.put(key, entity);
} else {
toUpdate.put(key, entity);
}
}
}
private void putInternal(FullEntity entity) {
Key key = entity.getKey();
toAdd.remove(key);
toUpdate.remove(key);
toDelete.remove(key);
toPut.put(key, entity);
}
@Override
public final Entity put(FullEntity entity) {
return DatastoreHelper.put(this, entity);
}
@SuppressWarnings("unchecked")
@Override
public final void putWithDeferredIdAllocation(FullEntity... entities) {
validateActive();
for (FullEntity entity : entities) {
IncompleteKey key = entity.getKey();
Preconditions.checkArgument(key != null, "Entity must have a key");
if (key instanceof Key) {
putInternal(Entity.convert((FullEntity) entity));
} else {
toAddAutoId.add((FullEntity) entity);
}
}
}
@SuppressWarnings("unchecked")
@Override
public final List put(FullEntity... entities) {
validateActive();
List incompleteKeys = Lists.newArrayListWithExpectedSize(entities.length);
for (FullEntity entity : entities) {
IncompleteKey key = entity.getKey();
Preconditions.checkArgument(key != null, "Entity must have a key");
if (!(key instanceof Key)) {
incompleteKeys.add(key);
}
}
Iterator allocated;
if (!incompleteKeys.isEmpty()) {
IncompleteKey[] toAllocate = Iterables.toArray(incompleteKeys, IncompleteKey.class);
allocated = getDatastore().allocateId(toAllocate).iterator();
} else {
allocated = Collections.emptyIterator();
}
List answer = Lists.newArrayListWithExpectedSize(entities.length);
for (FullEntity entity : entities) {
if (entity.getKey() instanceof Key) {
putInternal((FullEntity) entity);
answer.add(Entity.convert((FullEntity) entity));
} else {
Entity entityWithAllocatedId = Entity.newBuilder(allocated.next(), entity).build();
putInternal(entityWithAllocatedId);
answer.add(entityWithAllocatedId);
}
}
return answer;
}
@Override
public final void delete(Key... keys) {
validateActive();
for (Key key : keys) {
toAdd.remove(key);
toUpdate.remove(key);
toPut.remove(key);
toDelete.add(key);
}
}
@Override
public boolean isActive() {
return active;
}
protected String getName() {
return name;
}
protected Map> toAdd() {
return toAdd;
}
protected List> toAddAutoId() {
return toAddAutoId;
}
protected Map> toUpdate() {
return toUpdate;
}
protected Map> toPut() {
return toPut;
}
protected Set toDelete() {
return toDelete;
}
protected void deactivate() {
active = false;
}
protected void validateActive() {
if (!active) {
throw newInvalidRequest("%s is no longer active", name);
}
}
protected DatastoreException newInvalidRequest(String msg, Object... params) {
return DatastoreException.throwInvalidRequest(String.format(msg, params));
}
protected List toMutationPbList() {
List mutationsPb = new ArrayList<>();
for (FullEntity entity : toAddAutoId()) {
mutationsPb.add(
com.google.datastore.v1.Mutation.newBuilder().setInsert(entity.toPb()).build());
}
for (FullEntity entity : toAdd().values()) {
mutationsPb.add(
com.google.datastore.v1.Mutation.newBuilder().setInsert(entity.toPb()).build());
}
for (FullEntity entity : toUpdate().values()) {
mutationsPb.add(
com.google.datastore.v1.Mutation.newBuilder().setUpdate(entity.toPb()).build());
}
for (FullEntity entity : toPut().values()) {
mutationsPb.add(
com.google.datastore.v1.Mutation.newBuilder().setUpsert(entity.toPb()).build());
}
for (Key key : toDelete()) {
mutationsPb.add(com.google.datastore.v1.Mutation.newBuilder().setDelete(key.toPb()).build());
}
return mutationsPb;
}
protected abstract Datastore getDatastore();
}