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,159 @@
+/*
+ * 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 lombok.RequiredArgsConstructor;
+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() {
+        ImplementationVisitor implementationVisitor = new ImplementationVisitor();
+        FunctionalVisitor functionalVisitor = new FunctionalVisitor(implementationVisitor);
+        return Preconditions.check(new UsesType<>("org.springframework.data.domain.AuditorAware", true), new JavaIsoVisitor<ExecutionContext>() {
+            @Override
+            public @Nullable J visit(@Nullable Tree tree, ExecutionContext ctx) {
+                tree = implementationVisitor.visit(tree, ctx);
+                tree = functionalVisitor.visit(tree, ctx);
+                return (J) tree;
+            }
+        });
+    }
+
+    private static class ImplementationVisitor extends JavaIsoVisitor<ExecutionContext> {
+        @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;
+        }
+    }
+
+    @RequiredArgsConstructor
+    private static class FunctionalVisitor extends JavaIsoVisitor<ExecutionContext> {
+        private final JavaIsoVisitor<ExecutionContext> implementationVisitor;
+
+        @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) new MemberReferenceToMethodInvocation().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