Skip to content

Commit 26227ba

Browse files
committed
feat(database): add prompter for generating SQL script #80
Add a prompter to generate SQL script based on user selection. The prompter uses a template to render the SQL script and includes custom actions. The prompter is displayed in the chat window for the user to interact with.
1 parent 62d3e13 commit 26227ba

File tree

2 files changed

+39
-12
lines changed

2 files changed

+39
-12
lines changed

exts/database/src/main/kotlin/cc/unitmesh/database/actions/GenSqlScriptBySelection.kt

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
package cc.unitmesh.database.actions
22

33
import cc.unitmesh.devti.AutoDevBundle
4+
import cc.unitmesh.devti.gui.sendToChatWindow
45
import cc.unitmesh.devti.intentions.action.base.AbstractChatIntention
5-
import com.intellij.database.model.DasColumn
6+
import cc.unitmesh.devti.provider.ContextPrompter
7+
import cc.unitmesh.devti.template.TemplateRender
68
import com.intellij.database.model.DasTable
79
import com.intellij.database.model.ObjectKind
8-
import com.intellij.database.model.RawDataSource
9-
import com.intellij.database.psi.DbElement
1010
import com.intellij.database.psi.DbPsiFacade
1111
import com.intellij.database.util.DasUtil
12+
import com.intellij.openapi.diagnostic.logger
1213
import com.intellij.openapi.editor.Editor
1314
import com.intellij.openapi.project.Project
1415
import com.intellij.psi.PsiFile
@@ -23,11 +24,14 @@ class GenSqlScriptBySelection : AbstractChatIntention() {
2324

2425
override fun getText(): String = AutoDevBundle.message("migration.database.sql.generate")
2526

26-
override fun isAvailable(project: Project, editor: Editor?, psiFile: PsiFile?): Boolean {
27+
override fun isAvailable(project: Project, editor: Editor?, file: PsiFile?): Boolean {
28+
DbPsiFacade.getInstance(project).dataSources.firstOrNull() ?: return false
2729
return true
2830
}
2931

30-
override fun invoke(project: Project, editor: Editor?, psiFile: PsiFile?) {
32+
private val logger = logger<GenSqlScriptBySelection>()
33+
34+
override fun invoke(project: Project, editor: Editor?, file: PsiFile?) {
3135
val dbPsiFacade = DbPsiFacade.getInstance(project)
3236
val dataSource = dbPsiFacade.dataSources.firstOrNull() ?: return
3337

@@ -41,16 +45,37 @@ class GenSqlScriptBySelection : AbstractChatIntention() {
4145
tables.filter { table -> table.kind == ObjectKind.TABLE && table.dasParent?.name == schemaName }
4246
}.toList()
4347

44-
val tableColumns = DbContextProvider(dasTables).getTableColumns(dasTables.map { it.name })
45-
46-
DbContext(
48+
val dbContext = DbContext(
4749
requirement = selectedText ?: "",
4850
databaseVersion = databaseVersion.let {
4951
"name: ${it.name}, version: ${it.version}"
5052
},
5153
schemaName = schemaName,
5254
tableNames = dasTables.map { it.name },
5355
)
56+
57+
val actions = DbContextActionProvider(dasTables)
58+
val prompter = generateStepOnePrompt(dbContext, actions)
59+
60+
sendToChatWindow(project, getActionType()) { panel, service ->
61+
service.handlePromptAndResponse(panel, object : ContextPrompter() {
62+
override fun displayPrompt(): String = prompter
63+
override fun requestPrompt(): String = prompter
64+
}, null, false)
65+
}
66+
}
67+
68+
private fun generateStepOnePrompt(context: DbContext, actions: DbContextActionProvider): String {
69+
val templateRender = TemplateRender("genius/sql")
70+
val template = templateRender.getTemplate("sql-gen-clarify.vm")
71+
72+
templateRender.context = context
73+
templateRender.actions = actions
74+
75+
val prompter = templateRender.renderTemplate(template)
76+
77+
logger.info("Prompt: $prompter")
78+
return prompter
5479
}
5580
}
5681

@@ -61,10 +86,9 @@ data class DbContext(
6186
val tableNames: List<String>,
6287
// for step 2
6388
val tableInfos: List<String> = emptyList(),
64-
) {
65-
}
89+
)
6690

67-
data class DbContextProvider(val dasTables: List<DasTable>) {
91+
data class DbContextActionProvider(val dasTables: List<DasTable>) {
6892
/**
6993
* Retrieves the columns of the specified tables.
7094
*

src/main/kotlin/cc/unitmesh/devti/template/TemplateRender.kt

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@ class TemplateRender(val pathPrefix: String) {
1414
private val velocityContext = VelocityContext()
1515
private val splitter = TemplateRoleSplitter()
1616
var context: Any = ""
17-
17+
/**
18+
* Action is also a context, but with custom methods.
19+
*/
20+
var actions: Any = ""
1821

1922
/**
2023
* Retrieves the template for a given filename.

0 commit comments

Comments
 (0)