Skip to content

Commit 84f56c2

Browse files
GregBestlandolim7t
GregBestland
authored and
olim7t
committed
JAVA-2138: Allow query methods to customize statement attributes via annotations
1 parent f636b6d commit 84f56c2

18 files changed

Lines changed: 500 additions & 796 deletions

File tree

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

Lines changed: 198 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -21,63 +21,178 @@
2121
import static java.nio.charset.StandardCharsets.UTF_8;
2222
import static org.assertj.core.api.Assertions.assertThat;
2323

24-
import com.datastax.oss.driver.api.core.CqlIdentifier;
2524
import com.datastax.oss.driver.api.core.CqlSession;
2625
import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
27-
import com.datastax.oss.driver.api.mapper.StatementAttributes;
28-
import com.datastax.oss.driver.api.mapper.StatementAttributesBuilder;
26+
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
2927
import com.datastax.oss.driver.api.mapper.annotations.Dao;
3028
import com.datastax.oss.driver.api.mapper.annotations.DaoFactory;
31-
import com.datastax.oss.driver.api.mapper.annotations.DaoKeyspace;
3229
import com.datastax.oss.driver.api.mapper.annotations.Delete;
3330
import com.datastax.oss.driver.api.mapper.annotations.Entity;
3431
import com.datastax.oss.driver.api.mapper.annotations.Insert;
3532
import com.datastax.oss.driver.api.mapper.annotations.Mapper;
3633
import com.datastax.oss.driver.api.mapper.annotations.PartitionKey;
34+
import com.datastax.oss.driver.api.mapper.annotations.Query;
3735
import com.datastax.oss.driver.api.mapper.annotations.Select;
36+
import com.datastax.oss.driver.api.mapper.annotations.StatementAttributes;
37+
import com.datastax.oss.driver.api.mapper.annotations.Update;
3838
import com.datastax.oss.driver.api.testinfra.session.SessionRule;
39-
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
4039
import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule;
4140
import com.datastax.oss.protocol.internal.Message;
4241
import com.datastax.oss.protocol.internal.request.Execute;
4342
import com.datastax.oss.simulacron.common.cluster.ClusterQueryLogReport;
4443
import com.datastax.oss.simulacron.common.cluster.ClusterSpec;
4544
import com.datastax.oss.simulacron.common.cluster.QueryLog;
45+
import com.datastax.oss.simulacron.common.stubbing.PrimeDsl;
4646
import com.google.common.collect.ImmutableMap;
4747
import com.google.common.collect.Lists;
4848
import java.nio.ByteBuffer;
49-
import java.time.Duration;
5049
import java.util.Map;
5150
import java.util.Objects;
5251
import java.util.UUID;
5352
import java.util.concurrent.TimeUnit;
53+
import java.util.function.Function;
5454
import org.junit.Before;
55+
import org.junit.BeforeClass;
5556
import org.junit.ClassRule;
57+
import org.junit.Rule;
5658
import org.junit.Test;
59+
import org.junit.rules.ExpectedException;
60+
import org.junit.rules.RuleChain;
61+
import org.junit.rules.TestRule;
5762

5863
public class StatementAttributesIT {
5964

60-
@ClassRule
61-
public static SimulacronRule simulacronRule =
65+
private static SimulacronRule simulacronRule =
6266
new SimulacronRule(ClusterSpec.builder().withNodes(1));
63-
6467
private static SessionRule<CqlSession> sessionRule = SessionRule.builder(simulacronRule).build();
68+
@ClassRule public static TestRule chain = RuleChain.outerRule(simulacronRule).around(sessionRule);
6569

66-
private static SimpleDao simpleDao;
70+
@Rule public ExpectedException thrown = ExpectedException.none();
6771

6872
private static String PAGING_STATE = "paging_state";
6973
private static int PAGE_SIZE = 13;
7074

7175
private static final Simple simple = new Simple(UUID.randomUUID(), "DATA");
7276

77+
private static final Function<BoundStatementBuilder, BoundStatementBuilder> statementFunction =
78+
builder ->
79+
builder
80+
.setConsistencyLevel(DefaultConsistencyLevel.ANY)
81+
.setPageSize(PAGE_SIZE)
82+
.setSerialConsistencyLevel(DefaultConsistencyLevel.QUORUM)
83+
.setPagingState(ByteBuffer.wrap(PAGING_STATE.getBytes(UTF_8)));
84+
85+
private static final Function<BoundStatementBuilder, BoundStatementBuilder> badStatementFunction =
86+
builder -> {
87+
throw new IllegalStateException("mock error");
88+
};
89+
90+
private static SimpleDao dao;
91+
92+
@BeforeClass
93+
public static void setupClass() {
94+
primeDeleteQuery();
95+
primeInsertQuery();
96+
primeSelectQuery();
97+
primeCountQuery();
98+
primeUpdateQuery();
99+
100+
InventoryMapper inventoryMapper =
101+
new StatementAttributesIT_InventoryMapperBuilder(sessionRule.session()).build();
102+
dao = inventoryMapper.simpleDao();
103+
}
104+
73105
@Before
74106
public void setup() {
75-
simulacronRule.cluster().clearPrimes(true);
76107
simulacronRule.cluster().clearLogs();
77108
}
78109

79110
@Test
80-
public void should_honor_runtime_attributes_insert() {
111+
public void should_honor_runtime_attributes_on_insert() {
112+
dao.save(simple, statementFunction);
113+
114+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
115+
validateQueryOptions(report.getQueryLogs().get(0), true);
116+
}
117+
118+
@Test
119+
public void should_honor_annotation_attributes_on_insert() {
120+
dao.save2(simple);
121+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
122+
validateQueryOptions(report.getQueryLogs().get(0), false);
123+
}
124+
125+
@Test
126+
public void should_use_runtime_attributes_over_annotation_attributes() {
127+
dao.save3(simple, statementFunction);
128+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
129+
validateQueryOptions(report.getQueryLogs().get(0), false);
130+
}
131+
132+
@Test
133+
public void should_honor_runtime_attributes_on_delete() {
134+
dao.delete(simple, statementFunction);
135+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
136+
validateQueryOptions(report.getQueryLogs().get(0), true);
137+
}
138+
139+
@Test
140+
public void should_honor_annotation_attributes_on_delete() {
141+
dao.delete2(simple);
142+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
143+
validateQueryOptions(report.getQueryLogs().get(0), false);
144+
}
145+
146+
@Test
147+
public void should_honor_runtime_attributes_on_select() {
148+
dao.findByPk(simple.getPk(), statementFunction);
149+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
150+
validateQueryOptions(report.getQueryLogs().get(0), true);
151+
}
152+
153+
@Test
154+
public void should_honor_annotation_attributes_on_select() {
155+
dao.findByPk2(simple.getPk());
156+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
157+
validateQueryOptions(report.getQueryLogs().get(0), false);
158+
}
159+
160+
@Test
161+
public void should_honor_runtime_attributes_on_query() {
162+
dao.count(simple.getPk(), statementFunction);
163+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
164+
validateQueryOptions(report.getQueryLogs().get(0), true);
165+
}
166+
167+
@Test
168+
public void should_honor_annotation_attributes_on_query() {
169+
dao.count2(simple.getPk());
170+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
171+
validateQueryOptions(report.getQueryLogs().get(0), false);
172+
}
173+
174+
@Test
175+
public void should_honor_runtime_attributes_on_update() {
176+
dao.update(simple, statementFunction);
177+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
178+
validateQueryOptions(report.getQueryLogs().get(0), true);
179+
}
180+
181+
@Test
182+
public void should_honor_annotation_attributes_on_update() {
183+
dao.update2(simple);
184+
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
185+
validateQueryOptions(report.getQueryLogs().get(0), false);
186+
}
187+
188+
@Test
189+
public void should_fail_runtime_attributes_bad() {
190+
thrown.expect(IllegalStateException.class);
191+
thrown.expectMessage("mock error");
192+
dao.save(simple, badStatementFunction);
193+
}
194+
195+
private static void primeInsertQuery() {
81196
Map<String, Object> params = ImmutableMap.of("pk", simple.getPk(), "data", simple.getData());
82197
Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid", "data", "ascii");
83198
simulacronRule
@@ -91,19 +206,9 @@ public void should_honor_runtime_attributes_insert() {
91206
params,
92207
paramTypes))
93208
.then(noRows()));
94-
CqlSession session = SessionUtils.newSession(simulacronRule);
95-
InventoryMapper inventoryMapper =
96-
new StatementAttributesIT_InventoryMapperBuilder(session).build();
97-
simpleDao = inventoryMapper.simpleDao(sessionRule.keyspace());
98-
StatementAttributes attributes = buildRunTimeAttributes();
99-
simulacronRule.cluster().clearLogs();
100-
simpleDao.save(simple, attributes);
101-
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
102-
validateQueryOptions(report.getQueryLogs().get(0));
103209
}
104210

105-
@Test
106-
public void should_honor_runtime_attributes_delete() {
211+
private static void primeDeleteQuery() {
107212
Map<String, Object> params = ImmutableMap.of("pk", simple.getPk());
108213
Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid");
109214
simulacronRule
@@ -118,19 +223,9 @@ public void should_honor_runtime_attributes_delete() {
118223
paramTypes))
119224
.then(noRows())
120225
.delay(1, TimeUnit.MILLISECONDS));
121-
CqlSession session = SessionUtils.newSession(simulacronRule);
122-
InventoryMapper inventoryMapper =
123-
new StatementAttributesIT_InventoryMapperBuilder(session).build();
124-
simpleDao = inventoryMapper.simpleDao(sessionRule.keyspace());
125-
StatementAttributes attributes = buildRunTimeAttributes();
126-
simulacronRule.cluster().clearLogs();
127-
simpleDao.delete(simple, attributes);
128-
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
129-
validateQueryOptions(report.getQueryLogs().get(0));
130226
}
131227

132-
@Test
133-
public void should_honor_runtime_attributes_select() {
228+
private static void primeSelectQuery() {
134229
Map<String, Object> params = ImmutableMap.of("pk", simple.getPk());
135230
Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid");
136231
simulacronRule
@@ -145,35 +240,42 @@ public void should_honor_runtime_attributes_select() {
145240
paramTypes))
146241
.then(noRows())
147242
.delay(1, TimeUnit.MILLISECONDS));
148-
CqlSession session = SessionUtils.newSession(simulacronRule);
149-
InventoryMapper inventoryMapper =
150-
new StatementAttributesIT_InventoryMapperBuilder(session).build();
151-
simpleDao = inventoryMapper.simpleDao(sessionRule.keyspace());
243+
}
152244

153-
StatementAttributes attributes = buildRunTimeAttributes();
154-
simulacronRule.cluster().clearLogs();
155-
simpleDao.findByPk(simple.getPk(), attributes);
156-
ClusterQueryLogReport report = simulacronRule.cluster().getLogs();
157-
validateQueryOptions(report.getQueryLogs().get(0));
245+
private static void primeCountQuery() {
246+
Map<String, Object> params = ImmutableMap.of("pk", simple.getPk());
247+
Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid");
248+
simulacronRule
249+
.cluster()
250+
.prime(
251+
when(query(
252+
"SELECT count(*) FROM simple WHERE pk=:pk",
253+
Lists.newArrayList(
254+
com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ONE,
255+
com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ANY),
256+
params,
257+
paramTypes))
258+
.then(PrimeDsl.rows().row("count", 1L).columnTypes("count", "bigint").build())
259+
.delay(1, TimeUnit.MILLISECONDS));
158260
}
159261

160-
private StatementAttributes buildRunTimeAttributes() {
161-
StatementAttributes attributes =
162-
StatementAttributes.builder()
163-
.withTimeout(Duration.ofSeconds(3))
164-
.withConsistencyLevel(DefaultConsistencyLevel.QUORUM)
165-
.build();
166-
StatementAttributesBuilder builder = StatementAttributes.builder();
167-
168-
return builder
169-
.withConsistencyLevel(DefaultConsistencyLevel.ANY)
170-
.withPageSize(PAGE_SIZE)
171-
.withSerialConsistencyLevel(DefaultConsistencyLevel.QUORUM)
172-
.withPagingState(ByteBuffer.wrap(PAGING_STATE.getBytes(UTF_8)))
173-
.build();
262+
private static void primeUpdateQuery() {
263+
Map<String, Object> params = ImmutableMap.of("pk", simple.getPk(), "data", simple.getData());
264+
Map<String, String> paramTypes = ImmutableMap.of("pk", "uuid", "data", "ascii");
265+
simulacronRule
266+
.cluster()
267+
.prime(
268+
when(query(
269+
"UPDATE simple SET data=:data WHERE pk=:pk",
270+
Lists.newArrayList(
271+
com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ONE,
272+
com.datastax.oss.simulacron.common.codec.ConsistencyLevel.ANY),
273+
params,
274+
paramTypes))
275+
.then(noRows()));
174276
}
175277

176-
private void validateQueryOptions(QueryLog log) {
278+
private void validateQueryOptions(QueryLog log, boolean validatePageState) {
177279

178280
Message message = log.getFrame().message;
179281
assertThat(message).isInstanceOf(Execute.class);
@@ -183,26 +285,58 @@ private void validateQueryOptions(QueryLog log) {
183285
assertThat(queryExecute.options.serialConsistency)
184286
.isEqualTo(DefaultConsistencyLevel.QUORUM.getProtocolCode());
185287
assertThat(queryExecute.options.pageSize).isEqualTo(PAGE_SIZE);
186-
String pagingState = UTF_8.decode(queryExecute.options.pagingState).toString();
187-
assertThat(pagingState).isEqualTo(PAGING_STATE);
288+
if (validatePageState) {
289+
String pagingState = UTF_8.decode(queryExecute.options.pagingState).toString();
290+
assertThat(pagingState).isEqualTo(PAGING_STATE);
291+
}
188292
}
189293

190294
@Mapper
191295
public interface InventoryMapper {
192296
@DaoFactory
193-
StatementAttributesIT.SimpleDao simpleDao(@DaoKeyspace CqlIdentifier keyspace);
297+
StatementAttributesIT.SimpleDao simpleDao();
194298
}
195299

196300
@Dao
197301
public interface SimpleDao {
198302
@Insert
199-
void save(Simple simple, StatementAttributes attributes);
303+
void save(Simple simple, Function<BoundStatementBuilder, BoundStatementBuilder> function);
304+
305+
@Insert
306+
@StatementAttributes(consistencyLevel = "ANY", serialConsistencyLevel = "QUORUM", pageSize = 13)
307+
void save2(Simple simple);
308+
309+
@Insert
310+
@StatementAttributes(consistencyLevel = "ONE", pageSize = 500)
311+
void save3(Simple simple, Function<BoundStatementBuilder, BoundStatementBuilder> function);
312+
313+
@Delete
314+
void delete(Simple simple, Function<BoundStatementBuilder, BoundStatementBuilder> function);
200315

201316
@Delete
202-
void delete(Simple simple, StatementAttributes attributes);
317+
@StatementAttributes(consistencyLevel = "ANY", serialConsistencyLevel = "QUORUM", pageSize = 13)
318+
void delete2(Simple simple);
203319

204320
@Select
205-
Simple findByPk(UUID pk, StatementAttributes attributes);
321+
Simple findByPk(UUID pk, Function<BoundStatementBuilder, BoundStatementBuilder> function);
322+
323+
@Select
324+
@StatementAttributes(consistencyLevel = "ANY", serialConsistencyLevel = "QUORUM", pageSize = 13)
325+
Simple findByPk2(UUID pk);
326+
327+
@Query("SELECT count(*) FROM simple WHERE pk=:pk")
328+
long count(UUID pk, Function<BoundStatementBuilder, BoundStatementBuilder> function);
329+
330+
@Query("SELECT count(*) FROM simple WHERE pk=:pk")
331+
@StatementAttributes(consistencyLevel = "ANY", serialConsistencyLevel = "QUORUM", pageSize = 13)
332+
long count2(UUID pk);
333+
334+
@Update
335+
void update(Simple simple, Function<BoundStatementBuilder, BoundStatementBuilder> function);
336+
337+
@Update
338+
@StatementAttributes(consistencyLevel = "ANY", serialConsistencyLevel = "QUORUM", pageSize = 13)
339+
void update2(Simple simple);
206340
}
207341

208342
@Entity

0 commit comments

Comments
 (0)