fix(query-execution): 支持带前置注释的读查询结果识别

This commit is contained in:
tianqijiuyun-latiao
2026-03-06 16:05:34 +08:00
parent 0682baa14d
commit 0a0609c459
2 changed files with 62 additions and 11 deletions

View File

@@ -416,12 +416,7 @@ func (a *App) DBQueryWithCancel(config connection.ConnectionConfig, dbName strin
a.queryMu.Unlock()
}()
lowerQuery := strings.TrimSpace(strings.ToLower(query))
isReadQuery := strings.HasPrefix(lowerQuery, "select") || strings.HasPrefix(lowerQuery, "show") || strings.HasPrefix(lowerQuery, "describe") || strings.HasPrefix(lowerQuery, "explain")
// MongoDB JSON 命令中的 find/count/aggregate 也属于读查询
if !isReadQuery && strings.ToLower(strings.TrimSpace(runConfig.Type)) == "mongodb" && strings.HasPrefix(strings.TrimSpace(query), "{") {
isReadQuery = true
}
isReadQuery := isReadOnlySQLQuery(runConfig.Type, query)
runReadQuery := func(inst db.Database) ([]map[string]interface{}, []string, error) {
if q, ok := inst.(interface {
@@ -500,11 +495,7 @@ func (a *App) DBQueryIsolated(config connection.ConnectionConfig, dbName string,
ctx, cancel := utils.ContextWithTimeout(time.Duration(timeoutSeconds) * time.Second)
defer cancel()
lowerQuery := strings.TrimSpace(strings.ToLower(query))
isReadQuery := strings.HasPrefix(lowerQuery, "select") || strings.HasPrefix(lowerQuery, "show") || strings.HasPrefix(lowerQuery, "describe") || strings.HasPrefix(lowerQuery, "explain")
if !isReadQuery && strings.ToLower(strings.TrimSpace(runConfig.Type)) == "mongodb" && strings.HasPrefix(strings.TrimSpace(query), "{") {
isReadQuery = true
}
isReadQuery := isReadOnlySQLQuery(runConfig.Type, query)
if isReadQuery {
var data []map[string]interface{}

View File

@@ -5,6 +5,66 @@ import (
"unicode"
)
func leadingSQLKeyword(query string) string {
text := strings.TrimSpace(query)
for len(text) > 0 {
trimmed := strings.TrimLeft(text, " \t\r\n")
if trimmed == "" {
return ""
}
text = trimmed
switch {
case strings.HasPrefix(text, "--"):
if idx := strings.IndexByte(text, '\n'); idx >= 0 {
text = text[idx+1:]
continue
}
return ""
case strings.HasPrefix(text, "#"):
if idx := strings.IndexByte(text, '\n'); idx >= 0 {
text = text[idx+1:]
continue
}
return ""
case strings.HasPrefix(text, "/*"):
if idx := strings.Index(text, "*/"); idx >= 0 {
text = text[idx+2:]
continue
}
return ""
}
break
}
if text == "" {
return ""
}
for i, r := range text {
if unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' {
continue
}
if i == 0 {
return ""
}
return strings.ToLower(text[:i])
}
return strings.ToLower(text)
}
func isReadOnlySQLQuery(dbType string, query string) bool {
if strings.ToLower(strings.TrimSpace(dbType)) == "mongodb" && strings.HasPrefix(strings.TrimSpace(query), "{") {
return true
}
switch leadingSQLKeyword(query) {
case "select", "with", "show", "describe", "desc", "explain", "pragma", "values":
return true
default:
return false
}
}
func sanitizeSQLForPgLike(dbType string, query string) string {
switch strings.ToLower(strings.TrimSpace(dbType)) {
case "postgres", "kingbase", "highgo", "vastbase":