前言

其實，我們一直都很想，基於TiDB做一些很cool，很hacker的事情。比如我們團隊小夥伴發了一篇關於TiDB for Pg兼容Gitlab的一篇文章，具體文章可以參考鏈接：

TiDB4PG之兼容Gitlab - 知乎 (zhihu.com)

這篇文章我就來簡單聊聊實現兼容到Gitlab的艱苦過程。

我們採用了一個相對較笨的方式，將Gitlab的源碼通過編譯啓動的方式，連接到最開始的TiDB for PG，這樣肯定是報錯，不行的，畢竟很多東西沒有兼容。為了能快速實現兼容，我們決定採取抓包的方式，將Gitlab連接到TiDB For Pg的所執行的SQL語句找出來，進行粗略的分類整理，去看看有哪些SQL語句，去定製化的兼容開發。在兼容實現這些SQL語句時，難點之一，就有DML語句中的Returning關鍵字。

原理

在TiDB-Server裏面，執行一個SQL語句的流程，大致可以分為解析、編譯、優化、執行、返回結果這幾個階段。而實現一個關鍵字，同樣的需要在這幾個階段做一些文章。而對於Returning關鍵字而言，我們可以從DML語句中相對簡單的DELETE語句入手，所以接下來的改造過程，最終結果就是實現了DELETE RETURNING 句式。

改造實現過程

Parser

從SQL在TiDB中的流轉過程，邁入後續代碼的第一步，就是將客户端傳來的SQL語句解析為一個能夠被後續代碼認識的結構，也就是AST樹。而這一過程，主要就是在Parser這個模塊兒中實現。

在TiDB v5.0以前，Parser這個包是有一個專門的代碼倉庫的，通過go mod的方式導入到TiDB，而在5.0之後，TiDB將Parser包挪到TiDB源碼當中。TiDB for PG 的源碼也是基於TiDB v4.0.14改造的，這次我想嘗試一下，在TiDB 最新的源碼中實現RETURNING關鍵字，一個是為了hackathon的比賽作準備，另一個也是為了之後TiDB for PG向着TiDB新版本靠攏試試水。

Paser模塊主要靠Lexer & Yacc這兩個模塊共同構成。在解析的過程中，Lexer組件會將SQL文本轉換為一個又一個token傳給parser，而parser中最為重要的parser.go文件，則是goyacc工具根據parser.y文件生成的，根據文件中語法的定義，來決定lexer中傳過來的token能夠與什麼語法規則進行匹配，最終輸出AST結構樹，也就是parser/ast 中定義的各類stmt，而我們要實現的就是dml.go中的DeleteStmt。

// DeleteStmt is a statement to delete rows from table.
// See https://dev.mysql.com/doc/refman/5.7/en/delete.html
type DeleteStmt struct {
    dmlNode

    // TableRefs is used in both single table and multiple table delete statement.
    TableRefs *TableRefsClause
    // Tables is only used in multiple table delete statement.
    Tables       *DeleteTableList
    Where        ExprNode
    Order        *OrderByClause
    Limit        *Limit
    Priority     mysql.PriorityEnum
    IgnoreErr    bool
    Quick        bool
    IsMultiTable bool
    BeforeFrom   bool
    // TableHints represents the table level Optimizer Hint for join type.
    TableHints []*TableOptimizerHint
    With       *WithClause
    // 我們今天的主題，Returning 關鍵字
    Returning  *ReturningClause
}

type ReturningClause struct {
    node
    Fields *FieldList
}

func (n *ReturningClause) Restore(ctx *format.RestoreCtx) error {
    ctx.WriteKeyWord("Returning ")
    for i, item := range n.Fields.Fields {
        if i != 0 {
            ctx.WritePlain(",")
        }
        if err := item.Restore(ctx); err != nil {
            return errors.Annotatef(err, "An error occurred while restore ReturningClause.Fields[%d]", i)
        }
    }
    return nil
}

func (n *ReturningClause) Accept(v Visitor) (Node, bool) {
    newNode, skipChildren := v.Enter(n)
    if skipChildren {
        return v.Leave(newNode)
    }
    n = newNode.(*ReturningClause)

    if n.Fields != nil {
        node, ok := n.Fields.Accept(v)
        if !ok {
            return n, false
        }
        n.Fields = node.(*FieldList)
    }

    return v.Leave(n)
}

原諒篇幅有限，不能把所有代碼貼出來。這裏值得提一嘴的就是Accept()方法。在ast包中，幾乎所有的stmt結構都實現了ast.Node接口，這個接口中的Accept()方法，主要作用就是處理AST，通過Visitor模式遍歷所有的節點，並且對AST結構做一個轉換。而為了能正常將RETURNING關鍵字轉換成DeleteStmt，我們還需要在parser中去將RETURNING 關鍵字註冊為token。

在parser.y中definitions區域定義好RETURNING相關句式的token，比如RETURNING關鍵字，還有ReturningClause、ReturningOption句式。

關於parser的一些基礎知識可以參考文章：

TiDB Parser模塊的簡單解讀與改造方法 - 知乎 (zhihu.com)

TiDB 源碼閲讀系列文章（五）TiDB SQL Parser 的實現 | PingCAP

在做完這些之後，我們就能夠在parser.y的rule部分中，找到DELETE 句式，加入returning句式了，也就是ReturningOptional，接着在其中寫上簡單的邏輯。

/*******************************************************************
 *
 *  Delete Statement
 *
 *******************************************************************/
DeleteWithoutUsingStmt:
    "DELETE" TableOptimizerHintsOpt PriorityOpt QuickOptional IgnoreOptional "FROM" TableName PartitionNameListOpt TableAsNameOpt IndexHintListOpt WhereClauseOptional OrderByOptional LimitClause ReturningOptional
    {
        ... 此處省略 ...
        if $14 != nil {
            x.Returning = $14.(*ast.ReturningClause)
        }

        $$ = x
    }
|    "DELETE" TableOptimizerHintsOpt PriorityOpt QuickOptional IgnoreOptional TableAliasRefList "FROM" TableRefs WhereClauseOptional ReturningOptional
    {
        ... 此處省略 ...
        if $10 != nil {
            x.Returning = $10.(*ast.ReturningClause)
        }
        $$ = x
    }

DeleteWithUsingStmt:
    "DELETE" TableOptimizerHintsOpt PriorityOpt QuickOptional IgnoreOptional "FROM" TableAliasRefList "USING" TableRefs WhereClauseOptional ReturningOptional
    {
        ... 此處省略 ...
        if $11 != nil {
            x.Returning = $11.(*ast.ReturningClause)
        }
        $$ = x
    }

ReturningClause:
    "RETURNING" SelectStmtFieldList
    {
        $$ = &ast.ReturningClause{Fields: $2.(*ast.FieldList)}
    }

ReturningOptional:
    {
        $$ = nil
    }
|    ReturningClause
    {
        $$ = $1
    }

接着就能利用parser/bin/goyacc 工具，根據最新的paser.y生成最終的parser.go，進入parser包中，運行make all即可。

需要注意的是，對於關鍵字，在生成最新的parser.go之後，我們還需要在parser/misc.go中定義，這是由於lexer採用了字典樹技術進行token識別，而其實現代碼就是在其中，不然lexer會不認識這所謂的關鍵字。

改完之後的驗證其實很簡單，在parser包中找到parser_test.go的測試文件，寫一個delete returning的句式，運行一遍測試，過了，那就是OK了。

還可以啓動tidb源碼，用mysql客户端連上去，執行一個delete returning的句式，能夠成功返回，那麼説明，這個關鍵字同樣是兼容成功的。

簡單總結

到了這一步，初步關鍵字兼容已經實現了，注意，現在還只是初步兼容，而要使其生效，則需要進入到接下來的Plan制定以及執行器Exexutor執行的部分了。這一部分在TiDB v5.0的改造還在研究的過程中，畢竟相對於TiDB v4.0.14的計劃制定、優化有些許變動，還沒來得及去研究，我會在後續文章中詳細闡述。最後給大家看看TiDB for PG的returning兼容成果吧。