From 72fe382bbafe462957fe46de55731bc671da8dbe Mon Sep 17 00:00:00 2001
From: Thomas Darimont <tdarimont@gopivotal.com>
Date: Thu, 17 Apr 2014 09:44:05 +0200
Subject: [PATCH] DATAMONGO-913 - Improve DBRef handling in for
 LazyLoadingProxies.

We now use the captured DBRef of a given LazyLoadingProxy in MappingMongoConverter.toDBRef(..) in order to avoid a new DBRef creation that would fail for the proxy.

Original pull request: #174.
---
 .../core/convert/MappingMongoConverter.java   |  5 ++
 .../data/mongodb/core/MongoTemplateTests.java | 67 +++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java
index 2fcabf78c..fde98209d 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java
@@ -294,6 +294,11 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 			Assert.isTrue(annotation != null, "The referenced property has to be mapped with @DBRef!");
 		}
 
+		// @see DATAMONGO-913
+		if (object instanceof LazyLoadingProxy) {
+			return ((LazyLoadingProxy) object).toDBRef();
+		}
+
 		return createDBRef(object, referingProperty);
 	}
 
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java
index 12c65b337..01c8f9d91 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java
@@ -186,6 +186,8 @@ public class MongoTemplateTests {
 		template.dropCollection(DocumentWithCollectionOfSimpleType.class);
 		template.dropCollection(DocumentWithMultipleCollections.class);
 		template.dropCollection(DocumentWithDBRefCollection.class);
+		template.dropCollection(SomeContent.class);
+		template.dropCollection(SomeTemplate.class);
 	}
 
 	@Test
@@ -2635,6 +2637,62 @@ public class MongoTemplateTests {
 		assertThat(template.find(query, DoucmentWithNamedIdField.class), contains(two, one));
 	}
 
+	/**
+	 * @see DATAMONGO-913
+	 */
+	@Test
+	public void shouldRetrieveInitializedValueFromDbRefAssociationAfterLoad() {
+
+		SomeContent content = new SomeContent();
+		content.id = "content-1";
+		content.name = "Content 1";
+		content.text = "Some text";
+
+		template.save(content);
+
+		SomeTemplate tmpl = new SomeTemplate();
+		tmpl.id = "template-1";
+		tmpl.content = content;
+
+		template.save(tmpl);
+
+		SomeTemplate result = template.findOne(query(where("content").is(tmpl.getContent())), SomeTemplate.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.getContent(), is(notNullValue()));
+		assertThat(result.getContent().getId(), is(notNullValue()));
+		assertThat(result.getContent().getName(), is(notNullValue()));
+		assertThat(result.getContent().getText(), is(content.getText()));
+	}
+
+	/**
+	 * @see DATAMONGO-913
+	 */
+	@Test
+	public void shouldReuseExistingDBRefInQueryFromDbRefAssociationAfterLoad() {
+
+		SomeContent content = new SomeContent();
+		content.id = "content-1";
+		content.name = "Content 1";
+		content.text = "Some text";
+
+		template.save(content);
+
+		SomeTemplate tmpl = new SomeTemplate();
+		tmpl.id = "template-1";
+		tmpl.content = content;
+
+		template.save(tmpl);
+
+		SomeTemplate result = template.findOne(query(where("content").is(tmpl.getContent())), SomeTemplate.class);
+
+		// Use lazy-loading-proxy in query
+		result = template.findOne(query(where("content").is(result.getContent())), SomeTemplate.class);
+
+		assertNotNull(result.getContent().getName());
+		assertThat(result.getContent().getName(), is(content.getName()));
+	}
+
 	static class DoucmentWithNamedIdField {
 
 		@Id String someIdKey;
@@ -2876,6 +2934,15 @@ public class MongoTemplateTests {
 
 		String id;
 		String text;
+		String name;
+
+		public String getName() {
+			return name;
+		}
+
+		public String getId() {
+			return id;
+		}
 
 		public String getText() {
 			return text;