技术分享 | 使用 TiDB 的 SQL 解析器生成 SQL 指纹

什么是 SQL指纹

SQL 指纹指将一条 SQL 中的字面值替换成其他固定符号。可以用来做 SQL 脱敏或者 SQL 归类。例如：

select * from t1 where id = 100；

转换成：

select * from t1 where id = ?；

pt-fingerprint 的实现

从 pt-fingerprint 的代码实现看，它主要是通过正则匹配 SQL 字符串来替换对应字符。代码有 2 千多行，完全通过字符串解析会使得代码及其复杂而难以阅读，好处是无需关心 SQL 语义。

基于 TiDB SQL parser 的实现

TiDB SQL parser 的功能是把 SQL 语句按照 SQL 语法规则进行解析，将文本转换成抽象语法树，另外 TiDB SQL parser 支持将语法树转换成 SQL 文本，因此可以通过修改语法树结构达到修改 SQL 文本的目的。

1. 通过 TiDB SQL 解析器将 SQL 解析成语法树

解析出的语法树大致如下，其中"..." 代表之前存在多级。

&ast.SelectStmt {
    Fields:
        ... &ast.WildCard
    From: 
        ... &ast.TableName
            ... "t1"
    Where: &ast.BinaryOperationExpr
        L: &ast.ColumnNameExpr
            ... "id"
        R：&ast.ValueExpr
            ... 100
}                

2. 修改语法树上节点对应的值

TiDB 语法解析器代码实现了一套访问者的设计模式，可以通过实现一个Visitor 来遍历语法树。按照1中的语法树结构，我们只需要在遍历到ast.ValueExpr对象时将他的具体数值替换成?

Visitor 接口：

// Visitor visits a Node.
type Visitor interface {
 Enter(n Node) (node Node, skipChildren bool)
 Leave(n Node) (node Node, ok bool)
}

实现 Visitor 接口：

//此处省略N行代码

// 定义一个 FingerprintVisitor 使其实现 Visitor 接口
type FingerprintVisitor struct{}

func (f *FingerprintVisitor) Enter(n ast.Node) (node ast.Node, skipChildren bool) {
    // 当访问到ValueExpr 时，只需要将ValueExpr的值替换掉就行
 if v, ok := n.(*driver.ValueExpr); ok {
  v.Type.Charset = ""
  v.SetValue([]byte("?"))
 }
 return n, false
}

func (f *FingerprintVisitor) Leave(n ast.Node) (node ast.Node, ok bool) {
 return n, true
}

3. 将语法树还原成 SQL

TiDB SQL parser 从 v3 版本开始提供接口Restore(ctx *RestoreCtx) error 支持将语法树转化成 SQL 文本

完整代码

package main

import (
 "bytes"
 "fmt"

 "github.com/pingcap/parser"
 "github.com/pingcap/parser/ast"
 "github.com/pingcap/parser/format"
 driver "github.com/pingcap/tidb/types/parser_driver"
)

// 定义一个 FingerprintVisitor 使其实现 Visitor 接口
type FingerprintVisitor struct{}

func (f *FingerprintVisitor) Enter(n ast.Node) (node ast.Node, skipChildren bool) {
 // 当访问到ValueExpr 时，只需要将ValueExpr的值替换掉就行
 if v, ok := n.(*driver.ValueExpr); ok {
  v.Type.Charset = ""
  v.SetValue([]byte("?"))
 }
 return n, false
}

func (f *FingerprintVisitor) Leave(n ast.Node) (node ast.Node, ok bool) {
 return n, true
}

func main() {
 sql := "select * from t1 where id = 100;"
 p := parser.New()
 stmt, err := p.ParseOneStmt(sql, "", "")
 if err != nil {
  // 省略错误处理
  return
 }
 stmt.Accept(&FingerprintVisitor{})

 buf := new(bytes.Buffer)
 restoreCtx := format.NewRestoreCtx(format.RestoreKeyWordUppercase|format.RestoreNameBackQuotes, buf)
 err = stmt.Restore(restoreCtx)
 if nil != err {
  // 省略错误处理
  return
 }
 fmt.Println(buf.String())
 // SELECT * FROM `t1` WHERE `id`=?
}

总结

1. 使用 TiDB SQL parser 可以快速准确的实现 SQL 指纹，相比字符串解析降低了阅读的复杂度；
2. 额外的你需要花时间了解 TiDB 语法树的结构。