diff --git a/internal/app/methods_file.go b/internal/app/methods_file.go index e5bebf3..02e13e4 100644 --- a/internal/app/methods_file.go +++ b/internal/app/methods_file.go @@ -1700,8 +1700,36 @@ func exportInferObjectName(row map[string]interface{}) string { return "" } +func trimLeadingSQLComments(sql string) string { + trimmed := strings.TrimSpace(sql) + for trimmed != "" { + switch { + case strings.HasPrefix(trimmed, "--"): + if newline := strings.IndexByte(trimmed, '\n'); newline >= 0 { + trimmed = strings.TrimSpace(trimmed[newline+1:]) + continue + } + return "" + case strings.HasPrefix(trimmed, "#"): + if newline := strings.IndexByte(trimmed, '\n'); newline >= 0 { + trimmed = strings.TrimSpace(trimmed[newline+1:]) + continue + } + return "" + case strings.HasPrefix(trimmed, "/*"): + if end := strings.Index(trimmed, "*/"); end >= 0 { + trimmed = strings.TrimSpace(trimmed[end+2:]) + continue + } + return "" + } + break + } + return trimmed +} + func looksLikeSelectOrWith(sql string) bool { - trimmed := strings.TrimSpace(strings.TrimSuffix(sql, ";")) + trimmed := trimLeadingSQLComments(strings.TrimSuffix(sql, ";")) if trimmed == "" { return false } @@ -1947,8 +1975,7 @@ func (a *App) ExportQuery(config connection.ConnectionConfig, dbName string, que } query = sanitizeSQLForPgLike(runConfig.Type, query) - lowerQuery := strings.ToLower(strings.TrimSpace(query)) - if !(strings.HasPrefix(lowerQuery, "select") || strings.HasPrefix(lowerQuery, "with")) { + if !looksLikeSelectOrWith(query) { return connection.QueryResult{Success: false, Message: "仅支持 SELECT/WITH 查询导出"} } diff --git a/internal/app/methods_file_export_test.go b/internal/app/methods_file_export_test.go index f92bac7..7bd4c73 100644 --- a/internal/app/methods_file_export_test.go +++ b/internal/app/methods_file_export_test.go @@ -267,6 +267,20 @@ func TestGetExportQueryTimeout_CustomClickHouseUsesLongerMinimum(t *testing.T) { } } +func TestLooksLikeSelectOrWith_AllowsInnerJoinQueryAfterLeadingComments(t *testing.T) { + query := ` +-- query result export +/* generated by query editor */ +SELECT o.id, c.name +FROM orders o +INNER JOIN customers c ON c.id = o.customer_id +` + + if !looksLikeSelectOrWith(query) { + t.Fatalf("带前置注释的 INNER JOIN 查询应允许导出,query=%q", query) + } +} + func TestWriteRowsToFile_HTML_EscapeAndStyle(t *testing.T) { f, err := os.CreateTemp("", "gonavi-export-*.html") if err != nil {