使用 vitess-sqlparser 给 sql 语句添加 limit

背景

在写代码的时候遇到了这么一个需求。前端传过来的 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