Table annotation support specify supper class

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.Table` for interfaces, and `org.ktorm.schema.BaseTable` for classes.
+     */
+    val superClass: KClass<out BaseTable<out Any>> = Nothing::class
 )

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.superClassName.parameterizedBy(table.entityClass.toClassName()))
 
         addSuperclassConstructorParameter("%S", table.name)
         addSuperclassConstructorParameter("alias")

ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/parser/MetadataParser.kt

@@ -19,18 +19,26 @@ package org.ktorm.ksp.compiler.parser
 import com.google.devtools.ksp.*
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
-import com.google.devtools.ksp.symbol.*
 import com.google.devtools.ksp.symbol.ClassKind.*
+import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.KSPropertyDeclaration
+import com.google.devtools.ksp.symbol.KSType
+import com.squareup.kotlinpoet.ClassName
+import com.squareup.kotlinpoet.asClassName
+import com.squareup.kotlinpoet.ksp.KotlinPoetKspPreview
+import com.squareup.kotlinpoet.ksp.toClassName
 import org.ktorm.entity.Entity
 import org.ktorm.ksp.annotation.*
 import org.ktorm.ksp.compiler.util.*
 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.*
+import kotlin.reflect.KClass
 import kotlin.reflect.jvm.jvmName
 
 @OptIn(KspExperimental::class)
@@ -93,6 +101,20 @@ 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()
+
+        // due to the BUG of KSP, we cannot use `table.superClass` directly, so we use `parseAnnotationClassParameter` to get the value
+        // https://github.com/google/ksp/issues/1038
+        // TODO: remove the workaround after the bug is fixed
+        val annotationSuperClass = parseAnnotationClassParameter(table)
+        val superClass = if (annotationSuperClass != Nothing::class.asClassName()) {
+            annotationSuperClass
+        } else if(cls.classKind == INTERFACE) {
+            org.ktorm.schema.Table::class.asClassName()
+        }else {
+            BaseTable::class.asClassName()
+        }
+
+
         val tableMetadata = TableMetadata(
             entityClass = cls,
             name = table.name.ifEmpty { _databaseNamingStrategy.getTableName(cls) },
@@ -102,7 +124,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(),
+            superClassName = superClass
         )
 
         val columns = tableMetadata.columns as MutableList
@@ -296,4 +319,17 @@ internal class MetadataParser(resolver: Resolver, environment: SymbolProcessorEn
 
         stack.pop()
     }
+
+    @OptIn(KspExperimental::class, KotlinPoetKspPreview::class)
+    private fun parseAnnotationClassParameter(table: Table): ClassName {
+        return try {
+            ClassName(table.superClass.java.`package`.name, table.superClass.simpleName.orEmpty())
+        } catch (e: KSTypeNotPresentException) {
+            if (e.ksType.declaration is KSClassDeclaration){
+                return (e.ksType.declaration as KSClassDeclaration).toClassName()
+            }else {
+                throw e
+            }
+        }
+    }
 }

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
 

ktorm-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,29 @@ class TableClassGeneratorTest : BaseKspTest() {
             assert(Users.columns.map { it.name }.toSet() == setOf("id", "class", "operator"))
         }
     """.trimIndent())
+
+    @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.CstmTable", "kotlin.reflect.full.*"))
 }
+
+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
+)

ktorm-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")
 }

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 com.squareup.kotlinpoet.ClassName
 
 /**
  * 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 superClassName: ClassName
 )