22
22
import java .util .LinkedHashMap ;
23
23
import java .util .List ;
24
24
import java .util .Map ;
25
- import java .util .function .BiFunction ;
26
25
import java .util .function .Consumer ;
26
+ import java .util .function .Function ;
27
27
28
28
import javax .lang .model .element .Modifier ;
29
29
30
30
import org .apache .commons .logging .Log ;
31
31
import org .apache .commons .logging .LogFactory ;
32
32
import org .jspecify .annotations .Nullable ;
33
+
33
34
import org .springframework .aot .generate .ClassNameGenerator ;
34
35
import org .springframework .aot .generate .Generated ;
35
36
import org .springframework .data .projection .ProjectionFactory ;
43
44
import org .springframework .javapoet .JavaFile ;
44
45
import org .springframework .javapoet .TypeName ;
45
46
import org .springframework .javapoet .TypeSpec ;
47
+ import org .springframework .util .Assert ;
46
48
47
49
/**
48
50
* Builder for AOT repository fragments.
@@ -59,8 +61,8 @@ class AotRepositoryBuilder {
59
61
private final AotRepositoryFragmentMetadata generationMetadata ;
60
62
61
63
private @ Nullable Consumer <AotRepositoryConstructorBuilder > constructorCustomizer ;
62
- private @ Nullable BiFunction <Method , RepositoryInformation , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ;
63
- private ClassCustomizer customizer ;
64
+ private @ Nullable Function <Method , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ;
65
+ private Consumer < AotRepositoryClassBuilder > classCustomizer ;
64
66
65
67
private AotRepositoryBuilder (RepositoryInformation repositoryInformation , ProjectionFactory projectionFactory ) {
66
68
@@ -73,7 +75,7 @@ private AotRepositoryBuilder(RepositoryInformation repositoryInformation, Projec
73
75
.initializer ("$T.getLog($T.class)" , TypeName .get (LogFactory .class ), this .generationMetadata .getTargetTypeName ())
74
76
.build ());
75
77
76
- this .customizer = (info , builder ) -> {};
78
+ this .classCustomizer = (builder ) -> {};
77
79
}
78
80
79
81
public static <M extends QueryMethod > AotRepositoryBuilder forRepository (RepositoryInformation repositoryInformation ,
@@ -89,14 +91,14 @@ public AotRepositoryBuilder withConstructorCustomizer(
89
91
}
90
92
91
93
public AotRepositoryBuilder withQueryMethodContributor (
92
- BiFunction <Method , RepositoryInformation , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ) {
94
+ Function <Method , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ) {
93
95
this .methodContributorFunction = methodContributorFunction ;
94
96
return this ;
95
97
}
96
98
97
- public AotRepositoryBuilder withClassCustomizer (ClassCustomizer classCustomizer ) {
99
+ public AotRepositoryBuilder withClassCustomizer (Consumer < AotRepositoryClassBuilder > classCustomizer ) {
98
100
99
- this .customizer = classCustomizer ;
101
+ this .classCustomizer = classCustomizer ;
100
102
return this ;
101
103
}
102
104
@@ -110,7 +112,7 @@ public AotBundle build() {
110
112
repositoryInformation .getRepositoryInterface ());
111
113
112
114
// create the constructor
113
- AotRepositoryConstructorBuilder constructorBuilder = new AotRepositoryConstructorBuilder ( repositoryInformation ,
115
+ RepositoryConstructorBuilder constructorBuilder = new RepositoryConstructorBuilder (
114
116
generationMetadata );
115
117
if (constructorCustomizer != null ) {
116
118
constructorCustomizer .accept (constructorBuilder );
@@ -144,7 +146,12 @@ public AotBundle build() {
144
146
generationMetadata .getFields ().values ().forEach (builder ::addField );
145
147
146
148
// finally customize the file itself
147
- this .customizer .customize (repositoryInformation , builder );
149
+ this .classCustomizer .accept (customizer -> {
150
+
151
+ Assert .notNull (customizer , "ClassCustomizer must not be null" );
152
+ customizer .customize (builder );
153
+ });
154
+
148
155
JavaFile javaFile = JavaFile .builder (packageName (), builder .build ()).build ();
149
156
150
157
// TODO: module identifier
@@ -173,8 +180,7 @@ private void contributeMethod(Method method, RepositoryComposition repositoryCom
173
180
174
181
if (repositoryInformation .isQueryMethod (method ) && methodContributorFunction != null ) {
175
182
176
- MethodContributor <? extends QueryMethod > contributor = methodContributorFunction .apply (method ,
177
- repositoryInformation );
183
+ MethodContributor <? extends QueryMethod > contributor = methodContributorFunction .apply (method );
178
184
179
185
if (contributor != null ) {
180
186
@@ -234,20 +240,6 @@ public ProjectionFactory getProjectionFactory() {
234
240
return projectionFactory ;
235
241
}
236
242
237
- /**
238
- * Customizer interface to customize the AOT repository fragment class after it has been defined.
239
- */
240
- public interface ClassCustomizer {
241
-
242
- /**
243
- * Apply customization ot the AOT repository fragment class after it has been defined..
244
- *
245
- * @param information the repository information that is used for the AOT fragment.
246
- * @param builder the class builder to be customized.
247
- */
248
- void customize (RepositoryInformation information , TypeSpec .Builder builder );
249
-
250
- }
251
243
252
244
record AotBundle (JavaFile javaFile , JSONObject metadata ) {
253
245
}
0 commit comments