Skip to content

Commit 8f9ea98

Browse files
committed
feat(database): add safety check for dangerous SQL operations
- Implement a SQL safety check to detect and prevent potentially harmful operations - Add a list of dangerous SQL keywords (DELETE, DROP, TRUNCATE, etc.) - Analyze SQL statements to identify any dangerous operations - Return an error message if a dangerous operation is detected
1 parent 99dce5e commit 8f9ea98

File tree

1 file changed

+30
-0
lines changed
  • exts/ext-database/src/241/main/kotlin/cc/unitmesh/database/util

1 file changed

+30
-0
lines changed

exts/ext-database/src/241/main/kotlin/cc/unitmesh/database/util/SQLExecutor.kt

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,41 @@ import com.intellij.util.Consumer
2525
import java.util.concurrent.CompletableFuture
2626

2727
object SQLExecutor {
28+
private val DANGEROUS_OPERATIONS = setOf(
29+
"DELETE",
30+
"DROP",
31+
"TRUNCATE",
32+
"ALTER",
33+
"GRANT",
34+
"REVOKE"
35+
)
36+
37+
private fun checkSqlSafety(project: Project, psiFile: com.intellij.psi.PsiFile): Pair<Boolean, String> {
38+
val sqlFile = psiFile as? com.intellij.sql.psi.SqlFile ?: return Pair(true, "Not a SQL file")
39+
val statements = sqlFile.ddl
40+
41+
for (statement in statements) {
42+
val firstChild = statement.firstChild
43+
val operation = firstChild?.text?.uppercase() ?: continue
44+
45+
if (DANGEROUS_OPERATIONS.any { operation.contains(it) }) {
46+
return Pair(true, "Dangerous operation detected: $operation. Please confirm this operation.")
47+
}
48+
}
49+
50+
return Pair(false, "")
51+
}
52+
2853
fun executeSqlQuery(project: Project, sql: String): String {
2954
val file = LightVirtualFile("temp.sql", sql)
3055
val psiFile = runReadAction { PsiManager.getInstance(project).findFile(file) }
3156
?: return "ShireError[Database]: Can't find PSI file"
3257

58+
val (isDangerous, reason) = checkSqlSafety(project, psiFile)
59+
if (isDangerous) {
60+
return "ShireError[Database]: $reason"
61+
}
62+
3363
val dataSource = DatabaseSchemaAssistant.allRawDatasource(project).firstOrNull()
3464
?: throw IllegalArgumentException("ShireError[Database]: No database found")
3565

0 commit comments

Comments
 (0)