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

作者：孙健

爱可生研发工程师，负责高可用组建和 SQL 审核相关开发。

本文来源：原创投稿

*爱可生开源社区出品，原创内容未经授权不得随意使用，转载请联系小编并注明来源。

本文主要介绍如何借助 TiDB SQL 解析自定义生成 SQL 指纹，采用了一种有别于 pt-fingerprint(https://www.percona.com/doc/percona-toolkit/3.0/pt-fingerprint.html) 的方式。

什么是 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语法树的结构。