This commit is contained in:
Liang Ding 2023-05-04 10:11:29 +08:00
parent 8ae3b354ab
commit a48154ba84
No known key found for this signature in database
GPG key ID: 136F30F901A2231D
3 changed files with 91 additions and 8 deletions

View file

@ -19,6 +19,7 @@ package sql
import (
"bytes"
"database/sql"
"math"
"sort"
"strconv"
"strings"
@ -378,7 +379,45 @@ func QueryBookmarkLabels() (ret []string) {
return
}
func Query(stmt string) (ret []map[string]interface{}, err error) {
func QueryNoLimit(stmt string) (ret []map[string]interface{}, err error) {
return queryRawStmt(stmt, math.MaxInt)
}
func Query(stmt string, limit int) (ret []map[string]interface{}, err error) {
parsedStmt, err := sqlparser.Parse(stmt)
if nil != err {
return queryRawStmt(stmt, limit)
}
switch parsedStmt.(type) {
case *sqlparser.Select:
slct := parsedStmt.(*sqlparser.Select)
if nil == slct.Limit {
slct.Limit = &sqlparser.Limit{
Rowcount: &sqlparser.SQLVal{
Type: sqlparser.IntVal,
Val: []byte(strconv.Itoa(limit)),
},
}
} else {
if nil != slct.Limit.Rowcount && 0 < len(slct.Limit.Rowcount.(*sqlparser.SQLVal).Val) {
limit, _ = strconv.Atoi(string(slct.Limit.Rowcount.(*sqlparser.SQLVal).Val))
if 0 >= limit {
limit = 32
}
}
slct.Limit.Rowcount = &sqlparser.SQLVal{
Type: sqlparser.IntVal,
Val: []byte(strconv.Itoa(limit)),
}
}
stmt = sqlparser.String(slct)
default:
return
}
ret = []map[string]interface{}{}
rows, err := query(stmt)
if nil != err {
@ -413,6 +452,49 @@ func Query(stmt string) (ret []map[string]interface{}, err error) {
return
}
func queryRawStmt(stmt string, limit int) (ret []map[string]interface{}, err error) {
rows, err := query(stmt)
if nil != err {
if strings.Contains(err.Error(), "syntax error") {
return
}
return
}
defer rows.Close()
cols, err := rows.Columns()
if nil != err || nil == cols {
return
}
noLimit := !strings.Contains(strings.ToLower(stmt), " limit ")
var count, errCount int
for rows.Next() {
columns := make([]interface{}, len(cols))
columnPointers := make([]interface{}, len(cols))
for i := range columns {
columnPointers[i] = &columns[i]
}
if err = rows.Scan(columnPointers...); nil != err {
return
}
m := make(map[string]interface{})
for i, colName := range cols {
val := columnPointers[i].(*interface{})
m[colName] = *val
}
ret = append(ret, m)
count++
if (noLimit && limit < count) || 0 < errCount {
break
}
}
return
}
func SelectBlocksRawStmtNoParse(stmt string, limit int) (ret []*Block) {
return selectBlocksRawStmt(stmt, limit)
}
@ -491,7 +573,7 @@ func selectBlocksRawStmt(stmt string, limit int) (ret []*Block) {
}
defer rows.Close()
confLimit := !strings.Contains(strings.ToLower(stmt), " limit ")
noLimit := !strings.Contains(strings.ToLower(stmt), " limit ")
var count, errCount int
for rows.Next() {
count++
@ -502,7 +584,7 @@ func selectBlocksRawStmt(stmt string, limit int) (ret []*Block) {
errCount++
}
if (confLimit && limit < count) || 0 < errCount {
if (noLimit && limit < count) || 0 < errCount {
break
}
}