kotlin-orm · qumn · Oct 12, 2023 · Jun 12, 2024 · Jun 13, 2024
diff --git a/ktorm-core/src/main/kotlin/org/ktorm/schema/BaseTable.kt b/ktorm-core/src/main/kotlin/org/ktorm/schema/BaseTable.kt
@@ -372,7 +372,7 @@ public abstract class BaseTable<E : Any>(
     /**
      * Convert this table to a [TableExpression].
      */
-    public fun asExpression(): TableExpression {
+    public open fun asExpression(): TableExpression {
         return TableExpression(tableName, alias, catalog, schema)
     }
 

diff --git a/ktorm-ksp-annotations/src/main/kotlin/org/ktorm/ksp/annotation/Table.kt b/ktorm-ksp-annotations/src/main/kotlin/org/ktorm/ksp/annotation/Table.kt
@@ -16,6 +16,9 @@
 
 package org.ktorm.ksp.annotation
 
+import org.ktorm.schema.BaseTable
+import kotlin.reflect.KClass
+
 /**
  * Specify the table for an entity class.
  */
@@ -81,4 +84,12 @@ public annotation class Table(
      * Specify properties that should be ignored for generating column definitions.
      */
     val ignoreProperties: Array<String> = [],
+
+    /**
+     * The super class of the generated table class.
+     *
+     * If not specified, the super class will be determined based on the class kind of the entity:
+     * `org.ktorm.schema.BaseTable` for interfaces, and `org.ktorm.schema.Table` for classes.
+     */
+    val superClass: KClass<out BaseTable<out Any>> = Nothing::class
 )
diff --git a/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/generator/TableClassGenerator.kt b/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/generator/TableClassGenerator.kt
@@ -28,9 +28,7 @@ import org.ktorm.ksp.compiler.util._type
 import org.ktorm.ksp.compiler.util.getKotlinType
 import org.ktorm.ksp.compiler.util.getRegisteringCodeBlock
 import org.ktorm.ksp.spi.TableMetadata
-import org.ktorm.schema.BaseTable
 import org.ktorm.schema.Column
-import org.ktorm.schema.Table
 
 @OptIn(KotlinPoetKspPreview::class)
 internal object TableClassGenerator {
@@ -49,11 +47,7 @@ internal object TableClassGenerator {
     }
 
     private fun TypeSpec.Builder.configureSuperClass(table: TableMetadata): TypeSpec.Builder {
-        if (table.entityClass.classKind == ClassKind.INTERFACE) {
-            superclass(Table::class.asClassName().parameterizedBy(table.entityClass.toClassName()))
-        } else {
-            superclass(BaseTable::class.asClassName().parameterizedBy(table.entityClass.toClassName()))
-        }
+        superclass(table.superClass.asClassName().parameterizedBy(table.entityClass.toClassName()))
 
         addSuperclassConstructorParameter("%S", table.name)
         addSuperclassConstructorParameter("alias")

diff --git a/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/parser/MetadataParser.kt b/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/parser/MetadataParser.kt
@@ -28,6 +28,7 @@ import org.ktorm.ksp.spi.CodingNamingStrategy
 import org.ktorm.ksp.spi.ColumnMetadata
 import org.ktorm.ksp.spi.DatabaseNamingStrategy
 import org.ktorm.ksp.spi.TableMetadata
+import org.ktorm.schema.BaseTable
 import org.ktorm.schema.TypeReference
 import java.lang.reflect.InvocationTargetException
 import java.util.*
@@ -93,6 +94,9 @@ internal class MetadataParser(resolver: Resolver, environment: SymbolProcessorEn
 
         _logger.info("[ktorm-ksp-compiler] parse table metadata from entity: $className")
         val table = cls.getAnnotationsByType(Table::class).first()
+        val superClass = table.superClass.takeIf { it != Nothing::class }
+            ?: if (cls.classKind == INTERFACE) org.ktorm.schema.Table::class else BaseTable::class
+
         val tableMetadata = TableMetadata(
             entityClass = cls,
             name = table.name.ifEmpty { _databaseNamingStrategy.getTableName(cls) },
@@ -102,7 +106,8 @@ internal class MetadataParser(resolver: Resolver, environment: SymbolProcessorEn
             tableClassName = table.className.ifEmpty { _codingNamingStrategy.getTableClassName(cls) },
             entitySequenceName = table.entitySequenceName.ifEmpty { _codingNamingStrategy.getEntitySequenceName(cls) },
             ignoreProperties = table.ignoreProperties.toSet(),
-            columns = ArrayList()
+            columns = ArrayList(),
+            superClass = superClass
         )
 
         val columns = tableMetadata.columns as MutableList

diff --git a/ktorm-ksp-compiler/src/test/kotlin/org/ktorm/ksp/compiler/BaseKspTest.kt b/ktorm-ksp-compiler/src/test/kotlin/org/ktorm/ksp/compiler/BaseKspTest.kt
@@ -40,14 +40,14 @@ abstract class BaseKspTest {
     }
 
     protected fun kspFailing(message: String, @Language("kotlin") code: String, vararg options: Pair<String, String>) {
-        val result = compile(code, mapOf(*options))
+        val result = compile(code, emptyList(), mapOf(*options))
         assert(result.exitCode == KotlinCompilation.ExitCode.COMPILATION_ERROR)
         assert(result.messages.contains("e: Error occurred in KSP, check log for detail"))
         assert(result.messages.contains(message))
     }
 
-    protected fun runKotlin(@Language("kotlin") code: String, vararg options: Pair<String, String>) {
-        val result = compile(code, mapOf(*options))
+    protected fun runKotlin(@Language("kotlin") code: String, additionalImports: List<String> = emptyList(), vararg options: Pair<String, String>) {
+        val result = compile(code, additionalImports, mapOf(*options))
         assert(result.exitCode == KotlinCompilation.ExitCode.OK)
 
         try {
@@ -59,7 +59,7 @@ abstract class BaseKspTest {
         }
     }
 
-    private fun compile(@Language("kotlin") code: String, options: Map<String, String>): KotlinCompilation.Result {
+    private fun compile(@Language("kotlin") code: String, additionalImports: List<String>, options: Map<String, String>): KotlinCompilation.Result {
         @Language("kotlin")
         val header = """
             import java.math.*
@@ -72,6 +72,8 @@ abstract class BaseKspTest {
             import org.ktorm.dsl.*
             import org.ktorm.entity.*
             import org.ktorm.ksp.annotation.*
+
+            ${additionalImports.joinToString("\n") { "import $it" }}
 
             lateinit var database: Database
 

diff --git a/...-ksp-compiler/src/test/kotlin/org/ktorm/ksp/compiler/generator/TableClassGeneratorTest.kt b/...-ksp-compiler/src/test/kotlin/org/ktorm/ksp/compiler/generator/TableClassGeneratorTest.kt
@@ -1,7 +1,9 @@
 package org.ktorm.ksp.compiler.generator
 
 import org.junit.Test
+import org.ktorm.entity.Entity
 import org.ktorm.ksp.compiler.BaseKspTest
+import kotlin.reflect.KClass
 
 class TableClassGeneratorTest : BaseKspTest() {
 
@@ -101,7 +103,7 @@ class TableClassGeneratorTest : BaseKspTest() {
             assert(user.username == "jack")
             assert(user.phone == "12345")
         }
-    """.trimIndent(), "ktorm.allowReflection" to "true")
+    """.trimIndent(), emptyList(), "ktorm.allowReflection" to "true")
 
     @Test
     fun `ignore properties`() = runKotlin("""
@@ -199,4 +201,30 @@ class TableClassGeneratorTest : BaseKspTest() {
             assert(Users.columns.map { it.name }.toSet() == setOf("id", "class", "operator"))
         }
     """.trimIndent())
+
+    abstract class CstmTable<E: Entity<E>>(
+        tableName: String,
+        alias: String? = null,
+        catalog: String? = null,
+        schema: String? = null,
+        entityClass: KClass<E>? = null
+    ) : org.ktorm.schema.Table<E>(
+        tableName, alias, catalog, schema, entityClass
+    )
+    @Test
+    fun `super class`(){
+        runKotlin("""
+        @Table(superClass = CstmTable::class)
+        interface User: Entity<User> {
+            @PrimaryKey
+            var id: Int
+            var `class`: String
+            var operator: String
+        }
+
+        fun run() {
+            assert(CstmTable::class.isSubclassOf(CstmTable::class))
+        }
+    """.trimIndent(), listOf("org.ktorm.ksp.compiler.generator.TableClassGeneratorTest.CstmTable", "kotlin.reflect.full.*"))
+    }
 }
diff --git a/...-ksp-compiler/src/test/kotlin/org/ktorm/ksp/compiler/parser/DatabaseNamingStrategyTest.kt b/...-ksp-compiler/src/test/kotlin/org/ktorm/ksp/compiler/parser/DatabaseNamingStrategyTest.kt
@@ -101,7 +101,7 @@ class DatabaseNamingStrategyTest : BaseKspTest() {
             assert(UserProfiles.tableName == "USER_PROFILE")
             assert(UserProfiles.columns.map { it.name }.toSet() == setOf("ID", "PUBLIC_EMAIL", "PROFILE_PICTURE", "COMPANY_ID"))
         }
-    """.trimIndent(), "ktorm.dbNamingStrategy" to "upper-snake-case")
+    """.trimIndent(), emptyList(), "ktorm.dbNamingStrategy" to "upper-snake-case")
 
     @Test
     fun testUpperCamelCaseNamingByClassName() = runKotlin("""
@@ -128,5 +128,5 @@ class DatabaseNamingStrategyTest : BaseKspTest() {
             assert(UserProfiles.tableName == "USER_PROFILE")
             assert(UserProfiles.columns.map { it.name }.toSet() == setOf("ID", "PUBLIC_EMAIL", "PROFILE_PICTURE", "COMPANY_ID"))
         }
-    """.trimIndent(), "ktorm.dbNamingStrategy" to "org.ktorm.ksp.compiler.util.UpperSnakeCaseDatabaseNamingStrategy")
+    """.trimIndent(), emptyList(), "ktorm.dbNamingStrategy" to "org.ktorm.ksp.compiler.util.UpperSnakeCaseDatabaseNamingStrategy")
 }
diff --git a/ktorm-ksp-spi/src/main/kotlin/org/ktorm/ksp/spi/TableMetadata.kt b/ktorm-ksp-spi/src/main/kotlin/org/ktorm/ksp/spi/TableMetadata.kt
@@ -17,6 +17,7 @@
 package org.ktorm.ksp.spi
 
 import com.google.devtools.ksp.symbol.KSClassDeclaration
+import kotlin.reflect.KClass
 
 /**
  * Table definition metadata.
@@ -66,5 +67,10 @@ public data class TableMetadata(
     /**
      * Columns in the table.
      */
-    val columns: List<ColumnMetadata>
+    val columns: List<ColumnMetadata>,
+
+    /**
+     * The super class of the table class in the generated code.
+     */
+    val superClass: KClass<*>
 )