Migrate AuditorAware to an Optional<> for spring boot 1 to 2 migration

build.gradle.kts

@@ -48,6 +48,7 @@ configurations {
 }
 
 recipeDependencies {
+
     parserClasspath("javax.persistence:javax.persistence-api:2.+")
     parserClasspath("javax.validation:validation-api:2.0.1.Final")
     parserClasspath("org.junit.jupiter:junit-jupiter-api:latest.release")
@@ -83,6 +84,7 @@ recipeDependencies {
     parserClasspath("org.springframework:spring-webmvc:5.+")
 
     parserClasspath("org.springframework.data:spring-data-commons:2.+")
+    parserClasspath("org.springframework.data:spring-data-commons:1.+")
     parserClasspath("org.springframework.data:spring-data-jpa:2.+")
     parserClasspath("org.springframework.data:spring-data-jpa:2.3.+")
 

src/main/java/org/openrewrite/java/spring/data/MigrateAuditorAwareToOptional.java

@@ -0,0 +1,179 @@
+/*
+ * Copyright 2021 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.openrewrite.java.spring.data;
+
+import org.jspecify.annotations.Nullable;
+import org.openrewrite.*;
+import org.openrewrite.java.*;
+import org.openrewrite.java.search.UsesType;
+import org.openrewrite.java.spring.util.MemberReferenceToMethodInvocation;
+import org.openrewrite.java.tree.*;
+
+public class MigrateAuditorAwareToOptional extends Recipe {
+
+    private static final TypeMatcher isAuditorAware = new TypeMatcher("org.springframework.data.domain.AuditorAware", true);
+    private static final MethodMatcher isCurrentAuditor = new MethodMatcher("org.springframework.data.domain.AuditorAware getCurrentAuditor()", true);
+    private static final TypeMatcher isOptional = new TypeMatcher("java.util.Optional");
+
+    @Override
+    public String getDisplayName() {
+        return "Make AuditorAware.getCurrentAuditor return `Optional`";
+    }
+
+    @Override
+    public String getDescription() {
+        return "As of Spring boot 2.0, the `AuditorAware.getCurrentAuditor` method should return an `Optional`. " +
+               "This recipe will update the implementations of this method to return an `Optional` using the `ofNullable`.";
+    }
+
+    @Override
+    public TreeVisitor<?, ExecutionContext> getVisitor() {
+        JavaIsoVisitor<ExecutionContext> implementationVisitor = implementationVisitor();
+        JavaIsoVisitor<ExecutionContext> functionalVisitor = functionalVisitor(implementationVisitor);
+
+        return Preconditions.check(new UsesType<>("org.springframework.data.domain.AuditorAware", true), new TreeVisitor<Tree, ExecutionContext>() {
+
+            @Override
+            public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx, Cursor parent) {
+                if (!(tree instanceof SourceFile)) {
+                    return tree;
+                }
+
+                tree = implementationVisitor.visit(tree, ctx);
+                tree = functionalVisitor.visit(tree, ctx);
+                return tree;
+            }
+        });
+    }
+
+    private JavaIsoVisitor<ExecutionContext> implementationVisitor() {
+        return new JavaIsoVisitor<ExecutionContext>() {
+
+            @Override
+            public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDeclaration, ExecutionContext ctx) {
+                if (classDeclaration.getImplements() == null || classDeclaration.getImplements().stream().noneMatch(typeTree -> isAuditorAware.matches(typeTree.getType()))) {
+                    return classDeclaration;
+                }
+                return super.visitClassDeclaration(classDeclaration, ctx);
+            }
+
+            @Override
+            public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
+                TypeTree returnType = method.getReturnTypeExpression();
+                if (method.getMethodType() == null || !isCurrentAuditor.matches(method.getMethodType()) ||
+                    returnType == null || returnType.getType().toString().matches("java.util.Optional<.*>")) {
+                    return method;
+                }
+                Space space = returnType.getPrefix();
+                returnType = TypeTree.build("java.util.Optional<" + returnType.getType() + ">");
+                J.MethodDeclaration md = super.visitMethodDeclaration(method, ctx).withReturnTypeExpression(returnType.withPrefix(space));
+                doAfterVisit(ShortenFullyQualifiedTypeReferences.modifyOnly(md));
+                maybeAddImport("java.util.Optional");
+                return md;
+            }
+
+            @Override
+            public J.Return visitReturn(J.Return return_, ExecutionContext ctx) {
+                Expression expression = return_.getExpression();
+                if (expression == null) {
+                    return return_;
+                }
+                J.Return altered = JavaTemplate.builder("Optional.ofNullable(#{any()})")
+                        .imports("java.util.Optional")
+                        .build()
+                        .apply(getCursor(), expression.getCoordinates().replace(), expression);
+                if (altered == null) {
+                    return return_;
+                }
+                maybeAddImport("java.util.Optional");
+
+                return altered;
+            }
+        };
+    }
+
+    private JavaIsoVisitor<ExecutionContext> functionalVisitor(JavaIsoVisitor<ExecutionContext> implementationVisitor) {
+        MemberReferenceToMethodInvocation memberRefToInvocation = new MemberReferenceToMethodInvocation();
+        return new JavaIsoVisitor<ExecutionContext>() {
+            @Override
+            public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
+                if (!isAuditorAware.matches(method.getReturnTypeExpression()) || method.getBody() == null || method.getBody().getStatements().size() != 1) {
+                    return method;
+                }
+                Statement statement = method.getBody().getStatements().get(0);
+                if (!(statement instanceof J.Return)) {
+                    return method;
+                }
+
+                return super.visitMethodDeclaration(method, ctx);
+            }
+
+
+            @Override
+            public J.Return visitReturn(J.Return return_, ExecutionContext ctx) {
+
+                Expression expression = return_.getExpression();
+                if (expression instanceof J.MemberReference) {
+                    J.MemberReference memberReference = (J.MemberReference) expression;
+                    JavaType.Method methodType = memberReference.getMethodType();
+                    if (methodType == null || isOptional.matches(methodType.getReturnType())) {
+                        return return_;
+                    }
+
+                    expression = (Expression) memberRefToInvocation.visitNonNull(memberReference, ctx, new Cursor(getCursor(), expression).getParent());
+                }
+                if (expression instanceof J.Lambda) {
+                    J.Lambda lambda = ((J.Lambda) expression);
+                    J body = lambda.getBody();
+                    if (body instanceof J.MethodInvocation
+                            && (((J.MethodInvocation) body).getMethodType() != null && isOptional.matches(((J.MethodInvocation) body).getMethodType().getReturnType()))) {
+                        return return_;
+                    }
+                    if (body instanceof J.Literal || body instanceof J.MethodInvocation) {
+                        body = JavaTemplate.builder("Optional.ofNullable(#{any()})")
+                                .contextSensitive()
+                                .imports("java.util.Optional")
+                                .build()
+                                .apply(new Cursor(getCursor(), lambda), lambda.getCoordinates().replace(), body);
+                        JavaType.Method methodType = ((J.MethodInvocation) body).getMethodType();
+                        if (methodType != null) {
+                            methodType = methodType.withReturnType(JavaType.buildType("java.util.Optional"));
+                        }
+                        body = ((J.MethodInvocation) body).withMethodType(methodType);
+                        maybeAddImport("java.util.Optional");
+                        return return_.withExpression(lambda.withBody(body));
+                    }
+                    return super.visitReturn(return_, ctx);
+                }
+                if (expression instanceof J.MethodInvocation) {
+                    if (((J.MethodInvocation) expression).getMethodType() != null && isOptional.matches(((J.MethodInvocation) expression).getMethodType().getReturnType())) {
+                        return return_;
+                    }
+                    maybeAddImport("java.util.Optional");
+                    return return_.withExpression(JavaTemplate.builder("Optional.ofNullable(#{any()})")
+                            .imports("java.util.Optional")
+                            .build()
+                            .apply(new Cursor(getCursor(), expression), expression.getCoordinates().replace(), expression));
+                } else if (expression instanceof J.NewClass && isAuditorAware.matches(((J.NewClass) expression).getClazz().getType())) {
+                    implementationVisitor.setCursor(new Cursor(getCursor(), expression));
+                    maybeAddImport("java.util.Optional");
+                    return return_.withExpression(implementationVisitor.visitNewClass((J.NewClass) expression, ctx));
+                }
+                return return_;
+            }
+        };
+    }
+}

src/main/java/org/openrewrite/java/spring/util/concurrent/MemberReferenceToMethodInvocation.java → src/main/java/org/openrewrite/java/spring/util/MemberReferenceToMethodInvocation.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.openrewrite.java.spring.util.concurrent;
+package org.openrewrite.java.spring.util;
 
 import lombok.Value;
 import org.openrewrite.ExecutionContext;
@@ -28,7 +28,7 @@
 
 import static java.util.stream.Collectors.joining;
 
-class MemberReferenceToMethodInvocation extends JavaVisitor<ExecutionContext> {
+public class MemberReferenceToMethodInvocation extends JavaVisitor<ExecutionContext> {
     @Override
     public J visitMemberReference(J.MemberReference memberRef, ExecutionContext ctx) {
         J.MemberReference mr = (J.MemberReference) super.visitMemberReference(memberRef, ctx);

src/main/java/org/openrewrite/java/spring/util/concurrent/SuccessFailureCallbackToBiConsumerVisitor.java

@@ -19,6 +19,7 @@
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.JavaTemplate;
 import org.openrewrite.java.MethodMatcher;
+import org.openrewrite.java.spring.util.MemberReferenceToMethodInvocation;
 import org.openrewrite.java.tree.J;
 import org.openrewrite.staticanalysis.RemoveUnneededBlock;
 

src/main/java/org/openrewrite/java/spring/util/package-info.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2021 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.
+ */
+@NullMarked @NonNullFields
+package org.openrewrite.java.spring.util;
+
+import org.jspecify.annotations.NullMarked;
+import org.openrewrite.internal.lang.NonNullFields;

src/main/resources/META-INF/rewrite/spring-boot-20.yml

@@ -82,6 +82,7 @@ recipeList:
       newValue: "off"
   - org.openrewrite.java.spring.boot2.SpringBoot2BestPractices
   - org.openrewrite.apache.commons.lang.UpgradeApacheCommonsLang_2_3
+  - org.openrewrite.java.spring.data.MigrateAuditorAwareToOptional
 ---
 type: specs.openrewrite.org/v1beta/recipe
 name: org.openrewrite.java.spring.boot2.MigrateToWebServerFactoryCustomizer