使用 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

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值