Skip to content

Commit 1b242dd

Browse files
mp911demp911de
committed
Add support for Entity Graphs.
See #3830
1 parent c12bdb7 commit 1b242dd

File tree

5 files changed

+197
-10
lines changed

5 files changed

+197
-10
lines changed

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotEntityGraph.java

+31
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright 2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.jpa.repository.aot;
17+
18+
import java.util.List;
19+
20+
import org.jspecify.annotations.Nullable;
21+
22+
import org.springframework.data.jpa.repository.EntityGraph;
23+
24+
/**
25+
* AOT representation of an resolved entity graph. The graph can be either named or defined by attribute paths in case
26+
* the named entity graph cannot be looked up.
27+
*
28+
* @author Mark Paluch
29+
*/
30+
record AotEntityGraph(@Nullable String name, EntityGraph.EntityGraphType type, List<String> attributePaths) {
31+
}

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java

+51-6
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ static class QueryBlockBuilder {
7979
private String queryVariableName = "query";
8080
private @Nullable AotQueries queries;
8181
private MergedAnnotation<QueryHints> queryHints = MergedAnnotation.missing();
82+
private @Nullable AotEntityGraph entityGraph;
8283
private @Nullable String sqlResultSetMapping;
8384
private @Nullable Class<?> queryReturnType;
8485

@@ -112,6 +113,11 @@ public QueryBlockBuilder queryHints(MergedAnnotation<QueryHints> queryHints) {
112113
return this;
113114
}
114115

116+
public QueryBlockBuilder entityGraph(@Nullable AotEntityGraph entityGraph) {
117+
this.entityGraph = entityGraph;
118+
return this;
119+
}
120+
115121
public QueryBlockBuilder queryReturnType(@Nullable Class<?> queryReturnType) {
116122
this.queryReturnType = queryReturnType;
117123
return this;
@@ -162,7 +168,7 @@ public CodeBlock build() {
162168
}
163169

164170
builder.add(createQuery(queryVariableName, queryStringNameVariableName, queries.result(),
165-
this.sqlResultSetMapping, this.queryHints, this.queryReturnType));
171+
this.sqlResultSetMapping, this.queryHints, this.entityGraph, this.queryReturnType));
166172

167173
builder.add(applyLimits(queries.result().isExists()));
168174

@@ -173,7 +179,7 @@ public CodeBlock build() {
173179
boolean queryHints = this.queryHints.isPresent() && this.queryHints.getBoolean("forCounting");
174180

175181
builder.add(createQuery(countQueryVariableName, countQueryStringNameVariableName, queries.count(), null,
176-
queryHints ? this.queryHints : MergedAnnotation.missing(), Long.class));
182+
queryHints ? this.queryHints : MergedAnnotation.missing(), null, Long.class));
177183
builder.addStatement("return ($T) $L.getSingleResult()", Long.class, countQueryVariableName);
178184

179185
// end control flow does not work well with lambdas
@@ -190,8 +196,7 @@ private CodeBlock applySorting(String sort, String queryString, Class<?> actualR
190196
builder.beginControlFlow("if ($L.isSorted())", sort);
191197

192198
builder.addStatement("$T declaredQuery = $T.$L($L)", DeclaredQuery.class, DeclaredQuery.class,
193-
queries != null && queries.isNative() ? "nativeQuery" : "jpqlQuery",
194-
queryString);
199+
queries != null && queries.isNative() ? "nativeQuery" : "jpqlQuery", queryString);
195200

196201
builder.addStatement("$L = rewriteQuery(declaredQuery, $L, $T.class)", queryString, sort, actualReturnType);
197202
builder.endControlFlow();
@@ -238,13 +243,17 @@ private CodeBlock applyLimits(boolean exists) {
238243

239244
private CodeBlock createQuery(String queryVariableName, @Nullable String queryStringNameVariableName,
240245
AotQuery query, @Nullable String sqlResultSetMapping, MergedAnnotation<QueryHints> queryHints,
241-
@Nullable Class<?> queryReturnType) {
246+
@Nullable AotEntityGraph entityGraph, @Nullable Class<?> queryReturnType) {
242247

243248
Builder builder = CodeBlock.builder();
244249

245250
builder.add(
246251
doCreateQuery(queryVariableName, queryStringNameVariableName, query, sqlResultSetMapping, queryReturnType));
247252

253+
if (entityGraph != null) {
254+
builder.add(applyEntityGraph(entityGraph, queryVariableName));
255+
}
256+
248257
if (queryHints.isPresent()) {
249258
builder.add(applyHints(queryVariableName, queryHints));
250259
builder.add("\n");
@@ -363,6 +372,43 @@ private Object getParameter(ParameterBinding.ParameterOrigin origin) {
363372
throw new UnsupportedOperationException("Not supported yet");
364373
}
365374

375+
private CodeBlock applyEntityGraph(AotEntityGraph entityGraph, String queryVariableName) {
376+
377+
CodeBlock.Builder builder = CodeBlock.builder();
378+
379+
if (StringUtils.hasText(entityGraph.name())) {
380+
381+
builder.addStatement("$T<?> entityGraph = $L.getEntityGraph($S)", jakarta.persistence.EntityGraph.class,
382+
context.fieldNameOf(EntityManager.class), entityGraph.name());
383+
} else {
384+
385+
builder.addStatement("$T<$T> entityGraph = $L.createEntityGraph($T.class)",
386+
jakarta.persistence.EntityGraph.class, context.getActualReturnType().getType(),
387+
context.fieldNameOf(EntityManager.class), context.getActualReturnType().getType());
388+
389+
for (String attributePath : entityGraph.attributePaths()) {
390+
391+
String[] pathComponents = StringUtils.delimitedListToStringArray(attributePath, ".");
392+
393+
StringBuilder chain = new StringBuilder("entityGraph");
394+
for (int i = 0; i < pathComponents.length; i++) {
395+
396+
if (i < pathComponents.length - 1) {
397+
chain.append(".addSubgraph($S)");
398+
} else {
399+
chain.append(".addAttributeNodes($S)");
400+
}
401+
}
402+
403+
builder.addStatement(chain.toString(), (Object[]) pathComponents);
404+
}
405+
406+
builder.addStatement("$L.setHint($S, entityGraph)", queryVariableName, entityGraph.type().getKey());
407+
}
408+
409+
return builder.build();
410+
}
411+
366412
private CodeBlock applyHints(String queryVariableName, MergedAnnotation<QueryHints> queryHints) {
367413

368414
Builder hintsBuilder = CodeBlock.builder();
@@ -505,5 +551,4 @@ public static boolean returnsModifying(Class<?> returnType) {
505551

506552
}
507553

508-
509554
}

‎spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java

+71-1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package org.springframework.data.jpa.repository.aot;
1717

18+
import jakarta.persistence.Entity;
1819
import jakarta.persistence.EntityManager;
1920
import jakarta.persistence.EntityManagerFactory;
2021
import jakarta.persistence.Tuple;
@@ -23,16 +24,22 @@
2324

2425
import java.lang.reflect.Method;
2526
import java.util.Arrays;
27+
import java.util.Collection;
28+
import java.util.Collections;
29+
import java.util.LinkedHashSet;
2630
import java.util.List;
2731
import java.util.Map;
32+
import java.util.Set;
2833
import java.util.function.Function;
2934
import java.util.function.UnaryOperator;
3035

3136
import org.jspecify.annotations.Nullable;
3237

38+
import org.springframework.core.annotation.AnnotatedElementUtils;
3339
import org.springframework.core.annotation.MergedAnnotation;
3440
import org.springframework.data.jpa.provider.PersistenceProvider;
3541
import org.springframework.data.jpa.provider.QueryExtractor;
42+
import org.springframework.data.jpa.repository.EntityGraph;
3643
import org.springframework.data.jpa.repository.Modifying;
3744
import org.springframework.data.jpa.repository.NativeQuery;
3845
import org.springframework.data.jpa.repository.Query;
@@ -166,15 +173,17 @@ protected void customizeConstructor(AotRepositoryConstructorBuilder constructorB
166173
MergedAnnotation<Query> query = context.getAnnotation(Query.class);
167174
MergedAnnotation<NativeQuery> nativeQuery = context.getAnnotation(NativeQuery.class);
168175
MergedAnnotation<QueryHints> queryHints = context.getAnnotation(QueryHints.class);
176+
MergedAnnotation<EntityGraph> entityGraph = context.getAnnotation(EntityGraph.class);
169177
MergedAnnotation<Modifying> modifying = context.getAnnotation(Modifying.class);
170178

171179
body.add(context.codeBlocks().logDebug("invoking [%s]".formatted(context.getMethod().getName())));
172180

173181
AotQueries aotQueries = getQueries(context, query, selector, queryMethod, returnedType);
182+
AotEntityGraph aotEntityGraph = getAotEntityGraph(entityGraph, repositoryInformation, returnedType, queryMethod);
174183

175184
body.add(JpaCodeBlocks.queryBuilder(context, queryMethod).filter(aotQueries)
176185
.queryReturnType(getQueryReturnType(aotQueries.result(), returnedType, context)).nativeQuery(nativeQuery)
177-
.queryHints(queryHints).build());
186+
.queryHints(queryHints).entityGraph(aotEntityGraph).build());
178187

179188
body.add(
180189
JpaCodeBlocks.executionBuilder(context, queryMethod).modifying(modifying).query(aotQueries.result()).build());
@@ -360,4 +369,65 @@ private AotQuery createCountQuery(PartTree partTree, ReturnedType returnedType,
360369
return result;
361370
}
362371

372+
@SuppressWarnings("unchecked")
373+
private @Nullable AotEntityGraph getAotEntityGraph(MergedAnnotation<EntityGraph> entityGraph,
374+
RepositoryInformation information, ReturnedType returnedType, JpaQueryMethod queryMethod) {
375+
376+
if (!entityGraph.isPresent()) {
377+
return null;
378+
}
379+
380+
EntityGraph.EntityGraphType type = entityGraph.getEnum("type", EntityGraph.EntityGraphType.class);
381+
String[] attributePaths = entityGraph.getStringArray("attributePaths");
382+
Collection<String> entityGraphNames = getEntityGraphNames(entityGraph, information, queryMethod);
383+
List<Class<?>> candidates = Arrays.asList(returnedType.getDomainType(), returnedType.getReturnedType(),
384+
returnedType.getTypeToRead());
385+
386+
for (Class<?> candidate : candidates) {
387+
388+
Map<String, jakarta.persistence.EntityGraph<?>> namedEntityGraphs = emf
389+
.getNamedEntityGraphs(Class.class.cast(candidate));
390+
391+
if (namedEntityGraphs.isEmpty()) {
392+
continue;
393+
}
394+
395+
for (String entityGraphName : entityGraphNames) {
396+
if (namedEntityGraphs.containsKey(entityGraphName)) {
397+
return new AotEntityGraph(entityGraphName, type, Collections.emptyList());
398+
}
399+
}
400+
}
401+
402+
if (attributePaths.length > 0) {
403+
return new AotEntityGraph(null, type, Arrays.asList(attributePaths));
404+
}
405+
406+
return null;
407+
}
408+
409+
private Set<String> getEntityGraphNames(MergedAnnotation<EntityGraph> entityGraph, RepositoryInformation information,
410+
JpaQueryMethod queryMethod) {
411+
412+
Set<String> entityGraphNames = new LinkedHashSet<>();
413+
String value = entityGraph.getString("value");
414+
415+
if (StringUtils.hasText(value)) {
416+
entityGraphNames.add(value);
417+
}
418+
entityGraphNames.add(queryMethod.getNamedQueryName());
419+
entityGraphNames.add(getFallbackEntityGraphName(information, queryMethod));
420+
return entityGraphNames;
421+
}
422+
423+
private String getFallbackEntityGraphName(RepositoryInformation information, JpaQueryMethod queryMethod) {
424+
425+
Class<?> domainType = information.getDomainType();
426+
Entity entity = AnnotatedElementUtils.findMergedAnnotation(domainType, Entity.class);
427+
String entityName = entity != null && StringUtils.hasText(entity.name()) ? entity.name()
428+
: domainType.getSimpleName();
429+
430+
return entityName + "." + queryMethod.getName();
431+
}
432+
363433
}

‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java

+36-1
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.Optional;
2424
import java.util.stream.Stream;
2525

26+
import org.hibernate.proxy.HibernateProxy;
2627
import org.junit.jupiter.api.BeforeEach;
2728
import org.junit.jupiter.api.Test;
2829

@@ -33,6 +34,7 @@
3334
import org.springframework.data.domain.PageRequest;
3435
import org.springframework.data.domain.Slice;
3536
import org.springframework.data.domain.Sort;
37+
import org.springframework.data.jpa.domain.sample.Role;
3638
import org.springframework.data.jpa.domain.sample.User;
3739
import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
3840
import org.springframework.transaction.annotation.Transactional;
@@ -50,6 +52,7 @@ class JpaRepositoryContributorIntegrationTests {
5052
@Autowired UserRepository fragment;
5153
@Autowired EntityManager em;
5254
User luke, leia, han, chewbacca, yoda, vader, kylo;
55+
Role smuggler, jedi, imperium;
5356

5457
@Configuration
5558
static class JpaRepositoryContributorConfiguration extends AotFragmentTestConfigurationSupport {
@@ -62,17 +65,26 @@ public JpaRepositoryContributorConfiguration() {
6265
void beforeEach() {
6366

6467
em.createQuery("DELETE FROM %s".formatted(User.class.getName())).executeUpdate();
68+
em.createQuery("DELETE FROM %s".formatted(Role.class.getName())).executeUpdate();
69+
70+
smuggler = em.merge(new Role("Smuggler"));
71+
jedi = em.merge(new Role("Jedi"));
72+
imperium = em.merge(new Role("Imperium"));
6573

6674
luke = new User("Luke", "Skywalker", "[email protected]");
75+
luke.addRole(jedi);
6776
em.persist(luke);
6877

6978
leia = new User("Leia", "Organa", "[email protected]");
7079
em.persist(leia);
7180

7281
han = new User("Han", "Solo", "[email protected]");
82+
han.setManager(luke);
7383
em.persist(han);
7484

7585
chewbacca = new User("Chewbacca", "n/a", "[email protected]");
86+
chewbacca.setManager(han);
87+
chewbacca.addRole(smuggler);
7688
em.persist(chewbacca);
7789

7890
yoda = new User("Yoda", "n/a", "[email protected]");
@@ -83,6 +95,9 @@ void beforeEach() {
8395

8496
kylo = new User("Ben", "Solo", "[email protected]");
8597
em.persist(kylo);
98+
99+
em.flush();
100+
em.clear();
86101
}
87102

88103
@Test
@@ -388,6 +403,27 @@ void shouldApplyQueryHints() {
388403
.withMessageContaining("No enum constant jakarta.persistence.CacheStoreMode.foo");
389404
}
390405

406+
@Test
407+
void shouldApplyNamedEntityGraph() {
408+
409+
User chewie = fragment.findWithNamedEntityGraphByFirstname("Chewbacca");
410+
411+
assertThat(chewie.getManager()).isInstanceOf(HibernateProxy.class);
412+
assertThat(chewie.getRoles()).isNotInstanceOf(HibernateProxy.class);
413+
}
414+
415+
@Test
416+
void shouldApplyDeclaredEntityGraph() {
417+
418+
User chewie = fragment.findWithDeclaredEntityGraphByFirstname("Chewbacca");
419+
420+
assertThat(chewie.getRoles()).isNotInstanceOf(HibernateProxy.class);
421+
422+
User han = chewie.getManager();
423+
assertThat(han.getRoles()).isNotInstanceOf(HibernateProxy.class);
424+
assertThat(han.getManager()).isInstanceOf(HibernateProxy.class);
425+
}
426+
391427
@Test
392428
void testDerivedFinderReturningPageOfProjections() {
393429

@@ -464,7 +500,6 @@ void shouldApplySqlResultSetMapping() {
464500

465501
void todo() {
466502

467-
// entity graphs
468503
// interface projections
469504
// dynamic projections
470505
// class type parameter

‎spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java

+8-2
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import org.springframework.data.domain.Slice;
2828
import org.springframework.data.domain.Sort;
2929
import org.springframework.data.jpa.domain.sample.User;
30+
import org.springframework.data.jpa.repository.EntityGraph;
3031
import org.springframework.data.jpa.repository.Modifying;
3132
import org.springframework.data.jpa.repository.NativeQuery;
3233
import org.springframework.data.jpa.repository.Query;
@@ -111,7 +112,6 @@ interface UserRepository extends CrudRepository<User, Integer> {
111112
@Query("select u from User u where u.lastname like ?1%")
112113
Slice<User> findAnnotatedQuerySliceOfUsersByLastname(String lastname, Pageable pageable);
113114

114-
115115
// Value Expressions
116116

117117
@Query("select u from #{#entityName} u where u.emailAddress = ?1")
@@ -139,7 +139,7 @@ interface UserRepository extends CrudRepository<User, Integer> {
139139
// native queries
140140

141141
@Query(value = "SELECT firstname FROM SD_User ORDER BY UCASE(firstname)", countQuery = "SELECT count(*) FROM SD_User",
142-
nativeQuery = true)
142+
nativeQuery = true)
143143
Page<String> findByNativeQueryWithPageable(Pageable pageable);
144144

145145
// projections
@@ -158,6 +158,12 @@ interface UserRepository extends CrudRepository<User, Integer> {
158158
@QueryHints(value = { @QueryHint(name = "jakarta.persistence.cache.storeMode", value = "foo") }, forCounting = false)
159159
List<User> findHintedByLastname(String lastname);
160160

161+
@EntityGraph(type = EntityGraph.EntityGraphType.FETCH, value = "User.overview")
162+
User findWithNamedEntityGraphByFirstname(String firstname);
163+
164+
@EntityGraph(type = EntityGraph.EntityGraphType.FETCH, attributePaths = { "roles", "manager.roles" })
165+
User findWithDeclaredEntityGraphByFirstname(String firstname);
166+
161167
List<User> findByLastnameStartingWithOrderByFirstname(String lastname, Limit limit);
162168

163169
List<User> findByLastname(String lastname, Sort sort);

0 commit comments

Comments
 (0)