1616package org .springframework .data .repository .aot .generate ;
1717
1818import java .lang .reflect .Method ;
19+ import java .util .ArrayList ;
1920import java .util .Arrays ;
2021import java .util .Comparator ;
22+ import java .util .List ;
2123import java .util .Map ;
2224import java .util .function .BiFunction ;
2325import java .util .function .Consumer ;
2628
2729import org .apache .commons .logging .Log ;
2830import org .apache .commons .logging .LogFactory ;
31+ import org .jspecify .annotations .Nullable ;
2932
3033import org .springframework .aot .generate .ClassNameGenerator ;
3134import org .springframework .aot .generate .Generated ;
3235import org .springframework .data .projection .ProjectionFactory ;
36+ import org .springframework .data .repository .aot .generate .json .JSONException ;
37+ import org .springframework .data .repository .aot .generate .json .JSONObject ;
3338import org .springframework .data .repository .core .RepositoryInformation ;
39+ import org .springframework .data .repository .core .support .RepositoryComposition ;
40+ import org .springframework .data .repository .core .support .RepositoryFragment ;
3441import org .springframework .data .repository .query .QueryMethod ;
3542import org .springframework .javapoet .ClassName ;
3643import org .springframework .javapoet .FieldSpec ;
@@ -50,8 +57,8 @@ class AotRepositoryBuilder {
5057 private final ProjectionFactory projectionFactory ;
5158 private final AotRepositoryFragmentMetadata generationMetadata ;
5259
53- private Consumer <AotRepositoryConstructorBuilder > constructorCustomizer ;
54- private BiFunction <Method , RepositoryInformation , MethodContributor <? extends QueryMethod >> methodContributorFunction ;
60+ private @ Nullable Consumer <AotRepositoryConstructorBuilder > constructorCustomizer ;
61+ private @ Nullable BiFunction <Method , RepositoryInformation , MethodContributor <? extends QueryMethod >> methodContributorFunction ;
5562 private ClassCustomizer customizer ;
5663
5764 private AotRepositoryBuilder (RepositoryInformation repositoryInformation , ProjectionFactory projectionFactory ) {
@@ -92,7 +99,7 @@ public AotRepositoryBuilder withClassCustomizer(ClassCustomizer classCustomizer)
9299 return this ;
93100 }
94101
95- public JavaFile build () {
102+ public AotBundle build () {
96103
97104 // start creating the type
98105 TypeSpec .Builder builder = TypeSpec .classBuilder (this .generationMetadata .getTargetTypeName ()) //
@@ -104,51 +111,91 @@ public JavaFile build() {
104111 // create the constructor
105112 AotRepositoryConstructorBuilder constructorBuilder = new AotRepositoryConstructorBuilder (repositoryInformation ,
106113 generationMetadata );
107- constructorCustomizer .accept (constructorBuilder );
114+ if (constructorCustomizer != null ) {
115+ constructorCustomizer .accept (constructorBuilder );
116+ }
117+
108118 builder .addMethod (constructorBuilder .buildConstructor ());
109119
120+ List <AotRepositoryMethod > methodMetadata = new ArrayList <>();
121+ AotRepositoryMetadata .RepositoryType repositoryType = repositoryInformation .isReactiveRepository ()
122+ ? AotRepositoryMetadata .RepositoryType .REACTIVE
123+ : AotRepositoryMetadata .RepositoryType .IMPERATIVE ;
124+
125+ RepositoryComposition repositoryComposition = repositoryInformation .getRepositoryComposition ();
126+
110127 Arrays .stream (repositoryInformation .getRepositoryInterface ().getMethods ())
111128 .sorted (Comparator .<Method , String > comparing (it -> {
112129 return it .getDeclaringClass ().getName ();
113130 }).thenComparing (Method ::getName ).thenComparing (Method ::getParameterCount ).thenComparing (Method ::toString ))
114131 .forEach (method -> {
132+ contributeMethod (method , repositoryComposition , methodMetadata , builder );
133+ });
115134
116- if (repositoryInformation .isCustomMethod (method )) {
117- // TODO: fragment
118- return ;
119- }
135+ // write fields at the end so we make sure to capture things added by methods
136+ generationMetadata .getFields ().values ().forEach (builder ::addField );
120137
121- if (repositoryInformation .isBaseClassMethod (method )) {
122- // TODO: base
123- return ;
124- }
138+ // finally customize the file itself
139+ this .customizer .customize (repositoryInformation , generationMetadata , builder );
140+ JavaFile javaFile = JavaFile .builder (packageName (), builder .build ()).build ();
125141
126- if (method .isBridge () || method .isDefault () || java .lang .reflect .Modifier .isStatic (method .getModifiers ())) {
127- // TODO: report what we've skipped
128- return ;
129- }
142+ // TODO: module identifier
143+ AotRepositoryMetadata metadata = new AotRepositoryMetadata (repositoryInformation .getRepositoryInterface ().getName (),
144+ "" , repositoryType , methodMetadata );
130145
131- if (repositoryInformation .isQueryMethod (method )) {
146+ try {
147+ return new AotBundle (javaFile , metadata .toJson ());
148+ } catch (JSONException e ) {
149+ throw new IllegalStateException (e );
150+ }
151+ }
132152
133- MethodContributor <? extends QueryMethod > contributor = methodContributorFunction . apply ( method ,
134- repositoryInformation );
153+ private void contributeMethod ( Method method , RepositoryComposition repositoryComposition ,
154+ List < AotRepositoryMethod > methodMetadata , TypeSpec . Builder builder ) {
135155
136- if (contributor != null ) {
156+ if (repositoryInformation . isCustomMethod ( method ) || repositoryInformation . isBaseClassMethod ( method ) ) {
137157
138- AotQueryMethodGenerationContext context = new AotQueryMethodGenerationContext (repositoryInformation ,
139- method , contributor .getQueryMethod (), generationMetadata );
158+ RepositoryFragment <?> fragment = repositoryComposition .findFragment (method );
140159
141- builder .addMethod (contributor .contribute (context ));
142- }
143- }
144- });
160+ if (fragment != null ) {
161+ methodMetadata .add (getFragmentMetadata (method , fragment ));
162+ }
163+ return ;
164+ }
145165
146- // write fields at the end so we make sure to capture things added by methods
147- generationMetadata .getFields ().values ().forEach (builder ::addField );
166+ if (method .isBridge () || method .isDefault () || java .lang .reflect .Modifier .isStatic (method .getModifiers ())) {
167+ return ;
168+ }
148169
149- // finally customize the file itself
150- this .customizer .customize (repositoryInformation , generationMetadata , builder );
151- return JavaFile .builder (packageName (), builder .build ()).build ();
170+ if (repositoryInformation .isQueryMethod (method ) && methodContributorFunction != null ) {
171+
172+ MethodContributor <? extends QueryMethod > contributor = methodContributorFunction .apply (method ,
173+ repositoryInformation );
174+
175+ if (contributor != null ) {
176+
177+ if (contributor .contributesMethodSpec () && !repositoryInformation .isReactiveRepository ()) {
178+
179+ AotQueryMethodGenerationContext context = new AotQueryMethodGenerationContext (repositoryInformation , method ,
180+ contributor .getQueryMethod (), generationMetadata );
181+
182+ builder .addMethod (contributor .contribute (context ));
183+ }
184+
185+ methodMetadata
186+ .add (new AotRepositoryMethod (method .getName (), method .toGenericString (), contributor .getMetadata (), null ));
187+ }
188+ }
189+ }
190+
191+ private AotRepositoryMethod getFragmentMetadata (Method method , RepositoryFragment <?> fragment ) {
192+
193+ String signature = fragment .getSignatureContributor ().getName ();
194+ String implementation = fragment .getImplementation ().map (it -> it .getClass ().getName ()).orElse (null );
195+
196+ AotFragmentTarget fragmentTarget = new AotFragmentTarget (signature , implementation );
197+
198+ return new AotRepositoryMethod (method .getName (), method .toGenericString (), null , fragmentTarget );
152199 }
153200
154201 public AotRepositoryFragmentMetadata getGenerationMetadata () {
@@ -193,5 +240,10 @@ public interface ClassCustomizer {
193240 */
194241 void customize (RepositoryInformation information , AotRepositoryFragmentMetadata metadata ,
195242 TypeSpec .Builder builder );
243+
196244 }
245+
246+ record AotBundle (JavaFile javaFile , JSONObject metadata ) {
247+ }
248+
197249}
0 commit comments