From e38b219898eb943ef42bef584f1dab6feaca7b98 Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Wed, 2 Apr 2025 14:11:38 +0200 Subject: [PATCH] Add support for Entity Graphs. See #3830 --- .../jpa/repository/aot/AotEntityGraph.java | 31 ++++++++ .../jpa/repository/aot/JpaCodeBlocks.java | 57 +++++++++++++-- .../aot/JpaRepositoryContributor.java | 72 ++++++++++++++++++- ...RepositoryContributorIntegrationTests.java | 37 +++++++++- .../jpa/repository/aot/UserRepository.java | 10 ++- 5 files changed, 197 insertions(+), 10 deletions(-) create mode 100644 spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotEntityGraph.java diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotEntityGraph.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotEntityGraph.java new file mode 100644 index 000000000..388c041cb --- /dev/null +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/AotEntityGraph.java @@ -0,0 +1,31 @@ +/* + * Copyright 2025 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 + * + * https://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.jpa.repository.aot; + +import java.util.List; + +import org.jspecify.annotations.Nullable; + +import org.springframework.data.jpa.repository.EntityGraph; + +/** + * AOT representation of an resolved entity graph. The graph can be either named or defined by attribute paths in case + * the named entity graph cannot be looked up. + * + * @author Mark Paluch + */ +record AotEntityGraph(@Nullable String name, EntityGraph.EntityGraphType type, List attributePaths) { +} diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java index 75e74a78e..7b906e937 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaCodeBlocks.java @@ -79,6 +79,7 @@ class JpaCodeBlocks { private String queryVariableName = "query"; private @Nullable AotQueries queries; private MergedAnnotation queryHints = MergedAnnotation.missing(); + private @Nullable AotEntityGraph entityGraph; private @Nullable String sqlResultSetMapping; private @Nullable Class queryReturnType; @@ -112,6 +113,11 @@ class JpaCodeBlocks { return this; } + public QueryBlockBuilder entityGraph(@Nullable AotEntityGraph entityGraph) { + this.entityGraph = entityGraph; + return this; + } + public QueryBlockBuilder queryReturnType(@Nullable Class queryReturnType) { this.queryReturnType = queryReturnType; return this; @@ -162,7 +168,7 @@ class JpaCodeBlocks { } builder.add(createQuery(queryVariableName, queryStringNameVariableName, queries.result(), - this.sqlResultSetMapping, this.queryHints, this.queryReturnType)); + this.sqlResultSetMapping, this.queryHints, this.entityGraph, this.queryReturnType)); builder.add(applyLimits(queries.result().isExists())); @@ -173,7 +179,7 @@ class JpaCodeBlocks { boolean queryHints = this.queryHints.isPresent() && this.queryHints.getBoolean("forCounting"); builder.add(createQuery(countQueryVariableName, countQueryStringNameVariableName, queries.count(), null, - queryHints ? this.queryHints : MergedAnnotation.missing(), Long.class)); + queryHints ? this.queryHints : MergedAnnotation.missing(), null, Long.class)); builder.addStatement("return ($T) $L.getSingleResult()", Long.class, countQueryVariableName); // end control flow does not work well with lambdas @@ -190,8 +196,7 @@ class JpaCodeBlocks { builder.beginControlFlow("if ($L.isSorted())", sort); builder.addStatement("$T declaredQuery = $T.$L($L)", DeclaredQuery.class, DeclaredQuery.class, - queries != null && queries.isNative() ? "nativeQuery" : "jpqlQuery", - queryString); + queries != null && queries.isNative() ? "nativeQuery" : "jpqlQuery", queryString); builder.addStatement("$L = rewriteQuery(declaredQuery, $L, $T.class)", queryString, sort, actualReturnType); builder.endControlFlow(); @@ -238,13 +243,17 @@ class JpaCodeBlocks { private CodeBlock createQuery(String queryVariableName, @Nullable String queryStringNameVariableName, AotQuery query, @Nullable String sqlResultSetMapping, MergedAnnotation queryHints, - @Nullable Class queryReturnType) { + @Nullable AotEntityGraph entityGraph, @Nullable Class queryReturnType) { Builder builder = CodeBlock.builder(); builder.add( doCreateQuery(queryVariableName, queryStringNameVariableName, query, sqlResultSetMapping, queryReturnType)); + if (entityGraph != null) { + builder.add(applyEntityGraph(entityGraph, queryVariableName)); + } + if (queryHints.isPresent()) { builder.add(applyHints(queryVariableName, queryHints)); builder.add("\n"); @@ -363,6 +372,43 @@ class JpaCodeBlocks { throw new UnsupportedOperationException("Not supported yet"); } + private CodeBlock applyEntityGraph(AotEntityGraph entityGraph, String queryVariableName) { + + CodeBlock.Builder builder = CodeBlock.builder(); + + if (StringUtils.hasText(entityGraph.name())) { + + builder.addStatement("$T entityGraph = $L.getEntityGraph($S)", jakarta.persistence.EntityGraph.class, + context.fieldNameOf(EntityManager.class), entityGraph.name()); + } else { + + builder.addStatement("$T<$T> entityGraph = $L.createEntityGraph($T.class)", + jakarta.persistence.EntityGraph.class, context.getActualReturnType().getType(), + context.fieldNameOf(EntityManager.class), context.getActualReturnType().getType()); + + for (String attributePath : entityGraph.attributePaths()) { + + String[] pathComponents = StringUtils.delimitedListToStringArray(attributePath, "."); + + StringBuilder chain = new StringBuilder("entityGraph"); + for (int i = 0; i < pathComponents.length; i++) { + + if (i < pathComponents.length - 1) { + chain.append(".addSubgraph($S)"); + } else { + chain.append(".addAttributeNodes($S)"); + } + } + + builder.addStatement(chain.toString(), (Object[]) pathComponents); + } + + builder.addStatement("$L.setHint($S, entityGraph)", queryVariableName, entityGraph.type().getKey()); + } + + return builder.build(); + } + private CodeBlock applyHints(String queryVariableName, MergedAnnotation queryHints) { Builder hintsBuilder = CodeBlock.builder(); @@ -505,5 +551,4 @@ class JpaCodeBlocks { } - } diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java index 1cacad653..732441f3a 100644 --- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java +++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributor.java @@ -15,6 +15,7 @@ */ package org.springframework.data.jpa.repository.aot; +import jakarta.persistence.Entity; import jakarta.persistence.EntityManager; import jakarta.persistence.EntityManagerFactory; import jakarta.persistence.Tuple; @@ -23,16 +24,22 @@ import jakarta.persistence.metamodel.Metamodel; import java.lang.reflect.Method; import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Function; import java.util.function.UnaryOperator; import org.jspecify.annotations.Nullable; +import org.springframework.core.annotation.AnnotatedElementUtils; import org.springframework.core.annotation.MergedAnnotation; import org.springframework.data.jpa.provider.PersistenceProvider; import org.springframework.data.jpa.provider.QueryExtractor; +import org.springframework.data.jpa.repository.EntityGraph; import org.springframework.data.jpa.repository.Modifying; import org.springframework.data.jpa.repository.NativeQuery; import org.springframework.data.jpa.repository.Query; @@ -166,15 +173,17 @@ public class JpaRepositoryContributor extends RepositoryContributor { MergedAnnotation query = context.getAnnotation(Query.class); MergedAnnotation nativeQuery = context.getAnnotation(NativeQuery.class); MergedAnnotation queryHints = context.getAnnotation(QueryHints.class); + MergedAnnotation entityGraph = context.getAnnotation(EntityGraph.class); MergedAnnotation modifying = context.getAnnotation(Modifying.class); body.add(context.codeBlocks().logDebug("invoking [%s]".formatted(context.getMethod().getName()))); AotQueries aotQueries = getQueries(context, query, selector, queryMethod, returnedType); + AotEntityGraph aotEntityGraph = getAotEntityGraph(entityGraph, repositoryInformation, returnedType, queryMethod); body.add(JpaCodeBlocks.queryBuilder(context, queryMethod).filter(aotQueries) .queryReturnType(getQueryReturnType(aotQueries.result(), returnedType, context)).nativeQuery(nativeQuery) - .queryHints(queryHints).build()); + .queryHints(queryHints).entityGraph(aotEntityGraph).build()); body.add( JpaCodeBlocks.executionBuilder(context, queryMethod).modifying(modifying).query(aotQueries.result()).build()); @@ -360,4 +369,65 @@ public class JpaRepositoryContributor extends RepositoryContributor { return result; } + @SuppressWarnings("unchecked") + private @Nullable AotEntityGraph getAotEntityGraph(MergedAnnotation entityGraph, + RepositoryInformation information, ReturnedType returnedType, JpaQueryMethod queryMethod) { + + if (!entityGraph.isPresent()) { + return null; + } + + EntityGraph.EntityGraphType type = entityGraph.getEnum("type", EntityGraph.EntityGraphType.class); + String[] attributePaths = entityGraph.getStringArray("attributePaths"); + Collection entityGraphNames = getEntityGraphNames(entityGraph, information, queryMethod); + List> candidates = Arrays.asList(returnedType.getDomainType(), returnedType.getReturnedType(), + returnedType.getTypeToRead()); + + for (Class candidate : candidates) { + + Map> namedEntityGraphs = emf + .getNamedEntityGraphs(Class.class.cast(candidate)); + + if (namedEntityGraphs.isEmpty()) { + continue; + } + + for (String entityGraphName : entityGraphNames) { + if (namedEntityGraphs.containsKey(entityGraphName)) { + return new AotEntityGraph(entityGraphName, type, Collections.emptyList()); + } + } + } + + if (attributePaths.length > 0) { + return new AotEntityGraph(null, type, Arrays.asList(attributePaths)); + } + + return null; + } + + private Set getEntityGraphNames(MergedAnnotation entityGraph, RepositoryInformation information, + JpaQueryMethod queryMethod) { + + Set entityGraphNames = new LinkedHashSet<>(); + String value = entityGraph.getString("value"); + + if (StringUtils.hasText(value)) { + entityGraphNames.add(value); + } + entityGraphNames.add(queryMethod.getNamedQueryName()); + entityGraphNames.add(getFallbackEntityGraphName(information, queryMethod)); + return entityGraphNames; + } + + private String getFallbackEntityGraphName(RepositoryInformation information, JpaQueryMethod queryMethod) { + + Class domainType = information.getDomainType(); + Entity entity = AnnotatedElementUtils.findMergedAnnotation(domainType, Entity.class); + String entityName = entity != null && StringUtils.hasText(entity.name()) ? entity.name() + : domainType.getSimpleName(); + + return entityName + "." + queryMethod.getName(); + } + } diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java index d6e0edebb..5f609bad6 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/JpaRepositoryContributorIntegrationTests.java @@ -23,6 +23,7 @@ import java.util.List; import java.util.Optional; import java.util.stream.Stream; +import org.hibernate.proxy.HibernateProxy; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -33,6 +34,7 @@ import org.springframework.data.domain.Page; import org.springframework.data.domain.PageRequest; import org.springframework.data.domain.Slice; import org.springframework.data.domain.Sort; +import org.springframework.data.jpa.domain.sample.Role; import org.springframework.data.jpa.domain.sample.User; import org.springframework.test.context.junit.jupiter.SpringJUnitConfig; import org.springframework.transaction.annotation.Transactional; @@ -50,6 +52,7 @@ class JpaRepositoryContributorIntegrationTests { @Autowired UserRepository fragment; @Autowired EntityManager em; User luke, leia, han, chewbacca, yoda, vader, kylo; + Role smuggler, jedi, imperium; @Configuration static class JpaRepositoryContributorConfiguration extends AotFragmentTestConfigurationSupport { @@ -62,17 +65,26 @@ class JpaRepositoryContributorIntegrationTests { void beforeEach() { em.createQuery("DELETE FROM %s".formatted(User.class.getName())).executeUpdate(); + em.createQuery("DELETE FROM %s".formatted(Role.class.getName())).executeUpdate(); + + smuggler = em.merge(new Role("Smuggler")); + jedi = em.merge(new Role("Jedi")); + imperium = em.merge(new Role("Imperium")); luke = new User("Luke", "Skywalker", "luke@jedi.org"); + luke.addRole(jedi); em.persist(luke); leia = new User("Leia", "Organa", "leia@resistance.gov"); em.persist(leia); han = new User("Han", "Solo", "han@smuggler.net"); + han.setManager(luke); em.persist(han); chewbacca = new User("Chewbacca", "n/a", "chewie@smuggler.net"); + chewbacca.setManager(han); + chewbacca.addRole(smuggler); em.persist(chewbacca); yoda = new User("Yoda", "n/a", "yoda@jedi.org"); @@ -83,6 +95,9 @@ class JpaRepositoryContributorIntegrationTests { kylo = new User("Ben", "Solo", "kylo@new-empire.com"); em.persist(kylo); + + em.flush(); + em.clear(); } @Test @@ -388,6 +403,27 @@ class JpaRepositoryContributorIntegrationTests { .withMessageContaining("No enum constant jakarta.persistence.CacheStoreMode.foo"); } + @Test + void shouldApplyNamedEntityGraph() { + + User chewie = fragment.findWithNamedEntityGraphByFirstname("Chewbacca"); + + assertThat(chewie.getManager()).isInstanceOf(HibernateProxy.class); + assertThat(chewie.getRoles()).isNotInstanceOf(HibernateProxy.class); + } + + @Test + void shouldApplyDeclaredEntityGraph() { + + User chewie = fragment.findWithDeclaredEntityGraphByFirstname("Chewbacca"); + + assertThat(chewie.getRoles()).isNotInstanceOf(HibernateProxy.class); + + User han = chewie.getManager(); + assertThat(han.getRoles()).isNotInstanceOf(HibernateProxy.class); + assertThat(han.getManager()).isInstanceOf(HibernateProxy.class); + } + @Test void testDerivedFinderReturningPageOfProjections() { @@ -464,7 +500,6 @@ class JpaRepositoryContributorIntegrationTests { void todo() { - // entity graphs // interface projections // dynamic projections // class type parameter diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java index 9664faaea..1326960db 100644 --- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java +++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/UserRepository.java @@ -27,6 +27,7 @@ import org.springframework.data.domain.Pageable; import org.springframework.data.domain.Slice; import org.springframework.data.domain.Sort; import org.springframework.data.jpa.domain.sample.User; +import org.springframework.data.jpa.repository.EntityGraph; import org.springframework.data.jpa.repository.Modifying; import org.springframework.data.jpa.repository.NativeQuery; import org.springframework.data.jpa.repository.Query; @@ -111,7 +112,6 @@ interface UserRepository extends CrudRepository { @Query("select u from User u where u.lastname like ?1%") Slice findAnnotatedQuerySliceOfUsersByLastname(String lastname, Pageable pageable); - // Value Expressions @Query("select u from #{#entityName} u where u.emailAddress = ?1") @@ -139,7 +139,7 @@ interface UserRepository extends CrudRepository { // native queries @Query(value = "SELECT firstname FROM SD_User ORDER BY UCASE(firstname)", countQuery = "SELECT count(*) FROM SD_User", - nativeQuery = true) + nativeQuery = true) Page findByNativeQueryWithPageable(Pageable pageable); // projections @@ -158,6 +158,12 @@ interface UserRepository extends CrudRepository { @QueryHints(value = { @QueryHint(name = "jakarta.persistence.cache.storeMode", value = "foo") }, forCounting = false) List findHintedByLastname(String lastname); + @EntityGraph(type = EntityGraph.EntityGraphType.FETCH, value = "User.overview") + User findWithNamedEntityGraphByFirstname(String firstname); + + @EntityGraph(type = EntityGraph.EntityGraphType.FETCH, attributePaths = { "roles", "manager.roles" }) + User findWithDeclaredEntityGraphByFirstname(String firstname); + List findByLastnameStartingWithOrderByFirstname(String lastname, Limit limit); List findByLastname(String lastname, Sort sort);