Do not require typarray column when auto-registering extensions.

[closes #621]
pgjdbc · Dec 28, 2023 · c19ac40 · c19ac40
1 parent 29fd8dd
commit c19ac40
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 16 deletions.
diff --git a/src/main/java/io/r2dbc/postgresql/codec/BuiltinDynamicCodecs.java b/src/main/java/io/r2dbc/postgresql/codec/BuiltinDynamicCodecs.java
@@ -23,8 +23,10 @@
 import org.reactivestreams.Publisher;
 import reactor.util.annotation.Nullable;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
 import java.util.stream.Collectors;
 
 /**
@@ -62,7 +64,13 @@ public Iterable<Codec<?>> createCodec(ByteBufAllocator byteBufAllocator, int oid
                     return Collections.singletonList(new PostgisGeometryCodec(oid));
                 case VECTOR:
                     VectorCodec vectorCodec = new VectorCodec(byteBufAllocator, oid, typarray);
-                    return Arrays.asList(vectorCodec, new VectorCodec.VectorArrayCodec(byteBufAllocator, vectorCodec), new VectorFloatCodec(byteBufAllocator, oid));
+                    List<Codec<?>> codecs = new ArrayList<>(3);
+                    codecs.add(vectorCodec);
+                    if (typarray != PostgresTypes.NO_SUCH_TYPE) {
+                        codecs.add(new VectorCodec.VectorArrayCodec(byteBufAllocator, vectorCodec));
+                    }
+                    codecs.add(new VectorFloatCodec(byteBufAllocator, oid));
+                    return codecs;
                 default:
                     throw new UnsupportedOperationException(String.format("Codec %s for OID %d not supported", name(), oid));
             }
@@ -97,7 +105,7 @@ public Publisher<Void> register(PostgresqlConnection connection, ByteBufAllocato
             .flatMap(it -> it.map((row, rowMetadata) -> {
 
                     int oid = PostgresqlObjectId.toInt(row.get("oid", Long.class));
-                    int typarray = PostgresqlObjectId.toInt(row.get("typarray", Long.class));
+                    int typarray = rowMetadata.contains("typarray") ? PostgresqlObjectId.toInt(row.get("typarray", Long.class)) : PostgresTypes.NO_SUCH_TYPE;
                     String typname = row.get("typname", String.class);
 
                     BuiltinCodec lookup = BuiltinCodec.lookup(typname);

diff --git a/src/main/java/io/r2dbc/postgresql/codec/PostgresTypes.java b/src/main/java/io/r2dbc/postgresql/codec/PostgresTypes.java
@@ -18,6 +18,8 @@
 
 import io.r2dbc.postgresql.api.PostgresqlConnection;
 import io.r2dbc.postgresql.util.Assert;
+import io.r2dbc.spi.Row;
+import io.r2dbc.spi.RowMetadata;
 import io.r2dbc.spi.Type;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -35,6 +37,8 @@
  */
 public class PostgresTypes {
 
+    public final static int NO_SUCH_TYPE = -1;
+
     // parameterized with %s for the comparator (=, IN), %s for the actual criteria value and %s for a potential LIMIT 1 statement
     private static final String SELECT_PG_TYPE = "SELECT pg_type.oid, typarray, typname, typcategory "
         + "  FROM pg_catalog.pg_type "
@@ -74,13 +78,7 @@ public Mono<PostgresType> lookupType(String typeName) {
         }
 
         return this.connection.createStatement(String.format(SELECT_PG_TYPE, "=", "'" + typeName + "'", "LIMIT 1")).execute()
-            .flatMap(it -> it.map((row, rowMetadata) -> {
-
-                Long oid = row.get("oid", Long.class);
-                Long typarrayOid = row.get("typarray", Long.class);
-                return new PostgresType(PostgresqlObjectId.toInt(oid), oid.longValue(), PostgresqlObjectId.toInt(typarrayOid), typarrayOid, row.get("typname", String.class), row.get("typcategory",
-                    String.class));
-            })).singleOrEmpty();
+            .flatMap(it -> it.map(PostgresTypes::createType)).singleOrEmpty();
     }
 
     public Flux<PostgresType> lookupTypes(Iterable<String> typeNames) {
@@ -103,13 +101,18 @@ public Flux<PostgresType> lookupTypes(Iterable<String> typeNames) {
         }
 
         return this.connection.createStatement(String.format(SELECT_PG_TYPE, "IN", joiner, "")).execute()
-            .flatMap(it -> it.map((row, rowMetadata) -> {
+            .flatMap(it -> it.map(PostgresTypes::createType));
+    }
+
+    private static PostgresType createType(Row row, RowMetadata rowMetadata) {
+        Long oid = row.get("oid", Long.class);
+        String typname = row.get("typname", String.class);
+        String typcategory = row.get("typcategory", String.class);
+        Long typarrayOid = rowMetadata.contains("typarray") ? row.get("typarray", Long.class) : null;
 
-                Long oid = row.get("oid", Long.class);
-                Long typarrayOid = row.get("typarray", Long.class);
-                return new PostgresType(PostgresqlObjectId.toInt(oid), oid.longValue(), PostgresqlObjectId.toInt(typarrayOid), typarrayOid, row.get("typname", String.class), row.get("typcategory",
-                    String.class));
-            }));
+        long unsignedTyparray = typarrayOid != null ? typarrayOid : NO_SUCH_TYPE;
+        int typarray = typarrayOid != null ? PostgresqlObjectId.toInt(typarrayOid) : NO_SUCH_TYPE;
+        return new PostgresType(PostgresqlObjectId.toInt(oid), oid, typarray, unsignedTyparray, typname, typcategory);
     }
 
     public static class PostgresType implements Type, PostgresTypeIdentifier {

diff --git a/src/main/java/io/r2dbc/postgresql/codec/PostgresqlObjectId.java b/src/main/java/io/r2dbc/postgresql/codec/PostgresqlObjectId.java
@@ -584,7 +584,7 @@ public static int toInt(@Nullable Long oid) {
 
     public static int toInt(long oid) {
         if ((oid & 0xFFFFFFFF00000000L) != 0) {
-            throw new IllegalArgumentException("Value is not an OID:" + oid);
+            throw new IllegalArgumentException("Value is not an OID: " + oid);
         }
 
         return (int) oid;

diff --git a/src/test/java/io/r2dbc/postgresql/codec/EnumCodecUnitTests.java b/src/test/java/io/r2dbc/postgresql/codec/EnumCodecUnitTests.java
@@ -113,6 +113,35 @@ void shouldRegisterCodecAsFirst() {
         verify(mockCodecRegistry, never()).addLast(any(EnumCodec.class));
     }
 
+    @Test
+    void shouldRegisterCodecWithoutTyparray() {
+        CodecRegistrar codecRegistrar = EnumCodec
+            .builder()
+            .withRegistrationPriority(RegistrationPriority.FIRST)
+            .withEnum("foo", MyEnum.class)
+            .build();
+
+        ByteBufAllocator mockByteBufAllocator = mock(ByteBufAllocator.class);
+        CodecRegistry mockCodecRegistry = mock(CodecRegistry.class);
+
+        MockPostgresqlStatement mockPostgresqlStatement = MockPostgresqlStatement.builder()
+            .result(MockPostgresqlResult.builder()
+                .row(MockRow.builder()
+                    .identified("oid", Long.class, 42L)
+                    .identified("typname", String.class, "foo")
+                    .identified("typcategory", String.class, "E")
+                    .build())
+                .build())
+            .build();
+        MockPostgresqlConnection mockPostgresqlConnection = new MockPostgresqlConnection(mockPostgresqlStatement);
+
+        Publisher<Void> register = codecRegistrar.register(mockPostgresqlConnection, mockByteBufAllocator, mockCodecRegistry);
+        StepVerifier.create(register).verifyComplete();
+
+        verify(mockCodecRegistry, only()).addFirst(any(EnumCodec.class));
+        verify(mockCodecRegistry, never()).addLast(any(EnumCodec.class));
+    }
+
     @Test
     void shouldRegisterCodecAsLast() {
         CodecRegistrar codecRegistrar = EnumCodec