DATAMONGO-880 - Improved handling of persistence of lazy-loaded DBRefs.

Added LazyLoadingProxy interface that will be implemented by every LazyLoading-proxy that is created by the DefaultDbRefResolver. Clients can now cast those proxies to this interface and call it's methods initialize a proxy explicitly or to get the referenced DBRef if possible.

We now keep a reference to the DBRef that lead to the creation of a LazyLoadingProxy in order to be able to reuse it in case one assigns the proxy to a field that should be a DBRef. This avoids unnecessary conversion.

Previously saving of proxies wasn't possible since the mapping infrastructure did not know how to extract the entity information from the proxy. We now either store the DBRef backed by the proxy directly or we initialize the proxy first and use the result of LazyLoadingProxy.initialize().

Original pull request: #151.

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DbRefResolver.java

@@ -29,10 +29,11 @@ public interface DbRefResolver {
 
 	/**
 	 * @param property will never be {@literal null}.
+	 * @param dbref the {@link DBRef} to resolve.
 	 * @param callback will never be {@literal null}.
 	 * @return
 	 */
-	Object resolveDbRef(MongoPersistentProperty property, DbRefResolverCallback callback);
+	Object resolveDbRef(MongoPersistentProperty property, DBRef dbref, DbRefResolverCallback callback);
 
 	/**
 	 * Creates a {@link DBRef} instance for the given {@link org.springframework.data.mongodb.core.mapping.DBRef}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolver.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -78,13 +78,13 @@ public class DefaultDbRefResolver implements DbRefResolver {
 	 * @see org.springframework.data.mongodb.core.convert.DbRefResolver#resolveDbRef(org.springframework.data.mongodb.core.mapping.MongoPersistentProperty, org.springframework.data.mongodb.core.convert.DbRefResolverCallback)
 	 */
 	@Override
-	public Object resolveDbRef(MongoPersistentProperty property, DbRefResolverCallback callback) {
+	public Object resolveDbRef(MongoPersistentProperty property, DBRef dbref, DbRefResolverCallback callback) {
 
 		Assert.notNull(property, "Property must not be null!");
 		Assert.notNull(callback, "Callback must not be null!");
 
 		if (isLazyDbRef(property)) {
-			return createLazyLoadingProxy(property, callback);
+			return createLazyLoadingProxy(property, dbref, callback);
 		}
 
 		return callback.resolve(property);
@@ -109,10 +109,11 @@ public class DefaultDbRefResolver implements DbRefResolver {
 	 * eventually resolve the value of the property.
 	 * 
 	 * @param property must not be {@literal null}.
+	 * @param dbref
 	 * @param callback must not be {@literal null}.
 	 * @return
 	 */
-	private Object createLazyLoadingProxy(MongoPersistentProperty property, DbRefResolverCallback callback) {
+	private Object createLazyLoadingProxy(MongoPersistentProperty property, DBRef dbref, DbRefResolverCallback callback) {
 
 		ProxyFactory proxyFactory = new ProxyFactory();
 		Class<?> propertyType = property.getType();
@@ -121,7 +122,9 @@ public class DefaultDbRefResolver implements DbRefResolver {
 			proxyFactory.addInterface(type);
 		}
 
-		LazyLoadingInterceptor interceptor = new LazyLoadingInterceptor(property, exceptionTranslator, callback);
+		LazyLoadingInterceptor interceptor = new LazyLoadingInterceptor(property, dbref, exceptionTranslator, callback);
+
+		proxyFactory.addInterface(LazyLoadingProxy.class);
 
 		if (propertyType.isInterface()) {
 			proxyFactory.addInterface(propertyType);
@@ -158,27 +161,42 @@ public class DefaultDbRefResolver implements DbRefResolver {
 	static class LazyLoadingInterceptor implements MethodInterceptor, org.springframework.cglib.proxy.MethodInterceptor,
 			Serializable {
 
+		private static final Method initializeMethod;
+		private static final Method toDBRefMethod;
+
 		private final DbRefResolverCallback callback;
 		private final MongoPersistentProperty property;
 		private final PersistenceExceptionTranslator exceptionTranslator;
 
 		private volatile boolean resolved;
 		private Object result;
+		private DBRef dbref;
+
+		static {
+			try {
+				initializeMethod = LazyLoadingProxy.class.getMethod("initialize");
+				toDBRefMethod = LazyLoadingProxy.class.getMethod("toDBRef");
+			} catch (Exception e) {
+				throw new RuntimeException(e);
+			}
+		}
 
 		/**
 		 * Creates a new {@link LazyLoadingInterceptor} for the given {@link MongoPersistentProperty},
 		 * {@link PersistenceExceptionTranslator} and {@link DbRefResolverCallback}.
 		 * 
 		 * @param property must not be {@literal null}.
+		 * @param dbref
 		 * @param callback must not be {@literal null}.
 		 */
-		public LazyLoadingInterceptor(MongoPersistentProperty property, PersistenceExceptionTranslator exceptionTranslator,
-				DbRefResolverCallback callback) {
+		public LazyLoadingInterceptor(MongoPersistentProperty property, DBRef dbref,
+				PersistenceExceptionTranslator exceptionTranslator, DbRefResolverCallback callback) {
 
 			Assert.notNull(property, "Property must not be null!");
 			Assert.notNull(exceptionTranslator, "Exception translator must not be null!");
 			Assert.notNull(callback, "Callback must not be null!");
 
+			this.dbref = dbref;
 			this.callback = callback;
 			this.exceptionTranslator = exceptionTranslator;
 			this.property = property;
@@ -190,6 +208,15 @@ public class DefaultDbRefResolver implements DbRefResolver {
 		 */
 		@Override
 		public Object invoke(MethodInvocation invocation) throws Throwable {
+
+			if (invocation.getMethod().equals(initializeMethod)) {
+				return ensureResolved();
+			}
+
+			if (invocation.getMethod().equals(toDBRefMethod)) {
+				return this.dbref;
+			}
+
 			return intercept(invocation.getThis(), invocation.getMethod(), invocation.getArguments(), null);
 		}
 
@@ -199,6 +226,15 @@ public class DefaultDbRefResolver implements DbRefResolver {
 		 */
 		@Override
 		public Object intercept(Object obj, Method method, Object[] args, MethodProxy proxy) throws Throwable {
+
+			if (method.equals(initializeMethod)) {
+				return ensureResolved();
+			}
+
+			if (method.equals(toDBRefMethod)) {
+				return this.dbref;
+			}
+
 			return ReflectionUtils.isObjectMethod(method) ? method.invoke(obj, args) : method.invoke(ensureResolved(), args);
 		}
 
@@ -273,6 +309,7 @@ public class DefaultDbRefResolver implements DbRefResolver {
 			Enhancer enhancer = new Enhancer();
 			enhancer.setSuperclass(type);
 			enhancer.setCallbackType(org.springframework.cglib.proxy.MethodInterceptor.class);
+			enhancer.setInterfaces(new Class[] { LazyLoadingProxy.class });
 
 			Factory factory = (Factory) OBJENESIS.newInstance(enhancer.createClass());
 			factory.setCallbacks(new Callback[] { interceptor });

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/LazyLoadingProxy.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2014 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
+ *
+ *      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 org.springframework.data.mongodb.core.convert;
+
+import org.springframework.data.mongodb.core.convert.DefaultDbRefResolver.LazyLoadingInterceptor;
+
+import com.mongodb.DBRef;
+
+/**
+ * Allows direct interaction with the underlying {@link LazyLoadingInterceptor}.
+ * 
+ * @author Thomas Darimont
+ * @since 1.5
+ */
+public interface LazyLoadingProxy {
+
+	/**
+	 * Initializes the proxy and returns the wrapped value.
+	 * 
+	 * @return
+	 * @since 1.5
+	 */
+	Object initialize();
+
+	/**
+	 * Returns the {@link DBRef} represented by this {@link LazyLoadingProxy}, may be null.
+	 * 
+	 * @return
+	 * @since 1.5
+	 */
+	DBRef toDBRef();
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java

@@ -278,7 +278,9 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 
 				MongoPersistentProperty inverseProp = association.getInverse();
 
-				Object obj = dbRefResolver.resolveDbRef(inverseProp, new DbRefResolverCallback() {
+				Object value = dbo.get(inverseProp.getName());
+				DBRef dbref = value instanceof DBRef ? (DBRef) value : null;
+				Object obj = dbRefResolver.resolveDbRef(inverseProp, dbref, new DbRefResolverCallback() {
 
 					@Override
 					public Object resolve(MongoPersistentProperty property) {
@@ -403,6 +405,7 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 				Object propertyObj = wrapper.getProperty(prop, prop.getType(), fieldAccessOnly);
 
 				if (null != propertyObj) {
+
 					if (!conversions.isSimpleType(propertyObj.getClass())) {
 						writePropertyInternal(propertyObj, dbo, prop);
 					} else {
@@ -449,13 +452,31 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		}
 
 		if (prop.isDbReference()) {
-			DBRef dbRefObj = createDBRef(obj, prop);
+
+			DBRef dbRefObj = null;
+
+			/*
+			 * If we already have a LazyLoadingProxy, we use it's cached DBRef value instead of 
+			 * unnecessarily initializing it only to convert it to a DBRef a few instructions later.
+			 */
+			if (obj instanceof LazyLoadingProxy) {
+				dbRefObj = ((LazyLoadingProxy) obj).toDBRef();
+			}
+
+			dbRefObj = dbRefObj != null ? dbRefObj : createDBRef(obj, prop);
 			if (null != dbRefObj) {
 				accessor.put(prop, dbRefObj);
 				return;
 			}
 		}
 
+		/*
+		 * If we have a LazyLoadingProxy we make sure it is initialized first.
+		 */
+		if (obj instanceof LazyLoadingProxy) {
+			obj = ((LazyLoadingProxy) obj).initialize();
+		}
+
 		// Lookup potential custom target type
 		Class<?> basicTargetType = conversions.getCustomWriteTarget(obj.getClass(), null);
 

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java

@@ -2491,6 +2491,42 @@ public class MongoTemplateTests {
 		assertThat(result.get(0).dbRefProperty.field, is(sample.field));
 	}
 
+	/**
+	 * @see DATAMONGO-880
+	 */
+	@Test
+	public void savingAndReassigningLazyLoadingProxies() {
+
+		template.dropCollection(SomeTemplate.class);
+		template.dropCollection(SomeMessage.class);
+		template.dropCollection(SomeContent.class);
+
+		SomeContent content = new SomeContent();
+		content.id = "C1";
+		content.text = "BUBU";
+		template.save(content);
+
+		SomeTemplate tmpl = new SomeTemplate();
+		tmpl.id = "T1";
+		tmpl.content = content; // @DBRef(lazy=true) tmpl.content
+
+		template.save(tmpl);
+
+		SomeTemplate savedTmpl = template.findById(tmpl.id, SomeTemplate.class);
+
+		SomeMessage message = new SomeMessage();
+		message.id = "M1";
+		message.dbrefContent = savedTmpl.content; // @DBRef message.dbrefContent
+		message.normalContent = savedTmpl.content;
+
+		template.save(message);
+
+		SomeMessage savedMessage = template.findById(message.id, SomeMessage.class);
+
+		assertThat(savedMessage.dbrefContent.text, is(content.text));
+		assertThat(savedMessage.normalContent.text, is(content.text));
+	}
+
 	static class DocumentWithDBRefCollection {
 
 		@Id public String id;
@@ -2670,4 +2706,26 @@ public class MongoTemplateTests {
 		@Id String id;
 		EnumValue value;
 	}
+
+	static class SomeTemplate {
+
+		String id;
+		@org.springframework.data.mongodb.core.mapping.DBRef(lazy = true) SomeContent content;
+
+		public SomeContent getContent() {
+			return content;
+		}
+	}
+
+	static class SomeContent {
+
+		String id;
+		String text;
+	}
+
+	static class SomeMessage {
+		String id;
+		@org.springframework.data.mongodb.core.mapping.DBRef SomeContent dbrefContent;
+		SomeContent normalContent;
+	}
 }