Firestore - optimistic locking

docs/src/main/asciidoc/firestore.adoc

@@ -297,7 +297,7 @@ Now when you call the methods annotated with `@Transactional` on your service ob
 If an error occurs during the execution of a method annotated with `@Transactional`, the transaction will be rolled back. 
 If no error occurs, the transaction will be committed.
 
-==== Subcollections
+=== Subcollections
 A subcollection is a collection associated with a specific entity.
 Documents in subcollections can contain subcollections as well, allowing you to further nest data. You can nest data up to 100 levels deep.
 
@@ -321,6 +321,31 @@ include::{project-root}/spring-cloud-gcp-data-firestore/src/test/java/com/google
 
 ----
 
+=== Update Time and Optimistic Locking
+Firestore stores update time for every document.
+If you would like to retrieve it, you can add a field of `com.google.cloud.Timestamp` type to your entity and annotate it with `@UpdateTime` annotation.
+
+[source,java,indent=0]
+----
+@UpdateTime
+Timestamp updateTime;
+----
+
+===== Using update time for optimistic locking
+A field annotated with `@UpdateTime` can be used for optimistic locking.
+To enable that, you need to set `version` parameter to `true`:
+
+[source,java,indent=0]
+----
+@UpdateTime(version = true)
+Timestamp updateTime;
+----
+
+When you enable optimistic locking, a precondition will be automatically added to the write request to ensure that the document you are updating was not changed since your last read.
+It uses this field's value as a document version and checks that the version of the document you write is the same as the one you've read.
+
+If the field is empty, a precondition would check that the document with the same id does not exist to ensure you don't overwrite existing documents unintentionally.
+
 === Cloud Firestore Spring Boot Starter
 
 If you prefer using Firestore client only, Spring Cloud GCP provides a convenience starter which automatically configures authentication settings and client objects needed to begin using https://cloud.google.com/firestore/[Google Cloud Firestore] in native mode.

spring-cloud-gcp-autoconfigure/src/main/java/com/google/cloud/spring/autoconfigure/firestore/GcpFirestoreAutoConfiguration.java

@@ -131,8 +131,8 @@ public FirestoreMappingContext firestoreMappingContext() {
 
 		@Bean
 		@ConditionalOnMissingBean
-		public FirestoreClassMapper getClassMapper() {
-			return new FirestoreDefaultClassMapper();
+		public FirestoreClassMapper getClassMapper(FirestoreMappingContext mappingContext) {
+			return new FirestoreDefaultClassMapper(mappingContext);
 		}
 
 		@Bean
@@ -146,9 +146,9 @@ public FirestoreTemplate firestoreTemplate(FirestoreGrpc.FirestoreStub firestore
 		@Bean
 		@ConditionalOnMissingBean
 		public ReactiveFirestoreTransactionManager firestoreTransactionManager(
-				FirestoreGrpc.FirestoreStub firestoreStub) {
+				FirestoreGrpc.FirestoreStub firestoreStub, FirestoreClassMapper classMapper) {
 			return new ReactiveFirestoreTransactionManager(firestoreStub,
-					GcpFirestoreAutoConfiguration.this.firestoreRootPath);
+					GcpFirestoreAutoConfiguration.this.firestoreRootPath, classMapper);
 		}
 
 		@Bean

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/FirestoreTemplate.java

@@ -18,14 +18,17 @@
 
 import java.time.Duration;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
+import com.google.cloud.Timestamp;
 import com.google.cloud.spring.data.firestore.mapping.FirestoreClassMapper;
 import com.google.cloud.spring.data.firestore.mapping.FirestoreMappingContext;
 import com.google.cloud.spring.data.firestore.mapping.FirestorePersistentEntity;
 import com.google.cloud.spring.data.firestore.mapping.FirestorePersistentProperty;
+import com.google.cloud.spring.data.firestore.mapping.UpdateTime;
 import com.google.cloud.spring.data.firestore.transaction.ReactiveFirestoreResourceHolder;
 import com.google.cloud.spring.data.firestore.util.ObservableReactiveUtil;
 import com.google.cloud.spring.data.firestore.util.Util;
@@ -204,11 +207,13 @@ public <T> Flux<T> saveAll(Publisher<T> instances) {
 			if (transactionContext.isPresent()) {
 				ReactiveFirestoreResourceHolder holder = (ReactiveFirestoreResourceHolder) transactionContext.get()
 						.getResources().get(this.firestore);
-				List<Write> writes = holder.getWrites();
 				//In a transaction, all write operations should be sent in the commit request, so we just collect them
-				return Flux.from(instances).doOnNext(t -> writes.add(createUpdateWrite(t)));
+				return Flux.from(instances).doOnNext(t -> {
+					holder.getWrites().add(createUpdateWrite(t));
+					holder.getEntities().add(t);
+				});
 			}
-			return commitWrites(instances, this::createUpdateWrite);
+			return commitWrites(instances, this::createUpdateWrite, true);
 		});
 	}
 
@@ -288,11 +293,12 @@ private Flux<String> deleteDocumentsByName(Flux<String> documentNames) {
 				//In a transaction, all write operations should be sent in the commit request, so we just collect them
 				return Flux.from(documentNames).doOnNext(t -> writes.add(createDeleteWrite(t)));
 			}
-			return commitWrites(documentNames, this::createDeleteWrite);
+			return commitWrites(documentNames, this::createDeleteWrite, false);
 		});
 	}
 
-	private <T> Flux<T> commitWrites(Publisher<T> instances, Function<T, Write> converterToWrite) {
+	private <T> Flux<T> commitWrites(Publisher<T> instances, Function<T, Write> converterToWrite,
+			boolean setUpdateTime) {
 		return Flux.from(instances).bufferTimeout(this.writeBufferSize, this.writeBufferTimeout)
 				.flatMap(batch -> {
 					CommitRequest.Builder builder = CommitRequest.newBuilder()
@@ -302,7 +308,16 @@ private <T> Flux<T> commitWrites(Publisher<T> instances, Function<T, Write> conv
 
 					return ObservableReactiveUtil
 							.<CommitResponse>unaryCall(obs -> this.firestore.commit(builder.build(), obs))
-							.thenMany(Flux.fromIterable(batch));
+							.flatMapMany(
+									response -> {
+										if (setUpdateTime) {
+											for (T entity : batch) {
+												getClassMapper()
+														.setUpdateTime(entity, Timestamp.fromProto(response.getCommitTime()));
+											}
+										}
+										return Flux.fromIterable(batch);
+									});
 				});
 	}
 
@@ -376,6 +391,23 @@ private <T> Write createUpdateWrite(T entity) {
 		}
 		String resourceName = buildResourceName(entity);
 		Document document = getClassMapper().entityToDocument(entity, resourceName);
+		FirestorePersistentEntity<?> persistentEntity =
+				this.mappingContext.getPersistentEntity(entity.getClass());
+		FirestorePersistentProperty updateTimeProperty =
+				Objects.requireNonNull(persistentEntity).getUpdateTimeProperty();
+		if (updateTimeProperty != null
+				&& Objects.requireNonNull(updateTimeProperty.findAnnotation(UpdateTime.class)).version()) {
+			Object version = persistentEntity.getPropertyAccessor(entity).getProperty(updateTimeProperty);
+			if (version != null) {
+				builder.setCurrentDocument(
+						Precondition.newBuilder().setUpdateTime(((Timestamp) version).toProto()).build());
+			}
+			else {
+				//If an entity with an empty update time field is being saved, it must be new.
+				//Otherwise it will overwrite an existing document.
+				builder.setCurrentDocument(Precondition.newBuilder().setExists(false).build());
+			}
+		}
 		return builder.setUpdate(document).build();
 	}
 

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/mapping/FirestoreClassMapper.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spring.data.firestore.mapping;
 
+import com.google.cloud.Timestamp;
 import com.google.firestore.v1.Document;
 import com.google.firestore.v1.Value;
 
@@ -55,4 +56,7 @@ public interface FirestoreClassMapper {
 	 * @return the entity that the Firestore document was converted to
 	 */
 	<T> T documentToEntity(Document document, Class<T> clazz);
+
+	<T> T setUpdateTime(T entity, Timestamp updateTime);
+
 }

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/mapping/FirestoreDefaultClassMapper.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spring.data.firestore.mapping;
 
 import java.util.Map;
+import java.util.Objects;
 
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.DocumentSnapshot;
@@ -43,7 +44,10 @@ public final class FirestoreDefaultClassMapper implements FirestoreClassMapper {
 
 	private static final String NOT_USED_PATH = "/not/used/path";
 
-	public FirestoreDefaultClassMapper() {
+	private FirestoreMappingContext mappingContext;
+
+	public FirestoreDefaultClassMapper(FirestoreMappingContext mappingContext) {
+		this.mappingContext = mappingContext;
 	}
 
 	public <T> Value toFirestoreValue(T sourceValue) {
@@ -54,14 +58,39 @@ public <T> Value toFirestoreValue(T sourceValue) {
 
 	public <T> Document entityToDocument(T entity, String documentResourceName) {
 		DocumentSnapshot documentSnapshot = INTERNAL.snapshotFromObject(NOT_USED_PATH, entity);
-		Map<String, Value> valuesMap = INTERNAL.protoFromSnapshot(documentSnapshot);
 		return Document.newBuilder()
-				.putAllFields(valuesMap)
+				.putAllFields(removeUpdateTimestamp(INTERNAL.protoFromSnapshot(documentSnapshot), entity))
 				.setName(documentResourceName).build();
 	}
 
 	public <T> T documentToEntity(Document document, Class<T> clazz) {
 		DocumentSnapshot documentSnapshot = INTERNAL.snapshotFromProto(Timestamp.now(), document);
-		return documentSnapshot.toObject(clazz);
+		T entity = documentSnapshot.toObject(clazz);
+		return setUpdateTime(entity, documentSnapshot.getUpdateTime());
+	}
+
+	public  <T> T setUpdateTime(T entity, Timestamp updateTime) {
+		FirestorePersistentEntity<?> persistentEntity =
+				this.mappingContext.getPersistentEntity(entity.getClass());
+		FirestorePersistentProperty updateTimeProperty =
+				Objects.requireNonNull(persistentEntity).getUpdateTimeProperty();
+
+		if (updateTimeProperty != null) {
+			persistentEntity.getPropertyAccessor(entity).setProperty(updateTimeProperty, updateTime);
+		}
+
+		return entity;
+	}
+
+	private Map<String, Value> removeUpdateTimestamp(Map<String, Value> valuesMap, Object entity) {
+		FirestorePersistentEntity<?> persistentEntity =
+				this.mappingContext.getPersistentEntity(entity.getClass());
+		FirestorePersistentProperty updateTimeProperty =
+				Objects.requireNonNull(persistentEntity).getUpdateTimeProperty();
+		if (updateTimeProperty != null) {
+			valuesMap.remove(updateTimeProperty.getFieldName());
+		}
+		return valuesMap;
 	}
+
 }

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/mapping/FirestorePersistentEntity.java

@@ -44,4 +44,6 @@ public interface FirestorePersistentEntity<T> extends
 	 * @return the ID property.
 	 */
 	FirestorePersistentProperty getIdPropertyOrFail();
+
+	FirestorePersistentProperty getUpdateTimeProperty();
 }

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/mapping/FirestorePersistentEntityImpl.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spring.data.firestore.mapping;
 
+import com.google.cloud.Timestamp;
 import com.google.cloud.spring.data.firestore.Document;
 import com.google.cloud.spring.data.firestore.FirestoreDataException;
 
@@ -37,6 +38,8 @@ public class FirestorePersistentEntityImpl<T>
 
 	private final String collectionName;
 
+	private FirestorePersistentProperty updateTimeProperty;
+
 	public FirestorePersistentEntityImpl(TypeInformation<T> information) {
 		super(information);
 		this.collectionName = getEntityCollectionName(information);
@@ -62,6 +65,11 @@ public FirestorePersistentProperty getIdPropertyOrFail() {
 		return idProperty;
 	}
 
+	@Override
+	public FirestorePersistentProperty getUpdateTimeProperty() {
+		return updateTimeProperty;
+	}
+
 	private static <T> String getEntityCollectionName(TypeInformation<T> typeInformation) {
 		Document document = AnnotationUtils.findAnnotation(typeInformation.getType(), Document.class);
 		String collectionName = (String) AnnotationUtils.getValue(document, "collectionName");
@@ -74,4 +82,15 @@ private static <T> String getEntityCollectionName(TypeInformation<T> typeInforma
 			return collectionName;
 		}
 	}
+
+	@Override
+	public void addPersistentProperty(FirestorePersistentProperty property) {
+		super.addPersistentProperty(property);
+		if (property.findAnnotation(UpdateTime.class) != null) {
+			if (property.getActualType() != Timestamp.class) {
+				throw new FirestoreDataException("@UpdateTime annotated field should be of com.google.cloud.Timestamp type");
+			}
+			updateTimeProperty = property;
+		}
+	}
 }

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/mapping/UpdateTime.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2020-2020 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.spring.data.firestore.mapping;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Marks a field to be used for update time.
+ *
+ * @author Dmitry Solomakha
+ *
+ * @since 2.0.0
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ ElementType.METHOD, ElementType.FIELD })
+public @interface UpdateTime {
+	boolean version() default false;
+}

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/transaction/ReactiveFirestoreResourceHolder.java

@@ -34,6 +34,8 @@ public class ReactiveFirestoreResourceHolder extends ResourceHolderSupport {
 
 	private List<Write> writes = new ArrayList<>();
 
+	private List<Object> entities = new ArrayList<>();
+
 	public ReactiveFirestoreResourceHolder(ByteString transactionId) {
 		this.transactionId = transactionId;
 	}
@@ -45,4 +47,8 @@ public ByteString getTransactionId() {
 	public List<Write> getWrites() {
 		return this.writes;
 	}
+
+	public List<Object> getEntities() {
+		return entities;
+	}
 }

spring-cloud-gcp-data-firestore/src/main/java/com/google/cloud/spring/data/firestore/transaction/ReactiveFirestoreTransactionManager.java

@@ -16,13 +16,16 @@
 
 package com.google.cloud.spring.data.firestore.transaction;
 
+import com.google.cloud.Timestamp;
+import com.google.cloud.spring.data.firestore.mapping.FirestoreClassMapper;
 import com.google.cloud.spring.data.firestore.util.ObservableReactiveUtil;
 import com.google.cloud.spring.data.firestore.util.Util;
 import com.google.firestore.v1.BeginTransactionRequest;
 import com.google.firestore.v1.BeginTransactionResponse;
 import com.google.firestore.v1.CommitRequest;
 import com.google.firestore.v1.CommitResponse;
 import com.google.firestore.v1.FirestoreGrpc;
+import com.google.firestore.v1.FirestoreGrpc.FirestoreStub;
 import com.google.firestore.v1.RollbackRequest;
 import com.google.firestore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
@@ -50,16 +53,19 @@ public class ReactiveFirestoreTransactionManager extends AbstractReactiveTransac
 
 	private final String databasePath;
 
+	private FirestoreClassMapper classMapper;
+
 	/**
 	 * Constructor for ReactiveFirestoreTransactionManager.
 	 * @param firestore Firestore gRPC stub
 	 * @param parent the parent resource. For example:
 	 *     projects/{project_id}/databases/{database_id}/documents or
-	 *     projects/{project_id}/databases/{database_id}/documents/chatrooms/{chatroom_id}
+	 * @param classMapper Firestore class mapper
 	 */
-	public ReactiveFirestoreTransactionManager(FirestoreGrpc.FirestoreStub firestore, String parent) {
+	public ReactiveFirestoreTransactionManager(FirestoreStub firestore, String parent, FirestoreClassMapper classMapper) {
 		this.firestore = firestore;
 		this.databasePath = Util.extractDatabasePath(parent);
+		this.classMapper = classMapper;
 	}
 
 	@Override
@@ -101,10 +107,17 @@ protected Mono<Void> doCommit(TransactionSynchronizationManager transactionSynch
 
 			return ObservableReactiveUtil
 					.<CommitResponse>unaryCall(obs -> this.firestore.commit(builder.build(), obs))
-					.then();
+					.flatMap((response) -> {
+								for (Object entity : resourceHolder.getEntities()) {
+									this.classMapper.setUpdateTime(entity, Timestamp.fromProto(response.getCommitTime()));
+								}
+								return Mono.empty();
+							}
+					);
 		});
 	}
 
+
 	@Override
 	protected Mono<Void> doRollback(TransactionSynchronizationManager transactionSynchronizationManager,
 			GenericReactiveTransaction genericReactiveTransaction) throws TransactionException {