十分鐘成為 TiDB Contributor 系列 | 添加內建函數

01-27

最近我們對 TiDB 代碼做了些改進，大幅度簡化了添加內建函數的流程，這篇教程為大家分享如何為 TiDB 新增 builtin 函數。首先介紹一些必需的背景知識，然後介紹增加 builtin 函數的流程，最後會以一個函數作為示例。（作者：申礫）

背景知識

SQL 語句發送到 TiDB 後首先會經過 parser，從文本 parse 成為 AST（抽象語法樹），通過 Query Optimizer 生成執行計劃，得到一個可以執行的 plan，通過執行這個 plan 即可得到結果，這期間會涉及到如何獲取 table 中的數據，如何對數據進行過濾、計算、排序、聚合、濾重以及如何對表達式進行求值。

對於一個 builtin 函數，比較重要的是進行語法解析以及如何求值。其中語法解析部分需要了解如何寫 yacc 以及如何修改 TiDB 的詞法解析器，較為繁瑣，我們已經將這部分工作提前做好，大多數 builtin 函數的語法解析工作已經做完。

對 builtin 函數的求值需要在 TiDB 的表達式求值框架下完成，每個 builtin 函數被認為是一個表達式，用一個 ScalarFunction 來表示，每個 builtin 函數通過其函數名以及參數，獲取對應的函數類型以及函數簽名，然後通過函數簽名進行求值。

總體而言，上述流程對於不熟悉 TiDB 的朋友而言比較複雜，我們對這部分做了些工作，將一些流程性、較為繁瑣的工作做了統一處理，目前已經將大多數未實現的 buitlin 函數的語法解析以及尋找函數簽名的工作完成，但是函數實現部分留空。換句話說，只要找到留空的函數實現，將其補充完整，即可作為一個 PR。

添加 builtin 函數整體流程

找到未實現的函數
在 TiDB 源碼中的 expression 目錄下搜索 errFunctionNotExists，即可找到所有未實現的函數，從中選擇一個感興趣的函數，比如 SHA2 函數：

func (b *builtinSHA2Sig) eval(row []types.Datum) (d types.Datum, err error) {n return d, errFunctionNotExists.GenByArgs("SHA2")n}n

實現函數簽名接下來要做的事情就是實現 eval 方法，函數的功能請參考 MySQL 文檔，具體的實現方法可以參考目前已經實現函數。
在 typeinferer 中添加類型推導信息
在 plan/typeinferer.go 中的 handleFuncCallExpr() 裡面添加這個函數的返回結果類型，請保持和 MySQL 的結果一致。全部類型定義參見 MySQL Const。

* 注意大多數函數除了需要填寫返回值類型之外，還需要獲取返回值的長度。n

寫單元測試在 expression 目錄下，為函數的實現增加單元測試，同時也要在 plan/typeinferer_test.go 文件中添加 typeinferer 的單元測試
運行 make dev，確保所有的 test case 都能跑過

示例

這裡以新增 SHA1() 函數的 PR 為例，進行詳細說明

首先看 expression/builtin_encryption.go：

將 SHA1() 的求值方法補充完整

func (b *builtinSHA1Sig) eval(row []types.Datum) (d types.Datum, err error) {n // 首先對參數進行求值，這塊一般不用修改n args, err := b.evalArgs(row)n if err != nil {n return types.Datum{}, errors.Trace(err)n }n // 每個參數的意義請參考 MySQL 文檔n // SHA/SHA1 function only accept 1 parametern arg := args[0]n if arg.IsNull() {n return d, niln }n // 這裡對參數值做了一個類型轉換，函數的實現請參考 util/types/datum.gon bin, err := arg.ToBytes()n if err != nil {n return d, errors.Trace(err)n }n hasher := sha1.New()n hasher.Write(bin)n data := fmt.Sprintf("%x", hasher.Sum(nil))n // 設置返回值n d.SetString(data)n return d, niln}n

接下來給函數實現添加單元測試，參見 expression/builtin_encryption_test.go：

var shaCases = []struct {n origin interface{}n crypt stringn }{n {"test", "a94a8fe5ccb19ba61c4c0873d391e987982fbbd3"},n {"c4pt0r", "034923dcabf099fc4c8917c0ab91ffcd4c2578a6"},n {"pingcap", "73bf9ef43a44f42e2ea2894d62f0917af149a006"},n {"foobar", "8843d7f92416211de9ebb963ff4ce28125932878"},n {1024, "128351137a9c47206c4507dcf2e6fbeeca3a9079"},n {123.45, "22f8b438ad7e89300b51d88684f3f0b9fa1d7a32"},n }n n func (s *testEvaluatorSuite) TestShaEncrypt(c *C) {n defer testleak.AfterTest(c)() // 監測 goroutine 泄漏的工具，可以直接照搬n fc := funcs[ast.SHA]n for _, test := range shaCases {n in := types.NewDatum(test.origin)n f, _ := fc.getFunction(datumsToConstants([]types.Datum{in}), s.ctx)n crypt, err := f.eval(nil)n c.Assert(err, IsNil)n res, err := crypt.ToString()n c.Assert(err, IsNil)n c.Assert(res, Equals, test.crypt)n }n // test NULL input for shan var argNull types.Datumn f, _ := fc.getFunction(datumsToConstants([]types.Datum{argNull}), s.ctx)n crypt, err := f.eval(nil)n c.Assert(err, IsNil)n c.Assert(crypt.IsNull(), IsTrue)n}n* 注意，除了正常 case 之外，最好能添加一些異常的case，如輸入值為 nil，或者是多種類型的參數n

最後還需要添加類型推導信息以及 test case，參見 plan/typeinferer.go，plan/typeinferer_test.go：

case ast.SHA, ast.SHA1:n tp = types.NewFieldType(mysql.TypeVarString)n chs = v.defaultCharsetn tp.Flen = 40n

{`sha1(123)`, mysql.TypeVarString, "utf8"},n {`sha(123)`, mysql.TypeVarString, "utf8"},n

編輯按：添加 TiDB Robot 微信（ID：tidbai），加入 TiDB Contributor Club，無門檻參與開源項目，改變世界從這裡開始吧（萌萌噠）。

聽說現在成為 TiDB Contributor 就送限量版馬克杯呦~