JAVA-2228: Use CodeGeneratorFactory for DAO methods

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/CodeGeneratorFactory.java

@@ -18,6 +18,7 @@
 import com.datastax.oss.driver.api.mapper.annotations.Dao;
 import com.datastax.oss.driver.api.mapper.annotations.Entity;
 import com.datastax.oss.driver.api.mapper.annotations.Mapper;
+import com.datastax.oss.driver.internal.mapper.processor.dao.DaoImplementationSharedCode;
 import com.datastax.oss.driver.internal.mapper.processor.mapper.MapperImplementationSharedCode;
 import java.util.Optional;
 import javax.lang.model.element.ExecutableElement;
@@ -57,6 +58,22 @@ Optional<MethodGenerator> newMapperImplementationMethod(
   /** The builder associated to a {@link Mapper}-annotated interface. */
   CodeGenerator newMapperBuilder(TypeElement interfaceElement);
 
-  /** The implementation of a {@link Dao}-annotated interface. */
-  CodeGenerator newDao(TypeElement interfaceElement);
+  /**
+   * The implementation of a {@link Dao}-annotated interface.
+   *
+   * <p>The default code factory calls {@link #newDaoImplementationMethod(ExecutableElement,
+   * DaoImplementationSharedCode)} for each non-static, non-default method, but this is not a hard
+   * requirement.
+   */
+  CodeGenerator newDaoImplementation(TypeElement interfaceElement);
+
+  /**
+   * A method in the implementation of a {@link Dao}-annotated interface.
+   *
+   * @return empty if the processor can't determine what to generate. This will translate as a
+   *     compile-time error.
+   * @see #newDaoImplementation(TypeElement)
+   */
+  Optional<MethodGenerator> newDaoImplementationMethod(
+      ExecutableElement methodElement, DaoImplementationSharedCode enclosingClass);
 }

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/DefaultCodeGeneratorFactory.java

@@ -16,7 +16,14 @@
 package com.datastax.oss.driver.internal.mapper.processor;
 
 import com.datastax.oss.driver.api.mapper.annotations.DaoFactory;
+import com.datastax.oss.driver.api.mapper.annotations.GetEntity;
+import com.datastax.oss.driver.api.mapper.annotations.Insert;
+import com.datastax.oss.driver.api.mapper.annotations.SetEntity;
+import com.datastax.oss.driver.internal.mapper.processor.dao.DaoGetEntityMethodGenerator;
 import com.datastax.oss.driver.internal.mapper.processor.dao.DaoImplementationGenerator;
+import com.datastax.oss.driver.internal.mapper.processor.dao.DaoImplementationSharedCode;
+import com.datastax.oss.driver.internal.mapper.processor.dao.DaoInsertMethodGenerator;
+import com.datastax.oss.driver.internal.mapper.processor.dao.DaoSetEntityMethodGenerator;
 import com.datastax.oss.driver.internal.mapper.processor.entity.EntityHelperGenerator;
 import com.datastax.oss.driver.internal.mapper.processor.mapper.DaoFactoryMethodGenerator;
 import com.datastax.oss.driver.internal.mapper.processor.mapper.MapperBuilderGenerator;
@@ -66,7 +73,21 @@ public CodeGenerator newMapperBuilder(TypeElement interfaceElement) {
   }
 
   @Override
-  public CodeGenerator newDao(TypeElement interfaceElement) {
+  public CodeGenerator newDaoImplementation(TypeElement interfaceElement) {
     return new DaoImplementationGenerator(interfaceElement, context);
   }
+
+  @Override
+  public Optional<MethodGenerator> newDaoImplementationMethod(
+      ExecutableElement methodElement, DaoImplementationSharedCode enclosingClass) {
+    if (methodElement.getAnnotation(SetEntity.class) != null) {
+      return Optional.of(new DaoSetEntityMethodGenerator(methodElement, enclosingClass, context));
+    } else if (methodElement.getAnnotation(Insert.class) != null) {
+      return Optional.of(new DaoInsertMethodGenerator(methodElement, enclosingClass, context));
+    } else if (methodElement.getAnnotation(GetEntity.class) != null) {
+      return Optional.of(new DaoGetEntityMethodGenerator(methodElement, enclosingClass, context));
+    } else {
+      return Optional.empty();
+    }
+  }
 }

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/MapperProcessor.java

@@ -70,7 +70,7 @@ public boolean process(
     processAnnotatedTypes(
         roundEnvironment, Entity.class, ElementKind.CLASS, generatorFactory::newEntity);
     processAnnotatedTypes(
-        roundEnvironment, Dao.class, ElementKind.INTERFACE, generatorFactory::newDao);
+        roundEnvironment, Dao.class, ElementKind.INTERFACE, generatorFactory::newDaoImplementation);
     processAnnotatedTypes(
         roundEnvironment, Mapper.class, ElementKind.INTERFACE, generatorFactory::newMapper);
     return true;

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoGetEntityMethodGenerator.java

@@ -51,15 +51,15 @@ private enum Transformation {
   }
 
   private final ExecutableElement methodElement;
-  private final DaoImplementationGenerator daoImplementationGenerator;
+  private final DaoImplementationSharedCode daoImplementationGenerator;
   private final ProcessorContext context;
 
   public DaoGetEntityMethodGenerator(
       ExecutableElement methodElement,
-      DaoImplementationGenerator daoImplementationGenerator,
+      DaoImplementationSharedCode enclosingClass,
       ProcessorContext context) {
     this.methodElement = methodElement;
-    this.daoImplementationGenerator = daoImplementationGenerator;
+    this.daoImplementationGenerator = enclosingClass;
     this.context = context;
   }
 

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoImplementationGenerator.java

@@ -16,19 +16,16 @@
 package com.datastax.oss.driver.internal.mapper.processor.dao;
 
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
-import com.datastax.oss.driver.api.core.cql.SimpleStatement;
-import com.datastax.oss.driver.api.mapper.annotations.GetEntity;
-import com.datastax.oss.driver.api.mapper.annotations.Insert;
-import com.datastax.oss.driver.api.mapper.annotations.SetEntity;
 import com.datastax.oss.driver.internal.core.util.concurrent.BlockingOperation;
 import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
 import com.datastax.oss.driver.internal.mapper.DaoBase;
 import com.datastax.oss.driver.internal.mapper.MapperContext;
 import com.datastax.oss.driver.internal.mapper.processor.GeneratedNames;
+import com.datastax.oss.driver.internal.mapper.processor.MethodGenerator;
 import com.datastax.oss.driver.internal.mapper.processor.ProcessorContext;
 import com.datastax.oss.driver.internal.mapper.processor.SingleFileCodeGenerator;
+import com.datastax.oss.driver.internal.mapper.processor.SkipGenerationException;
 import com.datastax.oss.driver.internal.mapper.processor.util.NameIndex;
-import com.datastax.oss.driver.internal.mapper.processor.util.generation.BindableHandlingSharedCode;
 import com.datastax.oss.driver.internal.mapper.processor.util.generation.GenericTypeConstantGenerator;
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.CodeBlock;
@@ -43,6 +40,8 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.function.BiConsumer;
@@ -53,7 +52,7 @@
 import javax.lang.model.element.TypeElement;
 
 public class DaoImplementationGenerator extends SingleFileCodeGenerator
-    implements BindableHandlingSharedCode {
+    implements DaoImplementationSharedCode {
 
   private static final TypeName PREPARED_STATEMENT_STAGE =
       ParameterizedTypeName.get(CompletionStage.class, PreparedStatement.class);
@@ -93,17 +92,8 @@ public String addEntityHelperField(ClassName entityClassName) {
         });
   }
 
-  /**
-   * Requests the generation of a prepared statement in this DAO. It will be initialized in {@code
-   * initAsync}, and then passed to the constructor which will store it in a private field.
-   *
-   * @param methodElement the method that will be using this statement.
-   * @param simpleStatementGenerator a callback that generates code to create a {@link
-   *     SimpleStatement} local variable that will be used to create the statement. The first
-   *     parameter is the method to add to, and the second the name of the local variable.
-   * @return the name of the generated field that will hold the statement.
-   */
-  String addPreparedStatement(
+  @Override
+  public String addPreparedStatement(
       ExecutableElement methodElement,
       BiConsumer<MethodSpec.Builder, String> simpleStatementGenerator) {
     // Prepared statements are not shared between methods, so always generate a new name
@@ -121,20 +111,22 @@ protected String getFileName() {
 
   @Override
   protected JavaFile.Builder getContents() {
-
-    List<MethodSpec.Builder> methods = new ArrayList<>();
+    List<MethodGenerator> methodGenerators = new ArrayList<>();
     for (Element child : interfaceElement.getEnclosedElements()) {
       if (child.getKind() == ElementKind.METHOD) {
         ExecutableElement methodElement = (ExecutableElement) child;
-        if (methodElement.getAnnotation(SetEntity.class) != null) {
-          methods.add(new DaoSetEntityMethodGenerator(methodElement, this, context).generate());
-        } else if (methodElement.getAnnotation(Insert.class) != null) {
-          methods.add(new DaoInsertMethodGenerator(methodElement, this, context).generate());
-        }
-        if (methodElement.getAnnotation(GetEntity.class) != null) {
-          methods.add(new DaoGetEntityMethodGenerator(methodElement, this, context).generate());
+        Set<Modifier> modifiers = methodElement.getModifiers();
+        if (!modifiers.contains(Modifier.STATIC) && !modifiers.contains(Modifier.DEFAULT)) {
+          Optional<MethodGenerator> maybeGenerator =
+              context.getCodeGeneratorFactory().newDaoImplementationMethod(methodElement, this);
+          if (maybeGenerator.isPresent()) {
+            methodGenerators.add(maybeGenerator.get());
+          } else {
+            context
+                .getMessager()
+                .error(methodElement, "Don't know what implementation to generate for this method");
+          }
         }
-        // TODO handle other annotations
       }
     }
 
@@ -145,6 +137,13 @@ protected JavaFile.Builder getContents() {
             .superclass(DaoBase.class)
             .addSuperinterface(ClassName.get(interfaceElement));
 
+    for (MethodGenerator methodGenerator : methodGenerators) {
+      try {
+        classBuilder.addMethod(methodGenerator.generate().build());
+      } catch (SkipGenerationException ignored) {
+      }
+    }
+
     genericTypeConstantGenerator.generate(classBuilder);
 
     MethodSpec.Builder initAsyncBuilder = getInitAsyncContents();
@@ -188,10 +187,6 @@ protected JavaFile.Builder getContents() {
     classBuilder.addMethod(initBuilder.build());
     classBuilder.addMethod(constructorBuilder.build());
 
-    for (MethodSpec.Builder method : methods) {
-      classBuilder.addMethod(method.build());
-    }
-
     return JavaFile.builder(implementationName.packageName(), classBuilder.build());
   }
 

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoImplementationSharedCode.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * 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
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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 com.datastax.oss.driver.internal.mapper.processor.dao;
+
+import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.datastax.oss.driver.api.mapper.annotations.Dao;
+import com.datastax.oss.driver.internal.mapper.processor.util.generation.BindableHandlingSharedCode;
+import com.squareup.javapoet.MethodSpec;
+import java.util.function.BiConsumer;
+import javax.lang.model.element.ExecutableElement;
+
+/**
+ * Exposes callbacks that allow individual method generators for a {@link Dao}-annotated class to
+ * request the generation of class-level fields that they will use.
+ */
+public interface DaoImplementationSharedCode extends BindableHandlingSharedCode {
+
+  /**
+   * Requests the generation of a prepared statement in this DAO. It will be initialized in {@code
+   * initAsync}, and then passed to the constructor which will store it in a private field.
+   *
+   * @param methodElement the method that will be using this statement.
+   * @param simpleStatementGenerator a callback that generates code to create a {@link
+   *     SimpleStatement} local variable that will be used to create the statement. The first
+   *     parameter is the method to add to, and the second the name of the local variable.
+   * @return the name of the generated field that will hold the statement.
+   */
+  String addPreparedStatement(
+      ExecutableElement methodElement,
+      BiConsumer<MethodSpec.Builder, String> simpleStatementGenerator);
+}

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoInsertMethodGenerator.java

@@ -40,12 +40,12 @@
 public class DaoInsertMethodGenerator implements MethodGenerator {
 
   private final ExecutableElement methodElement;
-  private final DaoImplementationGenerator enclosingClass;
+  private final DaoImplementationSharedCode enclosingClass;
   private final ProcessorContext context;
 
   public DaoInsertMethodGenerator(
       ExecutableElement methodElement,
-      DaoImplementationGenerator enclosingClass,
+      DaoImplementationSharedCode enclosingClass,
       ProcessorContext context) {
     this.methodElement = methodElement;
     this.enclosingClass = enclosingClass;
@@ -117,7 +117,6 @@ public MethodSpec.Builder generate() {
       throw new SkipGenerationException();
     }
 
-
     // Generate the method:
     String helperFieldName = enclosingClass.addEntityHelperField(ClassName.get(entityElement));
     String statementName =

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoSetEntityMethodGenerator.java

@@ -36,15 +36,15 @@
 public class DaoSetEntityMethodGenerator implements MethodGenerator {
 
   private final ExecutableElement methodElement;
-  private final DaoImplementationGenerator daoImplementationGenerator;
+  private final DaoImplementationSharedCode enclosingClass;
   private final ProcessorContext context;
 
   public DaoSetEntityMethodGenerator(
       ExecutableElement methodElement,
-      DaoImplementationGenerator daoImplementationGenerator,
+      DaoImplementationSharedCode enclosingClass,
       ProcessorContext context) {
     this.methodElement = methodElement;
-    this.daoImplementationGenerator = daoImplementationGenerator;
+    this.enclosingClass = enclosingClass;
     this.context = context;
   }
 
@@ -117,8 +117,7 @@ public MethodSpec.Builder generate() {
     }
 
     // Generate the method:
-    String helperFieldName =
-        daoImplementationGenerator.addEntityHelperField(ClassName.get(entityElement));
+    String helperFieldName = enclosingClass.addEntityHelperField(ClassName.get(entityElement));
 
     // Forward to the base injector in the helper:
     return GeneratedCodePatterns.override(methodElement)