Proper handling of metadata from CrudMethodMetadataPostProcessor. (#1393)

Retrieve the data before any reactive lambda.
When the method returns, in can immediately be removed from
the currentInvocation, and the previous data restored.

Closes #1392.

Author: mikereiche

src/main/java/org/springframework/data/couchbase/repository/support/CrudMethodMetadataPostProcessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors
+ * Copyright 2012-2022 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.
@@ -168,10 +168,10 @@ public Object invoke(MethodInvocation invocation) throws Throwable {
 				try {
 					return invocation.proceed();
 				} finally {
-					// TransactionSynchronizationManager.unbindResource(method);
+					TransactionSynchronizationManager.unbindResource(method);
 				}
 			} finally {
-				// currentInvocation.set(oldInvocation);
+				currentInvocation.set(oldInvocation);
 			}
 		}
 	}
@@ -209,7 +209,7 @@ private static class DefaultCrudMethodMetadata implements CrudMethodMetadata {
 				return;
 			}
 
-			AnnotatedElement[] annotated = new AnnotatedElement[] { method, method.getDeclaringClass()};
+			AnnotatedElement[] annotated = new AnnotatedElement[] { method, method.getDeclaringClass() };
 			this.scanConsistency = OptionsBuilder.annotation(ScanConsistency.class, "query", QueryScanConsistency.NOT_BOUNDED,
 					annotated);
 			this.scope = OptionsBuilder.annotationString(Scope.class, CollectionIdentifier.DEFAULT_SCOPE, annotated);

src/main/java/org/springframework/data/couchbase/repository/support/SimpleReactiveCouchbaseRepository.java

@@ -68,9 +68,35 @@ public SimpleReactiveCouchbaseRepository(CouchbaseEntityInformation<T, String> e
 		this.operations = operations;
 	}
 
+	@Override
+	public Flux<T> findAll(Sort sort) {
+		return findAll(new Query().with(sort));
+	}
+
 	@SuppressWarnings("unchecked")
 	@Override
 	public <S extends T> Mono<S> save(S entity) {
+		return save(entity, getScope(), getCollection());
+	}
+
+	@Override
+	public <S extends T> Flux<S> saveAll(Iterable<S> entities) {
+		Assert.notNull(entities, "The given Iterable of entities must not be null!");
+		String scope = getScope();
+		String collection = getCollection();
+		return Flux.fromIterable(entities).flatMap(e -> save(e, scope, collection));
+	}
+
+	@Override
+	public <S extends T> Flux<S> saveAll(Publisher<S> entityStream) {
+		Assert.notNull(entityStream, "The given Iterable of entities must not be null!");
+		String scope = getScope();
+		String collection = getCollection();
+		return Flux.from(entityStream).flatMap(e -> save(e, scope, collection));
+	}
+
+	@SuppressWarnings("unchecked")
+	private <S extends T> Mono<S> save(S entity, String scope, String collection) {
 		Assert.notNull(entity, "Entity must not be null!");
 		Mono<S> result;
 		final CouchbasePersistentEntity<?> mapperEntity = operations.getConverter().getMappingContext()
@@ -83,97 +109,100 @@ public <S extends T> Mono<S> save(S entity) {
 
 		if (!versionPresent) { // the entity doesn't have a version property
 			// No version field - no cas
-			result = (Mono<S>) operations.upsertById(getJavaType()).inScope(getScope()).inCollection(getCollection())
-					.one(entity);
+			result = (Mono<S>) operations.upsertById(getJavaType()).inScope(scope).inCollection(collection).one(entity);
 		} else if (existingDocument) { // there is a version property, and it is non-zero
 			// Updating existing document with cas
-			result = (Mono<S>) operations.replaceById(getJavaType()).inScope(getScope()).inCollection(getCollection())
-					.one(entity);
+			result = (Mono<S>) operations.replaceById(getJavaType()).inScope(scope).inCollection(collection).one(entity);
 		} else { // there is a version property, but it's zero or not set.
 			// Creating new document
-			result = (Mono<S>) operations.insertById(getJavaType()).inScope(getScope()).inCollection(getCollection())
-					.one(entity);
+			result = (Mono<S>) operations.insertById(getJavaType()).inScope(scope).inCollection(collection).one(entity);
 		}
 		return result;
 	}
 
 	@Override
-	public Flux<T> findAll(Sort sort) {
-		return findAll(new Query().with(sort));
+	public Mono<T> findById(ID id) {
+		return findById(id, getScope(), getCollection());
 	}
 
 	@Override
-	public <S extends T> Flux<S> saveAll(Iterable<S> entities) {
-		Assert.notNull(entities, "The given Iterable of entities must not be null!");
-		return Flux.fromIterable(entities).flatMap(e -> save(e));
+	public Mono<T> findById(Publisher<ID> publisher) {
+		Assert.notNull(publisher, "The given Publisher must not be null!");
+		String scope = getScope();
+		String collection = getCollection();
+		return Mono.from(publisher).flatMap(id -> findById(id, scope, collection));
 	}
 
+	@SuppressWarnings("unchecked")
 	@Override
-	public <S extends T> Flux<S> saveAll(Publisher<S> entityStream) {
-		Assert.notNull(entityStream, "The given Iterable of entities must not be null!");
-		return Flux.from(entityStream).flatMap(this::save);
+	public Flux<T> findAllById(Iterable<ID> ids) {
+		Assert.notNull(ids, "The given Iterable of ids must not be null!");
+		List<String> convertedIds = Streamable.of(ids).stream().map(Objects::toString).collect(Collectors.toList());
+		return (Flux<T>) operations.findById(getJavaType()).inScope(getScope()).inCollection(getCollection())
+				.all(convertedIds);
 	}
 
 	@Override
-	public Mono<T> findById(ID id) {
-		return operations.findById(getJavaType()).inScope(getScope()).inCollection(getCollection()).one(id.toString());
+	public Flux<T> findAllById(Publisher<ID> entityStream) {
+		Assert.notNull(entityStream, "The given entityStream must not be null!");
+		String scope = getScope();
+		String collection = getCollection();
+		return Flux.from(entityStream).flatMap(id -> findById(id, scope, collection));
 	}
 
-	@Override
-	public Mono<T> findById(Publisher<ID> publisher) {
-		Assert.notNull(publisher, "The given Publisher must not be null!");
-		return Mono.from(publisher).flatMap(this::findById);
+	private Mono<T> findById(ID id, String scope, String collection) {
+		return operations.findById(getJavaType()).inScope(scope).inCollection(collection).one(id.toString());
 	}
 
 	@Override
 	public Mono<Boolean> existsById(ID id) {
 		Assert.notNull(id, "The given id must not be null!");
-		return operations.existsById(getJavaType()).inScope(getScope()).inCollection(getCollection()).one(id.toString());
+		return existsById(id, getScope(), getCollection());
+	}
+
+	private Mono<Boolean> existsById(ID id, String scope, String collection) {
+		Assert.notNull(id, "The given id must not be null!");
+		return operations.existsById(getJavaType()).inScope(scope).inCollection(collection).one(id.toString());
 	}
 
 	@Override
 	public Mono<Boolean> existsById(Publisher<ID> publisher) {
 		Assert.notNull(publisher, "The given Publisher must not be null!");
-		return Mono.from(publisher).flatMap(this::existsById);
+		String scope = getScope();
+		String collection = getCollection();
+		return Mono.from(publisher).flatMap(id -> existsById(id, scope, collection));
 	}
 
 	@Override
 	public Flux<T> findAll() {
 		return findAll(new Query());
 	}
 
-	@SuppressWarnings("unchecked")
-	@Override
-	public Flux<T> findAllById(Iterable<ID> ids) {
-		Assert.notNull(ids, "The given Iterable of ids must not be null!");
-		List<String> convertedIds = Streamable.of(ids).stream().map(Objects::toString).collect(Collectors.toList());
-		return (Flux<T>) operations.findById(getJavaType()).inScope(getScope()).inCollection(getCollection())
-				.all(convertedIds);
-	}
-
 	@Override
-	public Flux<T> findAllById(Publisher<ID> entityStream) {
-		Assert.notNull(entityStream, "The given entityStream must not be null!");
-		return Flux.from(entityStream).flatMap(this::findById);
+	public Mono<Void> deleteById(ID id) {
+		return deleteById(id, getScope(), getCollection());
 	}
 
-	@Override
-	public Mono<Void> deleteById(ID id) {
-		return operations.removeById(getJavaType()).inScope(getScope()).inCollection(getCollection()).one(id.toString())
-				.then();
+	private Mono<Void> deleteById(ID id, String scope, String collection) {
+		return operations.removeById(getJavaType()).inScope(scope).inCollection(collection).one(id.toString()).then();
 	}
 
 	@Override
 	public Mono<Void> deleteById(Publisher<ID> publisher) {
 		Assert.notNull(publisher, "The given id must not be null!");
-		return Mono.from(publisher).flatMap(this::deleteById);
+		String scope = getScope();
+		String collection = getCollection();
+		return Mono.from(publisher).flatMap(e -> deleteById(e, scope, collection));
 	}
 
 	@Override
 	public Mono<Void> delete(T entity) {
+		return delete(entity, getScope(), getCollection());
+	}
+
+	private Mono<Void> delete(T entity, String scope, String collection) {
 		Assert.notNull(entity, "Entity must not be null!");
-		return operations.removeById(getJavaType()).inScope(getScope()).inCollection(getCollection()).one(getId(entity))
-				.then();
+		return operations.removeById(getJavaType()).inScope(scope).inCollection(collection).one(getId(entity)).then();
 	}
 
 	@Override
@@ -191,13 +220,9 @@ public Mono<Void> deleteAll(Iterable<? extends T> entities) {
 	@Override
 	public Mono<Void> deleteAll(Publisher<? extends T> entityStream) {
 		Assert.notNull(entityStream, "The given publisher of entities must not be null!");
-		return Flux.from(entityStream).flatMap(this::delete).single();
-	}
-
-	@Override
-	public Mono<Long> count() {
-		return operations.findByQuery(getJavaType()).withConsistency(buildQueryScanConsistency()).inScope(getScope())
-				.inCollection(getCollection()).count();
+		String scope = getScope();
+		String collection = getCollection();
+		return Flux.from(entityStream).flatMap(e -> delete(e, scope, collection)).single();
 	}
 
 	@Override
@@ -206,6 +231,12 @@ public Mono<Void> deleteAll() {
 				.inCollection(getCollection()).all().then();
 	}
 
+	@Override
+	public Mono<Long> count() {
+		return operations.findByQuery(getJavaType()).withConsistency(buildQueryScanConsistency()).inScope(getScope())
+				.inCollection(getCollection()).count();
+	}
+
 	private Flux<T> findAll(Query query) {
 		return operations.findByQuery(getJavaType()).withConsistency(buildQueryScanConsistency()).inScope(getScope())
 				.inCollection(getCollection()).matching(query).all();