/* * Copyright 2016 Sai Pullabhotla. * * 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.jmethods.catatumbo.impl; import static com.jmethods.catatumbo.impl.DatastoreUtils.rollbackIfActive; import static com.jmethods.catatumbo.impl.DatastoreUtils.toEntities; import static com.jmethods.catatumbo.impl.DatastoreUtils.toNativeEntities; import static com.jmethods.catatumbo.impl.DatastoreUtils.toNativeFullEntities; import java.util.ArrayList; import java.util.List; import com.google.cloud.datastore.Batch; import com.google.cloud.datastore.Datastore; import com.google.cloud.datastore.DatastoreException; import com.google.cloud.datastore.DatastoreWriter; import com.google.cloud.datastore.Entity; import com.google.cloud.datastore.FullEntity; import com.google.cloud.datastore.Key; import com.google.cloud.datastore.Transaction; import com.jmethods.catatumbo.DatastoreKey; import com.jmethods.catatumbo.EntityManagerException; import com.jmethods.catatumbo.OptimisticLockException; import com.jmethods.catatumbo.impl.Marshaller.Intent; /** * Worker class for performing write operations on the Cloud Datastore. * * @author Sai Pullabhotla * */ public class DefaultDatastoreWriter { /** * A reference to the entity manager */ protected DefaultEntityManager entityManager; /** * Reference to the native DatastoreWriter for updating the Cloud Datastore. This could be the * {@link Datastore}, {@link Transaction} or {@link Batch}. */ protected DatastoreWriter nativeWriter; /** * A reference to the Datastore */ protected Datastore datastore; /** * Creates a new instance of <code>DefaultDatastoreWriter</code>. * * @param entityManager * a reference to the entity manager. */ public DefaultDatastoreWriter(DefaultEntityManager entityManager) { this.entityManager = entityManager; this.datastore = entityManager.getDatastore(); this.nativeWriter = datastore; } /** * Creates a new instance of <code>DefaultDatastoreWriter</code> for executing batch updates. * * @param batch * the {@link DefaultDatastoreBatch}. */ public DefaultDatastoreWriter(DefaultDatastoreBatch batch) { this.entityManager = batch.getEntityManager(); this.datastore = entityManager.getDatastore(); this.nativeWriter = batch.getNativeBatch(); } /** * Creates a new instance of <code>DefaultDatastoreWriter</code> for transactional updates. * * @param transaction * the {@link DefaultDatastoreTransaction}. */ public DefaultDatastoreWriter(DefaultDatastoreTransaction transaction) { this.entityManager = transaction.getEntityManager(); this.datastore = entityManager.getDatastore(); this.nativeWriter = transaction.getNativeTransaction(); } /** * Inserts the given entity into the Cloud Datastore. * * @param entity * the entity to insert * @return the inserted entity. The inserted entity will not be same as the passed in entity. For * example, the inserted entity may contain any generated ID, key, parent key, etc. * @throws EntityManagerException * if any error occurs while inserting. */ public <E> E insert(E entity) { try { entityManager.executeEntityListeners(CallbackType.PRE_INSERT, entity); FullEntity<?> nativeEntity = (FullEntity<?>) Marshaller.marshal(entityManager, entity, Intent.INSERT); Entity insertedNativeEntity = nativeWriter.add(nativeEntity); @SuppressWarnings("unchecked") E insertedEntity = (E) Unmarshaller.unmarshal(insertedNativeEntity, entity.getClass()); entityManager.executeEntityListeners(CallbackType.POST_INSERT, insertedEntity); return insertedEntity; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Inserts the given list of entities into the Cloud Datastore. * * @param entities * the entities to insert. * @return the inserted entities. The inserted entities will not be same as the passed in * entities. For example, the inserted entities may contain generated ID, key, parent key, * etc. * @throws EntityManagerException * if any error occurs while inserting. */ @SuppressWarnings("unchecked") public <E> List<E> insert(List<E> entities) { if (entities == null || entities.isEmpty()) { return new ArrayList<>(); } try { entityManager.executeEntityListeners(CallbackType.PRE_INSERT, entities); FullEntity<?>[] nativeEntities = toNativeFullEntities(entities, entityManager, Intent.INSERT); Class<?> entityClass = entities.get(0).getClass(); List<Entity> insertedNativeEntities = nativeWriter.add(nativeEntities); List<E> insertedEntities = (List<E>) toEntities(entityClass, insertedNativeEntities); entityManager.executeEntityListeners(CallbackType.POST_INSERT, insertedEntities); return insertedEntities; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Updates the given entity in the Cloud Datastore. The passed in Entity must have its ID set for * the update to work. * * @param entity * the entity to update * @return the updated entity. * @throws EntityManagerException * if any error occurs while updating. */ @SuppressWarnings("unchecked") public <E> E update(E entity) { try { entityManager.executeEntityListeners(CallbackType.PRE_UPDATE, entity); Intent intent = (nativeWriter instanceof Batch) ? Intent.BATCH_UPDATE : Intent.UPDATE; Entity nativeEntity = (Entity) Marshaller.marshal(entityManager, entity, intent); nativeWriter.update(nativeEntity); E updatedEntity = (E) Unmarshaller.unmarshal(nativeEntity, entity.getClass()); entityManager.executeEntityListeners(CallbackType.POST_UPDATE, updatedEntity); return updatedEntity; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Updates the given list of entities in the Cloud Datastore. * * @param entities * the entities to update. The passed in entities must have their ID set for the update * to work. * @return the updated entities * @throws EntityManagerException * if any error occurs while inserting. */ @SuppressWarnings("unchecked") public <E> List<E> update(List<E> entities) { if (entities == null || entities.isEmpty()) { return new ArrayList<>(); } try { Class<E> entityClass = (Class<E>) entities.get(0).getClass(); entityManager.executeEntityListeners(CallbackType.PRE_UPDATE, entities); Intent intent = (nativeWriter instanceof Batch) ? Intent.BATCH_UPDATE : Intent.UPDATE; Entity[] nativeEntities = toNativeEntities(entities, entityManager, intent); nativeWriter.update(nativeEntities); List<E> updatedEntities = toEntities(entityClass, nativeEntities); entityManager.executeEntityListeners(CallbackType.POST_UPDATE, updatedEntities); return updatedEntities; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Updates the given entity with optimistic locking, if the entity is set up to support optimistic * locking. Otherwise, a normal update is performed. * * @param entity * the entity to update * @return the updated entity which may be different than the given entity. */ public <E> E updateWithOptimisticLock(E entity) { PropertyMetadata versionMetadata = EntityIntrospector.getVersionMetadata(entity); if (versionMetadata == null) { return update(entity); } else { return updateWithOptimisticLockingInternal(entity, versionMetadata); } } /** * Updates the given list of entities using optimistic locking feature, if the entities are set up * to support optimistic locking. Otherwise, a normal update is performed. * * @param entities * the entities to update * @return the updated entities */ public <E> List<E> updateWithOptimisticLock(List<E> entities) { if (entities == null || entities.isEmpty()) { return new ArrayList<>(); } Class<?> entityClass = entities.get(0).getClass(); PropertyMetadata versionMetadata = EntityIntrospector.getVersionMetadata(entityClass); if (versionMetadata == null) { return update(entities); } else { return updateWithOptimisticLockInternal(entities, versionMetadata); } } /** * Worker method for updating the given entity with optimistic locking. * * @param entity * the entity to update * @param versionMetadata * the metadata for optimistic locking * @return the updated entity */ @SuppressWarnings("unchecked") protected <E> E updateWithOptimisticLockingInternal(E entity, PropertyMetadata versionMetadata) { Transaction transaction = null; try { entityManager.executeEntityListeners(CallbackType.PRE_UPDATE, entity); Entity nativeEntity = (Entity) Marshaller.marshal(entityManager, entity, Intent.UPDATE); transaction = datastore.newTransaction(); Entity storedNativeEntity = transaction.get(nativeEntity.getKey()); if (storedNativeEntity == null) { throw new OptimisticLockException( String.format("Entity does not exist: %s", nativeEntity.getKey())); } String versionPropertyName = versionMetadata.getMappedName(); long version = nativeEntity.getLong(versionPropertyName) - 1; long storedVersion = storedNativeEntity.getLong(versionPropertyName); if (version != storedVersion) { throw new OptimisticLockException( String.format("Expecting version %d, but found %d", version, storedVersion)); } transaction.update(nativeEntity); transaction.commit(); E updatedEntity = (E) Unmarshaller.unmarshal(nativeEntity, entity.getClass()); entityManager.executeEntityListeners(CallbackType.POST_UPDATE, updatedEntity); return updatedEntity; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } finally { rollbackIfActive(transaction); } } /** * Internal worker method for updating the entities using optimistic locking. * * @param entities * the entities to update * @param versionMetadata * the metadata of the version property * @return the updated entities */ @SuppressWarnings("unchecked") protected <E> List<E> updateWithOptimisticLockInternal(List<E> entities, PropertyMetadata versionMetadata) { Transaction transaction = null; try { entityManager.executeEntityListeners(CallbackType.PRE_UPDATE, entities); Entity[] nativeEntities = toNativeEntities(entities, entityManager, Intent.UPDATE); // The above native entities already have the version incremented by // the marshalling process Key[] nativeKeys = new Key[nativeEntities.length]; for (int i = 0; i < nativeEntities.length; i++) { nativeKeys[i] = nativeEntities[i].getKey(); } transaction = datastore.newTransaction(); List<Entity> storedNativeEntities = transaction.fetch(nativeKeys); String versionPropertyName = versionMetadata.getMappedName(); for (int i = 0; i < nativeEntities.length; i++) { long version = nativeEntities[i].getLong(versionPropertyName) - 1; Entity storedNativeEntity = storedNativeEntities.get(i); if (storedNativeEntity == null) { throw new OptimisticLockException( String.format("Entity does not exist: %s", nativeKeys[i])); } long storedVersion = storedNativeEntities.get(i).getLong(versionPropertyName); if (version != storedVersion) { throw new OptimisticLockException( String.format("Expecting version %d, but found %d", version, storedVersion)); } } transaction.update(nativeEntities); transaction.commit(); List<E> updatedEntities = (List<E>) toEntities(entities.get(0).getClass(), nativeEntities); entityManager.executeEntityListeners(CallbackType.POST_UPDATE, updatedEntities); return updatedEntities; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } finally { rollbackIfActive(transaction); } } /** * Updates or inserts the given entity in the Cloud Datastore. If the entity does not have an ID, * it may be generated. * * @param entity * the entity to update or insert * @return the updated/inserted entity. * @throws EntityManagerException * if any error occurs while saving. */ public <E> E upsert(E entity) { try { entityManager.executeEntityListeners(CallbackType.PRE_UPSERT, entity); FullEntity<?> nativeEntity = (FullEntity<?>) Marshaller.marshal(entityManager, entity, Intent.UPSERT); Entity upsertedNativeEntity = nativeWriter.put(nativeEntity); @SuppressWarnings("unchecked") E upsertedEntity = (E) Unmarshaller.unmarshal(upsertedNativeEntity, entity.getClass()); entityManager.executeEntityListeners(CallbackType.POST_UPSERT, upsertedEntity); return upsertedEntity; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Updates or inserts the given list of entities in the Cloud Datastore. If the entities do not * have a valid ID, IDs may be generated. * * @param entities * the entities to update/or insert. * @return the updated or inserted entities * @throws EntityManagerException * if any error occurs while saving. */ @SuppressWarnings("unchecked") public <E> List<E> upsert(List<E> entities) { if (entities == null || entities.isEmpty()) { return new ArrayList<>(); } try { entityManager.executeEntityListeners(CallbackType.PRE_UPSERT, entities); FullEntity<?>[] nativeEntities = toNativeFullEntities(entities, entityManager, Intent.UPSERT); Class<?> entityClass = entities.get(0).getClass(); List<Entity> upsertedNativeEntities = nativeWriter.put(nativeEntities); List<E> upsertedEntities = (List<E>) toEntities(entityClass, upsertedNativeEntities); entityManager.executeEntityListeners(CallbackType.POST_UPSERT, upsertedEntities); return upsertedEntities; } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes the given entity from the Cloud Datastore. * * @param entity * the entity to delete. The entity must have it ID set for the deletion to succeed. * @throws EntityManagerException * if any error occurs while deleting. */ public void delete(Object entity) { try { entityManager.executeEntityListeners(CallbackType.PRE_DELETE, entity); Key nativeKey = Marshaller.marshalKey(entityManager, entity); nativeWriter.delete(nativeKey); entityManager.executeEntityListeners(CallbackType.POST_DELETE, entity); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes the given entities from the Cloud Datastore. * * @param entities * the entities to delete. The entities must have it ID set for the deletion to succeed. * @throws EntityManagerException * if any error occurs while deleting. */ public void delete(List<?> entities) { try { entityManager.executeEntityListeners(CallbackType.PRE_DELETE, entities); Key[] nativeKeys = new Key[entities.size()]; for (int i = 0; i < entities.size(); i++) { nativeKeys[i] = Marshaller.marshalKey(entityManager, entities.get(i)); } nativeWriter.delete(nativeKeys); entityManager.executeEntityListeners(CallbackType.POST_DELETE, entities); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes the entity with the given ID. The entity is assumed to be a root entity (no parent). * The entity kind will be determined from the supplied entity class. * * @param entityClass * the entity class. * @param id * the ID of the entity. * @throws EntityManagerException * if any error occurs while inserting. */ public <E> void delete(Class<E> entityClass, long id) { try { EntityMetadata entityMetadata = EntityIntrospector.introspect(entityClass); Key nativeKey = entityManager.newNativeKeyFactory().setKind(entityMetadata.getKind()) .newKey(id); nativeWriter.delete(nativeKey); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes the entity with the given ID. The entity is assumed to be a root entity (no parent). * The entity kind will be determined from the supplied entity class. * * @param entityClass * the entity class. * @param id * the ID of the entity. * @throws EntityManagerException * if any error occurs while inserting. */ public <E> void delete(Class<E> entityClass, String id) { try { EntityMetadata entityMetadata = EntityIntrospector.introspect(entityClass); Key nativeKey = entityManager.newNativeKeyFactory().setKind(entityMetadata.getKind()) .newKey(id); nativeWriter.delete(nativeKey); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes the entity with the given ID and parent key. * * @param entityClass * the entity class. * @param parentKey * the parent key * @param id * the ID of the entity. * @throws EntityManagerException * if any error occurs while inserting. */ public <E> void delete(Class<E> entityClass, DatastoreKey parentKey, long id) { try { EntityMetadata entityMetadata = EntityIntrospector.introspect(entityClass); Key nativeKey = Key.newBuilder(parentKey.nativeKey(), entityMetadata.getKind(), id).build(); nativeWriter.delete(nativeKey); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes the entity with the given ID and parent key. * * @param entityClass * the entity class. * @param parentKey * the parent key * @param id * the ID of the entity. * @throws EntityManagerException * if any error occurs while inserting. */ public <E> void delete(Class<E> entityClass, DatastoreKey parentKey, String id) { try { EntityMetadata entityMetadata = EntityIntrospector.introspect(entityClass); Key nativeKey = Key.newBuilder(parentKey.nativeKey(), entityMetadata.getKind(), id).build(); nativeWriter.delete(nativeKey); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes an entity given its key. * * @param key * the entity's key * @throws EntityManagerException * if any error occurs while deleting. */ public void deleteByKey(DatastoreKey key) { try { nativeWriter.delete(key.nativeKey()); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } /** * Deletes the entities having the given keys. * * @param keys * the entities' keys * @throws EntityManagerException * if any error occurs while deleting. */ public void deleteByKey(List<DatastoreKey> keys) { try { Key[] nativeKeys = new Key[keys.size()]; for (int i = 0; i < keys.size(); i++) { nativeKeys[i] = keys.get(i).nativeKey(); } nativeWriter.delete(nativeKeys); } catch (DatastoreException exp) { throw DatastoreUtils.wrap(exp); } } }