JAVA-2139: Allow query methods to receive statement attributes via an…

… argument

integration-tests/src/test/java/com/datastax/oss/driver/mapper/StatementAttributesIT.java

@@ -0,0 +1,260 @@
+/*
+ * 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.mapper;
+
+import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.noRows;
+import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.query;
+import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.when;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.datastax.oss.driver.api.core.CqlIdentifier;
+import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
+import com.datastax.oss.driver.api.mapper.StatementAttributes;
+import com.datastax.oss.driver.api.mapper.StatementAttributesBuilder;
+import com.datastax.oss.driver.api.mapper.annotations.Dao;
+import com.datastax.oss.driver.api.mapper.annotations.DaoFactory;
+import com.datastax.oss.driver.api.mapper.annotations.DaoKeyspace;
+import com.datastax.oss.driver.api.mapper.annotations.Delete;
+import com.datastax.oss.driver.api.mapper.annotations.Entity;
+import com.datastax.oss.driver.api.mapper.annotations.Insert;
+import com.datastax.oss.driver.api.mapper.annotations.Mapper;
+import com.datastax.oss.driver.api.mapper.annotations.PartitionKey;
+import com.datastax.oss.driver.api.mapper.annotations.Select;
+import com.datastax.oss.driver.api.testinfra.session.SessionRule;
+import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
+import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule;
+import com.datastax.oss.protocol.internal.Message;
+import com.datastax.oss.protocol.internal.request.Execute;
+import com.datastax.oss.simulacron.common.cluster.ClusterQueryLogReport;
+import com.datastax.oss.simulacron.common.cluster.ClusterSpec;
+import com.datastax.oss.simulacron.common.cluster.QueryLog;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.util.Map;
+import java.util.Objects;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Test;
+
+public class StatementAttributesIT {
+
+  @ClassRule
+  public static SimulacronRule simulacronRule =
+      new SimulacronRule(ClusterSpec.builder().withNodes(1));
+
+  private static SessionRule<CqlSession> sessionRule = SessionRule.builder(simulacronRule).build();
+
+  private static SimpleDao simpleDao;
+
+  private static String PAGING_STATE = "paging_state";
+  private static int PAGE_SIZE = 13;
+
+  private static final Simple simple = new Simple(UUID.randomUUID(), "DATA");
+
+  @Before
+  public void setup() {
+    simulacronRule.cluster().clearPrimes(true);
+    simulacronRule.cluster().clearLogs();
+  }
+
+  @Test
+  public void should_honor_runtime_attributes_insert() {
+    Map<String, Object> params = ImmutableMap.of("pk", simple.getPk(), "data", simple.getData());
+    Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid", "data", "ascii");
+    simulacronRule
+        .cluster()
+        .prime(
+            when(query(
+                    "INSERT INTO simple (pk,data) VALUES (:pk,:data)",
+                    Lists.newArrayList(
+                        com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ONE,
+                        com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ANY),
+                    params,
+                    paramTypes))
+                .then(noRows()));
+    CqlSession session = SessionUtils.newSession(simulacronRule);
+    InventoryMapper inventoryMapper =
+        new StatementAttributesIT_InventoryMapperBuilder(session).build();
+    simpleDao = inventoryMapper.simpleDao(sessionRule.keyspace());
+    StatementAttributes attributes = buildRunTimeAttributes();
+    simulacronRule.cluster().clearLogs();
+    simpleDao.save(simple, attributes);
+    ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
+    validateQueryOptions(report.getQueryLogs().get(0));
+  }
+
+  @Test
+  public void should_honor_runtime_attributes_delete() {
+    Map<String, Object> params = ImmutableMap.of("pk", simple.getPk());
+    Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid");
+    simulacronRule
+        .cluster()
+        .prime(
+            when(query(
+                    "DELETE FROM simple WHERE pk=:pk",
+                    Lists.newArrayList(
+                        com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ONE,
+                        com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ANY),
+                    params,
+                    paramTypes))
+                .then(noRows())
+                .delay(1, TimeUnit.MILLISECONDS));
+    CqlSession session = SessionUtils.newSession(simulacronRule);
+    InventoryMapper inventoryMapper =
+        new StatementAttributesIT_InventoryMapperBuilder(session).build();
+    simpleDao = inventoryMapper.simpleDao(sessionRule.keyspace());
+    StatementAttributes attributes = buildRunTimeAttributes();
+    simulacronRule.cluster().clearLogs();
+    simpleDao.delete(simple, attributes);
+    ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
+    validateQueryOptions(report.getQueryLogs().get(0));
+  }
+
+  @Test
+  public void should_honor_runtime_attributes_select() {
+    Map<String, Object> params = ImmutableMap.of("pk", simple.getPk());
+    Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid");
+    simulacronRule
+        .cluster()
+        .prime(
+            when(query(
+                    "SELECT pk,data FROM simple WHERE pk=:pk",
+                    Lists.newArrayList(
+                        com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ONE,
+                        com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ANY),
+                    params,
+                    paramTypes))
+                .then(noRows())
+                .delay(1, TimeUnit.MILLISECONDS));
+    CqlSession session = SessionUtils.newSession(simulacronRule);
+    InventoryMapper inventoryMapper =
+        new StatementAttributesIT_InventoryMapperBuilder(session).build();
+    simpleDao = inventoryMapper.simpleDao(sessionRule.keyspace());
+
+    StatementAttributes attributes = buildRunTimeAttributes();
+    simulacronRule.cluster().clearLogs();
+    simpleDao.findByPk(simple.getPk(), attributes);
+    ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
+    validateQueryOptions(report.getQueryLogs().get(0));
+  }
+
+  private StatementAttributes buildRunTimeAttributes() {
+    StatementAttributes attributes =
+        StatementAttributes.builder()
+            .withTimeout(Duration.ofSeconds(3))
+            .withConsistencyLevel(DefaultConsistencyLevel.QUORUM)
+            .build();
+    StatementAttributesBuilder builder = StatementAttributes.builder();
+
+    return builder
+        .withConsistencyLevel(DefaultConsistencyLevel.ANY)
+        .withPageSize(PAGE_SIZE)
+        .withSerialConsistencyLevel(DefaultConsistencyLevel.QUORUM)
+        .withPagingState(ByteBuffer.wrap(PAGING_STATE.getBytes(UTF_8)))
+        .build();
+  }
+
+  private void validateQueryOptions(QueryLog log) {
+
+    Message message = log.getFrame().message;
+    assertThat(message).isInstanceOf(Execute.class);
+    Execute queryExecute = (Execute) message;
+    assertThat(queryExecute.options.consistency)
+        .isEqualTo(DefaultConsistencyLevel.ANY.getProtocolCode());
+    assertThat(queryExecute.options.serialConsistency)
+        .isEqualTo(DefaultConsistencyLevel.QUORUM.getProtocolCode());
+    assertThat(queryExecute.options.pageSize).isEqualTo(PAGE_SIZE);
+    String pagingState = UTF_8.decode(queryExecute.options.pagingState).toString();
+    assertThat(pagingState).isEqualTo(PAGING_STATE);
+  }
+
+  @Mapper
+  public interface InventoryMapper {
+    @DaoFactory
+    StatementAttributesIT.SimpleDao simpleDao(@DaoKeyspace CqlIdentifier keyspace);
+  }
+
+  @Dao
+  public interface SimpleDao {
+    @Insert
+    void save(Simple simple, StatementAttributes attributes);
+
+    @Delete
+    void delete(Simple simple, StatementAttributes attributes);
+
+    @Select
+    Simple findByPk(UUID pk, StatementAttributes attributes);
+  }
+
+  @Entity
+  public static class Simple {
+    @PartitionKey private UUID pk;
+    private String data;
+
+    public Simple() {}
+
+    public Simple(UUID pk, String data) {
+      this.pk = pk;
+      this.data = data;
+    }
+
+    public UUID getPk() {
+      return pk;
+    }
+
+    public String getData() {
+      return data;
+    }
+
+    public void setPk(UUID pk) {
+
+      this.pk = pk;
+    }
+
+    public void setData(String data) {
+      this.data = data;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof Simple)) {
+        return false;
+      }
+      Simple simple = (Simple) o;
+      return Objects.equals(pk, simple.pk) && Objects.equals(data, simple.data);
+    }
+
+    @Override
+    public int hashCode() {
+
+      return Objects.hash(pk, data);
+    }
+
+    @Override
+    public String toString() {
+      return "Simple{" + "pk=" + pk + ", data='" + data + '\'' + '}';
+    }
+  }
+}

manual/mapper/daos/.nav

@@ -6,4 +6,5 @@ queryprovider
 select
 setentity
 update
-null_saving
+null_saving
+statement_attributes

manual/mapper/daos/delete/README.md

@@ -44,6 +44,10 @@ the method must have corresponding parameters (same name, and a compatible Java
 void deleteIfDescriptionMatches(UUID productId, String expectedDescription);
 ```
 
+A [StatementAttributes](../statement_attributes/) can be added as the **last** parameter. This
+allows you to customize customize certain aspects of the request (page size, timeout, etc.) at
+runtime.
+
 ### Return type
 
 The method can return:

manual/mapper/daos/insert/README.md

@@ -25,6 +25,10 @@ void insertWithTtl(Product product, int ttl);
 The annotation can define a [null saving strategy](../null_saving/) that applies to the properties
 of the entity to insert.
 
+A [StatementAttributes](../statement_attributes/) can be added as the **last** parameter. This
+allows you to customize customize certain aspects of the request (page size, timeout, etc.) at
+runtime.
+
 ### Return type
 
 The method can return:

manual/mapper/daos/query/README.md

@@ -25,6 +25,10 @@ long countByIdAndYear(int id, int year);
 The annotation can define a [null saving strategy](../null_saving/) that applies to the method
 parameters.
 
+A [StatementAttributes](../statement_attributes/) can be added as the **last** parameter. This
+allows you to customize customize certain aspects of the request (page size, timeout, etc.) at
+runtime.
+
 ### Return type
 
 The method can return:

manual/mapper/daos/select/README.md

@@ -28,6 +28,10 @@ for each, with the same name and a compatible Java type.
 PagingIterable<Product> findByDescription(String searchString);
 ```
 
+A [StatementAttributes](../statement_attributes/) can be added as the **last** parameter. This
+allows you to customize customize certain aspects of the request (page size, timeout, etc.) at
+runtime.
+
 ### Return type
 
 In all cases, the method can return:

manual/mapper/daos/statement_attributes/README.md

@@ -0,0 +1,33 @@
+## Statement attributes
+
+The [@Delete](../delete/), [@Insert](../insert/), [@Query](../query/), [@Select](../select/) and
+[@Update](../update/) annotations allow you to control some aspects of the execution of the
+underlying statement, such as the consistency level, timeout, etc.
+
+### As a parameter
+
+If the **last** parameter of any of those methods is a [StatementAttributes], it will automatically
+be used to customize the statement:
+
+```java
+@Dao
+public interface ProductDao {
+  @Select
+  Product findById(int productId, StatementAttributes attributes);
+}
+
+StatementAttributes attributes =
+    StatementAttributes.builder()
+        .withTimeout(Duration.ofSeconds(3))
+        .withConsistencyLevel(DefaultConsistencyLevel.QUORUM)
+        .build();
+Product product = dao.findById(1, attributes);
+``` 
+
+Use this if you need to execute the same DAO methods with different configurations that can change
+dynamically.
+
+Note that the default implementation of [StatementAttributes] returned by the builder is immutable.
+If you use the same combinations often, you can store them as constants to reduce allocation costs. 
+
+[StatementAttributes]:  https://docs.datastax.com/en/drivers/java/4.0/com/datastax/oss/driver/api/mapper/StatementAttributes.html

manual/mapper/daos/update/README.md

@@ -73,6 +73,10 @@ template.setDescription("Coming soon"); // all other properties remain null
 dao.updateWhereIdIn(template, 42, 43);  // Will only update 'description' on the selected rows
 ```
 
+A [StatementAttributes](../statement_attributes/) can be added as the **last** parameter. This
+allows you to customize customize certain aspects of the request (page size, timeout, etc.) at
+runtime.
+
 ### Return type
 
 The method can return:

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

@@ -77,10 +77,15 @@ public Optional<MethodSpec> generate() {
 
     // Validate the arguments: either an entity instance, or the PK components (in the latter case,
     // the entity class has to be provided via the annotation).
+    // In either case, a StatementAttributes can be added in last position.
+    List<? extends VariableElement> parameters = methodElement.getParameters();
+    VariableElement statementAttributeParam = findStatementAttributesParam(methodElement);
+    if (statementAttributeParam != null) {
+      parameters = parameters.subList(0, parameters.size() - 1);
+    }
     TypeElement entityElement;
     EntityDefinition entityDefinition;
     boolean hasEntityParameter;
-    List<? extends VariableElement> parameters = methodElement.getParameters();
     if (parameters.isEmpty()) {
       context
           .getMessager()
@@ -164,7 +169,11 @@ public Optional<MethodSpec> generate() {
         "$T boundStatementBuilder = $L.boundStatementBuilder()",
         BoundStatementBuilder.class,
         statementName);
-
+    if (statementAttributeParam != null) {
+      deleteBuilder.addStatement(
+          "boundStatementBuilder = populateBoundStatementWithAttributes(boundStatementBuilder, $L)",
+          statementAttributeParam.getSimpleName().toString());
+    }
     int nextParameterIndex;
     if (hasEntityParameter) {
       // Bind entity's PK properties