Firestore - optimistic locking

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

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

@@ -22,10 +22,12 @@
 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;
@@ -208,7 +210,7 @@ public <T> Flux<T> saveAll(Publisher<T> instances) {
 				//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 commitWrites(instances, this::createUpdateWrite);
+			return commitWrites(instances, this::createUpdateWrite, true);
 		});
 	}
 
@@ -288,11 +290,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 +305,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 +388,21 @@ 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 = persistentEntity.getUpdateTimeProperty();
+		if (updateTimeProperty != null && 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

@@ -43,7 +43,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 +57,37 @@ 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 = 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 = 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.protobuf.Timestamp type");
+			}
+			updateTimeProperty = property;
+		}
+	}
 }

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

@@ -0,0 +1,34 @@
+/*
+ * 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;
+}