Skip to content

Commit 644537a

Browse files
christophstroblmp911de
christophstrobl
authored and
mp911de
committed
Add support for AOT generated repository implementations and wire fragments in BeanDefinition AOT code.
We now provide infrastructure to generate AOT repository method code that implements Query method behavior. No longer use spring.factories but write some custom bean config code so that one of the properties can provide an instance of the generated repository. Closes #3265
1 parent 9db2c4c commit 644537a

26 files changed

+1535
-11
lines changed

Diff for: ‎src/main/java/org/springframework/data/aot/AotContext.java

+7
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import org.springframework.beans.factory.config.BeanReference;
3131
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
3232
import org.springframework.beans.factory.support.RootBeanDefinition;
33+
import org.springframework.core.SpringProperties;
3334
import org.springframework.data.util.TypeScanner;
3435
import org.springframework.util.Assert;
3536

@@ -49,6 +50,12 @@
4950
*/
5051
public interface AotContext {
5152

53+
String GENERATED_REPOSITORIES_ENABLED = "spring.aot.repositories.enabled";
54+
55+
static boolean aotGeneratedRepositoriesEnabled() {
56+
return SpringProperties.getFlag(GENERATED_REPOSITORIES_ENABLED);
57+
}
58+
5259
/**
5360
* Create an {@link AotContext} backed by the given {@link BeanFactory}.
5461
*

Diff for: ‎src/main/java/org/springframework/data/repository/aot/generate/AotCodeContributor.java

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
/*
2+
* Copyright 2024 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.repository.aot.generate;
17+
18+
import org.springframework.aot.generate.GenerationContext;
19+
20+
/**
21+
* @author Christoph Strobl
22+
*/
23+
public interface AotCodeContributor {
24+
void contribute(GenerationContext generationContext);
25+
}

Diff for: ‎src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryBuilder.java

+169
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
/*
2+
* Copyright 2024 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.repository.aot.generate;
17+
18+
import java.time.YearMonth;
19+
import java.time.ZoneId;
20+
import java.time.temporal.ChronoField;
21+
import java.util.Map;
22+
import java.util.function.Consumer;
23+
import java.util.function.Function;
24+
25+
import javax.lang.model.element.Modifier;
26+
27+
import org.apache.commons.logging.Log;
28+
import org.apache.commons.logging.LogFactory;
29+
import org.springframework.aot.generate.ClassNameGenerator;
30+
import org.springframework.aot.generate.Generated;
31+
import org.springframework.data.repository.CrudRepository;
32+
import org.springframework.data.repository.core.RepositoryInformation;
33+
import org.springframework.javapoet.ClassName;
34+
import org.springframework.javapoet.FieldSpec;
35+
import org.springframework.javapoet.JavaFile;
36+
import org.springframework.javapoet.TypeName;
37+
import org.springframework.javapoet.TypeSpec;
38+
import org.springframework.stereotype.Component;
39+
import org.springframework.util.ReflectionUtils;
40+
41+
/**
42+
* @author Christoph Strobl
43+
*/
44+
public class AotRepositoryBuilder {
45+
46+
private final RepositoryInformation repositoryInformation;
47+
private final AotRepositoryImplementationMetadata generationMetadata;
48+
49+
private Consumer<AotRepositoryConstructorBuilder> constructorBuilderCustomizer;
50+
private Function<AotRepositoryMethodGenerationContext, AotRepositoryMethodBuilder> methodContextFunction;
51+
private RepositoryCustomizer customizer;
52+
53+
public static AotRepositoryBuilder forRepository(RepositoryInformation repositoryInformation) {
54+
return new AotRepositoryBuilder(repositoryInformation);
55+
}
56+
57+
AotRepositoryBuilder(RepositoryInformation repositoryInformation) {
58+
59+
this.repositoryInformation = repositoryInformation;
60+
this.generationMetadata = new AotRepositoryImplementationMetadata(className());
61+
this.generationMetadata.addField(FieldSpec
62+
.builder(TypeName.get(Log.class), "logger", Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
63+
.initializer("$T.getLog($T.class)", TypeName.get(LogFactory.class), this.generationMetadata.getTargetTypeName())
64+
.build());
65+
66+
this.customizer = (info, metadata, builder) -> {};
67+
}
68+
69+
public JavaFile javaFile() {
70+
71+
YearMonth creationDate = YearMonth.now(ZoneId.of("UTC"));
72+
73+
// start creating the type
74+
TypeSpec.Builder builder = TypeSpec.classBuilder(this.generationMetadata.getTargetTypeName()) //
75+
.addModifiers(Modifier.PUBLIC) //
76+
.addAnnotation(Generated.class) //
77+
.addJavadoc("AOT generated repository implementation for {@link $T}.\n",
78+
repositoryInformation.getRepositoryInterface()) //
79+
.addJavadoc("\n") //
80+
.addJavadoc("@since $L/$L\n", creationDate.get(ChronoField.YEAR), creationDate.get(ChronoField.MONTH_OF_YEAR)) //
81+
.addJavadoc("@author $L", "Spring Data"); // TODO: does System.getProperty("user.name") make sense here?
82+
83+
// TODO: we do not need that here
84+
// .addSuperinterface(repositoryInformation.getRepositoryInterface());
85+
86+
// create the constructor
87+
AotRepositoryConstructorBuilder constructorBuilder = new AotRepositoryConstructorBuilder(repositoryInformation,
88+
generationMetadata);
89+
constructorBuilderCustomizer.accept(constructorBuilder);
90+
builder.addMethod(constructorBuilder.buildConstructor());
91+
92+
// write methods
93+
// start with the derived ones
94+
ReflectionUtils.doWithMethods(repositoryInformation.getRepositoryInterface(), method -> {
95+
96+
AotRepositoryMethodGenerationContext context = new AotRepositoryMethodGenerationContext(method,
97+
repositoryInformation, generationMetadata);
98+
AotRepositoryMethodBuilder methodBuilder = methodContextFunction.apply(context);
99+
if (methodBuilder != null) {
100+
builder.addMethod(methodBuilder.buildMethod());
101+
}
102+
103+
}, it -> {
104+
105+
/*
106+
the isBaseClassMethod(it) check seems to have some issues.
107+
need to hard code it here
108+
*/
109+
110+
if (ReflectionUtils.findMethod(CrudRepository.class, it.getName(), it.getParameterTypes()) != null) {
111+
return false;
112+
}
113+
114+
return !repositoryInformation.isBaseClassMethod(it) && !repositoryInformation.isCustomMethod(it)
115+
&& !it.isDefault();
116+
});
117+
118+
// write fields at the end so we make sure to capture things added by methods
119+
generationMetadata.getFields().values().forEach(builder::addField);
120+
121+
// finally customize the file itself
122+
this.customizer.customize(repositoryInformation, generationMetadata, builder);
123+
return JavaFile.builder(packageName(), builder.build()).build();
124+
}
125+
126+
AotRepositoryBuilder withConstructorCustomizer(Consumer<AotRepositoryConstructorBuilder> constuctorBuilder) {
127+
128+
this.constructorBuilderCustomizer = constuctorBuilder;
129+
return this;
130+
}
131+
132+
AotRepositoryBuilder withDerivedMethodFunction(
133+
Function<AotRepositoryMethodGenerationContext, AotRepositoryMethodBuilder> methodContextFunction) {
134+
this.methodContextFunction = methodContextFunction;
135+
return this;
136+
}
137+
138+
AotRepositoryBuilder withFileCustomizer(RepositoryCustomizer repositoryCustomizer) {
139+
140+
this.customizer = repositoryCustomizer;
141+
return this;
142+
}
143+
144+
AotRepositoryImplementationMetadata getGenerationMetadata() {
145+
return generationMetadata;
146+
}
147+
148+
private ClassName className() {
149+
return new ClassNameGenerator(ClassName.get(packageName(), typeName())).generateClassName("Aot", null);
150+
}
151+
152+
private String packageName() {
153+
return repositoryInformation.getRepositoryInterface().getPackageName();
154+
}
155+
156+
private String typeName() {
157+
return "%sImpl".formatted(repositoryInformation.getRepositoryInterface().getSimpleName());
158+
}
159+
160+
Map<String, TypeName> getAutowireFields() {
161+
return generationMetadata.getConstructorArguments();
162+
}
163+
164+
public interface RepositoryCustomizer {
165+
166+
void customize(RepositoryInformation repositoryInformation, AotRepositoryImplementationMetadata metadata,
167+
TypeSpec.Builder builder);
168+
}
169+
}

Diff for: ‎src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryConstructorBuilder.java

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/*
2+
* Copyright 2024 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.repository.aot.generate;
17+
18+
import java.util.List;
19+
import java.util.Map.Entry;
20+
21+
import javax.lang.model.element.Modifier;
22+
23+
import org.springframework.core.ResolvableType;
24+
import org.springframework.data.repository.core.RepositoryInformation;
25+
import org.springframework.javapoet.MethodSpec;
26+
import org.springframework.javapoet.ParameterizedTypeName;
27+
import org.springframework.javapoet.TypeName;
28+
29+
/**
30+
* @author Christoph Strobl
31+
*/
32+
public class AotRepositoryConstructorBuilder {
33+
34+
private final RepositoryInformation repositoryInformation;
35+
private final AotRepositoryImplementationMetadata metadata;
36+
37+
private ConstructorCustomizer customizer = (info, builder) -> {};
38+
39+
AotRepositoryConstructorBuilder(RepositoryInformation repositoryInformation,
40+
AotRepositoryImplementationMetadata metadata) {
41+
42+
this.repositoryInformation = repositoryInformation;
43+
this.metadata = metadata;
44+
}
45+
46+
public void addParameter(String parameterName, Class<?> type) {
47+
48+
ResolvableType resolvableType = ResolvableType.forClass(type);
49+
if (!resolvableType.hasGenerics() || !resolvableType.hasResolvableGenerics()) {
50+
addParameter(parameterName, TypeName.get(type));
51+
return;
52+
}
53+
addParameter(parameterName, ParameterizedTypeName.get(type, resolvableType.resolveGenerics()));
54+
}
55+
56+
public void addParameter(String parameterName, TypeName type) {
57+
58+
this.metadata.addConstructorArgument(parameterName, type);
59+
this.metadata.addField(parameterName, type, Modifier.PRIVATE, Modifier.FINAL);
60+
}
61+
62+
public void customize(ConstructorCustomizer customizer) {
63+
this.customizer = customizer;
64+
}
65+
66+
MethodSpec buildConstructor() {
67+
68+
MethodSpec.Builder builder = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
69+
for (Entry<String, TypeName> parameter : this.metadata.getConstructorArguments().entrySet()) {
70+
builder.addParameter(parameter.getValue(), parameter.getKey()).addStatement("this.$N = $N", parameter.getKey(),
71+
parameter.getKey());
72+
}
73+
customizer.customize(repositoryInformation, builder);
74+
return builder.build();
75+
}
76+
77+
private static TypeName getDefaultStoreRepositoryImplementationType(RepositoryInformation repositoryInformation) {
78+
79+
ResolvableType resolvableType = ResolvableType.forClass(repositoryInformation.getRepositoryBaseClass());
80+
if (resolvableType.hasGenerics()) {
81+
List<Class<?>> generics = List.of();
82+
if (resolvableType.getGenerics().length == 2) { // TODO: Find some other way to resolve generics
83+
generics = List.of(repositoryInformation.getDomainType(), repositoryInformation.getIdType());
84+
}
85+
return ParameterizedTypeName.get(repositoryInformation.getRepositoryBaseClass(), generics.toArray(Class[]::new));
86+
}
87+
return TypeName.get(repositoryInformation.getRepositoryBaseClass());
88+
}
89+
90+
public interface ConstructorCustomizer {
91+
92+
void customize(RepositoryInformation repositoryInformation, MethodSpec.Builder builder);
93+
}
94+
}

0 commit comments

Comments
 (0)