十分鐘成爲 Contributor 系列 | 支持 AST 還原爲 SQL

做者：趙一霖

背景知識

SQL 語句發送到 TiDB 後首先會通過 parser，從文本 parse 成爲 AST（抽象語法樹），AST 節點與 SQL 文本結構是一一對應的，咱們經過遍歷整個 AST 樹就能夠拼接出一個與 AST 語義相同的 SQL 文本。

對 parser 不熟悉的小夥伴們能夠看 TiDB 源碼閱讀系列文章（五）TiDB SQL Parser 的實現。

爲了控制 SQL 文本的輸出格式，而且爲方便將來新功能的加入（例如在 SQL 文本中用 「*」 替代密碼），咱們引入了 RestoreFlags 並封裝了 RestoreCtx 結構（相關源碼）：

// `RestoreFlags` 中的互斥組:
// [RestoreStringSingleQuotes, RestoreStringDoubleQuotes]
// [RestoreKeyWordUppercase, RestoreKeyWordLowercase]
// [RestoreNameUppercase, RestoreNameLowercase]
// [RestoreNameDoubleQuotes, RestoreNameBackQuotes]
// 靠前的 flag 擁有更高的優先級。
const (
    RestoreStringSingleQuotes RestoreFlags = 1 << iota
    
    ...
)

// RestoreCtx is `Restore` context to hold flags and writer.
type RestoreCtx struct {
    Flags RestoreFlags
    In    io.Writer
}

// WriteKeyWord 用於向 `ctx` 中寫入關鍵字（例如：SELECT）。
// 它的大小寫受 `RestoreKeyWordUppercase`，`RestoreKeyWordLowercase` 控制
func (ctx *RestoreCtx) WriteKeyWord(keyWord string) {
    ...
}

// WriteString 用於向 `ctx` 中寫入字符串。
// 它是否被引號包裹及轉義規則受 `RestoreStringSingleQuotes`，`RestoreStringDoubleQuotes`，`RestoreStringEscapeBackslash` 控制。
func (ctx *RestoreCtx) WriteString(str string) {
    ...
}

// WriteName 用於向 `ctx` 中寫入名稱（庫名，表名，列名等）。
// 它是否被引號包裹及轉義規則受 `RestoreNameUppercase`，`RestoreNameLowercase`，`RestoreNameDoubleQuotes`，`RestoreNameBackQuotes` 控制。
func (ctx *RestoreCtx) WriteName(name string) {
    ...
}

// WriteName 用於向 `ctx` 中寫入普通文本。
// 它將被直接寫入不受 flag 影響。
func (ctx *RestoreCtx) WritePlain(plainText string) {
    ...
}

// WriteName 用於向 `ctx` 中寫入普通文本。
// 它將被直接寫入不受 flag 影響。
func (ctx *RestoreCtx) WritePlainf(format string, a ...interface{}) {
    ...
}

咱們在 ast.Node 接口中添加了一個 Restore(ctx *RestoreCtx) error 函數，這個函數將當前節點對應的 SQL 文本追加至參數 ctx 中，若是節點無效則返回 error。

type Node interface {
    // Restore AST to SQL text and append them to `ctx`.
    // return error when the AST is invalid.
    Restore(ctx *RestoreCtx) error
    
    ...
}

以 SQL 語句 SELECT column0 FROM table0 UNION SELECT column1 FROM table1 WHERE a = 1 爲例，以下圖所示，咱們經過遍歷整個 AST 樹，遞歸調用每一個節點的 Restore() 方法，便可拼接成一個完整的 SQL 文本。

值得注意的是，SQL 文本與 AST 是一個多對一的關係，咱們不可能從 AST 結構中還原出與原 SQL 徹底一致的文本，
所以咱們只要保證還原出的 SQL 文本與原 SQL 語義相同 便可。所謂語義相同，指的是由 AST 還原出的 SQL 文本再被解析爲 AST 後，兩個 AST 是相等的。

咱們已經完成了接口設計和測試框架，具體的Restore() 函數留空。所以只須要選擇一個留空的 Restore() 函數實現，並添加相應的測試數據，就能夠提交一個 PR 了！

實現 Restore() 函數的總體流程

1. 請先閱讀 Proposal、Issue

2. 在 Issue 中找到未實現的函數

   1. 在 Issue-pingcap/tidb#8532 中找到一個沒有被其餘貢獻者認領的任務，例如 ast/expressions.go: BetweenExpr。
   2. 在 pingcap/parser 中找到任務對應文件 ast/expressions.go。
   3. 在文件中找到 BetweenExpr 結構的 Restore 函數：

   // Restore implements Node interface.
   func (n *BetweenExpr) Restore(ctx *RestoreCtx) error {
       return errors.New("Not implemented")
   }

3. 實現 Restore() 函數

   根據 Node 節點結構和 SQL 語法實現函數功能。

   參考 MySQL 5.7 SQL Statement Syntax

4. 寫單元測試

   參考示例在相關文件下添加單元測試。

5. 運行 make test，確保全部的 test case 都能跑過。
6. 提交 PR

   PR 標題統一爲：parser: implement Restore for XXX
   請在 PR 中關聯 Issue: pingcap/tidb#8532

示例

這裏以實現 BetweenExpr 的 Restore 函數 PR 爲例，進行詳細說明：

1. 首先看 ast/expressions.go：

   1. 咱們要實現一個 ast.Node 結構的 Restore 函數，首先清楚該結構表明什麼短語，例如 BetweenExpr 表明 expr [NOT] BETWEEN expr AND expr （參見：MySQL 語法 - 比較函數和運算符）。
   2. 觀察 BetweenExpr 結構：

   // BetweenExpr is for "between and" or "not between and" expression.
   type BetweenExpr struct {
       exprNode
       // 被檢查的表達式
       Expr ExprNode
       // AND 左側的表達式
       Left ExprNode
       // AND 右側的表達式
       Right ExprNode
       // 是否有 NOT 關鍵字
       Not bool
   }

3. 實現 `BetweenExpr` 的 `Restore` 函數：

```
// Restore implements Node interface.
func (n *BetweenExpr) Restore(ctx *RestoreCtx) error {
    // 調用 Expr 的 Restore，向 ctx 寫入 Expr
    if err := n.Expr.Restore(ctx); err != nil {
        return errors.Annotate(err, "An error occurred while restore BetweenExpr.Expr")
    }
    // 判斷是否有 NOT，並寫入相應關鍵字
    if n.Not {
        ctx.WriteKeyWord(" NOT BETWEEN ")
    } else {
        ctx.WriteKeyWord(" BETWEEN ")
    }
    // 調用 Left 的 Restore
    if err := n.Left.Restore(ctx); err != nil {
        return errors.Annotate(err, "An error occurred while restore BetweenExpr.Left")
    }
    // 寫入 AND 關鍵字
    ctx.WriteKeyWord(" AND ")
    // 調用 Right 的 Restore
    if err := n.Right.Restore(ctx); err != nil {
        return errors.Annotate(err, "An error occurred while restore BetweenExpr.Right ")
    }
    return nil
}
```

1. 接下來給函數實現添加單元測試, ast/expressions_test.go：

   // 添加測試函數
   func (tc *testExpressionsSuite) TestBetweenExprRestore(c *C) {
       // 測試用例
       testCases := []NodeRestoreTestCase{
           {"b between 1 and 2", "`b` BETWEEN 1 AND 2"},
           {"b not between 1 and 2", "`b` NOT BETWEEN 1 AND 2"},
           {"b between a and b", "`b` BETWEEN `a` AND `b`"},
           {"b between '' and 'b'", "`b` BETWEEN '' AND 'b'"},
           {"b between '2018-11-01' and '2018-11-02'", "`b` BETWEEN '2018-11-01' AND '2018-11-02'"},
       }
       // 爲了避免依賴父節點實現，經過 extractNodeFunc 抽取待測節點
       extractNodeFunc := func(node Node) Node {
           return node.(*SelectStmt).Fields.Fields[0].Expr
       }
       // Run Test
       RunNodeRestoreTest(c, testCases, "select %s", extractNodeFunc)
   }

   至此 BetweenExpr 的 Restore 函數實現完成，能夠提交 PR 了。爲了更好的理解測試邏輯，下面咱們看 RunNodeRestoreTest：

   // 下面是測試邏輯，已經實現好了，不須要 contributor 實現
   func RunNodeRestoreTest(c *C, nodeTestCases []NodeRestoreTestCase, template string, extractNodeFunc func(node Node) Node) {
       parser := parser.New()
       for _, testCase := range nodeTestCases {
           // 經過 template 將測試用例拼接爲完整的 SQL
           sourceSQL := fmt.Sprintf(template, testCase.sourceSQL)
           expectSQL := fmt.Sprintf(template, testCase.expectSQL)
           stmt, err := parser.ParseOneStmt(sourceSQL, "", "")
           comment := Commentf("source %#v", testCase)
           c.Assert(err, IsNil, comment)
           var sb strings.Builder
           // 抽取指定節點並調用其 Restore 函數
           err = extractNodeFunc(stmt).Restore(NewRestoreCtx(DefaultRestoreFlags, &sb))
           c.Assert(err, IsNil, comment)
           // 經過 template 將 restore 結果拼接爲完整的 SQL
           restoreSql := fmt.Sprintf(template, sb.String())
           comment = Commentf("source %#v; restore %v", testCase, restoreSql)
           // 測試 restore 結果與預期一致
           c.Assert(restoreSql, Equals, expectSQL, comment)
           stmt2, err := parser.ParseOneStmt(restoreSql, "", "")
           c.Assert(err, IsNil, comment)
           CleanNodeText(stmt)
           CleanNodeText(stmt2)
           // 測試解析的 stmt 與原 stmt 一致
           c.Assert(stmt2, DeepEquals, stmt, comment)
       }
   }

**不過對於 ast.StmtNode（例如：ast.SelectStmt）測試方法有些不同，
因爲這類節點能夠還原爲一個完整的 SQL，所以直接在 parser_test.go 中測試。**

下面以實現 UseStmt 的 Restore 函數 PR 爲例，對測試進行說明：

1. Restore 函數實現過程略。

2. 給函數實現添加單元測試，參見 parser_test.go：

   在這個示例中，只添加了幾行測試數據就完成了測試：

   // 添加 testCase 結構的測試數據
   {"use `select`", true, "USE `select`"},
   {"use `sel``ect`", true, "USE `sel``ect`"},
   {"use select", false, "USE `select`"},

   咱們看 testCase 結構聲明：

   type testCase struct {
       // 原 SQL
       src     string
       // 是否能被正確 parse
       ok      bool
       // 預期的 restore SQL
       restore string
   }

   測試代碼會判斷原 SQL parse 出 AST 後再還原的 SQL 是否與預期的 restore SQL 相等，具體的測試邏輯在 parser_test.go 中 RunTest()、RunRestoreTest() 函數，邏輯與前例相似，此處再也不贅述。

---

加入 TiDB Contributor Club，無門檻參與開源項目，改變世界從這裏開始吧（萌萌噠）。