2222import java .util .LinkedHashMap ;
2323import java .util .List ;
2424import java .util .Map ;
25- import java .util .function .BiFunction ;
2625import java .util .function .Consumer ;
26+ import java .util .function .Function ;
2727
2828import javax .lang .model .element .Modifier ;
2929
3030import org .apache .commons .logging .Log ;
3131import org .apache .commons .logging .LogFactory ;
3232import org .jspecify .annotations .Nullable ;
33+
3334import org .springframework .aot .generate .ClassNameGenerator ;
3435import org .springframework .aot .generate .Generated ;
3536import org .springframework .data .projection .ProjectionFactory ;
4344import org .springframework .javapoet .JavaFile ;
4445import org .springframework .javapoet .TypeName ;
4546import org .springframework .javapoet .TypeSpec ;
47+ import org .springframework .util .Assert ;
4648
4749/**
4850 * Builder for AOT repository fragments.
@@ -59,8 +61,8 @@ class AotRepositoryBuilder {
5961 private final AotRepositoryFragmentMetadata generationMetadata ;
6062
6163 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 ;
6466
6567 private AotRepositoryBuilder (RepositoryInformation repositoryInformation , ProjectionFactory projectionFactory ) {
6668
@@ -73,7 +75,7 @@ private AotRepositoryBuilder(RepositoryInformation repositoryInformation, Projec
7375 .initializer ("$T.getLog($T.class)" , TypeName .get (LogFactory .class ), this .generationMetadata .getTargetTypeName ())
7476 .build ());
7577
76- this .customizer = (info , builder ) -> {};
78+ this .classCustomizer = (builder ) -> {};
7779 }
7880
7981 public static <M extends QueryMethod > AotRepositoryBuilder forRepository (RepositoryInformation repositoryInformation ,
@@ -89,14 +91,14 @@ public AotRepositoryBuilder withConstructorCustomizer(
8991 }
9092
9193 public AotRepositoryBuilder withQueryMethodContributor (
92- BiFunction <Method , RepositoryInformation , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ) {
94+ Function <Method , @ Nullable MethodContributor <? extends QueryMethod >> methodContributorFunction ) {
9395 this .methodContributorFunction = methodContributorFunction ;
9496 return this ;
9597 }
9698
97- public AotRepositoryBuilder withClassCustomizer (ClassCustomizer classCustomizer ) {
99+ public AotRepositoryBuilder withClassCustomizer (Consumer < AotRepositoryClassBuilder > classCustomizer ) {
98100
99- this .customizer = classCustomizer ;
101+ this .classCustomizer = classCustomizer ;
100102 return this ;
101103 }
102104
@@ -110,7 +112,7 @@ public AotBundle build() {
110112 repositoryInformation .getRepositoryInterface ());
111113
112114 // create the constructor
113- AotRepositoryConstructorBuilder constructorBuilder = new AotRepositoryConstructorBuilder ( repositoryInformation ,
115+ RepositoryConstructorBuilder constructorBuilder = new RepositoryConstructorBuilder (
114116 generationMetadata );
115117 if (constructorCustomizer != null ) {
116118 constructorCustomizer .accept (constructorBuilder );
@@ -144,7 +146,12 @@ public AotBundle build() {
144146 generationMetadata .getFields ().values ().forEach (builder ::addField );
145147
146148 // 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+
148155 JavaFile javaFile = JavaFile .builder (packageName (), builder .build ()).build ();
149156
150157 // TODO: module identifier
@@ -173,8 +180,7 @@ private void contributeMethod(Method method, RepositoryComposition repositoryCom
173180
174181 if (repositoryInformation .isQueryMethod (method ) && methodContributorFunction != null ) {
175182
176- MethodContributor <? extends QueryMethod > contributor = methodContributorFunction .apply (method ,
177- repositoryInformation );
183+ MethodContributor <? extends QueryMethod > contributor = methodContributorFunction .apply (method );
178184
179185 if (contributor != null ) {
180186
@@ -234,20 +240,6 @@ public ProjectionFactory getProjectionFactory() {
234240 return projectionFactory ;
235241 }
236242
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- }
251243
252244 record AotBundle (JavaFile javaFile , JSONObject metadata ) {
253245 }
0 commit comments