golang使用TiDB的SQL Parse对SQL指纹化

复制代码
//mkdir sqlaprse && cd sqlaprse
//export GOPROXY=https://proxy.golang.com.cn,direct
//go mod init sqlaprse && touch main.go
//go get -v github.com/pingcap/tidb/parser@40b72e7
//go get github.com/fatih/color
//go get github.com/pingcap/tidb/types/parser_driver@40b72e7
//go run main.go -sql="select * from t1 where id=1 and date>'2023-08-18'"
package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/fatih/color"
	"github.com/pingcap/tidb/parser"
	"github.com/pingcap/tidb/parser/ast"
	"github.com/pingcap/tidb/parser/format"
	driver "github.com/pingcap/tidb/types/parser_driver"
	"sync" // 导入 sync 包用于 WaitGroup
)

type FingerprintVisitor struct{}

func (f *FingerprintVisitor) Enter(n ast.Node) (node ast.Node, skipChildren bool) {
	if v, ok := n.(*driver.ValueExpr); ok {
		v.SetValue([]byte("?"))
	}
	return n, false
}

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

var (
	Success *color.Color
	Err     *color.Color
	sql     string
)

func processSQL(sql string, wg *sync.WaitGroup) {
	defer wg.Done() // 通知 WaitGroup 此 goroutine 已完成
	p := parser.New()
	stmt, err := p.ParseOneStmt(sql, "", "")
	if err != nil {
		Err.Println("解析错误:" + err.Error())
		return
	}
	stmt.Accept(&FingerprintVisitor{})

	buf := new(bytes.Buffer)
	restoreCtx := format.NewRestoreCtx(format.RestoreKeyWordUppercase|format.RestoreNameBackQuotes, buf)
	err = stmt.Restore(restoreCtx)
	if nil != err {
		Err.Println("解析错误:" + err.Error())
		return
	}
	Success.Println(buf.String())
}

func main() {
	Success = color.New(color.FgHiGreen, color.Bold, color.Underline)
	Err = color.New(color.FgHiRed, color.Bold)
	flag.StringVar(&sql, "sql", "", "SQL语句")
	flag.Parse()

	// 分割 SQL 语句
	statements := splitSQLStatements(sql)

	// 创建 WaitGroup 以等待所有 goroutine 完成
	var wg sync.WaitGroup

	for _, stmt := range statements {
		wg.Add(1) // 为每个 goroutine 增加 WaitGroup 计数
		go processSQL(stmt, &wg)
	}

	// 等待所有 goroutine 完成
	wg.Wait()
}

// 用于分割 SQL 语句的辅助函数
func splitSQLStatements(sql string) []string {
	// 在此实现分割 SQL 语句的逻辑
	// 例如，如果每个语句位于单独的行上并以分号结尾，则可以按分号分割
	// 并根据需要修改此函数。这是为了演示目的的简化示例。
	// 您可能需要处理边缘情况并改进分割逻辑。
	statements := []string{sql}
	return statements
}