Skip to content

Commit 4e08500

Browse files
committed
JAVA-2711: Fix handling of UDT keys in the mapper
1 parent 8beeb07 commit 4e08500

3 files changed

Lines changed: 219 additions & 12 deletions

File tree

changelog/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
### 4.6.0 (in progress)
66

7+
- [bug] JAVA-2711: Fix handling of UDT keys in the mapper
78
- [improvement] JAVA-2631: Add getIndex() shortcuts to TableMetadata
89
- [improvement] JAVA-2679: Add port information to QueryTrace and TraceEvent
910
- [improvement] JAVA-2184: Refactor DescribeIT to improve maintainability
Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
/*
2+
* Copyright DataStax, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.datastax.oss.driver.mapper;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
20+
import com.datastax.oss.driver.api.core.CqlIdentifier;
21+
import com.datastax.oss.driver.api.core.CqlSession;
22+
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
23+
import com.datastax.oss.driver.api.mapper.annotations.Dao;
24+
import com.datastax.oss.driver.api.mapper.annotations.DaoFactory;
25+
import com.datastax.oss.driver.api.mapper.annotations.DaoKeyspace;
26+
import com.datastax.oss.driver.api.mapper.annotations.Entity;
27+
import com.datastax.oss.driver.api.mapper.annotations.Insert;
28+
import com.datastax.oss.driver.api.mapper.annotations.Mapper;
29+
import com.datastax.oss.driver.api.mapper.annotations.PartitionKey;
30+
import com.datastax.oss.driver.api.mapper.annotations.Select;
31+
import com.datastax.oss.driver.api.testinfra.ccm.CcmRule;
32+
import com.datastax.oss.driver.api.testinfra.session.SessionRule;
33+
import com.datastax.oss.driver.categories.ParallelizableTests;
34+
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
35+
import java.util.List;
36+
import org.junit.BeforeClass;
37+
import org.junit.ClassRule;
38+
import org.junit.Test;
39+
import org.junit.experimental.categories.Category;
40+
import org.junit.rules.RuleChain;
41+
import org.junit.rules.TestRule;
42+
43+
@Category(ParallelizableTests.class)
44+
public class UdtKeyIT {
45+
46+
private static final CcmRule CCM_RULE = CcmRule.getInstance();
47+
48+
private static final SessionRule<CqlSession> SESSION_RULE = SessionRule.builder(CCM_RULE).build();
49+
50+
@ClassRule
51+
public static final TestRule CHAIN = RuleChain.outerRule(CCM_RULE).around(SESSION_RULE);
52+
53+
private static RecordDao dao;
54+
55+
@BeforeClass
56+
public static void setup() {
57+
CqlSession session = SESSION_RULE.session();
58+
for (String ddlQuery :
59+
ImmutableList.of(
60+
"CREATE TYPE key (value int)",
61+
"CREATE TABLE record(key frozen<key> PRIMARY KEY, value int)",
62+
"CREATE TABLE multi_key_record(key frozen<list<key>> PRIMARY KEY, value int)")) {
63+
session.execute(
64+
SimpleStatement.builder(ddlQuery)
65+
.setExecutionProfile(SESSION_RULE.slowProfile())
66+
.build());
67+
}
68+
69+
TestMapper mapper = new UdtKeyIT_TestMapperBuilder(SESSION_RULE.session()).build();
70+
dao = mapper.recordDao(SESSION_RULE.keyspace());
71+
}
72+
73+
@Test
74+
public void should_save_and_retrieve_entity_with_udt_pk() {
75+
// Given
76+
Key key = new Key(1);
77+
dao.save(new Record(key, 42));
78+
79+
// When
80+
Record record = dao.findByKey(key);
81+
82+
// Then
83+
assertThat(record.getValue()).isEqualTo(42);
84+
}
85+
86+
@Test
87+
public void should_save_and_retrieve_entity_with_udt_collection_pk() {
88+
// Given
89+
List<Key> key = ImmutableList.of(new Key(1), new Key(2));
90+
dao.saveMulti(new MultiKeyRecord(key, 42));
91+
92+
// When
93+
MultiKeyRecord record = dao.findMultiByKey(key);
94+
95+
// Then
96+
assertThat(record.getValue()).isEqualTo(42);
97+
}
98+
99+
@Entity
100+
public static class Key {
101+
private int value;
102+
103+
public Key() {}
104+
105+
public Key(int value) {
106+
this.value = value;
107+
}
108+
109+
public int getValue() {
110+
return value;
111+
}
112+
113+
public void setValue(int value) {
114+
this.value = value;
115+
}
116+
}
117+
118+
@Entity
119+
public static class Record {
120+
@PartitionKey private Key key;
121+
private int value;
122+
123+
public Record() {}
124+
125+
public Record(Key key, int value) {
126+
this.key = key;
127+
this.value = value;
128+
}
129+
130+
public Key getKey() {
131+
return key;
132+
}
133+
134+
public void setKey(Key key) {
135+
this.key = key;
136+
}
137+
138+
public int getValue() {
139+
return value;
140+
}
141+
142+
public void setValue(int value) {
143+
this.value = value;
144+
}
145+
}
146+
147+
@Entity
148+
public static class MultiKeyRecord {
149+
@PartitionKey private List<Key> key;
150+
private int value;
151+
152+
public MultiKeyRecord() {}
153+
154+
public MultiKeyRecord(List<Key> key, int value) {
155+
this.key = key;
156+
this.value = value;
157+
}
158+
159+
public List<Key> getKey() {
160+
return key;
161+
}
162+
163+
public void setKey(List<Key> key) {
164+
this.key = key;
165+
}
166+
167+
public int getValue() {
168+
return value;
169+
}
170+
171+
public void setValue(int value) {
172+
this.value = value;
173+
}
174+
}
175+
176+
@Dao
177+
interface RecordDao {
178+
@Select
179+
Record findByKey(Key key);
180+
181+
@Insert
182+
void save(Record record);
183+
184+
@Select
185+
MultiKeyRecord findMultiByKey(List<Key> key);
186+
187+
@Insert
188+
void saveMulti(MultiKeyRecord record);
189+
}
190+
191+
@Mapper
192+
interface TestMapper {
193+
@DaoFactory
194+
RecordDao recordDao(@DaoKeyspace CqlIdentifier keyspace);
195+
}
196+
}

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/util/generation/GeneratedCodePatterns.java

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -296,12 +296,17 @@ public static void setValue(
296296
udtValueName,
297297
NullSavingStrategy.class,
298298
NullSavingStrategy.DO_NOT_SET)
299-
.addStatement("$1L = $1L.setUdtValue($2L, $3L)", targetName, cqlName, udtValueName)
300-
.nextControlFlow(
301-
"else if ($L == $T.$L)",
302-
NULL_SAVING_STRATEGY,
303-
NullSavingStrategy.class,
304-
NullSavingStrategy.SET_TO_NULL)
299+
.addStatement("$1L = $1L.setUdtValue($2L, $3L)", targetName, cqlName, udtValueName);
300+
if (useNullSavingStrategy) {
301+
methodBuilder.nextControlFlow(
302+
"else if ($L == $T.$L)",
303+
NULL_SAVING_STRATEGY,
304+
NullSavingStrategy.class,
305+
NullSavingStrategy.SET_TO_NULL);
306+
} else {
307+
methodBuilder.nextControlFlow("else");
308+
}
309+
methodBuilder
305310
.addStatement("$1L = $1L.setUdtValue($2L, null)", targetName, cqlName)
306311
.endControlFlow();
307312
} else {
@@ -333,12 +338,17 @@ public static void setValue(
333338
targetName,
334339
cqlName,
335340
rawCollectionName,
336-
enclosingClass.addGenericTypeConstant(type.asRawTypeName()))
337-
.nextControlFlow(
338-
"else if ($L == $T.$L)",
339-
NULL_SAVING_STRATEGY,
340-
NullSavingStrategy.class,
341-
NullSavingStrategy.SET_TO_NULL)
341+
enclosingClass.addGenericTypeConstant(type.asRawTypeName()));
342+
if (useNullSavingStrategy) {
343+
methodBuilder.nextControlFlow(
344+
"else if ($L == $T.$L)",
345+
NULL_SAVING_STRATEGY,
346+
NullSavingStrategy.class,
347+
NullSavingStrategy.SET_TO_NULL);
348+
} else {
349+
methodBuilder.nextControlFlow("else");
350+
}
351+
methodBuilder
342352
.addStatement(
343353
"$1L = $1L.set($2L, null, $3L)",
344354
targetName,

0 commit comments

Comments
 (0)