Skip to content

Commit 7bbc06f

Browse files
vincentlauvlwjvincentlauvlwj
committed
refactor add function
1 parent fbe4ed1 commit 7bbc06f

File tree

3 files changed

+153
-62
lines changed

3 files changed

+153
-62
lines changed

‎ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/generator/AddFunctionGenerator.kt‎

Lines changed: 84 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import org.ktorm.dsl.AliasRemover
2424
import org.ktorm.entity.EntitySequence
2525
import org.ktorm.expression.ColumnAssignmentExpression
2626
import org.ktorm.expression.InsertExpression
27-
import org.ktorm.ksp.compiler.util._type
2827
import org.ktorm.ksp.spi.ColumnMetadata
2928
import org.ktorm.ksp.spi.TableMetadata
3029
import org.ktorm.schema.Column
@@ -34,44 +33,59 @@ internal object AddFunctionGenerator {
3433

3534
fun generate(table: TableMetadata): FunSpec {
3635
val primaryKeys = table.columns.filter { it.isPrimaryKey }
37-
val useGeneratedKey = primaryKeys.size == 1
38-
&& primaryKeys[0].entityProperty.isMutable
39-
&& primaryKeys[0].entityProperty._type.isMarkedNullable
40-
36+
val useGeneratedKey = primaryKeys.size == 1 && primaryKeys[0].entityProperty.isMutable
4137
val entityClass = table.entityClass.toClassName()
4238
val tableClass = ClassName(table.entityClass.packageName.asString(), table.tableClassName)
4339

4440
return FunSpec.builder("add")
4541
.addKdoc(kdoc(table, useGeneratedKey))
4642
.receiver(EntitySequence::class.asClassName().parameterizedBy(entityClass, tableClass))
47-
.addParameter("entity", entityClass)
48-
.addParameter(ParameterSpec.builder("isDynamic", typeNameOf<Boolean>()).defaultValue("false").build())
43+
.addParameters(parameters(entityClass, useGeneratedKey))
4944
.returns(Int::class.asClassName())
5045
.addCode(checkForDml())
51-
.addCode(addValFun())
52-
.addCode(addAssignments(table, useGeneratedKey))
46+
.addCode(addValFun(table, useGeneratedKey))
47+
.addCode(addAssignments(table))
5348
.addCode(createExpression())
5449
.addCode(executeUpdate(useGeneratedKey, primaryKeys))
5550
.build()
5651
}
5752

5853
private fun kdoc(table: TableMetadata, useGeneratedKey: Boolean): String {
59-
var kdoc = "" +
60-
"Insert the given entity into this sequence and return the affected record number. " +
61-
"If [isDynamic] is set to true, the generated SQL will include only the non-null columns. "
62-
6354
if (useGeneratedKey) {
6455
val pk = table.columns.single { it.isPrimaryKey }
6556
val pkName = table.entityClass.simpleName.asString() + "." + pk.entityProperty.simpleName.asString()
66-
67-
kdoc += "\n\n" +
68-
"Note that this function will obtain the generated primary key from the database and fill it into " +
69-
"the property [$pkName] after the insertion completes. But this requires us not to set " +
70-
"the primary key’s value beforehand, otherwise, if you do that, the given value will be " +
71-
"inserted into the database, and no keys generated."
57+
return """
58+
Insert the given entity into the table that the sequence object represents.
59+
60+
@param entity the entity to be inserted.
61+
@param isDynamic whether only non-null columns should be inserted.
62+
@param useGeneratedKey whether to obtain the generated primary key value and fill it into the property [$pkName] after insertion.
63+
@return the affected record number.
64+
""".trimIndent()
65+
} else {
66+
return """
67+
Insert the given entity into the table that the sequence object represents.
68+
69+
@param entity the entity to be inserted.
70+
@param isDynamic whether only non-null columns should be inserted.
71+
@return the affected record number.
72+
""".trimIndent()
7273
}
74+
}
7375

74-
return kdoc
76+
private fun parameters(entityClass: ClassName, useGeneratedKey: Boolean): List<ParameterSpec> {
77+
if (useGeneratedKey) {
78+
return listOf(
79+
ParameterSpec.builder("entity", entityClass).build(),
80+
ParameterSpec.builder("isDynamic", typeNameOf<Boolean>()).defaultValue("false").build(),
81+
ParameterSpec.builder("useGeneratedKey", typeNameOf<Boolean>()).defaultValue("false").build()
82+
)
83+
} else {
84+
return listOf(
85+
ParameterSpec.builder("entity", entityClass).build(),
86+
ParameterSpec.builder("isDynamic", typeNameOf<Boolean>()).defaultValue("false").build()
87+
)
88+
}
7589
}
7690

7791
internal fun checkForDml(): CodeBlock {
@@ -83,7 +97,6 @@ internal object AddFunctionGenerator {
8397
|| expression.orderBy.isNotEmpty()
8498
|| expression.offset != null
8599
|| expression.limit != null
86-
87100
if (isModified) {
88101
val msg = "" +
89102
"Entity manipulation functions are not supported by this sequence object. " +
@@ -97,38 +110,58 @@ internal object AddFunctionGenerator {
97110
return CodeBlock.of(code)
98111
}
99112

100-
internal fun addValFun(): CodeBlock {
101-
val code = """
102-
fun <T : Any> MutableList<%1T<*>>.addVal(column: %2T<T>, value: T?, isDynamic: Boolean) {
103-
if (!isDynamic || value != null) {
113+
private fun addValFun(table: TableMetadata, useGeneratedKey: Boolean): CodeBlock {
114+
if (useGeneratedKey) {
115+
val pk = table.columns.single { it.isPrimaryKey }
116+
val code = """
117+
fun <T : Any> MutableList<%1T<*>>.addVal(column: %2T<T>, value: T?) {
118+
if (useGeneratedKey && column === sourceTable.%3N) {
119+
return
120+
}
121+
122+
if (isDynamic && value == null) {
123+
return
124+
}
125+
104126
this += %1T(column.asExpression(), column.wrapArgument(value))
105127
}
106-
}
107-
128+
129+
130+
""".trimIndent()
131+
return CodeBlock.of(
132+
code,
133+
ColumnAssignmentExpression::class.asClassName(),
134+
Column::class.asClassName(),
135+
pk.columnPropertyName
136+
)
137+
} else {
138+
val code = """
139+
fun <T : Any> MutableList<%1T<*>>.addVal(column: %2T<T>, value: T?) {
140+
if (isDynamic && value == null) {
141+
return
142+
}
108143
109-
""".trimIndent()
110-
111-
return CodeBlock.of(code, ColumnAssignmentExpression::class.asClassName(), Column::class.asClassName())
144+
this += %1T(column.asExpression(), column.wrapArgument(value))
145+
}
146+
147+
148+
""".trimIndent()
149+
return CodeBlock.of(code, ColumnAssignmentExpression::class.asClassName(), Column::class.asClassName())
150+
}
112151
}
113152

114-
private fun addAssignments(table: TableMetadata, useGeneratedKey: Boolean): CodeBlock {
153+
private fun addAssignments(table: TableMetadata): CodeBlock {
115154
return buildCodeBlock {
116155
addStatement("val assignments = ArrayList<%T<*>>()", ColumnAssignmentExpression::class.asClassName())
117156

118157
for (column in table.columns) {
119-
val forceDynamic = useGeneratedKey
120-
&& column.isPrimaryKey && column.entityProperty._type.isMarkedNullable
121-
122158
addStatement(
123-
"assignments.addVal(sourceTable.%N, entity.%N, %L)",
159+
"assignments.addVal(sourceTable.%N, entity.%N)",
124160
column.columnPropertyName,
125-
column.entityProperty.simpleName.asString(),
126-
if (forceDynamic) "isDynamic·=·true" else "isDynamic"
161+
column.entityProperty.simpleName.asString()
127162
)
128163
}
129164

130-
add("\n")
131-
132165
beginControlFlow("if (assignments.isEmpty())")
133166
addStatement("return 0")
134167
endControlFlow()
@@ -138,33 +171,31 @@ internal object AddFunctionGenerator {
138171
}
139172

140173
private fun createExpression(): CodeBlock {
141-
val code = """
142-
val expression = database.dialect.createExpressionVisitor(%T).visit(
143-
%T(sourceTable.asExpression(), assignments)
174+
return buildCodeBlock {
175+
addStatement(
176+
"val visitor = database.dialect.createExpressionVisitor(%T)",
177+
AliasRemover::class.asClassName()
144178
)
145-
146-
147-
""".trimIndent()
148-
149-
return CodeBlock.of(code, AliasRemover::class.asClassName(), InsertExpression::class.asClassName())
179+
addStatement(
180+
"val expression = visitor.visit(%T(sourceTable.asExpression(), assignments))",
181+
InsertExpression::class.asClassName()
182+
)
183+
}
150184
}
151185

152186
private fun executeUpdate(useGeneratedKey: Boolean, primaryKeys: List<ColumnMetadata>): CodeBlock {
153187
return buildCodeBlock {
154188
if (!useGeneratedKey) {
155189
addStatement("return database.executeUpdate(expression)")
156190
} else {
157-
// If the primary key value is manually specified, not obtain the generated key.
158-
beginControlFlow("if (entity.%N != null)", primaryKeys[0].entityProperty.simpleName.asString())
191+
beginControlFlow("if (!useGeneratedKey)")
159192
addStatement("return database.executeUpdate(expression)")
160-
161-
// Else obtain the generated key value.
162193
nextControlFlow("else")
163194
addNamed(
164195
format = """
165196
val (effects, rowSet) = database.executeUpdateAndRetrieveKeys(expression)
166197
if (rowSet.next()) {
167-
val generatedKey = rowSet.%getGeneratedKey:M(sourceTable.%columnName:N)
198+
val generatedKey = rowSet.%getGeneratedKey:M(sourceTable.%columnPropertyName:N)
168199
if (generatedKey != null) {
169200
if (database.logger.isDebugEnabled()) {
170201
database.logger.debug("Generated Key: ${'$'}generatedKey")
@@ -179,12 +210,11 @@ internal object AddFunctionGenerator {
179210
""".trimIndent(),
180211

181212
arguments = mapOf(
182-
"columnName" to primaryKeys[0].columnPropertyName,
183213
"propertyName" to primaryKeys[0].entityProperty.simpleName.asString(),
214+
"columnPropertyName" to primaryKeys[0].columnPropertyName,
184215
"getGeneratedKey" to MemberName("org.ktorm.dsl", "getGeneratedKey", true)
185216
)
186217
)
187-
188218
endControlFlow()
189219
}
190220
}

‎ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/generator/UpdateFunctionGenerator.kt‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import org.ktorm.expression.ColumnAssignmentExpression
2626
import org.ktorm.expression.UpdateExpression
2727
import org.ktorm.ksp.compiler.util.*
2828
import org.ktorm.ksp.spi.TableMetadata
29+
import org.ktorm.schema.Column
2930

3031
@OptIn(KotlinPoetKspPreview::class)
3132
internal object UpdateFunctionGenerator {
@@ -45,14 +46,28 @@ internal object UpdateFunctionGenerator {
4546
.addParameter(ParameterSpec.builder("isDynamic", typeNameOf<Boolean>()).defaultValue("false").build())
4647
.returns(Int::class.asClassName())
4748
.addCode(AddFunctionGenerator.checkForDml())
48-
.addCode(AddFunctionGenerator.addValFun())
49+
.addCode(addValFun())
4950
.addCode(addAssignments(table))
5051
.addCode(buildConditions(table))
5152
.addCode(createExpression())
5253
.addStatement("return database.executeUpdate(expression)")
5354
.build()
5455
}
5556

57+
private fun addValFun(): CodeBlock {
58+
val code = """
59+
fun <T : Any> MutableList<%1T<*>>.addVal(column: %2T<T>, value: T?, isDynamic: Boolean) {
60+
if (!isDynamic || value != null) {
61+
this += %1T(column.asExpression(), column.wrapArgument(value))
62+
}
63+
}
64+
65+
66+
""".trimIndent()
67+
68+
return CodeBlock.of(code, ColumnAssignmentExpression::class.asClassName(), Column::class.asClassName())
69+
}
70+
5671
private fun addAssignments(table: TableMetadata): CodeBlock {
5772
return buildCodeBlock {
5873
addStatement("val assignments = ArrayList<%T<*>>()", ColumnAssignmentExpression::class.asClassName())

‎ktorm-ksp-compiler/src/test/kotlin/org/ktorm/ksp/compiler/generator/AddFunctionGeneratorTest.kt‎

Lines changed: 53 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,19 @@ import org.ktorm.ksp.compiler.BaseKspTest
66
class AddFunctionGeneratorTest : BaseKspTest() {
77

88
@Test
9-
fun `sequence add function`() = runKotlin("""
9+
fun testGenerateKey() = runKotlin("""
1010
@Table
1111
data class User(
1212
@PrimaryKey
13-
var id: Int?,
13+
var id: Int,
1414
var username: String,
1515
var age: Int,
1616
)
1717
1818
fun run() {
19-
database.users.add(User(null, "test", 100))
19+
val user = User(0, "test", 100)
20+
database.users.add(user, useGeneratedKey = true)
21+
assert(user.id == 4)
2022
2123
val users = database.users.toList()
2224
assert(users.size == 4)
@@ -28,19 +30,63 @@ class AddFunctionGeneratorTest : BaseKspTest() {
2830
""".trimIndent())
2931

3032
@Test
31-
fun `modified entity sequence call add fun`() = runKotlin("""
33+
fun testNoGenerateKey() = runKotlin("""
3234
@Table
3335
data class User(
3436
@PrimaryKey
35-
var id: Int?,
37+
var id: Int,
3638
var username: String,
37-
var age: Int
39+
var age: Int,
40+
)
41+
42+
fun run() {
43+
database.users.add(User(99, "test", 100))
44+
45+
val users = database.users.toList()
46+
assert(users.size == 4)
47+
assert(users[0] == User(id = 1, username = "jack", age = 20))
48+
assert(users[1] == User(id = 2, username = "lucy", age = 22))
49+
assert(users[2] == User(id = 3, username = "mike", age = 22))
50+
assert(users[3] == User(id = 99, username = "test", age = 100))
51+
}
52+
""".trimIndent())
53+
54+
@Test
55+
fun testNoGenerateKey1() = runKotlin("""
56+
@Table
57+
data class User(
58+
@PrimaryKey
59+
val id: Int,
60+
val username: String,
61+
val age: Int,
62+
)
63+
64+
fun run() {
65+
database.users.add(User(99, "test", 100))
66+
67+
val users = database.users.toList()
68+
assert(users.size == 4)
69+
assert(users[0] == User(id = 1, username = "jack", age = 20))
70+
assert(users[1] == User(id = 2, username = "lucy", age = 22))
71+
assert(users[2] == User(id = 3, username = "mike", age = 22))
72+
assert(users[3] == User(id = 99, username = "test", age = 100))
73+
}
74+
""".trimIndent())
75+
76+
@Test
77+
fun testSequenceModified() = runKotlin("""
78+
@Table
79+
data class User(
80+
@PrimaryKey
81+
val id: Int,
82+
val username: String,
83+
val age: Int,
3884
)
3985
4086
fun run() {
4187
try {
4288
val users = database.users.filter { it.id eq 1 }
43-
users.add(User(null, "lucy", 10))
89+
users.add(User(99, "lucy", 10))
4490
throw AssertionError("fail")
4591
} catch (_: UnsupportedOperationException) {
4692
}

0 commit comments

Comments
 (0)