Skip to content

Commit e4c774a

Browse files
jbellassaimp911de
jbellassai
authored and
mp911de
committed
Add codecs for working with enums as Strings
Automatically register an EnumStringCodec and an EnumStringArrayCodec for enums. [#454][resolves #429]
1 parent be4fe85 commit e4c774a

File tree

3 files changed

+154
-32
lines changed

3 files changed

+154
-32
lines changed

‎src/main/java/io/r2dbc/postgresql/codec/EnumCodec.java

+109-12
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,100 @@ public T decode(ByteBuf buffer, PostgresTypeIdentifier dataType, Format format,
155155

156156
}
157157

158+
static class EnumStringCodec implements Codec<String>, CodecMetadata {
159+
160+
private final ByteBufAllocator byteBufAllocator;
161+
162+
private final int oid;
163+
164+
EnumStringCodec(ByteBufAllocator byteBufAllocator, int oid) {
165+
this.oid = oid;
166+
this.byteBufAllocator = byteBufAllocator;
167+
}
168+
169+
@Override
170+
public boolean canDecode(int dataType, Format format, Class<?> type) {
171+
Assert.requireNonNull(type, "type must not be null");
172+
return type.isAssignableFrom(String.class) && dataType == this.oid;
173+
}
174+
175+
@Override
176+
public boolean canEncode(Object value) {
177+
Assert.requireNonNull(value, "value must not be null");
178+
return value instanceof String;
179+
}
180+
181+
@Override
182+
public boolean canEncodeNull(Class<?> type) {
183+
Assert.requireNonNull(type, "type must not be null");
184+
return String.class.equals(type);
185+
}
186+
187+
@Override
188+
public String decode(@Nullable ByteBuf buffer, int dataType, Format format, Class<? extends String> type) {
189+
Assert.requireNonNull(buffer, "byteBuf must not be null");
190+
return ByteBufUtils.decode(buffer);
191+
}
192+
193+
@Override
194+
public EncodedParameter encode(Object value) {
195+
return encode(value, this.oid);
196+
}
197+
198+
@Override
199+
public EncodedParameter encode(Object value, int dataType) {
200+
Assert.requireNonNull(value, "value must not be null");
201+
return new EncodedParameter(
202+
FORMAT_TEXT,
203+
dataType,
204+
Mono.fromSupplier(() -> ByteBufUtils.encode(this.byteBufAllocator, ((String) value))
205+
));
206+
}
207+
208+
@Override
209+
public EncodedParameter encodeNull() {
210+
return new EncodedParameter(Format.FORMAT_BINARY, this.oid, NULL_VALUE);
211+
}
212+
213+
@Override
214+
public Class<?> type() {
215+
return String.class;
216+
}
217+
218+
@Override
219+
public Iterable<PostgresTypeIdentifier> getDataTypes() {
220+
return Collections.singleton(AbstractCodec.getDataType(this.oid));
221+
}
222+
223+
}
224+
225+
static class EnumStringArrayCodec extends EnumStringCodec implements ArrayCodecDelegate<String> {
226+
227+
private final PostgresTypeIdentifier arrayType;
228+
229+
EnumStringArrayCodec(ByteBufAllocator byteBufAllocator, int oid, PostgresTypeIdentifier arrayType) {
230+
super(byteBufAllocator, oid);
231+
this.arrayType = arrayType;
232+
}
233+
234+
@Override
235+
public String encodeToText(String value) {
236+
Assert.requireNonNull(value, "value must not be null");
237+
return ArrayCodec.escapeArrayElement(value);
238+
}
239+
240+
@Override
241+
public PostgresTypeIdentifier getArrayDataType() {
242+
return arrayType;
243+
}
244+
245+
@Override
246+
public String decode(ByteBuf buffer, PostgresTypeIdentifier dataType, Format format, Class<? extends String> type) {
247+
return decode(buffer, dataType.getObjectId(), format, type);
248+
}
249+
250+
}
251+
158252
/**
159253
* Builder for {@link CodecRegistrar} to register {@link EnumCodec} for one or more enum type mappings.
160254
*/
@@ -228,20 +322,23 @@ public CodecRegistrar build() {
228322
missing.remove(it.getName());
229323
logger.debug("Registering codec for type '{}' with oid {} using Java enum type '{}'", it.getName(), it.getOid(), enumClass.getName());
230324

231-
if (this.registrationPriority == RegistrationPriority.LAST) {
232-
233-
if (it.getArrayObjectId() > 0) {
234-
registry.addLast(new ArrayCodec(allocator, new EnumArrayCodec(allocator, enumClass, it.getOid(), it.asArrayType()), enumClass));
235-
}
325+
EnumCodec enumCodec = new EnumCodec(allocator, enumClass, it.getOid());
326+
EnumStringCodec enumStringCodec = new EnumStringCodec(allocator, it.getOid());
327+
List<ArrayCodec> arrayCodecs = new ArrayList<>();
328+
if (it.getArrayObjectId() > 0) {
329+
PostgresTypes.PostgresType arrayType = it.asArrayType();
330+
arrayCodecs.add(new ArrayCodec(allocator, new EnumArrayCodec(allocator, enumClass, it.getOid(), arrayType), enumClass));
331+
arrayCodecs.add(new ArrayCodec(allocator, new EnumStringArrayCodec(allocator, it.getOid(), arrayType), String.class));
332+
}
236333

237-
registry.addLast(new EnumCodec(allocator, enumClass, it.getOid()));
334+
if (this.registrationPriority == RegistrationPriority.LAST) {
335+
arrayCodecs.forEach(registry::addLast);
336+
registry.addLast(enumCodec);
337+
registry.addLast(enumStringCodec);
238338
} else {
239-
240-
if (it.getArrayObjectId() > 0) {
241-
registry.addFirst(new ArrayCodec(allocator, new EnumArrayCodec(allocator, enumClass, it.getOid(), it.asArrayType()), enumClass));
242-
}
243-
244-
registry.addFirst(new EnumCodec(allocator, enumClass, it.getOid()));
339+
arrayCodecs.forEach(registry::addFirst);
340+
registry.addFirst(enumCodec);
341+
registry.addFirst(enumStringCodec);
245342
}
246343
}).doOnComplete(() -> {
247344

‎src/test/java/io/r2dbc/postgresql/codec/EnumCodecIntegrationTests.java

+37-18
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import io.r2dbc.postgresql.api.PostgresqlConnection;
2323
import io.r2dbc.postgresql.api.PostgresqlResult;
2424
import io.r2dbc.spi.Parameters;
25+
import org.junit.jupiter.api.BeforeAll;
2526
import org.junit.jupiter.api.Test;
2627
import org.springframework.dao.DataAccessException;
2728
import reactor.test.StepVerifier;
@@ -33,6 +34,15 @@
3334
*/
3435
final class EnumCodecIntegrationTests extends AbstractIntegrationTests {
3536

37+
@BeforeAll
38+
static void createEnum() {
39+
try {
40+
SERVER.getJdbcOperations().execute("CREATE TYPE my_enum AS ENUM ('HELLO', 'WORLD')");
41+
} catch (DataAccessException e) {
42+
// ignore duplicate types
43+
}
44+
}
45+
3646
@Override
3747
protected void customize(PostgresqlConnectionConfiguration.Builder builder) {
3848
builder.codecRegistrar(EnumCodec.builder().withEnum("my_enum", MyEnum.class).build());
@@ -59,12 +69,6 @@ void shouldReportUnresolvableTypes() {
5969
@Test
6070
void shouldBindEnumTypeAsString() {
6171

62-
try {
63-
SERVER.getJdbcOperations().execute("CREATE TYPE my_enum AS ENUM ('HELLO', 'WORLD')");
64-
} catch (DataAccessException e) {
65-
// ignore duplicate types
66-
}
67-
6872
SERVER.getJdbcOperations().execute("DROP TABLE IF EXISTS enum_test");
6973
SERVER.getJdbcOperations().execute("CREATE TABLE enum_test (the_value my_enum);");
7074

@@ -81,17 +85,29 @@ void shouldBindEnumTypeAsString() {
8185

8286
String result = SERVER.getJdbcOperations().queryForObject("SELECT the_value FROM enum_test", String.class);
8387
assertThat(result).isEqualTo("HELLO");
88+
89+
this.connection.createStatement("SELECT * FROM enum_test")
90+
.execute()
91+
.flatMap(it -> it.map(((row, rowMetadata) -> row.get(0, MyEnum.class))))
92+
.as(StepVerifier::create)
93+
.consumeNextWith(actual -> {
94+
assertThat(actual).isEqualTo(MyEnum.HELLO);
95+
})
96+
.verifyComplete();
97+
98+
this.connection.createStatement("SELECT * FROM enum_test")
99+
.execute()
100+
.flatMap(it -> it.map(((row, rowMetadata) -> row.get(0, String.class))))
101+
.as(StepVerifier::create)
102+
.consumeNextWith(actual -> {
103+
assertThat(actual).isEqualTo("HELLO");
104+
})
105+
.verifyComplete();
84106
}
85107

86108
@Test
87109
void shouldBindEnumArrayTypeAsString() {
88110

89-
try {
90-
SERVER.getJdbcOperations().execute("CREATE TYPE my_enum AS ENUM ('HELLO', 'WORLD')");
91-
} catch (DataAccessException e) {
92-
// ignore duplicate types
93-
}
94-
95111
SERVER.getJdbcOperations().execute("DROP TABLE IF EXISTS enum_test");
96112
SERVER.getJdbcOperations().execute("CREATE TABLE enum_test (the_value my_enum[]);");
97113

@@ -108,17 +124,20 @@ void shouldBindEnumArrayTypeAsString() {
108124

109125
String result = SERVER.getJdbcOperations().queryForObject("SELECT the_value FROM enum_test", String.class);
110126
assertThat(result).isEqualTo("{HELLO,WORLD}");
127+
128+
this.connection.createStatement("SELECT the_value FROM enum_test")
129+
.execute()
130+
.flatMap(it -> it.map(((row, rowMetadata) -> row.get(0, String[].class))))
131+
.as(StepVerifier::create)
132+
.consumeNextWith(actual -> {
133+
assertThat(actual).contains("HELLO", "WORLD");
134+
})
135+
.verifyComplete();
111136
}
112137

113138
@Test
114139
void shouldBindEnumArrayType() {
115140

116-
try {
117-
SERVER.getJdbcOperations().execute("CREATE TYPE my_enum AS ENUM ('HELLO', 'WORLD')");
118-
} catch (DataAccessException e) {
119-
// ignore duplicate types
120-
}
121-
122141
SERVER.getJdbcOperations().execute("DROP TABLE IF EXISTS enum_test");
123142
SERVER.getJdbcOperations().execute("CREATE TABLE enum_test (the_value my_enum[]);");
124143

‎src/test/java/io/r2dbc/postgresql/codec/EnumCodecUnitTests.java

+8-2
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@
4040
import static org.mockito.Mockito.any;
4141
import static org.mockito.Mockito.mock;
4242
import static org.mockito.Mockito.never;
43-
import static org.mockito.Mockito.only;
4443
import static org.mockito.Mockito.times;
4544
import static org.mockito.Mockito.verify;
4645

@@ -71,9 +70,14 @@ void canDecode() {
7170
assertThat(codec.canDecode(1, Format.FORMAT_TEXT, MyEnum.class)).isTrue();
7271
assertThat(codec.canDecode(1, FORMAT_BINARY, MyEnum.class)).isTrue();
7372
assertThat(codec.canDecode(1, FORMAT_BINARY, Object.class)).isTrue();
73+
7474
assertThat(codec.canDecode(VARCHAR.getObjectId(), FORMAT_BINARY, MyEnum.class)).isFalse();
7575
assertThat(codec.canDecode(JSON.getObjectId(), FORMAT_TEXT, MyEnum.class)).isFalse();
7676
assertThat(codec.canDecode(JSONB.getObjectId(), FORMAT_BINARY, MyEnum.class)).isFalse();
77+
78+
EnumCodec.EnumStringCodec stringCodec =
79+
new EnumCodec.EnumStringCodec(TestByteBufAllocator.TEST, 1);
80+
assertThat(stringCodec.canDecode(1, FORMAT_TEXT, String.class)).isTrue();
7781
}
7882

7983
@Test
@@ -109,7 +113,8 @@ void shouldRegisterCodecAsFirst() {
109113
Publisher<Void> register = codecRegistrar.register(mockPostgresqlConnection, mockByteBufAllocator, mockCodecRegistry);
110114
StepVerifier.create(register).verifyComplete();
111115

112-
verify(mockCodecRegistry, only()).addFirst(any(EnumCodec.class));
116+
verify(mockCodecRegistry).addFirst(any(EnumCodec.class));
117+
verify(mockCodecRegistry).addFirst(any(EnumCodec.EnumStringCodec.class));
113118
verify(mockCodecRegistry, never()).addLast(any(EnumCodec.class));
114119
}
115120

@@ -148,6 +153,7 @@ void shouldRegisterCodecAsLast() {
148153

149154
verify(mockCodecRegistry, never()).addFirst(any(EnumCodec.class));
150155
verify(mockCodecRegistry, times(2)).addLast(any(EnumCodec.class));
156+
verify(mockCodecRegistry, times(2)).addLast(any(EnumCodec.EnumStringCodec.class));
151157
}
152158

153159
enum MyEnum {

0 commit comments

Comments
 (0)