Improve extensibility of mapper internal classes

olim7t · olim7t · commit 3189a8032938 · 2019-06-25T06:38:10.000-07:00
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/DefaultCodeGeneratorFactory.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/DefaultCodeGeneratorFactory.java
@@ -50,7 +50,7 @@
 
 public class DefaultCodeGeneratorFactory implements CodeGeneratorFactory {
 
-  private final ProcessorContext context;
+  protected final ProcessorContext context;
   private final DaoReturnTypeParser daoReturnTypeParser;
 
   public DefaultCodeGeneratorFactory(ProcessorContext context) {
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/mapper/MapperBuilderGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/mapper/MapperBuilderGenerator.java
@@ -45,6 +45,10 @@ protected ClassName getPrincipalTypeName() {
     return builderName;
   }
 
+  protected Class<? extends CqlSession> getSessionClass() {
+    return CqlSession.class;
+  }
+
   @Override
   protected JavaFile.Builder getContents() {
     TypeSpec.Builder classContents =
@@ -55,14 +59,14 @@ protected JavaFile.Builder getContents() {
             .addJavadoc(
                 "Builds an instance of {@link $T} wrapping a driver {@link $T}.",
                 interfaceElement,
-                CqlSession.class)
+                getSessionClass())
             .addJavadoc(JAVADOC_PARAGRAPH_SEPARATOR)
             .addJavadoc(JAVADOC_GENERATED_WARNING)
             .addModifiers(Modifier.PUBLIC)
             .addMethod(
                 MethodSpec.constructorBuilder()
                     .addModifiers(Modifier.PUBLIC)
-                    .addParameter(CqlSession.class, "session")
+                    .addParameter(getSessionClass(), "session")
                     .addStatement("super(session)")
                     .build())
             .addMethod(
