JAVA-2700: Revisit execution profile parameters in mapper

olim7t · olim7t · commit 057fe5d2d58a · 2020-04-22T09:26:41.000-07:00
This amends JAVA-2633 (5e1e510).
diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/cql/StatementBuilder.java b/core/src/main/java/com/datastax/oss/driver/api/core/cql/StatementBuilder.java
@@ -91,14 +91,19 @@ protected StatementBuilder(StatementT template) {
   @NonNull
   public SelfT setExecutionProfileName(@Nullable String executionProfileName) {
     this.executionProfileName = executionProfileName;
+    if (executionProfileName != null) {
+      this.executionProfile = null;
+    }
     return self;
   }
 
   /** @see Statement#setExecutionProfile(DriverExecutionProfile) */
   @NonNull
   public SelfT setExecutionProfile(@Nullable DriverExecutionProfile executionProfile) {
     this.executionProfile = executionProfile;
-    this.executionProfileName = null;
+    if (executionProfile != null) {
+      this.executionProfileName = null;
+    }
     return self;
   }
 
diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/mapper/ProfileIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/mapper/ProfileIT.java
diff --git a/manual/mapper/mapper/README.md b/manual/mapper/mapper/README.md
@@ -111,6 +111,8 @@ assert dao1 == dao2;
 
 ### DAO parameterization
 
+#### Keyspace and table
+
 The mapper allows you to reuse the same DAO interface for different tables. For example, given the
 following definitions:
 
@@ -159,6 +161,24 @@ ProductDao dao3 = inventoryMapper.productDao("keyspace3", "table3");
 The DAO's keyspace and table can also be injected into custom query strings; see [Query
 methods](../daos/query/).
 
+#### Execution profile
+
+Similarly, a DAO can be parameterized to use a particular [configuration
+profile](../../core/configuration/#execution-profiles):
+
+```java
+@Mapper
+public interface InventoryMapper {
+  @DaoFactory
+  ProductDao productDao(@DaoProfile String profileName);
+
+  @DaoFactory
+  ProductDao productDao(@DaoProfile DriverExecutionProfile profile);
+}
+```
+
+The mapper will call `setExecutionProfileName` / `setExecutionProfile` on every generated statement.
+
 ### Schema validation
 
 The mapper validates entity mappings against the database schema at runtime. This check is performed
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoDeleteMethodGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoDeleteMethodGenerator.java
@@ -223,7 +223,6 @@ public Optional<MethodSpec> generate() {
         statementName);
     populateBuilderWithStatementAttributes(methodBodyBuilder, methodElement);
     populateBuilderWithFunction(methodBodyBuilder, boundStatementFunction);
-    populateBuilderWithProfile(methodBodyBuilder);
 
     int nextParameterIndex = 0;
     if (hasEntityParameter) {
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoInsertMethodGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoInsertMethodGenerator.java
@@ -140,7 +140,6 @@ public Optional<MethodSpec> generate() {
 
     populateBuilderWithStatementAttributes(methodBodyBuilder, methodElement);
     populateBuilderWithFunction(methodBodyBuilder, boundStatementFunction);
-    populateBuilderWithProfile(methodBodyBuilder);
 
     warnIfCqlNamePresent(parameters.subList(0, 1));
     String entityParameterName = parameters.get(0).getSimpleName().toString();
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoMethodGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoMethodGenerator.java
@@ -146,13 +146,6 @@ protected void populateBuilderWithFunction(
     }
   }
 
-  protected void populateBuilderWithProfile(CodeBlock.Builder builder) {
-    builder.beginControlFlow("if(context.getExecutionProfileName() != null)");
-    builder.addStatement(
-        "boundStatementBuilder = boundStatementBuilder.setExecutionProfileName(context.getExecutionProfileName())");
-    builder.endControlFlow();
-  }
-
   protected void populateBuilderWithStatementAttributes(
       CodeBlock.Builder builder, ExecutableElement methodElement) {
     StatementAttributes statementAttributes =
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoQueryMethodGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoQueryMethodGenerator.java
@@ -109,7 +109,6 @@ public Optional<MethodSpec> generate() {
 
     populateBuilderWithStatementAttributes(methodBodyBuilder, methodElement);
     populateBuilderWithFunction(methodBodyBuilder, boundStatementFunction);
-    populateBuilderWithProfile(methodBodyBuilder);
 
     if (validateCqlNamesPresent(parameters)) {
       GeneratedCodePatterns.bindParameters(
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoSelectMethodGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoSelectMethodGenerator.java
@@ -161,7 +161,6 @@ public Optional<MethodSpec> generate() {
         statementName);
     populateBuilderWithStatementAttributes(methodBodyBuilder, methodElement);
     populateBuilderWithFunction(methodBodyBuilder, boundStatementFunction);
-    populateBuilderWithProfile(methodBodyBuilder);
 
     if (!primaryKeyParameters.isEmpty()) {
       List<CodeBlock> primaryKeyNames =
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoUpdateMethodGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoUpdateMethodGenerator.java
@@ -126,7 +126,6 @@ public Optional<MethodSpec> generate() {
 
     populateBuilderWithStatementAttributes(methodBodyBuilder, methodElement);
     populateBuilderWithFunction(methodBodyBuilder, boundStatementFunction);
-    populateBuilderWithProfile(methodBodyBuilder);
 
     String entityParameterName = parameters.get(0).getSimpleName().toString();
 
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
@@ -75,7 +75,8 @@ protected JavaFile.Builder getContents() {
                     .addAnnotation(Override.class)
                     .returns(ClassName.get(interfaceElement))
                     .addStatement(
-                        "$1T context = new $1T(session, defaultKeyspaceId, customState)",
+                        "$1T context = new $1T(session, defaultKeyspaceId, "
+                            + "defaultExecutionProfileName, defaultExecutionProfile, customState)",
                         DefaultMapperContext.class)
                     .addStatement(
                         "return new $T(context)",
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/mapper/MapperDaoFactoryMethodGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/mapper/MapperDaoFactoryMethodGenerator.java
@@ -105,8 +105,8 @@ public Optional<MethodSpec> generate() {
     // Validate the arguments
     String keyspaceArgumentName = null;
     String tableArgumentName = null;
-    String executionProfileName = null;
-    boolean executionProfileIsClass = false;
+    String profileArgumentName = null;
+    boolean profileIsClass = false;
 
     for (VariableElement parameterElement : methodElement.getParameters()) {
       if (parameterElement.getAnnotation(DaoKeyspace.class) != null) {
@@ -125,15 +125,11 @@ public Optional<MethodSpec> generate() {
         }
       } else if (parameterElement.getAnnotation(DaoProfile.class) != null) {
 
-        executionProfileName =
-            validateExecutionProfile(
-                parameterElement, executionProfileName, DaoProfile.class, context);
-        if (context
-            .getClassUtils()
-            .isSame(parameterElement.asType(), DriverExecutionProfile.class)) {
-          executionProfileIsClass = true;
-        }
-        if (executionProfileName == null) {
+        profileArgumentName =
+            validateExecutionProfile(parameterElement, profileArgumentName, context);
+        profileIsClass =
+            context.getClassUtils().isSame(parameterElement.asType(), DriverExecutionProfile.class);
+        if (profileArgumentName == null) {
           return Optional.empty();
         }
       } else {
@@ -151,22 +147,22 @@ public Optional<MethodSpec> generate() {
         return Optional.empty();
       }
     }
-    boolean isCachedByKeyspaceAndTable =
-        (keyspaceArgumentName != null || tableArgumentName != null || executionProfileName != null);
+    boolean isCachedByMethodArguments =
+        (keyspaceArgumentName != null || tableArgumentName != null || profileArgumentName != null);
 
     TypeName returnTypeName = ClassName.get(methodElement.getReturnType());
     String suggestedFieldName = methodElement.getSimpleName() + "Cache";
     String fieldName =
-        isCachedByKeyspaceAndTable
+        isCachedByMethodArguments
             ? enclosingClass.addDaoMapField(suggestedFieldName, returnTypeName)
             : enclosingClass.addDaoSimpleField(
                 suggestedFieldName, returnTypeName, daoImplementationName, isAsync);
 
     MethodSpec.Builder overridingMethodBuilder = GeneratedCodePatterns.override(methodElement);
 
-    if (isCachedByKeyspaceAndTable) {
-      // DaoCacheKey key = new DaoCacheKey(x, y)
-      // where x, y is either the name of the parameter or "(CqlIdentifier)null"
+    if (isCachedByMethodArguments) {
+      // DaoCacheKey key = new DaoCacheKey(<ks>, <table>, <profileName>, <profile>)
+      // where <ks>, <table> is either the name of the parameter or "(CqlIdentifier)null"
       overridingMethodBuilder.addCode("$1T key = new $1T(", DaoCacheKey.class);
       if (keyspaceArgumentName == null) {
         overridingMethodBuilder.addCode("($T)null", CqlIdentifier.class);
@@ -180,28 +176,23 @@ public Optional<MethodSpec> generate() {
         overridingMethodBuilder.addCode("$L", tableArgumentName);
       }
       overridingMethodBuilder.addCode(", ");
-      if (executionProfileName == null) {
-        overridingMethodBuilder.addCode("($T)null", String.class);
+      if (profileArgumentName == null) {
+        overridingMethodBuilder.addCode("null, null);\n");
       } else {
-
-        if (!executionProfileIsClass) {
-          overridingMethodBuilder.addCode("$L", executionProfileName);
+        if (profileIsClass) {
+          overridingMethodBuilder.addCode("null, $L);\n", profileArgumentName);
         } else {
-          overridingMethodBuilder.addCode("$L.getName()", executionProfileName);
+          overridingMethodBuilder.addCode("$L, null);\n", profileArgumentName);
         }
       }
-      overridingMethodBuilder.addCode(");\n");
 
-      overridingMethodBuilder.addCode(
+      overridingMethodBuilder.addStatement(
           "return $L.computeIfAbsent(key, "
-              + "k -> $T.$L(context.withKeyspaceAndTable(k.getKeyspaceId(), k.getTableId())",
+              + "k -> $T.$L(context.withDaoParameters(k.getKeyspaceId(), k.getTableId(), "
+              + "k.getExecutionProfileName(), k.getExecutionProfile())))",
           fieldName,
           daoImplementationName,
           isAsync ? "initAsync" : "init");
-      if (executionProfileName != null) {
-        overridingMethodBuilder.addCode(".withExecutionProfile($L)", executionProfileName);
-      }
-      overridingMethodBuilder.addCode("));\n");
     } else {
       overridingMethodBuilder.addStatement("return $L.get()", fieldName);
     }
@@ -232,8 +223,8 @@ private String validateKeyspaceOrTableParameter(
   }
 
   private String validateExecutionProfile(
-      VariableElement candidate, String previous, Class<?> annotation, ProcessorContext context) {
-    if (!isSingleAnnotation(candidate, previous, annotation, context)) {
+      VariableElement candidate, String previous, ProcessorContext context) {
+    if (!isSingleAnnotation(candidate, previous, DaoProfile.class, context)) {
       return null;
     }
     TypeMirror type = candidate.asType();
@@ -245,7 +236,7 @@ private String validateExecutionProfile(
               candidate,
               processedType,
               "Invalid parameter type: @%s-annotated parameter of %s methods must be of type %s or %s ",
-              annotation.getSimpleName(),
+              DaoProfile.class.getSimpleName(),
               DaoFactory.class.getSimpleName(),
               String.class.getSimpleName(),
               DriverExecutionProfile.class.getSimpleName());
diff --git a/mapper-runtime/revapi.json b/mapper-runtime/revapi.json
@@ -60,8 +60,13 @@
       {
         "code": "java.method.addedToInterface",
         "new": "method java.lang.String com.datastax.oss.driver.api.mapper.MapperContext::getExecutionProfileName()",
-        "justification": "JAVA-2633 Adding execution profile to mapper"
+        "justification": "JAVA-2633: Add execution profile argument to DAO factory method (accept API break -- it's unlikely that MapperContext will be implemented outside of the driver)"
 
+      },
+      {
+        "code": "java.method.addedToInterface",
+        "new": "method com.datastax.oss.driver.api.core.config.DriverExecutionProfile com.datastax.oss.driver.api.mapper.MapperContext::getExecutionProfile()",
+        "justification": "JAVA-2633: Add execution profile argument to DAO factory method (accept API break -- it's unlikely that MapperContext will be implemented outside of the driver)"
       }
     ]
   }
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperBuilder.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperBuilder.java
@@ -17,6 +17,7 @@
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.mapper.annotations.DaoFactory;
 import com.datastax.oss.driver.api.mapper.annotations.Mapper;
 import com.datastax.oss.driver.api.mapper.annotations.QueryProvider;
@@ -38,6 +39,8 @@ public abstract class MapperBuilder<MapperT> {
   protected final CqlSession session;
   protected CqlIdentifier defaultKeyspaceId;
   protected Map<Object, Object> customState;
+  protected String defaultExecutionProfileName;
+  protected DriverExecutionProfile defaultExecutionProfile;
 
   protected MapperBuilder(CqlSession session) {
     this.session = session;
@@ -95,6 +98,50 @@ public MapperBuilder<MapperT> withDefaultKeyspace(@Nullable String keyspaceName)
     return withDefaultKeyspace(keyspaceName == null ? null : CqlIdentifier.fromCql(keyspaceName));
   }
 
+  /**
+   * Specifies a default execution profile name that will be used for all DAOs built with this
+   * mapper (unless they specify their own execution profile).
+   *
+   * <p>This works the same way as the {@linkplain #withDefaultKeyspace(CqlIdentifier) default
+   * keyspace}.
+   *
+   * <p>Note that if you had already set a profile with #withDefaultExecutionProfile, this method
+   * erases it.
+   *
+   * @see DaoFactory
+   */
+  @NonNull
+  public MapperBuilder<MapperT> withDefaultExecutionProfileName(
+      @Nullable String executionProfileName) {
+    this.defaultExecutionProfileName = executionProfileName;
+    if (executionProfileName != null) {
+      this.defaultExecutionProfile = null;
+    }
+    return this;
+  }
+
+  /**
+   * Specifies a default execution profile name that will be used for all DAOs built with this
+   * mapper (unless they specify their own execution profile).
+   *
+   * <p>This works the same way as the {@linkplain #withDefaultKeyspace(CqlIdentifier) default
+   * keyspace}.
+   *
+   * <p>Note that if you had already set a profile name with #withDefaultExecutionProfileName, this
+   * method erases it.
+   *
+   * @see DaoFactory
+   */
+  @NonNull
+  public MapperBuilder<MapperT> withDefaultExecutionProfile(
+      @Nullable DriverExecutionProfile executionProfile) {
+    this.defaultExecutionProfile = executionProfile;
+    if (executionProfile != null) {
+      this.defaultExecutionProfileName = null;
+    }
+    return this;
+  }
+
   /**
    * When the new instance of a class annotated with {@code @Dao} is created an automatic check for
    * schema validation is performed. It verifies if all {@code @Dao} entity fields are present in
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperContext.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperContext.java
@@ -17,6 +17,7 @@
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.mapper.entity.naming.NameConverter;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
@@ -46,12 +47,25 @@ public interface MapperContext {
   CqlIdentifier getTableId();
 
   /**
-   * If this context belongs to a DAO that was built with a execution-profile-parameterized mapper
-   * method, the value of that parameter. Otherwise null.
+   * If this context belongs to a DAO that was built with a method that takes an execution profile
+   * name as parameter, the value of that parameter. Otherwise null.
+   *
+   * <p>Note that this is mutually exclusive with {@link #getExecutionProfile()}: at most one of the
+   * two methods returns a non-null value (or both return null if no profile was provided).
    */
   @Nullable
   String getExecutionProfileName();
 
+  /**
+   * If this context belongs to a DAO that was built with a method that takes an execution profile
+   * as parameter, the value of that parameter. Otherwise null.
+   *
+   * <p>Note that this is mutually exclusive with {@link #getExecutionProfileName()}: at most one of
+   * the two methods returns a non-null value (or both return null if no profile was provided).
+   */
+  @Nullable
+  DriverExecutionProfile getExecutionProfile();
+
   /**
    * Returns an instance of the given converter class.
    *
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/annotations/DaoProfile.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/annotations/DaoProfile.java
@@ -15,6 +15,7 @@
  */
 package com.datastax.oss.driver.api.mapper.annotations;
 
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -26,18 +27,16 @@
  *
  * <p>Example:
  *
- * <p>*
- *
  * <pre>
  *  * &#64;Mapper
  *  * public interface InventoryMapper {
- *  *   ProductDao productDao(@DaoTable String executionProfile);
+ *  *   ProductDao productDao(@DaoProfile String executionProfile);
  *  * }
  *  * </pre>
  *
- * The annotated parameter can be a {@link String}. If it is present, the value will be injected in
- * the DAO instance, where it will be used in generated queries. This allows you to reuse the same
- * DAO for different execution profiles.
+ * The annotated parameter can be a {@link String} or {@link DriverExecutionProfile}. If it is
+ * present, the value will be injected in the DAO instance, where it will be used in generated
+ * queries. This allows you to reuse the same DAO for different execution profiles.
  *
  * @see DaoFactory
  */
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DaoBase.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DaoBase.java
@@ -60,6 +60,11 @@ public class DaoBase {
 
   protected static CompletionStage<PreparedStatement> prepare(
       SimpleStatement statement, MapperContext context) {
+    if (context.getExecutionProfileName() != null) {
+      statement = statement.setExecutionProfileName(context.getExecutionProfileName());
+    } else if (context.getExecutionProfile() != null) {
+      statement = statement.setExecutionProfile(context.getExecutionProfile());
+    }
     return context.getSession().prepareAsync(statement);
   }
 
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DaoCacheKey.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DaoCacheKey.java
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DefaultMapperContext.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DefaultMapperContext.java

Original file line number	Diff line number	Diff line change
`@@ -146,13 +146,6 @@ protected void populateBuilderWithFunction(`
`146`	`146`	`}`
`147`	`147`	`}`
`148`	`148`
`149`		`- protected void populateBuilderWithProfile(CodeBlock.Builder builder) {`
`150`		`- builder.beginControlFlow("if(context.getExecutionProfileName() != null)");`
`151`		`- builder.addStatement(`
`152`		`- "boundStatementBuilder = boundStatementBuilder.setExecutionProfileName(context.getExecutionProfileName())");`
`153`		`- builder.endControlFlow();`
`154`		`- }`
`155`		`-`
`156`	`149`	`protected void populateBuilderWithStatementAttributes(`
`157`	`150`	`CodeBlock.Builder builder, ExecutableElement methodElement) {`
`158`	`151`	`StatementAttributes statementAttributes =`