Support CTE in raw query method matching (#3236)

micronaut-projects · Nov 20, 2024 · 2050ffe · 2050ffe
1 parent 023449a
commit 2050ffe
Show file tree

Hide file tree

Showing 5 changed files with 132 additions and 10 deletions.
diff --git a/...sor/src/main/java/io/micronaut/data/processor/visitors/finders/RawQueryMethodMatcher.java b/...sor/src/main/java/io/micronaut/data/processor/visitors/finders/RawQueryMethodMatcher.java
@@ -61,12 +61,11 @@
  */
 public class RawQueryMethodMatcher implements MethodMatcher {
 
-    private static final String SELECT = "select";
-    private static final String DELETE = "delete";
-    private static final String UPDATE = "update";
-    private static final String INSERT = "insert";
-
+    private static final Pattern UPDATE_PATTERN = Pattern.compile(".*\\bupdate\\b.*");
+    private static final Pattern DELETE_PATTERN = Pattern.compile(".*\\bdelete\\b.*");
+    private static final Pattern INSERT_PATTERN = Pattern.compile(".*\\binsert\\b.*");
     private static final Pattern RETURNING_PATTERN = Pattern.compile(".*\\breturning\\b.*");
+
     private static final Pattern VARIABLE_PATTERN = Pattern.compile("([^:\\\\]*)((?<![:]):([a-zA-Z0-9]+))([^:]*)");
 
     @Override
@@ -173,22 +172,21 @@ private boolean isValidReturnType(ClassElement returnType, DataMethod.OperationT
 
     private DataMethod.OperationType findOperationType(String methodName, String query, boolean readOnly) {
         query = query.trim().toLowerCase(Locale.ENGLISH);
-        if (query.startsWith(SELECT)) {
-            return DataMethod.OperationType.QUERY;
-        } else if (query.startsWith(DELETE)) {
+
+        if (DELETE_PATTERN.matcher(query).find()) {
             if (RETURNING_PATTERN.matcher(query).find()) {
                 return DataMethod.OperationType.DELETE_RETURNING;
             }
             return DataMethod.OperationType.DELETE;
-        } else if (query.startsWith(UPDATE)) {
+        } else if (UPDATE_PATTERN.matcher(query).find()) {
             if (RETURNING_PATTERN.matcher(query).find()) {
                 return DataMethod.OperationType.UPDATE_RETURNING;
             }
             if (DeleteMethodMatcher.METHOD_PATTERN.matcher(methodName.toLowerCase(Locale.ENGLISH)).matches()) {
                 return DataMethod.OperationType.DELETE;
             }
             return DataMethod.OperationType.UPDATE;
-        } else if (query.startsWith(INSERT)) {
+        } else if (INSERT_PATTERN.matcher(query).find()) {
             if (RETURNING_PATTERN.matcher(query).find()) {
                 return DataMethod.OperationType.INSERT_RETURNING;
             }

diff --git a/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildDeleteSpec.groovy b/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildDeleteSpec.groovy
@@ -15,6 +15,7 @@
  */
 package io.micronaut.data.processor.sql
 
+import io.micronaut.data.intercept.DeleteAllInterceptor
 import io.micronaut.data.intercept.DeleteReturningManyInterceptor
 import io.micronaut.data.intercept.DeleteReturningOneInterceptor
 import io.micronaut.data.intercept.annotation.DataMethod
@@ -557,4 +558,33 @@ interface PersonRepository extends CrudRepository<Person, Long> {
         'java.util.List<Long>'                        | DeleteReturningManyInterceptor
         'Long'                                        | DeleteReturningOneInterceptor
     }
+
+    @Unroll
+    void "test build delete with CTE"() {
+        given:
+        def repository = buildRepository('test.PersonRepository', """
+import io.micronaut.data.jdbc.annotation.JdbcRepository;
+import io.micronaut.data.model.query.builder.sql.Dialect;
+import io.micronaut.data.tck.entities.Person;
+
+@JdbcRepository(dialect= Dialect.MYSQL)
+@io.micronaut.context.annotation.Executable
+interface PersonRepository extends CrudRepository<Person, Long> {
+
+    @Query(\"""
+            WITH ids AS (SELECT id FROM person)
+            DELETE FROM person
+            WHERE id = :id
+            \""")
+    void customDelete(Long id);
+
+}
+""")
+        def method = repository.findPossibleMethods("customDelete").findFirst().get()
+        def deleteQuery = getQuery(method)
+
+        expect:
+        deleteQuery.replace('\n', ' ') == "WITH ids AS (SELECT id FROM person) DELETE FROM person WHERE id = :id "
+        method.classValue(DataMethod, "interceptor").get() == DeleteAllInterceptor
+    }
 }
diff --git a/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildInsertSpec.groovy b/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildInsertSpec.groovy
@@ -15,6 +15,7 @@
  */
 package io.micronaut.data.processor.sql
 
+import io.micronaut.data.intercept.SaveEntityInterceptor
 import io.micronaut.data.intercept.annotation.DataMethod
 import io.micronaut.data.model.DataType
 import io.micronaut.data.model.entities.Person
@@ -589,4 +590,33 @@ interface AccountRepository extends CrudRepository<Account, Long> {
             getResultDataType(saveMethod) == DataType.ENTITY
             getOperationType(saveMethod) == DataMethod.OperationType.INSERT
     }
+
+    @Unroll
+    void "test build insert with CTE"() {
+        given:
+        def repository = buildRepository('test.PersonRepository', """
+import io.micronaut.data.jdbc.annotation.JdbcRepository;
+import io.micronaut.data.model.query.builder.sql.Dialect;
+import io.micronaut.data.tck.entities.Person;
+
+@JdbcRepository(dialect= Dialect.MYSQL)
+@io.micronaut.context.annotation.Executable
+interface PersonRepository extends CrudRepository<Person, Long> {
+
+    @Query(\"""
+            WITH ids AS (SELECT id FROM person)
+            INSERT INTO person(name, age, enabled)
+            VALUES (:name, :age, TRUE)
+            \""")
+    void customInsert(Person person);
+
+}
+""")
+        def method = repository.findPossibleMethods("customInsert").findFirst().get()
+        def insertQuery = getQuery(method)
+
+        expect:
+        insertQuery.replace('\n', ' ') == "WITH ids AS (SELECT id FROM person) INSERT INTO person(name, age, enabled) VALUES (:name, :age, TRUE) "
+        method.classValue(DataMethod, "interceptor").get() == SaveEntityInterceptor
+    }
 }
diff --git a/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildQuerySpec.groovy b/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildQuerySpec.groovy
@@ -17,6 +17,8 @@ package io.micronaut.data.processor.sql
 
 import io.micronaut.core.annotation.AnnotationMetadata
 import io.micronaut.data.annotation.Join
+import io.micronaut.data.intercept.FindAllInterceptor
+import io.micronaut.data.intercept.FindOneInterceptor
 import io.micronaut.data.intercept.annotation.DataMethod
 import io.micronaut.data.model.CursoredPageable
 import io.micronaut.data.model.DataType
@@ -2077,4 +2079,37 @@ interface OtherRepository extends GenericRepository<BookEntity, Long> {
         then:
         noExceptionThrown()
     }
+
+    @Unroll
+    void "test build select with CTE for type #type"() {
+        given:
+        def repository = buildRepository('test.PersonRepository', """
+import io.micronaut.data.jdbc.annotation.JdbcRepository;
+import io.micronaut.data.model.query.builder.sql.Dialect;
+import io.micronaut.data.tck.entities.Person;
+
+@JdbcRepository(dialect= Dialect.MYSQL)
+@io.micronaut.context.annotation.Executable
+interface PersonRepository extends CrudRepository<Person, Long> {
+
+    @Query(\"""
+            WITH ids AS (SELECT id FROM person)
+            SELECT * FROM person
+            \""")
+    $type customSelect(Long id);
+
+}
+""")
+        def method = repository.findPossibleMethods("customSelect").findFirst().get()
+        def selectQuery = getQuery(method)
+
+        expect:
+        selectQuery.replace('\n', ' ') == "WITH ids AS (SELECT id FROM person) SELECT * FROM person "
+        method.classValue(DataMethod, "interceptor").get() == interceptor
+
+        where:
+        type                                          | interceptor
+        'java.util.List<Person>'                      | FindAllInterceptor
+        'Person'                                      | FindOneInterceptor
+    }
 }
diff --git a/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildUpdateSpec.groovy b/data-processor/src/test/groovy/io/micronaut/data/processor/sql/BuildUpdateSpec.groovy
@@ -197,6 +197,35 @@ interface PersonRepository extends CrudRepository<Person, Long> {
         'Long'                                        | UpdateReturningOneInterceptor
     }
 
+    @Unroll
+    void "test build update with CTE"() {
+        given:
+        def repository = buildRepository('test.PersonRepository', """
+import io.micronaut.data.jdbc.annotation.JdbcRepository;
+import io.micronaut.data.model.query.builder.sql.Dialect;
+import io.micronaut.data.tck.entities.Person;
+
+@JdbcRepository(dialect= Dialect.MYSQL)
+@io.micronaut.context.annotation.Executable
+interface PersonRepository extends CrudRepository<Person, Long> {
+
+    @Query(\"""
+            WITH ids AS (SELECT id FROM person)
+            UPDATE person SET name = 'test'
+            WHERE id = :id
+            \""")
+    void customUpdate(Long id);
+
+}
+""")
+        def method = repository.findPossibleMethods("customUpdate").findFirst().get()
+        def updateQuery = getQuery(method)
+
+        expect:
+        updateQuery.replace('\n', ' ') == "WITH ids AS (SELECT id FROM person) UPDATE person SET name = 'test' WHERE id = :id "
+        method.classValue(DataMethod, "interceptor").get() == UpdateInterceptor
+    }
+
     @Unroll
     void "test build update with datasource set"() {
         given: