Hackathon 實用指南丨快速給 TiDB 新增一個功能

語言: CN / TW / HK
時間 2022-10-11 06:15:32
主題:

TiDB Hackathon 2022 火熱報名中！你報名了嗎（還沒報名看這裡）？你有 idea 了嗎（沒有 idea 看這裡）？

有了 idea，但是不夠了解 TiDB，不知道如何動手實踐？本文將通過 step-by-step 的方式，介紹如何快速給 TiDB 新增一個功能，讓沒有太多知識背景的人也能快速上手。

ps：參加 TiDB 產品組的小夥伴，想給 TiDB 元件增加新功能的，快來圍觀！

假設我們想要將 SST 檔案匯入 TiDB 中，通過新增 LOAD SST FILE <file_path> 語法來實現。

TiDB 資料庫在收到一條 SQL 請求後，大概的執行流程是 生成 AST 語法樹 -> 生成執行計劃 -> 構造 Executor 並執行。我們先來實現語法。

語法實現

要如何實現語法呢？我們可以照葫蘆畫瓢，找一個類似的 LOAD DATA 語法作為葫蘆，然後開始畫瓢。

Step-1: 新增 AST 語法樹

LOAD DATA 語法是用 ast.LoadDataStmt 表示的，我們照葫蘆畫瓢在 tidb/parser/ast/dml.go 中新增一個 LoadSSTFileStmt AST 語法樹：

```sql // LoadSSTFileStmt is a statement to load sst file. type LoadSSTFileStmt struct { dmlNode

Path string }

// Restore implements Node interface. func (n LoadSSTFileStmt) Restore(ctx format.RestoreCtx) error { ctx.WriteKeyWord("LOAD SST FILE ") ctx.WriteString(n.Path) return nil }

// Accept implements Node Accept interface. func (n *LoadSSTFileStmt) Accept(v Visitor) (Node, bool) { newNode, _ := v.Enter(n) return v.Leave(newNode) } ```

Restore 方法用來根據 AST 語法樹還原出對應的 SQL 語句。 Accept 方法是方便其他工具遍歷這個 AST 語法樹，例如 TiDB 在預處理是會通過 AST 語法樹的 Accept 方法來遍歷 AST 語法樹中的所有節點。

Step-2：新增語法

LOAD DATA 語法是通過 LoadDataStmt 實現的，我們也照葫蘆畫瓢，在 tidb/parser/parser.y 中，新增 LoadSSTFileStmt 語法，這裡需要修改好幾處地方，下面用 git diff 展示修改：

```sql diff --git a/parser/parser.y b/parser/parser.y index 1539bb13db..079859e8a9 100644 --- a/parser/parser.y +++ b/parser/parser.y @@ -243,6 +243,7 @@ import ( sqlCalcFoundRows "SQL_CALC_FOUND_ROWS" sqlSmallResult "SQL_SMALL_RESULT" ssl "SSL" + sst "SST" starting "STARTING" statsExtended "STATS_EXTENDED" straightJoin "STRAIGHT_JOIN" @@ -908,6 +909,7 @@ import ( IndexAdviseStmt "INDEX ADVISE statement" KillStmt "Kill statement" LoadDataStmt "Load data statement" + LoadSSTFileStmt "Load sst file statement" LoadStatsStmt "Load statistic statement" LockTablesStmt "Lock tables statement" NonTransactionalDeleteStmt "Non-transactional delete statement" @@ -11324,6 +11326,7 @@ Statement: | IndexAdviseStmt | KillStmt | LoadDataStmt +| LoadSSTFileStmt | LoadStatsStmt | PlanReplayerStmt | PreparedStmt @@ -13496,6 +13499,14 @@ LoadDataStmt: $ = x }

+LoadSSTFileStmt: + "LOAD" "SST" "FILE" stringLit + { + $ = &ast.LoadSSTFileStmt{ + Path: $4, + } + } + ```

上面的修改中：

  • 第 9 行是因為語法中 SST 是一個新的關鍵字，所以需要註冊一個新的關鍵字。
  • 第 17 行 和 25 行是註冊一個新語法叫 LoadSSTFileStmt 。
  • 第 33 - 40 行是定義 LoadSSTFileStmt 語法結構為： LOAD SST FILE <file_path> ，這裡前 3 個關鍵字都是固定的，所以直接定義 "LOAD" "SST" "FILE" 即可，第 4 個是檔案路徑，一個變數值，我們用 stringLit 來提取這個變數的值，然後再用這個的值來初始化 ast.LoadSSTFileStmt ，其中 $4 是指第 4 個變數 stringLit 的值。

因為引入了新的關鍵字 SST ，所以還需要在 tidb/parser/misc.go 中新增這個關鍵字：

c++ diff --git a/parser/misc.go b/parser/misc.go index 140619bb07..418e9dd6a4 100644 --- a/parser/misc.go +++ b/parser/misc.go @@ -669,6 +669,7 @@ var tokenMap = map[string]int{ "SQL_TSI_YEAR": sqlTsiYear, "SQL": sql, "SSL": ssl, + "SST": sst, "STALENESS": staleness, "START": start, "STARTING": starting,

Step-3：編譯和測試

編譯生成新的 parser 檔案。

sql cd parser make fmt #格式化程式碼 make # 編譯生成新的 parser 檔案

我們可以在 tidb/parser/parser_test.go 檔案中的 TestDMLStmt 中新增一個測試，來驗證我們新增的語法生效了，下面是 git diff 展示的修改：

``c++ diff --git a/parser/parser_test.go b/parser/parser_test.go index 7093c3889f..d2c75c4c59 100644 --- a/parser/parser_test.go +++ b/parser/parser_test.go @@ -666,6 +666,9 @@ func TestDMLStmt(t *testing.T) { {"LOAD DATA LOCAL INFILE '/tmp/t.csv' IGNORE INTO TABLE t1 FIELDS TERMINATED BY ',' LINES TERMINATED BY '\n';", true, "LOAD DATA LOCAL INFILE '/tmp/t.csv' IGNORE INTO TABLEt1FIELDS TERMINATED BY ','"}, {"LOAD DATA LOCAL INFILE '/tmp/t.csv' REPLACE INTO TABLE t1 FIELDS TERMINATED BY ',' LINES TERMINATED BY '\n';", true, "LOAD DATA LOCAL INFILE '/tmp/t.csv' REPLACE INTO TABLEt1` FIELDS TERMINATED BY ','"},

  • // load sst file test
  • {"load sst file 'table0.sst'", true, "LOAD SST FILE 'table0.sst'"}, + ```

然後跑測試：

sql cd parser make test #跑 parser 的所有測試，快速驗證可以用 go test -run="TestDMLStmt" 命令只跑修改的 TestDMLStmt 測試

生成執行計劃

TiDB 在生成 AST 語法樹後，需要生成對應的執行計劃。我們需要先定義 LOAD SST FILE 的執行計劃。同樣的照葫蘆畫瓢，我們先在 tidb/planner/core/common_plans.go 檔案中找到 LOAD DATA 的執行計劃 LoadData , 然後開始畫瓢定義 LoadSSTFile 執行計劃：

```sql // LoadSSTFile represents a load sst file plan. type LoadSSTFile struct { baseSchemaProducer

    Path        string

} ```

為了讓 TiDB 能更具 ast.LoadSSTFileStmt 語法樹生成對應的 LoadSSTFile 執行計劃，

需要在 tidb/planner/core/planbuilder.go 檔案中，參考 buildLoadData 方法，來實現我們的 buildLoadSSTFile 方法，用來生成執行計劃, 下面是 git diff 展示修改內容：

```sql diff --git a/planner/core/planbuilder.go b/planner/core/planbuilder.go index ad7ce64748..c68e992b35 100644 --- a/planner/core/planbuilder.go +++ b/planner/core/planbuilder.go @@ -734,6 +734,8 @@ func (b PlanBuilder) Build(ctx context.Context, node ast.Node) (Plan, error) { return b.buildInsert(ctx, x) case ast.LoadDataStmt: return b.buildLoadData(ctx, x) + case ast.LoadSSTFileStmt: + return b.buildLoadSSTFile(x) @@ -3979,6 +3981,13 @@ func (b PlanBuilder) buildLoadData(ctx context.Context, ld *ast.LoadDataStmt) ( return p, nil }

+func (b PlanBuilder) buildLoadSSTFile(ld ast.LoadSSTFileStmt) (Plan, error) { + p := &LoadSSTFile{ + Path: ld.Path, + } + return p, nil +} + ```

構造 Executor 並執行

生成執行計劃之後，就需要構造對應的 Executor 然後執行了。TiDB 是用 Volcano 執行引擎，你可以將相關的初始化工作放在 Open 方法中，將主要功能的實現都放在 Next 方法中，以及執行完成後，在 Close 方法中執行相關的清理和釋放資源的操作。

我們需要先定義 LOAD SST FILE 的 Executor，並讓其實現 executor.Executor 介面，可以把相關定義放到 tidb/executor/executor.go 檔案中：

```sql // LoadSSTFileExec represents a load sst file executor. type LoadSSTFileExec struct { baseExecutor

path string done bool }

// Open implements the Executor Open interface. func (e *LoadSSTFileExec) Open(ctx context.Context) error { logutil.BgLogger().Warn("----- load sst file open, you can initialize some resource here") return nil }

// Next implements the Executor Next interface. func (e LoadSSTFileExec) Next(ctx context.Context, req chunk.Chunk) error { req.Reset() if e.done { return nil } e.done = true

logutil.BgLogger().Warn("----- load sst file exec", zap.String("file", e.path)) return nil }

// Close implements the Executor Close interface. func (e *LoadSSTFileExec) Close() error { logutil.BgLogger().Warn("----- load sst file close, you can release some resource here") return nil } ```

如果沒有初始化工作和清理工作，你也可以不用實現 Open 和 Close 方法，因為 baseExecutor 已經實現過了。

這裡為了簡化教程在 LoadSSTFileExec Executor 中僅僅是輸出了幾條 Log，你需要將自己功能具體實現的程式碼放在這裡。

然後為了讓 TiDB 能夠根據 LoadSSTFile 執行計劃來生成 LoadSSTFileExec Executor, 需要修改 tidb/executor/builder.go 檔案，下面是用 git diff 展示的修改：

```sql diff --git a/executor/builder.go b/executor/builder.go index 1154633bd5..4f0478daa6 100644 --- a/executor/builder.go +++ b/executor/builder.go @@ -199,6 +199,8 @@ func (b executorBuilder) build(p plannercore.Plan) Executor { return b.buildInsert(v) case plannercore.LoadData: return b.buildLoadData(v) + case plannercore.LoadSSTFile: + return b.buildLoadSSTFile(v) case plannercore.LoadStats: return b.buildLoadStats(v) case plannercore.IndexAdvise: @@ -944,6 +946,14 @@ func (b executorBuilder) buildLoadData(v *plannercore.LoadData) Executor { return loadDataExec }

+func (b executorBuilder) buildLoadSSTFile(v plannercore.LoadSSTFile) Executor { + e := &LoadSSTFileExec{ + baseExecutor: newBaseExecutor(b.ctx, nil, v.ID()), + path: v.Path, + } + return e +} + ```

驗證

到此，我們已經成功的在 TiDB 中新增了一個 “功能”， 我們可以編譯 TiDB 並啟動後驗證下：

sql make #編譯 TiDB server bin/tidb-server # 啟動一個 TiDB server

然後新起一個終端，用 mysql 客戶端連上去試試新功能：

```sql ▶ mysql -u root -h 127.0.0.1 -P 4000

mysql> load sst file 'table0.sst'; Query OK, 0 rows affected (0.00 sec) ```

可以看到執行成功了，並且在 tidb-server 的輸出日誌中，可以看到我們這個功能的 Executor 執行時的日誌輸出：

c++ [2022/09/19 15:24:02.745 +08:00] [WARN] [executor.go:2213] ["----- load sst file open, you can initialize some resource here"] [2022/09/19 15:24:02.745 +08:00] [WARN] [executor.go:2225] ["----- load sst file exec"] [file=table0.sst] [2022/09/19 15:24:02.745 +08:00] [WARN] [executor.go:2231] ["----- load sst file close, you can release some resource here"]

總結 本文的程式碼示例：https://github.com/pingcap/tidb/pull/37936/files

本文通過“照葫蘆畫瓢” 的方式，教你如何在 TiDB 中新增一個功能，但也忽略了一些細節，例如許可權檢查，新增完備的測試等等，希望能對讀者有所幫助。如果想要了解更多的知識背景和細節，推薦閱讀 TiDB Development Guide和 TiDB 原始碼閱讀部落格。

「」