add enum array support

klahap · Mar 4, 2025 · 508ca8b · 508ca8b
1 parent c233dd8
commit 508ca8b
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 9 deletions.
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -11,7 +11,4 @@ services:
     ports:
       - "5438:5432"
     volumes:
-      - db-data:/var/lib/postgresql/default_code
-
-volumes:
-  db-data:
+      - ./test-schema.sql:/docker-entrypoint-initdb.d/init.sql
diff --git a/src/main/kotlin/util/codegen/CodeGenContext.kt b/src/main/kotlin/util/codegen/CodeGenContext.kt
@@ -11,6 +11,8 @@ data class CodeGenContext(
 
     val getArrayColumnType
         get() = ClassName(packageCustomColumn.name, "getArrayColumnType")
+    val customEnumerationArray
+        get() = ClassName(packageCustomColumn.name, "customEnumerationArray")
     val typeNameMultiRange
         get() = ClassName(packageCustomColumn.name, "MultiRange")
     val defaultJsonColumnType

diff --git a/src/main/kotlin/util/codegen/Column.kt b/src/main/kotlin/util/codegen/Column.kt
@@ -42,10 +42,31 @@ fun PropertySpec.Builder.initializer(column: Table.Column, postfix: String, post
     }
 
     when (val type = column.type) {
-        is Table.Column.Type.NonPrimitive.Array -> initializer(
-            "array<%T>(name = %S)$postfix",
-            type.getTypeName(), columnName, *postArgs
-        )
+        is Table.Column.Type.NonPrimitive.Array -> {
+            when (val elementType = type.elementType) {
+                is Table.Column.Type.NonPrimitive.Enum -> initializer(
+                    """
+                    |%T<%T>(
+                    |    name = %S,
+                    |    sql = %S,
+                    |    fromDb = { %T<%T>(it as String) },
+                    |    toDb = { it.toPgObject() },
+                    |)$postfix""".trimMargin(),
+                    customEnumerationArray,
+                    type.getTypeName(),
+                    columnName,
+                    "${elementType.name.schema.schemaName}.${elementType.name.name}",
+                    typeNameGetPgEnumByLabel,
+                    elementType.name.typeName,
+                    *postArgs,
+                )
+
+                else -> initializer(
+                    "array<%T>(name = %S)$postfix",
+                    type.getTypeName(), columnName, *postArgs
+                )
+            }
+        }
 
         is Table.Column.Type.NonPrimitive.Enum -> initializer(
             """
@@ -56,7 +77,7 @@ fun PropertySpec.Builder.initializer(column: Table.Column, postfix: String, post
                 toDb = { it.toPgObject() },
             )$postfix""".trimIndent(),
             columnName,
-            type.name.name,
+            "${type.name.schema.schemaName}.${type.name.name}",
             typeNameGetPgEnumByLabel,
             type.name.typeName,
             *postArgs

diff --git a/src/main/resources/default_code/column_type/Util.kt b/src/main/resources/default_code/column_type/Util.kt
@@ -4,7 +4,10 @@ import default_code.column_type.RawRange.Empty
 import default_code.column_type.RawRange.Normal
 import default_code.column_type.RawRangeBorder.Infinity
 import org.jetbrains.exposed.sql.ColumnType
+import org.jetbrains.exposed.sql.Column
 import org.jetbrains.exposed.sql.ArrayColumnType
+import org.jetbrains.exposed.sql.CustomEnumerationColumnType
+import org.jetbrains.exposed.sql.Table
 import org.postgresql.util.PGobject
 import kotlin.enums.enumEntries
 
@@ -28,6 +31,21 @@ inline fun <reified T> getPgEnumByLabel(label: String): T
         ?: error("enum with label '$label' not found in '${T::class.qualifiedName}'")
 }
 
+fun <T : Enum<T>> Table.customEnumerationArray(
+    name: String,
+    sql: String?,
+    fromDb: (Any) -> T,
+    toDb: (T) -> Any
+): Column<List<T>> {
+    val enumColumnType = CustomEnumerationColumnType(
+        name = "${name}_element",
+        sql = sql,
+        fromDb = fromDb,
+        toDb = toDb,
+    )
+    return array(name = name, columnType = enumColumnType)
+}
+
 internal fun List<RawRange>.toInt4MultiRange(): MultiRange<Int> = MultiRange(map { it.toInt4Range() }.toSet())
 internal fun List<RawRange>.toInt8MultiRange(): MultiRange<Long> = MultiRange(map { it.toInt8Range() }.toSet())
 internal fun String.parseMultiRange(): List<RawRange> = trimStart('{').trimEnd('}')