TiDBのSQLパーサーを使ってみました

Go言語で開発をしていると少なからずデータの出し入れ作業があるので自動化したいなと思うことが時々あります。

OASやORMあたりのライブラリを利用すると入出力まわりはたいていあるのですが、自分で作ってみたいなと思いMySqlのCreate文から構造を取得した際にTiDBのパーサーをつかってみたので簡単に紹介させていただきます。

TiDBはGo言語で実装され、mysqlとの高い互換性がありパーサーの利用方法もシンプルでした。
TiDBのパーサーの利用方法はquickstartを実施してもらうとすぐに利用できるようになりますが、TableのCreate文から必要な要素を抽出した処理を簡単にですが紹介させていただきます。

### TiDB紹介

- TiDB
https://github.com/pingcap/tidb/blob/master/README.md

TiDB ("Ti" stands for Titanium) is an open-source NewSQL database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads. It is MySQL compatible and features horizontal scalability, strong consistency, and high availability.

- QuickStart
https://github.com/pingcap/tidb/blob/master/parser/docs/quickstart.md

## SQLパース

quickstart.md ではselect文のパースを行なっていますが、下記のcreate文のテーブル定義をGoの構造体に入れる例となります。
quickstart.md の`Parse SQL text`のコード (https://github.com/pingcap/tidb/blob/master/parser/docs/quickstart.md#parse-sql-text) に追加していくことで動くようになっています。
`Parse SQL text` をローカル環境で動くようにしておくと動作確認できるようになっています。

```
var crateTableStmt = `
CREATE TABLE users (
	id BIGINT NOT NULL AUTO_INCREMENT COMMENT 'ID',
	name VARCHAR(10) NULL COMMENT '名前',
	birthday DATE NULL COMMENT '誕生日',
	PRIMARY KEY(id)
)COMMENT 'ユーザ情報';
`
```

SQLをパースする実装はQuickStartのサンプルを引用させていただきます。関数にSQLを渡すとパース結果を取得することができます。
複数のSQLをまとめてパースすることができ配列で返却されます。今回は検証で1個のSQLのため配列の先頭を返します。
```
func parse(sql string) (*ast.StmtNode, error) {
	p := parser.New()
	// 第2引数：charset 第3引数：collation
	stmtNodes, _, err := p.Parse(sql, "", "") 
	if err != nil {
		return nil, err
	}
	return &stmtNodes[0], nil 
}
```

StmtNodeとしてcreate文の要素が保持されています。そこから必要な定義を取得するには Visitorインターフェースを実装します。(https://pkg.go.dev/github.com/pingcap/tidb/parser/ast#Visitor)
```
type Visitor interface {
	Enter(n Node) (node Node, skipChildren bool)
	Leave(n Node) (node Node, ok bool)
}
```

今回はCreate Tableの要素を取得するインターフェース実装となります。以下の構造体に対してEnter/Leaveを実装します。
```
// CraateTable DDLからテーブル名とカラムの情報を保持する構造体
type CreateTable struct {
	TableName string     // テーブル名
	Columns   []Column   // カラム
	CurrentColumn Column // 分析中のカラムの情報を保持する
}
// カラムの情報を保持する構造体
type Column struct {
	Name     string // カラム名
	Type     string // 型
	Nullable bool   // Null可
	Comment  string
}
```

CreateTableの構造体を渡すことでパースされた定義をセットしていくようなメソッドを実装します。
データの構成に沿ってEnter/Leaveメソッドが呼ばれます。今回のSQLだと以下のようなネストした順番となります。（実際は記載していない複数のノードの呼び出しがあります）

1\. Enter(テーブル:users)
2\. Enter(カラム:id) ... Enter(コメント:ID) Leave(コメント:ID) Leave(カラム:id)
3\. Enter(カラム:name) ... Leave(カラム:name)
4\. Enter(カラム:birthday) ... Leave(カラム:birthday)
5\. Leave(テーブル:users)

```
// Enterメソッド. 開始時に呼ばれる
func (v *CreateTable) Enter(node ast.Node) (ast.Node, bool) {
	// tableの抽出。型アサーションでマッチするNodeを判定します
	if tab, ok := node.(*ast.CreateTableStmt); ok {
		// .O はOriginal、.L はLower caseの文字列を返します
		v.TableName = tab.Table.Name.O 
	}
	// columnの抽出
	if col, ok := node.(*ast.ColumnDef); ok {
		v.CurrentColumn = Column{}
		// users.birthdayの場合 col.Name.Name.O=>birthday, col.Name.OrigColName()=>users.birthday
		v.CurrentColumn.Name = col.Name.Name.O
		// カラムの型
 		v.CurrentColumn.Type = col.Tp.String() 
	}
	// null/not nullの抽出
	if opt, ok := node.(*ast.ColumnOption); ok {
		if opt.Tp == ast.ColumnOptionNull {
			v.CurrentColumn.Nullable = true
		} else if opt.Tp == ast.ColumnOptionNotNull {
			v.CurrentColumn.Nullable = false
		}
	}
	// commentの抽出. ValueExprにキャストして値を取得します
	if opt, ok := node.(*ast.ColumnOption); ok && opt.Tp == ast.ColumnOptionComment {
		if val, ok := opt.Expr.(ast.ValueExpr); ok {
			v.CurrentColumn.Comment = val.GetDatumString()
		}
	}
	// bool は skipChildren でtrueを渡すことで子要素の処理をスキップすることができます。
	return node, false 
}

// Leaveメソッド. 終了時に呼ばれる
func (v *CreateTable) Leave(node ast.Node) (ast.Node, bool) {
	if _, ok := node.(*ast.ColumnDef); ok {
		// カラム定義が終了したら抽出した CurrentColumn を配列に追加
		v.Columns = append(v.Columns, v.CurrentColumn)
	}
	return node, true
}
```

では実行です。
```
func main(){
	stmtNode, err := parse(crateTableStmt) // SQLをパース
	if err != nil {
		panic(err)
	}
	createTable := CreateTable{} // Visitorインターフェースを実装している構造体
	createTable.Columns = []Column{}
	// Accept で Enter/Leaveが処理され、構造体に定義情報が設定されます。
	(*stmtNode).Accept(&createTable)
	fmt.Printf("%v\n", createTable) // 設定を出力
}
```

TiDBのパーサーの利用することでシンプルな実装でテーブル名、カラム名、型、null可、コメントを構造体に抽出することができました。

SQLを複数パースさせることもできます。
```
var crateTableStmt = `
CREATE TABLE users (
	id BIGINT NOT NULL AUTO_INCREMENT COMMENT 'ID'
)COMMENT 'ユーザ情報';

CREATE TABLE items (
	id BIGINT NOT NULL AUTO_INCREMENT COMMENT 'ID'
)COMMENT 'アイテム情報';
`
```
quickstart.md の`Parse SQL text`では最初のSQLを返すようにしているのでここを配列で返すようにすると複数のSQLを処理することもできます。
```
func parse(sql string) (*[]ast.StmtNode, error) {
	... // 省略
	return &stmtNodes, nil
}
```

## さいごに

SQLパーサーのライブラリは多くありますが、mysql互換で使ったことのないものを選んでみました。ライブラリを読むことで実装方法など学べプロダクトにも興味をもつことができました。