背景
在写代码的时候遇到了这么一个需求。前端传过来的 sql 语句要先判断有没有 limit,如果没有就在后面加上一个默认的 limit,防止返回的数据量太大。例如
select * from student
// 转换为
select * from student limit 10
方案1-strings.Index
最开始直接用的是 strings 的搜索,先匹配 select,再匹配 limit。如果没匹配到 limit 就在 sql 语句后面 append 一个。
func SetDefaultLimit(sql string, limit int64) string {
// sql 转为小写
sql = strings.ToLower(sql)
// 如果是非查询语句则直接跳过
if strings.Index(sql, "select") < 0 {
return sql
}
// 如果包含limit 则直接跳过
if strings.Index(sql, "limit") >= 0 {
return sql
}
// 加 limit
if strings.HasSuffix(sql, ";") {
sql = strings.Replace(sql, ";", fmt.Sprintf(" limit %d", limit), -1)
} else {
sql = sql + fmt.Sprintf(" limit %d ;", limit)
}
return sql
}
这种做法的缺点是,在最开始将 sql 语句转换为小写,这对于数据库表名大小写敏感的情况很不友好,会造成 sql 语句执行失败。
如果转换为小写,那么情况就会变为很多,LIMIT 、limit 大写的,小写的,大小写混合的情况多到无法判断。
为了解决这个问题,想到了可不可以先转换成 ast 树,然后在 ast 上面加 limit,再转换回 sql 语句,所以有了方案2。
方案2-vitess-sqlparser
“vitess.io/vitess/go/vt/sqlparser”
这个库的原页面已经打不开了,可以在这看到 api
sqlparser package - vitess.io/vitess/go/vt/sqlparser - Go Packages
这个库是支持 sql 解析的一个库,里面功能也比较全。
// SetDefaultLimit 设置 sql 语句默认的 limit ,防止数据量过大
func SetDefaultLimit(ctx context.Context, sql string, limit int64) (string, error) {
stmt, err := sqlparser.Parse(sql)
if err != nil {
return err
}
switch t := stmt.(type) {
case *sqlparser.Select:
if t.Limit == nil {
t.SetLimit(&sqlparser.Limit{
Offset: nil,
Rowcount: &sqlparser.Literal{
Type: sqlparser.IntVal,
Val: strconv.FormatInt(limit, 10),
},
})
return sqlparser.String(t), nil
}
default:
return sql, nil
}
return sql, nil
}
这种通过解析 sql 来加 limit 的参数就规避了使用 strings 的那种情况。
额外
这个库先解析 sql 语句为 ast,所以有很多其他的用途,可以做一些 sql 语句的校验,比如建表语句有没有主键,添加的索引是否超过了4列等等。
并且还可以做一些 sql 优化的操作,比如对SELECT补全*指代的列名,同一列的 OR 过滤条件使用 IN() 替代,如果值有相等的会进行合并等等。详见 XiaoMi/soar