Generate UDT writers with specific PGTypes

impossibl · Apr 10, 2020 · 724730e · 724730e
1 parent 3674915
commit 724730e
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 9 deletions.
diff --git a/driver/src/main/java/com/impossibl/postgres/api/jdbc/PGAnyType.java b/driver/src/main/java/com/impossibl/postgres/api/jdbc/PGAnyType.java
@@ -33,7 +33,7 @@
 import java.sql.SQLType;
 
 
-public interface PGAnyType extends SQLType  {
+public interface PGAnyType extends SQLType {
 
   String VENDOR_NAME = "PostgreSQL";
 

diff --git a/udt-gen/src/main/kotlin/UDTGenerator.kt b/udt-gen/src/main/kotlin/UDTGenerator.kt
@@ -2,6 +2,7 @@ package com.impossibl.postgres.tools
 
 import com.impossibl.postgres.api.jdbc.PGAnyType
 import com.impossibl.postgres.api.jdbc.PGConnection
+import com.impossibl.postgres.api.jdbc.PGType
 import com.impossibl.postgres.types.QualifiedName
 import com.squareup.javapoet.*
 import com.xenomachina.argparser.*
@@ -364,11 +365,14 @@ class UDTGenerator(
            attrSqlType.javaType.writerTypeName == "Object" ->
              CodeBlock.of("out.writeObject(this.\$L, \$T.\$L);\n", attrPropName, JDBCType::class.java, "STRUCT")
 
+           attrTypeName.box().isBoxedPrimitive ->
+             CodeBlock.of("out.writeObject(this.\$L, \$T.\$L);\n", attrPropName, JDBCType::class.java, attrTypeName.primitiveJDBCType)
+
+           PGType.values().any { it.vendorTypeNumber == attrSqlType.vendorTypeNumber } ->
+             CodeBlock.of("out.writeObject(this.\$L, \$T.\$L);\n", attrPropName, PGType::class.java, PGType.valueOf(attrSqlType.vendorTypeNumber).name)
+
            else ->
-             if (attrTypeName.box().isBoxedPrimitive)
-               CodeBlock.of("out.writeObject(this.\$L, \$T.\$L);\n", attrPropName, JDBCType::class.java, attrTypeName.primitiveJDBCType)
-             else
-               CodeBlock.of("out.write\$L(this.\$L);\n", attrSqlType.javaType.readerTypeName, attrPropName)
+             CodeBlock.of("out.write\$L(this.\$L);\n", attrSqlType.javaType.readerTypeName, attrPropName)
          }
       )
 

diff --git a/udt-gen/src/test/kotlin/UDTGeneratorTest.kt b/udt-gen/src/test/kotlin/UDTGeneratorTest.kt
@@ -11,6 +11,7 @@ import org.junit.jupiter.api.Test
 import java.io.File
 import java.sql.DriverManager
 import java.util.*
+import javax.tools.Diagnostic
 import kotlin.random.Random
 
 class UDTGeneratorTest {
@@ -31,7 +32,7 @@ class UDTGeneratorTest {
 
   @Test
   fun testCompile() {
-    System.out.println("Generating code from $url")
+    println("Generating code from $url")
 
     DriverManager.getConnection(url, props).use { connection ->
 
@@ -41,7 +42,7 @@ class UDTGeneratorTest {
         connection.createStatement().use { stmt ->
           stmt.execute("CREATE SCHEMA $schemaName")
           stmt.execute("CREATE TYPE $schemaName.title as enum ('mr', 'mrs', 'ms', 'dr')")
-          stmt.execute("CREATE TYPE $schemaName.address as (street text, city text, state char(2), zip char(5))")
+          stmt.execute("CREATE TYPE $schemaName.address as (street varchar, city text, state char(2), zip char(5))")
           stmt.execute("CREATE TYPE $schemaName.v_card as (id int4, name text, title $schemaName.title, addresses $schemaName.address[])")
           stmt.execute("SET SEARCH_PATH = public, $schemaName")
         }
@@ -55,6 +56,7 @@ class UDTGeneratorTest {
         val result = javac()
            .compile(files + JavaFileObjects.forResource("VCardTest.java"))
 
+        assertThat(result.errors(), equalTo(emptyList<Diagnostic<*>>()))
         assertThat(result.status(), equalTo(Compilation.Status.SUCCESS))
 
       }
@@ -69,7 +71,7 @@ class UDTGeneratorTest {
 
   @Test
   fun testGenerate() {
-    System.out.println("Generating code from $url")
+    println("Generating code from $url")
 
     DriverManager.getConnection(url, props).use { connection ->
 
@@ -93,7 +95,7 @@ class UDTGeneratorTest {
            .generate(outDirectory)
 
         val pkgFileNames = File(outDirectory, pkgName.replace('.', '/'))
-           .listFiles()
+           .listFiles()!!
            .map { it.name }
 
         assertThat(pkgFileNames.size, equalTo(3))