From fadca10f62183c992e063f0929dd38c6939fc5c5 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@vmware.com>
Date: Mon, 30 May 2022 11:34:08 +0200
Subject: [PATCH] Support `@DocumentReference` via Querydsl.

Closes #4037
Original pull request: #4069.
---
 .../support/SpringDataMongodbSerializer.java  | 17 ++++++++++++---
 .../SpringDataMongodbSerializerUnitTests.java | 21 +++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializer.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializer.java
index eebd5e952..1d7c41278 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializer.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializer.java
@@ -22,7 +22,6 @@ import java.util.Set;
 import java.util.regex.Pattern;
 
 import org.bson.Document;
-
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
@@ -158,8 +157,20 @@ class SpringDataMongodbSerializer extends MongodbDocumentSerializer {
 
 		MongoPersistentProperty property = getPropertyFor(path);
 
-		return property.isIdProperty() ? asReference(constant.getConstant(), path.getMetadata().getParent())
-				: asReference(constant.getConstant(), path);
+		if (property.isDocumentReference()) {
+			return converter.toDocumentPointer(constant.getConstant(), property).getPointer();
+		}
+
+		if (property.isIdProperty()) {
+
+			MongoPersistentProperty propertyForPotentialDbRef = getPropertyForPotentialDbRef(path);
+			if (propertyForPotentialDbRef != null && propertyForPotentialDbRef.isDocumentReference()) {
+				return converter.toDocumentPointer(constant.getConstant(), propertyForPotentialDbRef).getPointer();
+			}
+			return asReference(constant.getConstant(), path.getMetadata().getParent());
+		}
+
+		return asReference(constant.getConstant(), path);
 	}
 
 	@Nullable
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializerUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializerUnitTests.java
index 49c2d3a79..46fe10e3a 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializerUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/support/SpringDataMongodbSerializerUnitTests.java
@@ -50,6 +50,7 @@ import com.querydsl.core.types.dsl.BooleanOperation;
 import com.querydsl.core.types.dsl.PathBuilder;
 import com.querydsl.core.types.dsl.SimplePath;
 import com.querydsl.core.types.dsl.StringPath;
+import org.springframework.data.mongodb.repository.User;
 
 /**
  * Unit tests for {@link SpringDataMongodbSerializer}.
@@ -225,6 +226,26 @@ public class SpringDataMongodbSerializerUnitTests {
 		assertThat(serializer.handle(testExpression)).isEqualTo(expected);
 	}
 
+	@Test // GH-4037
+	void parsesDocumentReference() {
+
+		User user = new User();
+		user.setId("007");
+		Predicate predicate = QPerson.person.spiritAnimal.eq(user);
+
+		assertThat(serializer.handle(predicate)).isEqualTo(Document.parse("{ 'spiritAnimal' : '007' }"));
+	}
+
+	@Test // GH-4037
+	void parsesDocumentReferenceOnId() {
+
+		User user = new User();
+		user.setId("007");
+		Predicate predicate = QPerson.person.spiritAnimal.id.eq("007");
+
+		assertThat(serializer.handle(predicate)).isEqualTo(Document.parse("{ 'spiritAnimal' : '007' }"));
+	}
+
 	class Address {
 		String id;
 		String street;